"""Irradiance-volume capture: desktop (Vulkan) backend for ``IrradianceVolume3D``.
Bakes a grid of diffuse-GI probes (design D10/RM-E9). Per budgeted probe it
renders a small radiance cubemap from the probe position (reusing the
reflection-probe face-render machinery: ``_FaceCamera`` + a
``GameViewportRenderer`` face target + a source cube), then dispatches the
``irradiance_sh_reduce.comp`` compute to project that cube onto SH-L1 and write
the 12 coefficients into the volume SSBO (forward set0 binding 18) at the probe's
slice. The uber shader trilinearly blends the 8 surrounding probes and evaluates
their SH per fragment normal (``irradiance_volume_diffuse`` in cube_textured.frag),
so a dynamic object inside the box picks up soft coloured indirect light.
All work records into the PRIMARY frame command buffer (the ``cmd`` from the
app's ``pre_render`` hook), exactly like :class:`ReflectionProbePass`: face
renders, the SH-reduce dispatch, and the write->read barrier are inline, so the
main pass later in the same cmd reads this frame's freshly-baked SH with no queue
stall. ``update_budget`` probes bake per frame so a large grid amortises.
Zero-cost when unused: the pass is lazily created the first frame a volume is
present, and the uber's ``irradiance_volume_enabled`` gate stays 0 (SSBO never
read) otherwise, so a scene with no volume is byte-identical.
"""
from __future__ import annotations
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.pipeline_compute import create_compute_pipeline
from .buffer_manager import (
IRRADIANCE_VOLUME_BUFFER_SIZE,
IRRADIANCE_VOLUME_HEADER_SIZE,
)
from .game_viewport import GameViewportRenderer
from .reflection_probe_pass import _CUBE_FORMAT, _FACE_DIRS, _FaceCamera
__all__ = ["IrradianceVolumePass"]
log = logging.getLogger(__name__)
# Radiance-cube face resolution. Irradiance is very low frequency (SH-L1 has 4
# bands), so a tiny face captures plenty of directional signal cheaply.
FACE_SIZE = 32
# Sphere samples per probe in the SH reduce (multiple of the compute's 64-wide
# workgroup). 256 is ample for the 4 L1 coefficients.
SAMPLE_COUNT = 256
_PC_SIZE = 16 # uvec4 (probe_index, sample_count, pad, pad)
[docs]
class IrradianceVolumePass:
"""Owns the SH-reduce compute + capture scratch; bakes IrradianceVolume3D grids."""
def __init__(self, engine: Any) -> None:
self._engine = engine
self._ready = False
self._face_target: GameViewportRenderer | None = None
self._src_image: Any = None
self._src_memory: Any = None
self._src_view: Any = None
self._src_sampler: Any = None
self._src_initialised = False
self._desc_layout: Any = None
self._desc_pool: Any = None
self._desc_set: Any = None
self._pipeline: Any = None
self._pipe_layout: Any = None
self._module: Any = None
# Last uploaded header hash (skip redundant header uploads).
self._header_hash: int | None = None
# ------------------------------------------------------------------ setup
[docs]
def setup(self) -> None:
"""Create the SH-reduce compute pipeline + descriptor set (scratch is lazy)."""
device = self._engine.ctx.device
cis = vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
sb = vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER
cs = vk.VK_SHADER_STAGE_COMPUTE_BIT
self._desc_layout = create_descriptor_set_layout(
device,
[
(0, cis, cs, 1), # source radiance cube
(1, sb, cs, 1), # volume SH SSBO (binding 18 buffer, written here)
],
)
self._desc_pool = create_pool_for_types(device, {cis: 1, sb: 1}, max_sets=1)
self._desc_set = allocate_descriptor_set(device, self._desc_pool, self._desc_layout)
self._pipeline, self._pipe_layout, self._module = create_compute_pipeline(
device, self._engine.shader_dir / "irradiance_sh_reduce.comp", [self._desc_layout], _PC_SIZE
)
self._ready = True
log.debug("IrradianceVolumePass initialised")
def _ensure_scratch(self) -> None:
"""Lazily build the face target + source cube + descriptor writes."""
if self._face_target is not None:
return
e = self._engine
self._face_target = GameViewportRenderer(e)
self._face_target.create(FACE_SIZE, FACE_SIZE)
self._create_source_cube()
self._src_sampler = vk.vkCreateSampler(
e.ctx.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,
),
None,
)
# binding 1 = the forward set's volume SSBO (buffer_manager owns it).
vol_buf = e.renderer._buffers.irradiance_volume_buf
with DescriptorWriteBatch(e.ctx.device) as b:
b.image(self._desc_set, 0, self._src_view, self._src_sampler)
b.ssbo(self._desc_set, 1, vol_buf, IRRADIANCE_VOLUME_BUFFER_SIZE)
def _create_source_cube(self) -> None:
from ..gpu.memory import _find_memory_type
device = self._engine.ctx.device
phys = self._engine.ctx.physical_device
ffi = vk.ffi
ci = ffi.new("VkImageCreateInfo*")
ci.sType = vk.VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO
ci.imageType = vk.VK_IMAGE_TYPE_2D
ci.format = _CUBE_FORMAT
ci.extent.width = FACE_SIZE
ci.extent.height = FACE_SIZE
ci.extent.depth = 1
ci.mipLevels = 1
ci.arrayLayers = 6
ci.samples = vk.VK_SAMPLE_COUNT_1_BIT
ci.tiling = vk.VK_IMAGE_TILING_OPTIMAL
ci.usage = vk.VK_IMAGE_USAGE_SAMPLED_BIT | vk.VK_IMAGE_USAGE_TRANSFER_DST_BIT
ci.sharingMode = vk.VK_SHARING_MODE_EXCLUSIVE
ci.initialLayout = vk.VK_IMAGE_LAYOUT_UNDEFINED
ci.flags = vk.VK_IMAGE_CREATE_CUBE_COMPATIBLE_BIT
img_out = ffi.new("VkImage*")
if vk._vulkan._callApi(vk._vulkan.lib.vkCreateImage, device, ci, ffi.NULL, img_out) != vk.VK_SUCCESS:
raise RuntimeError("vkCreateImage (irradiance source cube) failed")
self._src_image = img_out[0]
req = vk.vkGetImageMemoryRequirements(device, self._src_image)
self._src_memory = vk.vkAllocateMemory(
device,
vk.VkMemoryAllocateInfo(
allocationSize=req.size,
memoryTypeIndex=_find_memory_type(phys, req.memoryTypeBits, vk.VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT),
),
None,
)
vk.vkBindImageMemory(device, self._src_image, self._src_memory, 0)
self._src_view = vk.vkCreateImageView(
device,
vk.VkImageViewCreateInfo(
image=self._src_image,
viewType=vk.VK_IMAGE_VIEW_TYPE_CUBE,
format=_CUBE_FORMAT,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=6,
),
),
None,
)
# ----------------------------------------------------------------- capture
[docs]
def update_volumes(self, cmd: Any, adapter: Any, tree: Any, volumes: list) -> bool:
"""Bake up to ``update_budget`` probes of the active volume into *cmd*.
Only the first volume in tree order is baked (single-volume support;
multiple volumes are a later phase). Returns ``True`` if a probe rendered
this frame (the face render used the shared offscreen target).
"""
if not self._ready or adapter is None or tree is None or not volumes:
return False
vol = volumes[0]
mode = getattr(vol, "bake_mode", "once")
if mode == "disabled":
return False
n = int(vol.probe_count())
if n <= 0:
return False
budget = max(1, int(getattr(vol, "update_budget", 4)))
# Header upload (grid + bounds + intensity): hash-gated so a static volume
# uploads once. Only the 64-byte header, never the SH region (GPU-written).
self._upload_header(vol, n)
# Schedule which probes to bake this frame.
if mode == "always":
start = getattr(vol, "_bake_cursor", 0) % n
indices = [(start + k) % n for k in range(min(budget, n))]
vol._bake_cursor = (start + len(indices)) % n
else: # "once": bake forward from the cursor until the whole grid is done.
if getattr(vol, "_baked", False):
return False
start = int(getattr(vol, "_bake_cursor", 0))
indices = list(range(start, min(start + budget, n)))
vol._bake_cursor = start + len(indices)
if vol._bake_cursor >= n:
vol._baked = True
vol._version = int(getattr(vol, "_version", 0)) + 1
if not indices:
return False
self._ensure_scratch()
positions = vol.probe_positions()
cull = 0xFFFFFFFF
for idx in indices:
self._bake_probe(cmd, adapter, tree, positions[idx], int(idx), cull)
# Make the compute SH writes visible to the main pass's fragment reads.
self._buffer_barrier(cmd)
return True
def _bake_probe(self, cmd: Any, adapter: Any, tree: Any, eye: np.ndarray, idx: int, cull: int) -> None:
"""Render 6 radiance faces from ``eye`` into the source cube, then SH-reduce into slice ``idx``."""
eye = np.asarray(eye, dtype=np.float32)
near, far = 0.05, 200.0
# Source cube: UNDEFINED (or SHADER_READ from a prior probe) -> TRANSFER_DST.
old = vk.VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL if self._src_initialised else vk.VK_IMAGE_LAYOUT_UNDEFINED
self._img_barrier(
cmd,
self._src_image,
6,
old,
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,
)
self._src_initialised = True
for face, (fwd, up) in enumerate(_FACE_DIRS):
cam = _FaceCamera(eye, fwd, up, near, far, cull)
adapter.render_to_target(
cmd, self._face_target, tree, camera=cam, sru_id=0x49560000 ^ (idx << 3) ^ (face + 1)
)
ct = self._face_target._target.colour_image
self._img_barrier(
cmd,
ct,
1,
vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
vk.VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
vk.VK_ACCESS_SHADER_READ_BIT,
vk.VK_ACCESS_TRANSFER_READ_BIT,
vk.VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
)
self._copy_face(cmd, ct, self._src_image, face)
self._img_barrier(
cmd,
ct,
1,
vk.VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
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,
)
# Source cube -> SHADER_READ for the SH-reduce compute to sample it.
self._img_barrier(
cmd,
self._src_image,
6,
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,
)
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._pipeline)
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._pipe_layout, 0, 1, [self._desc_set], 0, None
)
pc = np.array([idx, SAMPLE_COUNT, 0, 0], dtype=np.uint32).tobytes()
cbuf = vk.ffi.new("char[]", pc)
vk._vulkan.lib.vkCmdPushConstants(cmd, self._pipe_layout, vk.VK_SHADER_STAGE_COMPUTE_BIT, 0, _PC_SIZE, cbuf)
vk.vkCmdDispatch(cmd, 1, 1, 1)
# Return the cube to TRANSFER_DST for the next probe's faces.
self._img_barrier(
cmd,
self._src_image,
6,
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,
)
# ----------------------------------------------------------------- payload
def _upload_header(self, vol: Any, count: int) -> None:
"""Build + upload the 64-byte volume header (grid + bounds + intensity)."""
nx, ny, nz = vol.grid_dims()
lo, hi = vol.bounds_world()
intensity = float(getattr(vol, "intensity", 1.0))
header = bytearray(IRRADIANCE_VOLUME_HEADER_SIZE)
header[0:16] = np.array([nx, ny, nz, count], dtype=np.int32).tobytes()
header[16:32] = np.array([float(lo.x), float(lo.y), float(lo.z), 0.0], dtype=np.float32).tobytes()
header[32:48] = np.array([float(hi.x), float(hi.y), float(hi.z), 0.0], dtype=np.float32).tobytes()
header[48:64] = np.array([intensity, 0.0, 0.0, 0.0], dtype=np.float32).tobytes()
h = hash(bytes(header))
if h == self._header_hash:
return
self._header_hash = h
self._engine.renderer._buffers.write_irradiance_volume_header(np.frombuffer(bytes(header), dtype=np.uint8))
# --------------------------------------------------------- Vulkan helpers
def _copy_face(self, cmd: Any, src_image: Any, dst_cube: Any, dst_layer: int) -> None:
region = vk.VkImageCopy(
srcSubresource=vk.VkImageSubresourceLayers(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT, mipLevel=0, baseArrayLayer=0, layerCount=1
),
srcOffset=vk.VkOffset3D(x=0, y=0, z=0),
dstSubresource=vk.VkImageSubresourceLayers(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT, mipLevel=0, baseArrayLayer=dst_layer, layerCount=1
),
dstOffset=vk.VkOffset3D(x=0, y=0, z=0),
extent=vk.VkExtent3D(width=FACE_SIZE, height=FACE_SIZE, depth=1),
)
vk.vkCmdCopyImage(
cmd,
src_image,
vk.VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
dst_cube,
vk.VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
1,
[region],
)
def _img_barrier(
self,
cmd: Any,
image: Any,
layers: int,
old: int,
new: int,
src_access: int,
dst_access: int,
src_stage: int,
dst_stage: int,
) -> None:
barrier = 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=layers,
),
)
vk.vkCmdPipelineBarrier(cmd, src_stage, dst_stage, 0, 0, None, 0, None, 1, [barrier])
def _buffer_barrier(self, cmd: Any) -> None:
"""Compute SH writes -> fragment (uber) reads of the volume SSBO."""
b = vk.VkBufferMemoryBarrier(
srcAccessMask=vk.VK_ACCESS_SHADER_WRITE_BIT,
dstAccessMask=vk.VK_ACCESS_SHADER_READ_BIT,
srcQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
dstQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
buffer=self._engine.renderer._buffers.irradiance_volume_buf,
offset=0,
size=IRRADIANCE_VOLUME_BUFFER_SIZE,
)
vk.vkCmdPipelineBarrier(
cmd,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
vk.VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
0,
0,
None,
1,
[b],
0,
None,
)
# ----------------------------------------------------------------- cleanup
[docs]
def cleanup(self) -> None:
if not self._ready:
return
device = self._engine.ctx.device
vk.vkDeviceWaitIdle(device)
if self._face_target is not None:
self._face_target.destroy()
self._face_target = None
if self._src_view:
vk.vkDestroyImageView(device, self._src_view, None)
if self._src_image:
vk.vkDestroyImage(device, self._src_image, None)
if self._src_memory:
vk.vkFreeMemory(device, self._src_memory, None)
if self._src_sampler:
vk.vkDestroySampler(device, self._src_sampler, None)
if self._pipeline:
vk.vkDestroyPipeline(device, self._pipeline, None)
if self._pipe_layout:
vk.vkDestroyPipelineLayout(device, self._pipe_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)
self._ready = False