"""Screen-space reflections pass (design D7, RM-E4).
Hi-Z traced, half-resolution SSR that writes reflected radiance to the design-D8
indirect-specular hook (set0 b16) so it never edits the uber shading math. Each
frame, after the HDR pass:
1. the HDR colour is copied into a small mip chain (A12 mipgen) so the trace can
read a roughness-blurred reflection;
2. a MIN-Z depth pyramid is built with the :class:`HiZPass` machinery
(``ssr_hiz.comp``) for conservative empty-space skipping;
3. a compute trace (``ssr_trace.comp``) marches the reflected view ray through
screen space, reading the thin G-buffer (A9: world normal + roughness) and
the pyramid, and writes the reflected colour (with screen-edge + distance
fade) to a half-res RGBA16F target.
The uber shader samples that target at b16 and applies the Fresnel weight; a ray
miss writes 0 there so the probe/IBL/flat ambient (D8 fallback) is untouched.
Because the trace reads the HDR colour produced by the forward pass, the
reflection is one frame late; it converges within a couple of frames for a
static view (the golden captures a settled frame).
Zero-cost when unused: the whole pass is lazily created the first frame SSR is
active (mirroring SSAO/velocity), and the uber's ``indirect_specular_enabled``
gate stays 0 otherwise, so b16 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, create_image, mip_chain_length
from ..gpu.pipeline_compute import create_compute_pipeline
from .hiz_pass import HiZPass
__all__ = ["SSRPass"]
log = logging.getLogger(__name__)
# Push constant: mat4 proj(64) + mat4 inv_proj(64) + mat4 view(64) + vec4 params(16)
# + vec4 res(16) = 224 bytes.
_PC_SIZE = 224
# HDR colour format (matches the HDR target / scene-copy colour): required for the
# vkCmdCopyImage into mip0 and the linear-blit mip generation.
_COLOUR_FORMAT = vk.VK_FORMAT_R16G16B16A16_SFLOAT
_SSR_FORMAT = vk.VK_FORMAT_R16G16B16A16_SFLOAT
# Scene-colour mip chain is capped at 5 levels (LOD 0..4), matching the IBL
# prefilter's ``roughness * 4`` addressing (design D7).
_MAX_SCENE_MIPS = 5
[docs]
class SSRPass:
"""Half-res Hi-Z screen-space reflections into the D8 indirect-specular hook."""
def __init__(self, engine: Any) -> None:
self._engine = engine
self._ready = False
# MIN-Z depth pyramid (reuses the HiZ machinery with ssr_hiz.comp).
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 # full-chain sampled view
self._scene_mips: int = 1
# Half-res SSR output (rgba16f, GENERAL layout, sampled at b16).
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._output_sampler: Any = None
# Trace pipeline + descriptors.
self._pipeline: Any = None
self._layout: Any = None
self._module: Any = None
self._desc_pool: Any = None
self._desc_layout: Any = None
self._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
# Dimensions (internal HDR-chain extent).
self._width = 0
self._height = 0
# Public tuning (driven by env-sync / QualitySettings).
self.enabled: bool = True
self.intensity: float = 1.0
self.max_distance: float = 40.0
self.roughness_cutoff: float = 0.6
self.thickness: float = 1.0
# ------------------------------------------------------------------ views
[docs]
@property
def output_view(self) -> Any:
"""The half-res reflection target view (bound to set0 b16 by the renderer)."""
return self._out_view
[docs]
@property
def output_sampler(self) -> Any:
"""Linear/clamp sampler for reading the reflection target at b16."""
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-colour mip chain, output target and pipeline."""
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_output(width, height)
self._create_descriptors()
self._create_pipeline()
self._ready = True
log.debug("SSR pass initialised (%dx%d, half-res trace, %d scene mips)", width, height, self._scene_mips)
def _create_samplers(self) -> None:
device = self._engine.ctx.device
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=vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
addressModeV=vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
addressModeW=vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
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=vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
addressModeV=vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
addressModeW=vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
), 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=vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
addressModeV=vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
addressModeW=vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
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 _create_output(self, width: int, height: int) -> None:
device = self._engine.ctx.device
phys = self._engine.ctx.physical_device
hw, hh = max(1, width // 2), max(1, height // 2)
self._out_image, self._out_memory = create_image(
device, phys, hw, hh, _SSR_FORMAT,
vk.VK_IMAGE_USAGE_STORAGE_BIT | vk.VK_IMAGE_USAGE_SAMPLED_BIT,
)
self._out_view = vk.vkCreateImageView(device, vk.VkImageViewCreateInfo(
image=self._out_image, viewType=vk.VK_IMAGE_VIEW_TYPE_2D, format=_SSR_FORMAT,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0, levelCount=1, baseArrayLayer=0, layerCount=1,
),
), None)
from ..gpu.memory import transition_image_layout
# GENERAL permanently: written by the trace (storage) and sampled at b16
# (the SSAO ao_view precedent samples a GENERAL storage image directly).
transition_image_layout(
device, self._engine.ctx.graphics_queue, self._engine.ctx.command_pool,
self._out_image, vk.VK_IMAGE_LAYOUT_UNDEFINED, vk.VK_IMAGE_LAYOUT_GENERAL,
)
def _create_descriptors(self) -> None:
device = self._engine.ctx.device
cs = vk.VK_SHADER_STAGE_COMPUTE_BIT
self._desc_layout = create_descriptor_set_layout(device, [
(0, vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, cs, 1), # min-Z pyramid
(1, vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, cs, 1), # gbuffer
(2, vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, cs, 1), # scene colour mips
(3, vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, cs, 1), # ssr output
])
self._desc_pool = create_pool_for_types(device, {
vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER: 3,
vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE: 1,
})
self._desc_set = allocate_descriptor_set(device, self._desc_pool, self._desc_layout)
self._write_descriptors()
def _write_descriptors(self) -> None:
with DescriptorWriteBatch(self._engine.ctx.device) as b:
b.image(
self._desc_set, 0, self._hiz.sampled_view, self._hiz.sampler,
image_layout=vk.VK_IMAGE_LAYOUT_GENERAL,
)
b.image(self._desc_set, 1, self._gbuffer_view, self._gbuffer_sampler)
b.image(self._desc_set, 2, self._scene_view, self._scene_sampler)
b.storage_image(self._desc_set, 3, self._out_view)
def _create_pipeline(self) -> None:
e = self._engine
self._pipeline, self._layout, self._module = create_compute_pipeline(
e.ctx.device, e.shader_dir / "ssr_trace.comp", [self._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, and trace reflections.
Call after ``end_hdr_pass`` (HDR colour SHADER_READ, depth DEPTH_RO,
gbuffer SHADER_READ), outside any render pass.
"""
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)
def _copy_scene_colour(self, cmd: Any) -> None:
"""Copy the HDR colour into mip0 and generate the reflection mip chain."""
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
# HDR colour -> TRANSFER_SRC; whole scene chain (contents discarded) -> TRANSFER_DST.
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)])
# Generate the reflection mip chain (mip0 filled, all levels TRANSFER_DST
# on entry; leaves every level SHADER_READ_ONLY).
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 _trace(self, cmd: Any, proj_matrix: np.ndarray, view_matrix: np.ndarray) -> None:
ffi = vk.ffi
hw, hh = max(1, self._width // 2), max(1, self._height // 2)
# Make the gbuffer's colour-attachment write visible to the compute read
# (layout already SHADER_READ from the HDR pass; execution/memory only).
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._pipeline)
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._layout, 0, 1, [self._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, self.roughness_cutoff, 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,
)
pc_data = proj_t.tobytes() + inv_proj_t.tobytes() + view_t.tobytes() + params.tobytes() + res.tobytes()
cbuf = ffi.new("char[]", pc_data)
vk._vulkan.lib.vkCmdPushConstants(cmd, self._layout, vk.VK_SHADER_STAGE_COMPUTE_BIT, 0, _PC_SIZE, cbuf)
vk.vkCmdDispatch(cmd, (hw + 7) // 8, (hh + 7) // 8, 1)
# Trace write -> fragment-shader read (next frame's uber b16 sample).
out_barrier = vk.VkImageMemoryBarrier(
srcAccessMask=vk.VK_ACCESS_SHADER_WRITE_BIT, dstAccessMask=vk.VK_ACCESS_SHADER_READ_BIT,
oldLayout=vk.VK_IMAGE_LAYOUT_GENERAL, newLayout=vk.VK_IMAGE_LAYOUT_GENERAL,
srcQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED, dstQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
image=self._out_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_COMPUTE_SHADER_BIT, vk.VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
0, 0, None, 0, None, 1, [out_barrier])
# ----------------------------------------------------------------- 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 + output + 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_output(width, height)
self._write_descriptors()
[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._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._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
if self._pipeline:
vk.vkDestroyPipeline(device, self._pipeline, None)
if self._layout:
vk.vkDestroyPipelineLayout(device, self._layout, None)
if self._module:
vk.vkDestroyShaderModule(device, self._module, None)
if self._desc_pool:
vk.vkDestroyDescriptorPool(device, self._desc_pool, None)
if self._desc_layout:
vk.vkDestroyDescriptorSetLayout(device, self._desc_layout, None)
for s in (self._scene_sampler, self._gbuffer_sampler, self._output_sampler):
if s:
vk.vkDestroySampler(device, s, None)
self._scene_sampler = self._gbuffer_sampler = self._output_sampler = None
self._destroy_targets()
self._ready = False