"""Screen-space global illumination pass (design D8, RM-E8).
Half-res, Hi-Z traced one-bounce indirect DIFFUSE that writes the design-D8
indirect-diffuse hook (set0 b17) so it never edits the uber shading math. It is
the diffuse sibling of :class:`SSRPass` (RM-E4) and deliberately reuses its
screen-trace machinery: the same MIN-Z depth pyramid (``ssr_hiz.comp`` via
:class:`HiZPass`), the same scene-colour mip chain (A12 mipgen), and the same
thin G-buffer read (A9). Each frame, after the HDR pass:
1. the HDR colour is copied into a small mip chain so the gather reads a slightly
blurred bounce (shared idiom with SSR);
2. a MIN-Z depth pyramid is built with the :class:`HiZPass` machinery;
3. a compute trace (``ssgi_trace.comp``) casts a cosine hemisphere of rays per
half-res pixel, marches each through screen space, and gathers the on-screen
scene colour it hits: one bounce of indirect diffuse, written noisy to a
half-res raw target;
4. a temporal filter (``ssgi_temporal.comp``) reprojects the accumulated history
through the previous view-projection and blends the fresh trace in (TAA-style
denoise), writing the settled result to the b17 target; the driver copies that
back into the history image for next frame.
The uber shader samples b17 and composites the bounce OVER the flat/IBL ambient
(RM-D8fix); a ray miss / open direction contributes 0 there so the ambient
fallback is untouched. Like SSR the result is one frame late and converges within
a few frames for a static view (the golden captures a settled frame).
Zero-cost when unused: the whole pass is lazily created the first frame SSGI is
active (mirroring SSAO/SSR/velocity), and the uber's ``indirect_diffuse_enabled``
gate stays 0 otherwise, so b17 is never sampled and feature-off frames are
byte-identical.
"""
import logging
from typing import Any
import numpy as np
import vulkan as vk
from ..gpu.descriptors import (
DescriptorWriteBatch,
allocate_descriptor_set,
create_descriptor_set_layout,
create_pool_for_types,
)
from ..gpu.memory import (
_record_mip_chain,
begin_single_time_commands,
create_image,
end_single_time_commands,
mip_chain_length,
)
from ..gpu.pipeline_compute import create_compute_pipeline
from .hiz_pass import HiZPass
__all__ = ["SSGIPass"]
log = logging.getLogger(__name__)
# Trace PC: mat4 proj(64) + inv_proj(64) + view(64) + vec4 params(16) + vec4 res(16).
# Temporal PC: mat4 inv_proj(64) + inv_view(64) + prev_view_proj(64) + params(16) + res(16).
# Both are 224 bytes; one push-constant range size covers both pipelines.
_PC_SIZE = 224
_COLOUR_FORMAT = vk.VK_FORMAT_R16G16B16A16_SFLOAT
_GI_FORMAT = vk.VK_FORMAT_R16G16B16A16_SFLOAT
# Scene-colour mip chain capped at 5 levels (LOD 0..4), like SSR (design D7).
_MAX_SCENE_MIPS = 5
# History feedback cap: the steady-state blend weight once the running-average
# ramp is over. Higher = smoother but slower to forget a perturbation, which the
# GI feedback loop (the trace reads the post-GI HDR colour) amplifies into slow
# cross-run drift. 0.8 forgets a perturbation in a handful of frames (0.8^N decays
# fast) so a settled view converges deterministically, at the cost of a little
# more per-frame grain (denoised by NUM_RAYS). Ghosting under motion accepted (D8).
_FEEDBACK_MAX = 0.8
[docs]
class SSGIPass:
"""Half-res Hi-Z screen-space GI into the D8 indirect-diffuse hook."""
def __init__(self, engine: Any) -> None:
self._engine = engine
self._ready = False
self._hiz: HiZPass | None = None
# Scene-colour mip chain (copied from the HDR colour each frame).
self._scene_image: Any = None
self._scene_memory: Any = None
self._scene_view: Any = None
self._scene_mips: int = 1
# Half-res targets (all rgba16f): raw trace, accumulated history, output.
self._raw_image: Any = None
self._raw_memory: Any = None
self._raw_view: Any = None
self._hist_image: Any = None
self._hist_memory: Any = None
self._hist_view: Any = None
self._out_image: Any = None
self._out_memory: Any = None
self._out_view: Any = None
# Samplers.
self._scene_sampler: Any = None
self._gbuffer_sampler: Any = None
self._hist_sampler: Any = None
self._output_sampler: Any = None
# Trace pipeline + descriptors.
self._trace_pipeline: Any = None
self._trace_layout: Any = None
self._trace_module: Any = None
# Temporal pipeline + descriptors.
self._temporal_pipeline: Any = None
self._temporal_layout: Any = None
self._temporal_module: Any = None
self._desc_pool: Any = None
self._trace_desc_layout: Any = None
self._trace_desc_set: Any = None
self._temporal_desc_layout: Any = None
self._temporal_desc_set: Any = None
# External references (owned by the HDR target).
self._hdr_colour_image: Any = None
self._depth_view: Any = None
self._depth_image: Any = None
self._gbuffer_view: Any = None
self._gbuffer_image: Any = None
self._width = 0
self._height = 0
# Temporal state: previous world->clip and accumulated valid-frame count.
self._prev_view_proj: np.ndarray | None = None
self._frame_count: int = 0
# Public tuning (driven by env-sync / QualitySettings).
self.enabled: bool = True
self.intensity: float = 1.0
self.max_distance: float = 8.0
self.thickness: float = 0.6
# ------------------------------------------------------------------ views
[docs]
@property
def output_view(self) -> Any:
"""The half-res GI target view (bound to set0 b17 by the renderer)."""
return self._out_view
[docs]
@property
def output_sampler(self) -> Any:
"""Linear/clamp sampler for reading the GI target at b17."""
return self._output_sampler
# ------------------------------------------------------------------ setup
[docs]
def setup(
self,
width: int,
height: int,
hdr_colour_image: Any,
depth_view: Any,
depth_image: Any,
gbuffer_view: Any,
gbuffer_image: Any = None,
) -> None:
"""Allocate the pyramid, scene chain, raw/history/output targets and pipelines."""
self._width = width
self._height = height
self._hdr_colour_image = hdr_colour_image
self._depth_view = depth_view
self._depth_image = depth_image
self._gbuffer_view = gbuffer_view
self._gbuffer_image = gbuffer_image
self._hiz = HiZPass(self._engine)
self._hiz.setup(width, height, depth_view, depth_image, shader="ssr_hiz.comp")
self._create_samplers()
self._create_scene_chain(width, height)
self._create_targets(width, height)
self._create_descriptors()
self._create_pipelines()
self._prev_view_proj = None
self._frame_count = 0
self._ready = True
log.debug("SSGI pass initialised (%dx%d, half-res trace + temporal)", width, height)
def _create_samplers(self) -> None:
device = self._engine.ctx.device
clamp = vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE
self._scene_sampler = vk.vkCreateSampler(
device,
vk.VkSamplerCreateInfo(
magFilter=vk.VK_FILTER_LINEAR,
minFilter=vk.VK_FILTER_LINEAR,
mipmapMode=vk.VK_SAMPLER_MIPMAP_MODE_LINEAR,
addressModeU=clamp,
addressModeV=clamp,
addressModeW=clamp,
minLod=0.0,
maxLod=float(_MAX_SCENE_MIPS),
),
None,
)
self._gbuffer_sampler = vk.vkCreateSampler(
device,
vk.VkSamplerCreateInfo(
magFilter=vk.VK_FILTER_NEAREST,
minFilter=vk.VK_FILTER_NEAREST,
mipmapMode=vk.VK_SAMPLER_MIPMAP_MODE_NEAREST,
addressModeU=clamp,
addressModeV=clamp,
addressModeW=clamp,
),
None,
)
self._hist_sampler = vk.vkCreateSampler(
device,
vk.VkSamplerCreateInfo(
magFilter=vk.VK_FILTER_LINEAR,
minFilter=vk.VK_FILTER_LINEAR,
mipmapMode=vk.VK_SAMPLER_MIPMAP_MODE_NEAREST,
addressModeU=clamp,
addressModeV=clamp,
addressModeW=clamp,
minLod=0.0,
maxLod=0.0,
),
None,
)
self._output_sampler = vk.vkCreateSampler(
device,
vk.VkSamplerCreateInfo(
magFilter=vk.VK_FILTER_LINEAR,
minFilter=vk.VK_FILTER_LINEAR,
mipmapMode=vk.VK_SAMPLER_MIPMAP_MODE_NEAREST,
addressModeU=clamp,
addressModeV=clamp,
addressModeW=clamp,
minLod=0.0,
maxLod=0.0,
),
None,
)
def _create_scene_chain(self, width: int, height: int) -> None:
device = self._engine.ctx.device
phys = self._engine.ctx.physical_device
self._scene_mips = max(1, min(_MAX_SCENE_MIPS, mip_chain_length(width, height)))
self._scene_image, self._scene_memory = create_image(
device,
phys,
width,
height,
_COLOUR_FORMAT,
vk.VK_IMAGE_USAGE_SAMPLED_BIT | vk.VK_IMAGE_USAGE_TRANSFER_DST_BIT | vk.VK_IMAGE_USAGE_TRANSFER_SRC_BIT,
mip_levels=self._scene_mips,
)
self._scene_view = vk.vkCreateImageView(
device,
vk.VkImageViewCreateInfo(
image=self._scene_image,
viewType=vk.VK_IMAGE_VIEW_TYPE_2D,
format=_COLOUR_FORMAT,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=self._scene_mips,
baseArrayLayer=0,
layerCount=1,
),
),
None,
)
def _make_target(self, width: int, height: int, usage: int, layout: int) -> tuple[Any, Any, Any]:
"""Create a half-res GI target, zero-clear it (so no pass ever reads
undefined float16 memory: an unwritten history texel is a defined 0, not
a NaN that ``mix``/reprojection would propagate), and leave it in *layout*.
Requires TRANSFER_DST usage, which all three GI targets carry."""
device = self._engine.ctx.device
phys = self._engine.ctx.physical_device
image, memory = create_image(
device,
phys,
width,
height,
_GI_FORMAT,
usage | vk.VK_IMAGE_USAGE_TRANSFER_DST_BIT,
)
view = vk.vkCreateImageView(
device,
vk.VkImageViewCreateInfo(
image=image,
viewType=vk.VK_IMAGE_VIEW_TYPE_2D,
format=_GI_FORMAT,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=1,
),
),
None,
)
rng = vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=1,
)
cmd = begin_single_time_commands(device, self._engine.ctx.command_pool)
def barrier(old, new, src, dst, src_stage, dst_stage):
b = vk.VkImageMemoryBarrier(
srcAccessMask=src,
dstAccessMask=dst,
oldLayout=old,
newLayout=new,
srcQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
dstQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
image=image,
subresourceRange=rng,
)
vk.vkCmdPipelineBarrier(cmd, src_stage, dst_stage, 0, 0, None, 0, None, 1, [b])
barrier(
vk.VK_IMAGE_LAYOUT_UNDEFINED,
vk.VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
0,
vk.VK_ACCESS_TRANSFER_WRITE_BIT,
vk.VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
)
vk.vkCmdClearColorImage(
cmd,
image,
vk.VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
vk.VkClearColorValue(float32=[0.0, 0.0, 0.0, 0.0]),
1,
[rng],
)
barrier(
vk.VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
layout,
vk.VK_ACCESS_TRANSFER_WRITE_BIT,
vk.VK_ACCESS_SHADER_READ_BIT,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
vk.VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,
)
end_single_time_commands(device, self._engine.ctx.graphics_queue, self._engine.ctx.command_pool, cmd)
return image, memory, view
def _create_targets(self, width: int, height: int) -> None:
hw, hh = max(1, width // 2), max(1, height // 2)
# Raw trace target: storage-written by the trace, storage-read by temporal.
self._raw_image, self._raw_memory, self._raw_view = self._make_target(
hw,
hh,
vk.VK_IMAGE_USAGE_STORAGE_BIT,
vk.VK_IMAGE_LAYOUT_GENERAL,
)
# History: sampled by temporal (reprojected), copy-dst from the output.
self._hist_image, self._hist_memory, self._hist_view = self._make_target(
hw,
hh,
vk.VK_IMAGE_USAGE_SAMPLED_BIT | vk.VK_IMAGE_USAGE_TRANSFER_DST_BIT,
vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
)
# Output (b17): storage-written by temporal, sampled by the uber (GENERAL,
# the SSR/SSAO precedent), copy-src back into history.
self._out_image, self._out_memory, self._out_view = self._make_target(
hw,
hh,
vk.VK_IMAGE_USAGE_STORAGE_BIT | vk.VK_IMAGE_USAGE_SAMPLED_BIT | vk.VK_IMAGE_USAGE_TRANSFER_SRC_BIT,
vk.VK_IMAGE_LAYOUT_GENERAL,
)
def _create_descriptors(self) -> None:
device = self._engine.ctx.device
cs = vk.VK_SHADER_STAGE_COMPUTE_BIT
cis = vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
si = vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE
self._trace_desc_layout = create_descriptor_set_layout(
device,
[
(0, cis, cs, 1), # min-Z pyramid
(1, cis, cs, 1), # gbuffer
(2, cis, cs, 1), # scene colour mips
(3, si, cs, 1), # raw output
],
)
self._temporal_desc_layout = create_descriptor_set_layout(
device,
[
(0, si, cs, 1), # raw trace (readonly)
(1, cis, cs, 1), # history
(2, cis, cs, 1), # current depth (min-Z pyramid LOD0)
(3, si, cs, 1), # b17 output
],
)
self._desc_pool = create_pool_for_types(device, {cis: 5, si: 3}, max_sets=2)
self._trace_desc_set = allocate_descriptor_set(device, self._desc_pool, self._trace_desc_layout)
self._temporal_desc_set = allocate_descriptor_set(device, self._desc_pool, self._temporal_desc_layout)
self._write_descriptors()
def _write_descriptors(self) -> None:
gen = vk.VK_IMAGE_LAYOUT_GENERAL
with DescriptorWriteBatch(self._engine.ctx.device) as b:
b.image(self._trace_desc_set, 0, self._hiz.sampled_view, self._hiz.sampler, image_layout=gen)
b.image(self._trace_desc_set, 1, self._gbuffer_view, self._gbuffer_sampler)
b.image(self._trace_desc_set, 2, self._scene_view, self._scene_sampler)
b.storage_image(self._trace_desc_set, 3, self._raw_view)
b.storage_image(self._temporal_desc_set, 0, self._raw_view)
b.image(self._temporal_desc_set, 1, self._hist_view, self._hist_sampler)
b.image(self._temporal_desc_set, 2, self._hiz.sampled_view, self._hiz.sampler, image_layout=gen)
b.storage_image(self._temporal_desc_set, 3, self._out_view)
def _create_pipelines(self) -> None:
e = self._engine
self._trace_pipeline, self._trace_layout, self._trace_module = create_compute_pipeline(
e.ctx.device,
e.shader_dir / "ssgi_trace.comp",
[self._trace_desc_layout],
_PC_SIZE,
)
self._temporal_pipeline, self._temporal_layout, self._temporal_module = create_compute_pipeline(
e.ctx.device,
e.shader_dir / "ssgi_temporal.comp",
[self._temporal_desc_layout],
_PC_SIZE,
)
# ----------------------------------------------------------------- render
[docs]
def render(self, cmd: Any, proj_matrix: np.ndarray, view_matrix: np.ndarray) -> None:
"""Copy scene colour, build the min-Z pyramid, trace + temporally denoise.
Call after ``end_hdr_pass`` (HDR colour SHADER_READ, depth DEPTH_RO,
gbuffer SHADER_READ), outside any render pass. Mirrors SSRPass.render.
"""
if not self._ready or not self.enabled:
return
self._copy_scene_colour(cmd)
self._hiz.render(cmd)
self._trace(cmd, proj_matrix, view_matrix)
self._temporal(cmd, proj_matrix, view_matrix)
# Advance temporal state for next frame.
self._prev_view_proj = np.ascontiguousarray(proj_matrix @ view_matrix, dtype=np.float32)
self._frame_count += 1
def _copy_scene_colour(self, cmd: Any) -> None:
"""Copy the HDR colour into mip0 and generate the gather mip chain (SSR idiom)."""
w, h = self._width, self._height
c_aspect = vk.VK_IMAGE_ASPECT_COLOR_BIT
RO = vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
TSRC = vk.VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL
TDST = vk.VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
UNDEF = vk.VK_IMAGE_LAYOUT_UNDEFINED
def barrier(image, old, new, src_access, dst_access, base_mip=0, mips=1):
return vk.VkImageMemoryBarrier(
srcAccessMask=src_access,
dstAccessMask=dst_access,
oldLayout=old,
newLayout=new,
srcQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
dstQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
image=image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=c_aspect,
baseMipLevel=base_mip,
levelCount=mips,
baseArrayLayer=0,
layerCount=1,
),
)
SR = vk.VK_ACCESS_SHADER_READ_BIT
TR = vk.VK_ACCESS_TRANSFER_READ_BIT
TW = vk.VK_ACCESS_TRANSFER_WRITE_BIT
vk.vkCmdPipelineBarrier(
cmd,
vk.VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
0,
0,
None,
0,
None,
2,
[
barrier(self._hdr_colour_image, RO, TSRC, SR, TR),
barrier(self._scene_image, UNDEF, TDST, 0, TW, mips=self._scene_mips),
],
)
layers = vk.VkImageSubresourceLayers(aspectMask=c_aspect, mipLevel=0, baseArrayLayer=0, layerCount=1)
region = vk.VkImageCopy(
srcSubresource=layers,
srcOffset=vk.VkOffset3D(x=0, y=0, z=0),
dstSubresource=layers,
dstOffset=vk.VkOffset3D(x=0, y=0, z=0),
extent=vk.VkExtent3D(width=w, height=h, depth=1),
)
vk.vkCmdCopyImage(cmd, self._hdr_colour_image, TSRC, self._scene_image, TDST, 1, [region])
# Restore HDR colour for the downstream bloom/TAA/tonemap samplers.
vk.vkCmdPipelineBarrier(
cmd,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
vk.VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
0,
0,
None,
0,
None,
1,
[barrier(self._hdr_colour_image, TSRC, RO, TR, SR)],
)
if self._scene_mips > 1:
_record_mip_chain(cmd, self._scene_image, w, h, self._scene_mips)
else:
vk.vkCmdPipelineBarrier(
cmd,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
0,
0,
None,
0,
None,
1,
[barrier(self._scene_image, TDST, RO, TW, SR)],
)
def _push(self, cmd: Any, layout: Any, data: bytes) -> None:
cbuf = vk.ffi.new("char[]", data)
vk._vulkan.lib.vkCmdPushConstants(cmd, layout, vk.VK_SHADER_STAGE_COMPUTE_BIT, 0, _PC_SIZE, cbuf)
def _trace(self, cmd: Any, proj_matrix: np.ndarray, view_matrix: np.ndarray) -> None:
hw, hh = max(1, self._width // 2), max(1, self._height // 2)
# Gbuffer colour-attachment write -> compute read (layout already SHADER_READ).
if self._gbuffer_image is not None:
gbar = vk.VkImageMemoryBarrier(
srcAccessMask=vk.VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT,
dstAccessMask=vk.VK_ACCESS_SHADER_READ_BIT,
oldLayout=vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
newLayout=vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
srcQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
dstQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
image=self._gbuffer_image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=1,
),
)
vk.vkCmdPipelineBarrier(
cmd,
vk.VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
0,
0,
None,
0,
None,
1,
[gbar],
)
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._trace_pipeline)
vk.vkCmdBindDescriptorSets(
cmd,
vk.VK_PIPELINE_BIND_POINT_COMPUTE,
self._trace_layout,
0,
1,
[self._trace_desc_set],
0,
None,
)
proj_t = np.ascontiguousarray(proj_matrix.T, dtype=np.float32)
inv_proj_t = np.ascontiguousarray(np.linalg.inv(proj_matrix).T, dtype=np.float32)
view_t = np.ascontiguousarray(view_matrix.T, dtype=np.float32)
params = np.array(
[self.intensity, self.max_distance, float(self._frame_count & 1023), self.thickness],
dtype=np.float32,
)
res = np.array(
[float(self._width), float(self._height), 1.0 / self._width, 1.0 / self._height],
dtype=np.float32,
)
self._push(
cmd,
self._trace_layout,
proj_t.tobytes() + inv_proj_t.tobytes() + view_t.tobytes() + params.tobytes() + res.tobytes(),
)
vk.vkCmdDispatch(cmd, (hw + 7) // 8, (hh + 7) // 8, 1)
# Trace write -> temporal read.
self._image_barrier(
cmd,
self._raw_image,
vk.VK_IMAGE_LAYOUT_GENERAL,
vk.VK_IMAGE_LAYOUT_GENERAL,
vk.VK_ACCESS_SHADER_WRITE_BIT,
vk.VK_ACCESS_SHADER_READ_BIT,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
)
def _temporal(self, cmd: Any, proj_matrix: np.ndarray, view_matrix: np.ndarray) -> None:
hw, hh = max(1, self._width // 2), max(1, self._height // 2)
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._temporal_pipeline)
vk.vkCmdBindDescriptorSets(
cmd,
vk.VK_PIPELINE_BIND_POINT_COMPUTE,
self._temporal_layout,
0,
1,
[self._temporal_desc_set],
0,
None,
)
inv_proj_t = np.ascontiguousarray(np.linalg.inv(proj_matrix).T, dtype=np.float32)
inv_view_t = np.ascontiguousarray(np.linalg.inv(view_matrix).T, dtype=np.float32)
if self._prev_view_proj is None:
prev_vp = proj_matrix @ view_matrix
frame_count = 0.0
else:
prev_vp = self._prev_view_proj
frame_count = float(self._frame_count)
prev_vp_t = np.ascontiguousarray(prev_vp.T, dtype=np.float32)
params = np.array([frame_count, _FEEDBACK_MAX, 0.0, 0.0], dtype=np.float32)
res = np.array(
[float(self._width), float(self._height), 1.0 / self._width, 1.0 / self._height],
dtype=np.float32,
)
self._push(
cmd,
self._temporal_layout,
inv_proj_t.tobytes() + inv_view_t.tobytes() + prev_vp_t.tobytes() + params.tobytes() + res.tobytes(),
)
vk.vkCmdDispatch(cmd, (hw + 7) // 8, (hh + 7) // 8, 1)
# Temporal write -> (a) fragment read at b17 next frame, (b) copy src now.
self._image_barrier(
cmd,
self._out_image,
vk.VK_IMAGE_LAYOUT_GENERAL,
vk.VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
vk.VK_ACCESS_SHADER_WRITE_BIT,
vk.VK_ACCESS_TRANSFER_READ_BIT,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
)
self._image_barrier(
cmd,
self._hist_image,
vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
vk.VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
vk.VK_ACCESS_SHADER_READ_BIT,
vk.VK_ACCESS_TRANSFER_WRITE_BIT,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
)
# Copy the settled output into history for next frame's reprojection.
layers = vk.VkImageSubresourceLayers(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
mipLevel=0,
baseArrayLayer=0,
layerCount=1,
)
region = vk.VkImageCopy(
srcSubresource=layers,
srcOffset=vk.VkOffset3D(x=0, y=0, z=0),
dstSubresource=layers,
dstOffset=vk.VkOffset3D(x=0, y=0, z=0),
extent=vk.VkExtent3D(width=hw, height=hh, depth=1),
)
vk.vkCmdCopyImage(
cmd,
self._out_image,
vk.VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
self._hist_image,
vk.VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
1,
[region],
)
# Restore layouts: output -> GENERAL for the uber b17 sample; history ->
# SHADER_READ for next frame's temporal reprojection.
self._image_barrier(
cmd,
self._out_image,
vk.VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
vk.VK_IMAGE_LAYOUT_GENERAL,
vk.VK_ACCESS_TRANSFER_READ_BIT,
vk.VK_ACCESS_SHADER_READ_BIT,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
vk.VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
)
self._image_barrier(
cmd,
self._hist_image,
vk.VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
vk.VK_ACCESS_TRANSFER_WRITE_BIT,
vk.VK_ACCESS_SHADER_READ_BIT,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
)
@staticmethod
def _image_barrier(cmd, image, old, new, src_access, dst_access, src_stage, dst_stage) -> None:
bar = vk.VkImageMemoryBarrier(
srcAccessMask=src_access,
dstAccessMask=dst_access,
oldLayout=old,
newLayout=new,
srcQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
dstQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
image=image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=1,
),
)
vk.vkCmdPipelineBarrier(cmd, src_stage, dst_stage, 0, 0, None, 0, None, 1, [bar])
# ----------------------------------------------------------------- resize
[docs]
def resize(
self,
width: int,
height: int,
hdr_colour_image: Any,
depth_view: Any,
depth_image: Any,
gbuffer_view: Any,
gbuffer_image: Any = None,
) -> None:
"""Recreate the scene chain + targets + pyramid for a new internal extent."""
if not self._ready:
return
self._width = width
self._height = height
self._hdr_colour_image = hdr_colour_image
self._depth_view = depth_view
self._depth_image = depth_image
self._gbuffer_view = gbuffer_view
self._gbuffer_image = gbuffer_image
self._destroy_targets()
self._hiz.resize(width, height, depth_view, depth_image)
self._create_scene_chain(width, height)
self._create_targets(width, height)
self._write_descriptors()
self._prev_view_proj = None
self._frame_count = 0
[docs]
def set_gbuffer(self, gbuffer_view: Any, gbuffer_image: Any) -> None:
"""Re-point the G-buffer inputs (after a G-buffer toggle rebuild)."""
self._gbuffer_view = gbuffer_view
self._gbuffer_image = gbuffer_image
if self._ready:
self._write_descriptors()
# ---------------------------------------------------------------- cleanup
def _destroy_targets(self) -> None:
device = self._engine.ctx.device
for view, img, mem in (
(self._scene_view, self._scene_image, self._scene_memory),
(self._raw_view, self._raw_image, self._raw_memory),
(self._hist_view, self._hist_image, self._hist_memory),
(self._out_view, self._out_image, self._out_memory),
):
if view is not None:
vk.vkDestroyImageView(device, view, None)
if img is not None:
vk.vkDestroyImage(device, img, None)
if mem is not None:
vk.vkFreeMemory(device, mem, None)
self._scene_view = self._scene_image = self._scene_memory = None
self._raw_view = self._raw_image = self._raw_memory = None
self._hist_view = self._hist_image = self._hist_memory = None
self._out_view = self._out_image = self._out_memory = None
[docs]
def cleanup(self) -> None:
if not self._ready:
return
device = self._engine.ctx.device
if self._hiz is not None:
self._hiz.cleanup()
self._hiz = None
for pipe in (self._trace_pipeline, self._temporal_pipeline):
if pipe:
vk.vkDestroyPipeline(device, pipe, None)
for lay in (self._trace_layout, self._temporal_layout):
if lay:
vk.vkDestroyPipelineLayout(device, lay, None)
for mod in (self._trace_module, self._temporal_module):
if mod:
vk.vkDestroyShaderModule(device, mod, None)
if self._desc_pool:
vk.vkDestroyDescriptorPool(device, self._desc_pool, None)
for dl in (self._trace_desc_layout, self._temporal_desc_layout):
if dl:
vk.vkDestroyDescriptorSetLayout(device, dl, None)
for s in (self._scene_sampler, self._gbuffer_sampler, self._hist_sampler, self._output_sampler):
if s:
vk.vkDestroySampler(device, s, None)
self._scene_sampler = self._gbuffer_sampler = self._hist_sampler = self._output_sampler = None
self._destroy_targets()
self._ready = False