"""Reflection-probe capture: desktop (Vulkan) backend for ``ReflectionProbe3D``.
Captures a local environment cubemap per probe and feeds the engine's
existing split-sum IBL precompute (irradiance + GGX-prefiltered specular,
:class:`IBLPass`) so meshes inside a probe's influence box pick up the local
room reflection instead of the global skybox IBL.
Pipeline (per probe, run **once** on enter + on explicit ``request_update()``):
1. Render the scene six times from the probe's ``capture_position`` into a
square offscreen target (``GameViewportRenderer``), one per cube face, with a
90° FOV perspective camera. After each face, ``vkCmdCopyImage`` the colour
image into face *f* of a per-probe source cubemap.
2. Record :meth:`IBLPass.record_convolution` on that source cube (using a
per-probe descriptor bundle) into the same cmd → irradiance + prefiltered
specular for this probe (reusing the global IBL compute shaders).
3. ``vkCmdCopyImage`` the probe's irradiance cube (6 layers) and prefiltered
cube (6 layers × N mips) into the shared cubemap **arrays** at the probe's
array slice (slice = ``probe_index * 6``).
The shared arrays + a probe-box SSBO are bound to the forward set (bindings
10/11/12). The fragment shader tests each fragment's world position against the
probe boxes (bounded loop, ``MAX_PROBES``) and, when inside, samples that
probe's maps (optionally box-projected) instead of the global IBL.
Capture records into the PRIMARY frame command buffer (the ``cmd`` passed to the
app's ``pre_render`` hook): the six face renders, the per-probe IBL convolution,
and the cube-array copy are all recorded inline with intra-cmd image barriers,
so an always-mode probe no longer submits-and-waits its own buffer every frame.
Each probe owns its source cubemap **and** its IBL descriptor bundle, so several
always-mode probes convolved in one frame never cross-contaminate (the
convolution reads each probe's own cube; ``vkUpdateDescriptorSets`` mutates a set
at record time, so a shared bundle's last write would otherwise win for all).
"""
from __future__ import annotations
import logging
from typing import Any
import numpy as np
import vulkan as vk
from .buffer_manager import MAX_PROBES, PROBE_BUFFER_SIZE
from .game_viewport import GameViewportRenderer
from .ibl_pass import IRRADIANCE_SIZE, PREFILTER_MIP_LEVELS, PREFILTER_SIZE, IBLPass
from .render_target import RenderTarget # noqa: F401 (referenced in docstring / type intent)
__all__ = ["ReflectionProbePass"]
log = logging.getLogger(__name__)
# Default per-face capture resolution. Small: the IBL convolution downsamples
# to 32×32 irradiance / 128×128 prefilter anyway, so a high-res face buys
# little. Per-pass override via ``WorldEnvironment.probe_face_size``.
FACE_SIZE = 128
_CUBE_FORMAT = vk.VK_FORMAT_R16G16B16A16_SFLOAT
# Six cube-face look directions (+X, -X, +Y, -Y, +Z, -Z) and up vectors,
# matching the standard Vulkan/GL cubemap face ordering used by the skybox
# and IBL shaders. forward/up are world-space.
_FACE_DIRS = [
((1.0, 0.0, 0.0), (0.0, -1.0, 0.0)), # +X
((-1.0, 0.0, 0.0), (0.0, -1.0, 0.0)), # -X
((0.0, 1.0, 0.0), (0.0, 0.0, 1.0)), # +Y
((0.0, -1.0, 0.0), (0.0, 0.0, -1.0)), # -Y
((0.0, 0.0, 1.0), (0.0, -1.0, 0.0)), # +Z
((0.0, 0.0, -1.0), (0.0, -1.0, 0.0)), # -Z
]
class _FaceCamera:
"""Lightweight duck-typed camera for one cube face.
Satisfies the slice of the camera protocol that ``SceneAdapter.submit_scene``
reads: ``view_matrix`` (property), ``projection_matrix(aspect)``,
``cull_mask``, and ``_visible_in_hierarchy``. Built per-face from the probe's
world capture position so we never mutate a real scene camera.
"""
def __init__(self, eye: np.ndarray, forward: tuple, up: tuple, near: float, far: float, cull_mask: int):
from simvx.core.math.matrices import look_at, perspective
self._eye = np.asarray(eye, dtype=np.float32)
self._fwd = np.asarray(forward, dtype=np.float32)
self._up = np.asarray(up, dtype=np.float32)
self._near = near
self._far = far
self.cull_mask = cull_mask
self._visible_in_hierarchy = True
self._look_at = look_at
self._perspective = perspective
@property
def view_matrix(self) -> np.ndarray:
return self._look_at(self._eye, self._eye + self._fwd, self._up)
def projection_matrix(self, aspect: float = 1.0) -> np.ndarray:
import math
proj = self._perspective(math.radians(90.0), aspect, self._near, self._far)
proj[1, 1] *= -1 # Vulkan Y-flip (matches Camera3D.projection_matrix)
return proj
[docs]
class ReflectionProbePass:
"""Owns the shared probe cubemap arrays + box SSBO and drives per-probe capture."""
def __init__(self, engine: Any) -> None:
self._engine = engine
self._ready = False
# Shared cubemap arrays (MAX_PROBES * 6 layers each).
self._irr_image: Any = None
self._irr_memory: Any = None
self._irr_view: Any = None
self._pre_image: Any = None
self._pre_memory: Any = None
self._pre_view: Any = None
self._sampler: Any = None
# Shared cube sampler (read-only) + shared offscreen face target.
self._src_sampler: Any = None
self._face_target: GameViewportRenderer | None = None
self._ibl: IBLPass | None = None
# Probe id -> assigned array slot (0..MAX_PROBES-1).
self._slots: dict[int, int] = {}
# Last uploaded probe-box payload hash (skip redundant SSBO uploads).
self._box_hash: int | None = None
# Per-probe capture state keyed by id(probe): each probe owns a persistent
# source cubemap (so its convolution reads its OWN faces and never mixes
# with another always-probe captured the same frame) plus a face cursor
# (0..5) for round-robin time-slicing. Full-capture probes write all six
# faces in one frame; round-robin writes one face per frame. Persistent so
# the cube stays in TRANSFER_DST between a round-robin cycle's faces.
self._probe_state: dict[int, dict[str, Any]] = {}
# Per-face capture resolution. The default
# matches the historical constant; ``set_face_size`` applies a
# WorldEnvironment.probe_face_size override. Only the capture scratch
# (face target + per-probe source cubes) depends on it: the shared
# irradiance/prefilter arrays are fixed-size convolution outputs.
self._face_size = FACE_SIZE
# Set when the capture scratch was rebuilt at a new face size: every
# active probe must recapture so no slot keeps an old-resolution bake.
self._recapture_all = False
# ------------------------------------------------------------------ setup
[docs]
def setup(self) -> None:
"""Allocate the shared cubemap arrays + capture scratch resources."""
e = self._engine
device = e.ctx.device
phys = e.ctx.physical_device
self._irr_image, self._irr_memory, self._irr_view = self._create_cube_array(
device, phys, IRRADIANCE_SIZE, 1,
)
self._pre_image, self._pre_memory, self._pre_view = self._create_cube_array(
device, phys, PREFILTER_SIZE, PREFILTER_MIP_LEVELS,
)
self._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(PREFILTER_MIP_LEVELS),
), None)
self._src_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,
), None)
# Transition both arrays to SHADER_READ_ONLY so the descriptor is valid
# before any probe captures (the shader's probe_count gate skips them).
self._transition_array(self._irr_image, 1,
vk.VK_IMAGE_LAYOUT_UNDEFINED,
vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
self._transition_array(self._pre_image, PREFILTER_MIP_LEVELS,
vk.VK_IMAGE_LAYOUT_UNDEFINED,
vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
self._ready = True
log.debug("ReflectionProbePass initialised (max_probes=%d)", MAX_PROBES)
def _create_cube_array(self, device: Any, phys: Any, size: int, mips: int) -> tuple[Any, Any, Any]:
"""Create a CUBE_ARRAY image (MAX_PROBES*6 layers) + a CUBE_ARRAY view."""
from ..gpu.memory import _find_memory_type
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 = size
ci.extent.height = size
ci.extent.depth = 1
ci.mipLevels = mips
ci.arrayLayers = MAX_PROBES * 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 (probe cube array) failed")
image = img_out[0]
req = vk.vkGetImageMemoryRequirements(device, image)
mem = 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, image, mem, 0)
view = vk.vkCreateImageView(device, vk.VkImageViewCreateInfo(
image=image,
viewType=vk.VK_IMAGE_VIEW_TYPE_CUBE_ARRAY,
format=_CUBE_FORMAT,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0, levelCount=mips,
baseArrayLayer=0, layerCount=MAX_PROBES * 6,
),
), None)
return image, mem, view
# --------------------------------------------------------------- accessors
[docs]
def get_irradiance_array_view(self) -> Any:
return self._irr_view
[docs]
def get_prefilter_array_view(self) -> Any:
return self._pre_view
[docs]
def get_sampler(self) -> Any:
return self._sampler
# ------------------------------------------------------------ quality dials
[docs]
def set_face_size(self, size: int | None) -> None:
"""Set the per-face capture resolution.
``None`` means the built-in default (:data:`FACE_SIZE`), so an unset
``WorldEnvironment.probe_face_size`` keeps the exact historical path.
A change tears down the capture scratch (shared face target + per-probe
source cubes; the fixed-size IBL machinery survives) and schedules a
full recapture of every active probe. Rare transition: device-wait-idle
is acceptable.
"""
size = FACE_SIZE if size is None else max(16, int(size))
if size == self._face_size:
return
self._face_size = size
if self._face_target is None and not self._probe_state:
return # nothing captured yet: the lazy scratch picks the size up
vk.vkDeviceWaitIdle(self._engine.ctx.device)
for pid in list(self._probe_state):
self._destroy_probe_state(pid)
if self._face_target is not None:
self._face_target.destroy()
self._face_target = None
self._recapture_all = True
[docs]
def disable_probes(self) -> None:
"""Zero the probe-box SSBO so the forward blend loop is a no-op.
Called per frame while ``WorldEnvironment.probe_blend_count == 0``
(capture is skipped upstream). Hash-gated: after the first call this
is a dictionary lookup per frame.
"""
if self._ready:
self._upload_boxes([])
# ----------------------------------------------------------------- capture
[docs]
def update_probes(self, cmd: Any, adapter: Any, tree: Any, probes: list) -> bool:
"""Capture any probe that is new or has ``_update_requested`` set, into *cmd*.
All capture work (face renders, IBL convolution, array copy) records into
the primary frame command buffer *cmd*: no separate submit/wait. Returns
``True`` if at least one probe captured this frame (the caller must then
re-submit the main scene, since face rendering clobbers the renderer's
per-frame submission lists: same contract as SubViewports).
"""
if not self._ready or adapter is None or tree is None:
return False
# Bound the active set to MAX_PROBES: no closest-to-origin priority;
# the first MAX_PROBES in tree order win.
active = probes[:MAX_PROBES]
# (Re)assign slots for the active set; drop stale ones.
live_ids = {id(p) for p in active}
self._slots = {pid: s for pid, s in self._slots.items() if pid in live_ids}
for p in active:
if id(p) not in self._slots:
self._slots[id(p)] = self._next_free_slot()
# Drop per-probe staging (source cube + IBL bundle) for departed probes.
for pid in [pid for pid in self._probe_state if pid not in live_ids]:
self._destroy_probe_state(pid)
# A face-size change (set_face_size) invalidated every existing bake:
# force one full recapture per active probe this frame.
recapture_all, self._recapture_all = self._recapture_all, False
captured = False
for probe in active:
slot = self._slots[id(probe)]
mode = getattr(probe, "capture_mode", "once")
never_captured = getattr(probe, "_cubemap_version", 0) == 0
requested = getattr(probe, "_update_requested", False) or recapture_all
if mode == "always" and not requested and not never_captured:
# Steady-state "always" refresh: time-sliced or full per frame.
slicing = getattr(probe, "time_slicing", "round_robin")
if slicing == "round_robin":
captured |= self._capture_round_robin(cmd, adapter, tree, probe, slot)
else: # "every_frame"
self._capture_probe(cmd, adapter, tree, probe, slot)
self._mark_updated(probe)
captured = True
continue
# First capture (never captured) or explicit request_update():
# always a full 6-face capture + convolve, regardless of mode.
if requested or never_captured:
self._capture_probe(cmd, adapter, tree, probe, slot)
self._mark_updated(probe)
captured = True
self._upload_boxes(active)
return captured
def _mark_updated(self, probe: Any) -> None:
"""Clear the request flag, bump the version, and emit cubemap_updated."""
probe._update_requested = False
probe._cubemap_version = getattr(probe, "_cubemap_version", 0) + 1
sig = getattr(probe, "cubemap_updated", None)
if sig is not None:
sig.emit()
def _next_free_slot(self) -> int:
used = set(self._slots.values())
for s in range(MAX_PROBES):
if s not in used:
return s
return 0 # full: reuse slot 0 (bounded by MAX_PROBES anyway)
def _capture_probe(self, cmd: Any, adapter: Any, tree: Any, probe: Any, slot: int) -> None:
"""Render 6 faces from *probe*, convolve, and copy into array *slot*: all in *cmd*."""
self._ensure_capture_scratch()
st = self._probe_state_for(cmd, probe)
src_image, src_view = st["image"], st["view"]
eye = np.asarray(probe.capture_position, dtype=np.float32)
near, far = 0.05, max(2.0, float(max(probe.size)) * 4.0 + 50.0)
cull = int(getattr(probe, "cull_mask", 0xFFFFFFFF))
# --- 1. Render the six faces into this probe's own source cube ---
# The cube lives in TRANSFER_DST between captures (set on creation and
# restored after each convolve), so no transition is needed here for a
# re-capture; the initial creation transition handles the first use.
for face, (fwd, up) in enumerate(_FACE_DIRS):
cam = _FaceCamera(eye, fwd, up, near, far, cull)
# Stable per-probe-face SRU id keys the visibility cache distinctly.
adapter.render_to_target(cmd, self._face_target, tree, camera=cam,
sru_id=(id(probe) ^ (face + 1)))
# The offscreen pass leaves colour in SHADER_READ_ONLY; move it to
# TRANSFER_SRC, copy into source-cube face `face`, then back so the
# next face render's pass (initialLayout=UNDEFINED) is happy.
ct = self._face_target._target.colour_image
self._barrier(cmd, ct, 1, 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, src_image, dst_layer=face, size=self._face_size)
self._barrier(cmd, ct, 1, 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_ONLY so the convolution (COMPUTE) can sample it.
self._barrier(cmd, src_image, 6, 1,
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)
# --- 2. IBL convolution into the shared outputs, then copy to the array ---
self._convolve_and_store(cmd, st, slot, src_view)
# Return the cube to TRANSFER_DST for the next capture's face copies.
self._barrier(cmd, src_image, 6, 1,
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)
log.debug("Reflection probe captured into slot %d", slot)
def _convolve_and_store(self, cmd: Any, st: dict[str, Any], slot: int, src_view: Any) -> None:
"""Convolve this probe's source cube (its OWN IBL bundle) + copy into the array slot.
Records into *cmd*. The per-probe bundle guarantees the convolution reads
this probe's cube even if another always-probe convolved earlier in the
same cmd; the array copy is recorded right after so the shared IBL output
images are consumed before the next probe's convolution overwrites them.
"""
# Serialise this refresh against the immediately-prior frame. The
# convolution rewrites the shared IBL output images and the array copy
# overwrites this probe's destination slice, which the prior frame's
# forward pass is still sampling (and, for back-to-back convolutions, its
# copy-out is still reading the shared images). The per-slot fences only
# gate FRAMES_IN_FLIGHT frames back, never one frame back, so without this
# the refresh races the prior frame and corrupts the probe map under load.
# Refreshes are rare (once for once-mode, every six frames for
# round-robin), so the wait is cheap; only an every-frame probe pays it
# each refresh, matching that mode's documented cost.
self._wait_prior_frame()
# Pick this probe's bundle for the current frame-in-flight so updating it
# cannot race the previous frame's still-executing convolution. Keyed by the
# bounded array slot (not id(probe)): slots are unique per live probe and
# reused, so the bundle pool stays bounded (no leak/exhaustion under churn);
# record_convolution rebinds the source cube each call, so a reused slot is
# correctly repointed at the new probe's cube.
frame_slot = getattr(self._engine._sync, "current_frame", 0)
bundle = self._ibl.acquire_bundle(slot, frame_slot)
self._ibl.record_convolution(cmd, src_view, self._src_sampler, bundle)
base = slot * 6
self._copy_cube_into_array(cmd, self._ibl.get_irradiance_image(), self._irr_image,
IRRADIANCE_SIZE, 1, base)
self._copy_cube_into_array(cmd, self._ibl.get_prefiltered_image(), self._pre_image,
PREFILTER_SIZE, PREFILTER_MIP_LEVELS, base)
def _wait_prior_frame(self) -> None:
"""Block until the immediately-prior in-flight frame has fully retired.
See :meth:`_convolve_and_store`. A no-op when there is no GPU frame-sync
(CPU-only tests) or no pipelining (``frame_count <= 1``: frames are
already fully serialised, and the prior slot would be this frame's own
not-yet-submitted fence). Never resets the fence: the frame loop owns
the reset/signal lifecycle, we only observe completion.
"""
sync = getattr(self._engine, "_sync", None)
frame_count = getattr(sync, "frame_count", 0)
fences = getattr(sync, "fences", None)
if not fences or frame_count <= 1:
return
prior = (getattr(sync, "current_frame", 0) - 1) % frame_count
vk.vkWaitForFences(self._engine.ctx.device, 1, [fences[prior]], vk.VK_TRUE, 2_000_000_000)
# ---------------------------------------------------- round-robin capture
def _capture_round_robin(self, cmd: Any, adapter: Any, tree: Any, probe: Any, slot: int) -> bool:
"""Capture ONE cube face per call into the probe's persistent source cube.
Records into the primary frame *cmd*. Advances a per-probe face cursor
0..5; once the 6th face lands, runs the IBL convolution + copies into the
shared array slot (a full refresh every 6 frames). Returns ``True``
whenever a face was rendered this frame: even a mid-cycle face uses the
shared offscreen face target and clobbers the renderer's per-frame
submission lists, so the caller must re-submit the main scene every
round-robin frame (same contract as full capture).
"""
self._ensure_capture_scratch()
st = self._probe_state_for(cmd, probe)
face = st["cursor"]
eye = np.asarray(probe.capture_position, dtype=np.float32)
near, far = 0.05, max(2.0, float(max(probe.size)) * 4.0 + 50.0)
cull = int(getattr(probe, "cull_mask", 0xFFFFFFFF))
fwd, up = _FACE_DIRS[face]
# --- Render + copy just this one face into the persistent cube layer ---
cam = _FaceCamera(eye, fwd, up, near, far, cull)
adapter.render_to_target(cmd, self._face_target, tree, camera=cam,
sru_id=(id(probe) ^ (face + 1)))
ct = self._face_target._target.colour_image
self._barrier(cmd, ct, 1, 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, st["image"], dst_layer=face, size=self._face_size)
self._barrier(cmd, ct, 1, 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)
st["cursor"] = face + 1
if st["cursor"] < 6:
# Mid-cycle: array slice unchanged, but a face was rendered into the
# shared offscreen target, so the main scene must be re-submitted.
return True
# 6th face landed: convolve the full cube + copy into the shared slot, all
# in this cmd. The cube has accumulated all six faces across prior frames
# and lives in TRANSFER_DST; move it to SHADER_READ for the convolution.
st["cursor"] = 0
self._barrier(cmd, st["image"], 6, 1,
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)
self._convolve_and_store(cmd, st, slot, st["view"])
# Return the persistent cube to TRANSFER_DST for the next cycle's faces.
self._barrier(cmd, st["image"], 6, 1,
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)
self._mark_updated(probe)
log.debug("Reflection probe (round-robin) refreshed slot %d", slot)
return True
def _probe_state_for(self, cmd: Any, probe: Any) -> dict[str, Any]:
"""Return *probe*'s persistent capture state, creating it on first use.
State = its own source cubemap + its own IBL descriptor bundle + a
round-robin face cursor. Per probe so two always-mode probes captured in
the same frame each convolve from their OWN cube via their OWN bundle. On
first creation the cube's UNDEFINED -> TRANSFER_DST transition is recorded
into the frame *cmd* (no separate submit/wait), so even the per-probe init
adds zero queue stalls.
"""
st = self._probe_state.get(id(probe))
if st is None:
img, mem, view = self._create_source_cube()
# Record the one-time init transition into the frame cmd: the cube must
# be in TRANSFER_DST before the first face copy targets it.
self._barrier(cmd, img, 6, 1,
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)
st = {"image": img, "memory": mem, "view": view, "cursor": 0, "id": id(probe)}
self._probe_state[id(probe)] = st
return st
def _destroy_probe_state(self, pid: int) -> None:
"""Free a probe's per-probe source cube + IBL bundle (probe left the active set)."""
st = self._probe_state.pop(pid, None)
if st is None:
return
device = self._engine.ctx.device
# Bundles are keyed by the bounded array slot and reused across probe
# lifetimes (see _convolve_and_store), so there is nothing per-probe to free here.
if st.get("view"):
vk.vkDestroyImageView(device, st["view"], None)
if st.get("image"):
vk.vkDestroyImage(device, st["image"], None)
if st.get("memory"):
vk.vkFreeMemory(device, st["memory"], None)
def _ensure_capture_scratch(self) -> None:
"""Lazily create the offscreen face target and IBLPass (per-probe cubes are lazy)."""
if self._face_target is not None:
return
e = self._engine
self._face_target = GameViewportRenderer(e)
self._face_target.create(self._face_size, self._face_size)
self._ibl = IBLPass(e)
self._ibl.setup()
def _create_source_cube(self) -> tuple[Any, Any, Any]:
"""Create a 6-layer, 1-mip CUBE source image + view (transfer-dst, sampled)."""
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 = self._face_size
ci.extent.height = self._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 (probe source cube) failed")
image = img_out[0]
req = vk.vkGetImageMemoryRequirements(device, image)
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, image, memory, 0)
view = vk.vkCreateImageView(device, vk.VkImageViewCreateInfo(
image=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)
return image, memory, view
# ------------------------------------------------------------- box SSBO
def _upload_boxes(self, probes: list) -> None:
"""Build + upload the probe-box SSBO (count header + Probe array)."""
data = bytearray(PROBE_BUFFER_SIZE)
count = min(len(probes), MAX_PROBES)
data[0:4] = np.array([count], dtype=np.uint32).tobytes()
for probe in probes[:count]:
slot = self._slots[id(probe)]
centre = np.asarray(probe.capture_position, dtype=np.float32)
half = np.asarray(probe.size, dtype=np.float32)
intensity = float(getattr(probe, "intensity", 1.0))
box_proj = 1.0 if getattr(probe, "box_projection", False) else 0.0
blend = float(getattr(probe, "blend_distance", 1.0))
off = 16 + slot * 48
# vec4 centre_slice (xyz centre, w = array slice as float)
data[off:off + 16] = np.array([centre[0], centre[1], centre[2], float(slot * 6)],
dtype=np.float32).tobytes()
# vec4 half_extent_intensity (xyz half extents, w = intensity)
data[off + 16:off + 32] = np.array([half[0], half[1], half[2], intensity],
dtype=np.float32).tobytes()
# vec4 flags (x = box_projection, y = blend_distance, zw reserved)
data[off + 32:off + 48] = np.array([box_proj, blend, 0.0, 0.0], dtype=np.float32).tobytes()
payload = np.frombuffer(bytes(data), dtype=np.uint8)
h = hash(payload.tobytes())
if h == self._box_hash:
return
self._box_hash = h
self._engine.renderer._buffers.write_probe_buffer(payload)
# --------------------------------------------------------- Vulkan helpers
def _copy_face(self, cmd: Any, src_image: Any, dst_cube: Any, dst_layer: int, size: int) -> None:
"""Copy a 2D colour image into one layer of a cube image."""
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=size, height=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 _copy_cube_into_array(self, cmd: Any, src_cube: Any, dst_array: Any,
size: int, mips: int, base_layer: int) -> None:
"""Copy a 6-layer cube (all mips) into the array at ``base_layer``.
IBLPass leaves its outputs in SHADER_READ_ONLY; transition src->TRANSFER_SRC
and the destination slice->TRANSFER_DST, copy each mip, then restore both.
"""
self._barrier(cmd, src_cube, 6, mips,
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._barrier(cmd, dst_array, 6, mips,
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_FRAGMENT_SHADER_BIT, vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
base_layer=base_layer)
regions = []
for mip in range(mips):
msize = max(1, size >> mip)
regions.append(vk.VkImageCopy(
srcSubresource=vk.VkImageSubresourceLayers(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT, mipLevel=mip, baseArrayLayer=0, layerCount=6),
srcOffset=vk.VkOffset3D(x=0, y=0, z=0),
dstSubresource=vk.VkImageSubresourceLayers(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT, mipLevel=mip,
baseArrayLayer=base_layer, layerCount=6),
dstOffset=vk.VkOffset3D(x=0, y=0, z=0),
extent=vk.VkExtent3D(width=msize, height=msize, depth=1),
))
vk.vkCmdCopyImage(
cmd,
src_cube, vk.VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
dst_array, vk.VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
len(regions), regions,
)
self._barrier(cmd, src_cube, 6, mips,
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)
self._barrier(cmd, dst_array, 6, mips,
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_FRAGMENT_SHADER_BIT,
base_layer=base_layer)
def _barrier(self, cmd: Any, image: Any, layer_count: int, mips: int,
old: int, new: int, src_access: int, dst_access: int,
src_stage: int, dst_stage: int, base_layer: int = 0) -> 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=mips,
baseArrayLayer=base_layer, layerCount=layer_count,
),
)
vk.vkCmdPipelineBarrier(cmd, src_stage, dst_stage, 0, 0, None, 0, None, 1, [barrier])
def _transition_array(self, image: Any, mips: int, old: int, new: int) -> None:
"""One-shot full-array layout transition (init-only, at setup; not steady state)."""
from ..gpu.memory import begin_single_time_commands, end_single_time_commands
e = self._engine
cmd = begin_single_time_commands(e.ctx.device, e.ctx.command_pool)
self._barrier(cmd, image, MAX_PROBES * 6, mips, old, new,
0, vk.VK_ACCESS_SHADER_READ_BIT,
vk.VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, vk.VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
end_single_time_commands(e.ctx.device, e.ctx.graphics_queue, e.ctx.command_pool, cmd)
# ----------------------------------------------------------------- cleanup
[docs]
def cleanup(self) -> None:
if not self._ready:
return
device = self._engine.ctx.device
vk.vkDeviceWaitIdle(device)
for pid in list(self._probe_state):
self._destroy_probe_state(pid)
if self._ibl is not None:
self._ibl.cleanup()
self._ibl = None
if self._face_target is not None:
self._face_target.destroy()
self._face_target = None
for view, img, mem in (
(self._irr_view, self._irr_image, self._irr_memory),
(self._pre_view, self._pre_image, self._pre_memory),
):
if view:
vk.vkDestroyImageView(device, view, None)
if img:
vk.vkDestroyImage(device, img, None)
if mem:
vk.vkFreeMemory(device, mem, None)
for sampler in (self._sampler, self._src_sampler):
if sampler:
vk.vkDestroySampler(device, sampler, None)
self._ready = False