"""BufferManager: owns the forward renderer's SSBOs and descriptor sets.
Extracted from Renderer so transform/material/light/shadow/joint
buffers and the Forward+ tile-culling placeholders live in one place. The
descriptor set layout and cubemap placeholder are owned here too so the
renderer can swap the IBL cubemap in via ``write_cubemap_descriptor``.
"""
import logging
from typing import Any
import numpy as np
import vulkan as vk
from .. import frame_globals as fg
from ..gpu.descriptors import (
allocate_descriptor_set,
create_descriptor_pool,
create_ssbo_layout,
write_image_descriptor,
write_ssbo_descriptor,
write_ubo_descriptor,
)
from ..gpu.memory import create_buffer, upload_image_data, upload_numpy
from ..types import AABB_DTYPE, FRAMES_IN_FLIGHT, LIGHT_DTYPE, MATERIAL_DTYPE, TRANSFORM_DTYPE
from .decal_pack import DECAL_BUFFER_SIZE
__all__ = ["BufferManager", "SHADOW_DATA_SIZE"]
log = logging.getLogger(__name__)
# Shadow SSBO total size: must match the ShadowBuffer struct in cube_textured.frag.
# Layout: cascade_vps[3](192) + cascade_splits(16) + flags/indices(32) +
# point_light_pos_range(16) + spot_vp(64) + spot_light_pos_range(16) + ambient_colour(16) +
# ambient_mode/indirect-hook flags (16, zeros == default behaviour) +
# debug_view + pad (16, zeros == debug off) = 384
SHADOW_DATA_SIZE = 384
# Default ambient colour (cool grey fill) written at offset 336 when no WorldEnvironment overrides it.
_DEFAULT_AMBIENT = np.array([0.15, 0.15, 0.2, 1.0], dtype=np.float32)
_HOST_FLAGS = vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
_SSBO_USAGE = vk.VK_BUFFER_USAGE_STORAGE_BUFFER_BIT
# Frame-arena grow policy: when the high-water mark of slots reserved in a frame
# crosses this fraction of capacity, schedule a proactive grow before an actual
# overflow forces clamping. The transform SSBO grows to the next power of two
# that fits the demand.
_GROW_THRESHOLD = 0.8
def _next_pow2(n: int) -> int:
"""Smallest power of two >= ``n`` (>= 1)."""
if n <= 1:
return 1
return 1 << (n - 1).bit_length()
def _normal_matrices(model_mats: np.ndarray) -> np.ndarray:
"""Per-instance normal matrices packed (N,4,4), row-major for the column-major GPU read.
The shared transform SSBO is row-major numpy read column-major on the GPU, so the
GPU sees ``transpose(stored)``. The general normal matrix is ``transpose(inverse(M3))``,
so we store ``inverse(M3)`` (the GPU then reads ``transpose(inverse(M3))``).
Uniform-scale (similarity) instances skip the inverse entirely: storing ``M3ᵀ`` makes
the GPU read ``M3 = s·R``, and every lit fragment shader normalizes the interpolated
normal, so the uniform factor ``s`` washes out and the rotation ``R`` is applied
correctly. Only genuinely non-uniform-scale instances pay ``np.linalg.inv`` — and only
on that subset. Assumes TRS composition (no shear), which is what node/MultiMesh
transforms produce; a sheared instance with coincidentally equal-length axes would get
a slightly-off normal (never catastrophic, falls under normalize()).
"""
count = model_mats.shape[0]
m3 = model_mats[:, :3, :3]
col_sq = np.sum(m3 * m3, axis=1) # (N,3): squared length of each column (axis scale²)
mx = col_sq.max(axis=1)
mn = col_sq.min(axis=1)
uniform = mx <= mn * 1.0001 + 1e-12 # equal-length axes -> uniform scale
normal3 = np.empty((count, 3, 3), dtype=np.float32)
normal3[uniform] = m3[uniform].transpose(0, 2, 1) # M3ᵀ -> GPU reads M3 (∝ R), normalized in shader
nonuniform = ~uniform
if nonuniform.any():
sub = m3[nonuniform]
try:
normal3[nonuniform] = np.linalg.inv(sub)
except np.linalg.LinAlgError:
out = np.empty_like(sub)
for j in range(sub.shape[0]):
try:
out[j] = np.linalg.inv(sub[j])
except np.linalg.LinAlgError:
out[j] = sub[j].T
normal3[nonuniform] = out
normal4 = np.zeros((count, 4, 4), dtype=np.float32)
normal4[:, 3, 3] = 1.0
normal4[:, :3, :3] = normal3
return normal4
# Reflection probes: fixed small cap (bounded VRAM + bounded per-fragment loop).
# Each probe contributes 6 cube faces to two shared cubemap arrays. The probe
# box SSBO carries one std430 ``Probe`` per slot: three vec4s = 48 bytes, plus a
# 16-byte count header. Must match the ``ProbeBuffer`` struct in cube_textured.frag.
MAX_PROBES = 8
_PROBE_STRIDE = 48
PROBE_BUFFER_SIZE = 16 + MAX_PROBES * _PROBE_STRIDE
# Irradiance-volume SH SSBO (binding 18). Header = 4 vec4
# (grid ivec4 + bounds_min + bounds_max + vol_params = 64 B), then one ProbeSH
# (3 vec4 = 48 B) per probe. Must match ``IrradianceVolumeBuffer`` in
# cube_textured.frag and irradiance_sh_reduce.comp, and the probe cap in
# ``simvx.core.irradiance_volume.MAX_PROBES``.
IRRADIANCE_VOLUME_HEADER_SIZE = 64
IRRADIANCE_VOLUME_MAX_PROBES = 256
_IRRADIANCE_VOLUME_STRIDE = 48
IRRADIANCE_VOLUME_BUFFER_SIZE = (
IRRADIANCE_VOLUME_HEADER_SIZE + IRRADIANCE_VOLUME_MAX_PROBES * _IRRADIANCE_VOLUME_STRIDE
)
# Decal SSBO (binding 19): 16-byte count header + one std430
# Decal record (mat4 + 3 vec4 = 112 B) per projector. ``DECAL_BUFFER_SIZE`` and
# the record layout live in ``renderer.decal_pack`` (mirrored by ``DecalBuffer``
# in cube_textured.frag).
[docs]
class BufferManager:
"""Owns the renderer's SSBOs and descriptor sets.
Main descriptor set (``ssbo_set``) exposes eighteen bindings:
0: transforms, 1: materials, 2: lights, 3: shadow,
4: IBL cubemap sampler, 5: tile light indices, 6: tile info,
7/8/9: IBL irradiance cube / prefilter cube / BRDF 2D LUT,
10/11: reflection-probe irradiance / prefilter cube arrays,
12: reflection-probe box SSBO, 13: FrameGlobals UBO,
14: scene colour copy, 15: scene depth copy (written by
the renderer's :class:`SceneCopyTargets`),
16/17: indirect specular / diffuse hook textures (1x1
black until an SSR / SSGI producer writes them).
Joint descriptor set (``joint_set``) is set 2 binding 0 for skinned meshes.
"""
def __init__(
self, engine: Any, max_objects: int, max_materials: int = 1024, max_lights: int = 256, max_joints: int = 256
) -> None:
self._engine = engine
self.max_objects = max_objects
self.max_materials = max_materials
self.max_lights = max_lights
self.max_joints = max_joints
# SSBO resources.
#
# transform / aabb / shadow / frame_globals hold data that CHANGES every
# frame (object + camera motion), so they are ringed across
# FRAMES_IN_FLIGHT: a single copy would be host-overwritten while the
# previous in-flight frame's GPU passes still read it (a cross-frame WAR
# race). Each is stored as an N-slot list; the ``*_buf`` / ``*_mem``
# properties below return the current frame's slot, so every upload and
# bind site follows the ring with no change. material / light are
# dirty-tracked and mostly static, so they stay single (gated, not ringed).
self._transform_bufs: list[Any] = []
self._transform_mems: list[Any] = []
self.material_buf: Any = None
self.material_mem: Any = None
self.light_buf: Any = None
self.light_mem: Any = None
self._shadow_bufs: list[Any] = []
self._shadow_mems: list[Any] = []
self.tile_light_idx_buf: Any = None
self.tile_light_idx_mem: Any = None
self.tile_info_buf: Any = None
self.tile_info_mem: Any = None
# Per-instance LOCAL-AABB SSBO, slot-aligned 1:1 with the transform SSBO
# (row i == _instances[i]). Read-only by the GPU occlusion-cull compute
# (phase O3). Always allocated (small, 32 bytes/slot) but only ever read
# when occlusion culling is enabled; it carries no fragment/vertex binding.
self._aabb_bufs: list[Any] = []
self._aabb_mems: list[Any] = []
# Descriptors. The main SSBO set is ringed across FRAMES_IN_FLIGHT: set
# ``i`` binds the per-frame slot ``i`` of transform / shadow / frame_globals
# (all other bindings are the same shared buffer in every set). The
# ``ssbo_set`` property returns the current frame's set, so every bind site
# follows the ring unchanged; ``ssbo_sets`` exposes the whole list for the
# few external writers that must fan a shared binding out to all sets.
self.ssbo_layout: Any = None
self.ssbo_pool: Any = None
self.ssbo_sets: list[Any] = []
# Joint (bone palette) SSBO: host-written every skinned frame with the
# animated palette, so it is ringed across FRAMES_IN_FLIGHT like transform.
self.joint_layout: Any = None
self.joint_pool: Any = None
self._joint_bufs: list[Any] = []
self._joint_mems: list[Any] = []
self._joint_sets: list[Any] = []
# IBL cubemap placeholder (owned for lifetime of manager)
self.placeholder_cubemap_view: Any = None
self.placeholder_cubemap_sampler: Any = None
self.placeholder_cubemap_img: Any = None
self.placeholder_cubemap_mem: Any = None
# 1×1 2D placeholder for the BRDF-LUT slot (binding 9) when no sky is set.
self.placeholder_lut_view: Any = None
self.placeholder_lut_img: Any = None
self.placeholder_lut_mem: Any = None
# Reflection-probe cubemap-array placeholder (bindings 10/11): a 1-layer
# cube array bound when no probe has captured yet. Replaced by
# ``write_probe_descriptors`` once ReflectionProbePass owns real arrays.
self.placeholder_cubearray_view: Any = None
self.placeholder_cubearray_img: Any = None
self.placeholder_cubearray_mem: Any = None
self.placeholder_cubearray_sampler: Any = None
# Reflection-probe box SSBO (binding 12): count header + MAX_PROBES boxes.
self.probe_buf: Any = None
self.probe_mem: Any = None
# Irradiance-volume SH SSBO (binding 18).
self.irradiance_volume_buf: Any = None
self.irradiance_volume_mem: Any = None
# Decal SSBO (binding 19; grows for the screen-tile masks).
self.decal_buf: Any = None
self.decal_mem: Any = None
self.decal_buf_size: int = 0
# FrameGlobals UBO (binding 13): one 256-byte std140 block of per-frame
# globals (time / wind / wetness / render extents). Host-visible, written
# once with zeros at setup (never-unbound fallback) and repacked per frame
# by ``update_frame_globals``; read by the water / ocean / cube shaders.
# Ringed across FRAMES_IN_FLIGHT (per-frame content).
self._frame_globals_bufs: list[Any] = []
self._frame_globals_mems: list[Any] = []
# Per-frame counter + previous scene time, used to derive frame_index and
# delta for the block without a renderer-wide frame clock.
self._frame_globals_index: int = 0
self._frame_globals_prev_now: float = 0.0
# Dirty-tracking: skip redundant GPU uploads when data hasn't changed
self._materials_hash: int = 0
self._lights_hash: int = 0
# ---- Frame-arena bump allocator over the shared transform SSBO ----
# Capacity (in slots) == ``max_objects``. Each scene-render-unit reserves
# a contiguous slot range [base, base+count); the cursor resets at the
# start of each real frame. Tracks the per-frame high-water mark so a
# grow can be scheduled BEFORE an actual overflow forces clamping.
self._frame_cursor: int = 0
# Peak total slots reserved in any single frame seen so far (high-water).
self.transform_high_water: int = 0
# Capacity to grow to at the next frame boundary (0 == no grow pending).
self._pending_grow: int = 0
# Latch so the overflow WARNING is logged once per overflow episode (not
# every frame while saturated); cleared once a grow lands.
self._overflow_warned: bool = False
# Joint (bone palette) SSBO grow scheduling, mirroring the transform arena:
# the skinned pass clamps the concatenated palette to current capacity and
# schedules a grow for the next frame boundary so it never writes OOB.
self._pending_joint_grow: int = 0
self._joint_overflow_warned: bool = False
# ---------------------------------------------- per-frame ring accessors
@property
def _frame(self) -> int:
"""The frame-in-flight slot currently recording (0..FRAMES_IN_FLIGHT-1)."""
return self._engine.current_frame
[docs]
@property
def aabb_buf(self) -> Any:
return self._aabb_bufs[self._frame] if self._aabb_bufs else None
[docs]
@property
def aabb_mem(self) -> Any:
return self._aabb_mems[self._frame] if self._aabb_mems else None
[docs]
@property
def shadow_buf(self) -> Any:
return self._shadow_bufs[self._frame] if self._shadow_bufs else None
[docs]
@property
def shadow_mem(self) -> Any:
return self._shadow_mems[self._frame] if self._shadow_mems else None
[docs]
@property
def frame_globals_buf(self) -> Any:
return self._frame_globals_bufs[self._frame] if self._frame_globals_bufs else None
[docs]
@property
def frame_globals_mem(self) -> Any:
return self._frame_globals_mems[self._frame] if self._frame_globals_mems else None
[docs]
@property
def ssbo_set(self) -> Any:
"""The current frame's main SSBO descriptor set (what every pass binds)."""
return self.ssbo_sets[self._frame] if self.ssbo_sets else None
[docs]
@property
def joint_buf(self) -> Any:
return self._joint_bufs[self._frame] if self._joint_bufs else None
[docs]
@property
def joint_mem(self) -> Any:
return self._joint_mems[self._frame] if self._joint_mems else None
[docs]
@property
def joint_set(self) -> Any:
"""The current frame's joint descriptor set (bound for skinned meshes)."""
return self._joint_sets[self._frame] if self._joint_sets else None
def _write_ssbo_all(self, binding: int, buf: Any, size: int) -> None:
"""Write an SSBO binding shared by every frame into all N sets."""
device = self._engine.ctx.device
for s in self.ssbo_sets:
write_ssbo_descriptor(device, s, binding, buf, size)
def _write_ubo_all(self, binding: int, buf: Any, size: int) -> None:
device = self._engine.ctx.device
for s in self.ssbo_sets:
write_ubo_descriptor(device, s, binding, buf, size)
def _write_image_all(self, binding: int, view: Any, sampler: Any) -> None:
device = self._engine.ctx.device
for s in self.ssbo_sets:
write_image_descriptor(device, s, binding, view, sampler)
# ------------------------------------------------------------------ setup
[docs]
def setup(self) -> None:
"""Allocate all SSBOs and descriptor sets."""
e = self._engine
device = e.ctx.device
phys = e.ctx.physical_device
transform_size = self.max_objects * TRANSFORM_DTYPE.itemsize
material_size = self.max_materials * MATERIAL_DTYPE.itemsize
light_size = self.max_lights * LIGHT_DTYPE.itemsize
joint_buf_size = self.max_joints * 64 # mat4 = 64 bytes
aabb_size = self.max_objects * AABB_DTYPE.itemsize
# Per-frame (ringed) buffers: one copy per frame in flight.
for _ in range(FRAMES_IN_FLIGHT):
tbuf, tmem = create_buffer(device, phys, transform_size, _SSBO_USAGE, _HOST_FLAGS)
abuf, amem = create_buffer(device, phys, aabb_size, _SSBO_USAGE, _HOST_FLAGS)
sbuf, smem = create_buffer(device, phys, SHADOW_DATA_SIZE, _SSBO_USAGE, _HOST_FLAGS)
self._transform_bufs.append(tbuf)
self._transform_mems.append(tmem)
self._aabb_bufs.append(abuf)
self._aabb_mems.append(amem)
self._shadow_bufs.append(sbuf)
self._shadow_mems.append(smem)
# Shared (single, dirty-tracked or static) buffers.
self.material_buf, self.material_mem = create_buffer(device, phys, material_size, _SSBO_USAGE, _HOST_FLAGS)
self.light_buf, self.light_mem = create_buffer(device, phys, light_size, _SSBO_USAGE, _HOST_FLAGS)
self.tile_light_idx_buf, self.tile_light_idx_mem = create_buffer(device, phys, 16, _SSBO_USAGE, _HOST_FLAGS)
self.tile_info_buf, self.tile_info_mem = create_buffer(device, phys, 16, _SSBO_USAGE, _HOST_FLAGS)
for _ in range(FRAMES_IN_FLIGHT):
jbuf, jmem = create_buffer(device, phys, joint_buf_size, _SSBO_USAGE, _HOST_FLAGS)
self._joint_bufs.append(jbuf)
self._joint_mems.append(jmem)
# Main SSBO set: 4 SSBOs + 1 cubemap sampler (4) + 2 trailing SSBOs (5-6)
# + 3 IBL samplers (7=irradiance cube, 8=prefilter cube, 9=BRDF 2D LUT)
# + 2 reflection-probe cubemap arrays (10=irradiance array, 11=prefilter
# array) + 1 reflection-probe box SSBO (12).
# The forward set's image-sampler bindings (cubemap, IBL maps, reflection-
# probe arrays) are written after the set is bound into the recording frame
# command buffer (skybox install in pre_render, per-capture probe updates),
# so the set is UPDATE_AFTER_BIND. Removes the old "sync every descriptor
# before the first bind" ordering constraint.
# ``tail_ubos=1`` appends binding 13, the FrameGlobals UBO.
# ``tail3_samplers=2`` appends bindings 14/15, the scene colour / depth
# copy samplers; they sit after the UBO tail and default
# to a 1x1 fallback so they are never unbound when the split is inactive.
# ``tail4_samplers=2`` appends bindings 16/17, the pluggable-ambient
# indirect specular / diffuse hooks; 1x1 black fallbacks
# gated off by the shadow SSBO's ``indirect_*_enabled`` flags.
self.ssbo_layout = create_ssbo_layout(
device,
binding_count=4,
extra_samplers=1,
trailing_ssbos=2,
extra_samplers_tail=3,
tail2_samplers=2,
tail2_ssbos=1,
tail_ubos=1,
tail3_samplers=2,
tail4_samplers=2,
tail5_ssbos=1,
tail6_ssbos=1,
update_after_bind=True,
)
# One set per frame in flight; pool sizes scale with FRAMES_IN_FLIGHT.
self.ssbo_pool = create_descriptor_pool(
device,
max_sets=FRAMES_IN_FLIGHT,
extra_samplers=10 * FRAMES_IN_FLIGHT,
ssbo_count=9 * FRAMES_IN_FLIGHT,
ubo_count=1 * FRAMES_IN_FLIGHT,
update_after_bind=True,
)
self.ssbo_sets = [
allocate_descriptor_set(device, self.ssbo_pool, self.ssbo_layout) for _ in range(FRAMES_IN_FLIGHT)
]
for i, s in enumerate(self.ssbo_sets):
# Per-frame bindings point at slot i; shared bindings are identical.
write_ssbo_descriptor(device, s, 0, self._transform_bufs[i], transform_size)
write_ssbo_descriptor(device, s, 3, self._shadow_bufs[i], SHADOW_DATA_SIZE)
self._write_ssbo_all(1, self.material_buf, material_size)
self._write_ssbo_all(2, self.light_buf, light_size)
self._write_ssbo_all(5, self.tile_light_idx_buf, 16)
self._write_ssbo_all(6, self.tile_info_buf, 16)
# Joint SSBO set (set 2, binding 0): one per frame in flight, each bound to
# its own ringed joint buffer.
self.joint_layout = create_ssbo_layout(device, binding_count=1)
self.joint_pool = create_descriptor_pool(device, max_sets=FRAMES_IN_FLIGHT)
for i in range(FRAMES_IN_FLIGHT):
js = allocate_descriptor_set(device, self.joint_pool, self.joint_layout)
self._joint_sets.append(js)
write_ssbo_descriptor(device, js, 0, self._joint_bufs[i], joint_buf_size)
# Shadow SSBO defaults: no-shadow sentinels + ambient colour
init_shadow = np.zeros(SHADOW_DATA_SIZE, dtype=np.uint8)
sentinel = np.array([0xFF, 0xFF, 0xFF, 0xFF], dtype=np.uint8)
init_shadow[208:212] = sentinel
init_shadow[220:224] = sentinel
init_shadow[224:228] = sentinel
init_shadow[336:352] = _DEFAULT_AMBIENT.view(np.uint8)
for smem in self._shadow_mems:
upload_numpy(device, smem, init_shadow)
# IBL cubemap placeholder (replaced by Renderer.set_skybox)
from ..assets.cubemap_loader import load_cubemap
(
self.placeholder_cubemap_view,
self.placeholder_cubemap_sampler,
self.placeholder_cubemap_img,
self.placeholder_cubemap_mem,
) = load_cubemap(device, phys, e.ctx.graphics_queue, e.ctx.command_pool, colour=(0.0, 0.0, 0.0))
self._write_image_all(4, self.placeholder_cubemap_view, self.placeholder_cubemap_sampler)
# 1×1 black 2D texture for the BRDF-LUT slot when no sky is set. Never
# sampled while ``ibl_enabled == 0`` (the shader branches it out), but
# the descriptor must be a valid sampler2D.
self.placeholder_lut_img, self.placeholder_lut_mem = upload_image_data(
device,
phys,
e.ctx.graphics_queue,
e.ctx.command_pool,
np.zeros((1, 1, 4), dtype=np.uint8),
1,
1,
)
self.placeholder_lut_view = vk.vkCreateImageView(
device,
vk.VkImageViewCreateInfo(
image=self.placeholder_lut_img,
viewType=vk.VK_IMAGE_VIEW_TYPE_2D,
format=vk.VK_FORMAT_R8G8B8A8_UNORM,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=1,
),
),
None,
)
# Bind the fallbacks to 7/8 (cube) and 9 (2D). Replaced by
# ``write_ibl_descriptors`` when a skybox + IBL precompute is installed.
self._write_ibl_fallback()
# Pluggable-ambient hook fallbacks: bind the shared 1x1
# black 2D placeholder to 16 (indirect specular) / 17 (indirect diffuse)
# so the descriptors are never unbound. The shadow SSBO's
# ``indirect_*_enabled`` flags (0 until an SSR / SSGI producer flips
# them) branch the samples out entirely, so these stay unread today.
s = self.placeholder_cubemap_sampler
self._write_image_all(16, self.placeholder_lut_view, s)
self._write_image_all(17, self.placeholder_lut_view, s)
# Reflection-probe box SSBO (binding 12): defaults to zero probes.
self.probe_buf, self.probe_mem = create_buffer(
device,
phys,
PROBE_BUFFER_SIZE,
_SSBO_USAGE,
_HOST_FLAGS,
)
upload_numpy(device, self.probe_mem, np.zeros(PROBE_BUFFER_SIZE, dtype=np.uint8))
self._write_ssbo_all(12, self.probe_buf, PROBE_BUFFER_SIZE)
# Reflection-probe cubemap-array placeholders (bindings 10/11). A single
# black 1×1×6 cube layer bound as a CUBE_ARRAY so the descriptor is valid
# before any probe captures (the shader's ``probe_count == 0`` gate skips
# sampling). Replaced by ``write_probe_descriptors`` once a probe is live.
self._create_cubearray_placeholder(device, phys)
self._write_probe_array_fallback()
# FrameGlobals UBO (binding 13): host-visible 256-byte block, zero-filled
# so the descriptor is never unbound before the first per-frame repack.
# Ringed: set i binds slot i (content is repacked every frame).
zeroed = fg.zeroed()
for s in self.ssbo_sets:
buf, mem = create_buffer(
device, phys, fg.FRAME_GLOBALS_SIZE, vk.VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, _HOST_FLAGS
)
upload_numpy(device, mem, zeroed)
self._frame_globals_bufs.append(buf)
self._frame_globals_mems.append(mem)
write_ubo_descriptor(device, s, 13, buf, fg.FRAME_GLOBALS_SIZE)
# Irradiance-volume SH SSBO (binding 18): host-visible,
# zero-filled so ``grid.w`` (probe count) is 0 and the uber's
# ``irradiance_volume_enabled`` gate never reads it before a volume bakes.
# The IrradianceVolumePass SH-reduce compute writes probe slices into it.
self.irradiance_volume_buf, self.irradiance_volume_mem = create_buffer(
device, phys, IRRADIANCE_VOLUME_BUFFER_SIZE, _SSBO_USAGE, _HOST_FLAGS
)
upload_numpy(device, self.irradiance_volume_mem, np.zeros(IRRADIANCE_VOLUME_BUFFER_SIZE, dtype=np.uint8))
self._write_ssbo_all(18, self.irradiance_volume_buf, IRRADIANCE_VOLUME_BUFFER_SIZE)
# Decal SSBO (binding 19): host-visible, zero-filled
# so ``decal_count`` is 0 and the uber's decal projection loop never runs
# before a Decal3D is packed. Rewritten each frame by ``write_decals``
# from ``forward._pack_decals`` (a scene with no decal keeps it zeroed, so
# the frame stays byte-identical: zero cost when unused).
self.decal_buf, self.decal_mem = create_buffer(
device, phys, DECAL_BUFFER_SIZE, _SSBO_USAGE, _HOST_FLAGS
)
self.decal_buf_size = DECAL_BUFFER_SIZE
upload_numpy(device, self.decal_mem, np.zeros(DECAL_BUFFER_SIZE, dtype=np.uint8))
self._write_ssbo_all(19, self.decal_buf, DECAL_BUFFER_SIZE)
# -------------------------------------------------------- frame arena
[docs]
def begin_frame_arena(self) -> None:
"""Reset the per-frame bump cursor. Call once at the start of every frame."""
self._frame_cursor = 0
[docs]
def reserve_slots(self, count: int) -> int:
"""Reserve a contiguous range of ``count`` transform slots; return its base.
Bump-allocates from the shared SSBO arena. On overflow (the running cursor
would pass capacity) the request is clamped to what fits, a WARNING is
logged once, and a grow to ``next_pow2(needed)`` is scheduled for the next
frame boundary. The high-water mark is updated either way so a proactive
grow can fire before an actual overflow recurs.
Returns the base slot. The caller records draws with
``first_instance = base + local_index`` and writes its transforms into the
SSBO at byte offset ``base * TRANSFORM_DTYPE.itemsize``.
"""
if count < 0:
count = 0
base = self._frame_cursor
cap = self.max_objects
needed = base + count
if needed > cap:
granted = max(0, cap - base)
if not self._overflow_warned:
log.warning(
"Transform SSBO arena overflow: frame needs %d slots, capacity %d; "
"rendering %d, growing buffer before next frame",
needed,
cap,
base + granted,
)
self._overflow_warned = True
self._pending_grow = max(self._pending_grow, _next_pow2(needed))
self._frame_cursor = cap
self.transform_high_water = max(self.transform_high_water, needed)
return base
self._frame_cursor = needed
self.transform_high_water = max(self.transform_high_water, needed)
return base
[docs]
def request_joint_capacity(self, joint_count: int) -> int:
"""Reserve room for ``joint_count`` bone matrices in the joint SSBO.
Returns the number that fit in the CURRENT buffer (so the caller clamps
its upload and never writes past the allocation). When the request exceeds
capacity a grow is scheduled for the next frame boundary (mirrors the
transform arena), so the dropped skeletons appear one frame later instead
of corrupting device memory. Symmetric with the web ``_ensureBoneBuf``.
"""
if joint_count <= self.max_joints:
return joint_count
if not self._joint_overflow_warned:
log.warning(
"Joint SSBO overflow: frame needs %d bone matrices, capacity %d; "
"rendering %d, growing buffer before next frame",
joint_count,
self.max_joints,
self.max_joints,
)
self._joint_overflow_warned = True
self._pending_joint_grow = max(self._pending_joint_grow, _next_pow2(joint_count))
return self.max_joints
[docs]
def maybe_grow(self) -> bool:
"""Grow the transform/AABB and joint SSBOs if scheduled. Frame-boundary only.
Reallocates when a grow is pending (an overflow clamped this frame) OR the
high-water mark has crossed ``_GROW_THRESHOLD`` of capacity (proactive).
Must run OUTSIDE command recording (begin_frame or post-submit) so the
reallocated buffer + descriptor rewrite never race the GPU. Returns True
when a grow happened.
"""
grew = False
target = self._pending_grow
if target <= self.max_objects and self.transform_high_water > _GROW_THRESHOLD * self.max_objects:
target = max(target, _next_pow2(self.max_objects * 2))
if target > self.max_objects:
self._grow_transform_buffer(target)
grew = True
if self._pending_joint_grow > self.max_joints:
self._grow_joint_buffer(self._pending_joint_grow)
grew = True
return grew
def _grow_joint_buffer(self, new_capacity: int) -> None:
"""Reallocate the joint SSBO to ``new_capacity`` mat4s and repoint its descriptor.
Run only at a frame boundary (via :meth:`maybe_grow`), so no GPU work
references the old buffer when it is freed.
"""
e = self._engine
device = e.ctx.device
phys = e.ctx.physical_device
old_cap = self.max_joints
joint_buf_size = new_capacity * 64
old = list(zip(self._joint_bufs, self._joint_mems, strict=True))
self._joint_bufs, self._joint_mems = [], []
for i in range(FRAMES_IN_FLIGHT):
jbuf, jmem = create_buffer(device, phys, joint_buf_size, _SSBO_USAGE, _HOST_FLAGS)
self._joint_bufs.append(jbuf)
self._joint_mems.append(jmem)
write_ssbo_descriptor(device, self._joint_sets[i], 0, jbuf, joint_buf_size)
self.max_joints = new_capacity
for old_joint_buf, old_joint_mem in old:
if old_joint_buf:
vk.vkDestroyBuffer(device, old_joint_buf, None)
if old_joint_mem:
vk.vkFreeMemory(device, old_joint_mem, None)
self._pending_joint_grow = 0
self._joint_overflow_warned = False
log.info("Joint SSBO grown: %d -> %d bone matrices", old_cap, new_capacity)
def _grow_transform_buffer(self, new_capacity: int) -> None:
"""Reallocate the transform + AABB SSBO ring to ``new_capacity`` slots.
Grows all FRAMES_IN_FLIGHT copies in lock-step and rewrites the transform
descriptor (binding 0) of every set. The AABB SSBO is slot-aligned 1:1 and
read by the occlusion pass, so it is grown alongside; the occlusion pass's
cached descriptor sets captured the OLD buffers, so its cache is dropped.
Runs at a frame boundary (no GPU work in flight) so the frees are safe.
"""
e = self._engine
device = e.ctx.device
phys = e.ctx.physical_device
old_cap = self.max_objects
transform_size = new_capacity * TRANSFORM_DTYPE.itemsize
aabb_size = new_capacity * AABB_DTYPE.itemsize
old = list(zip(self._transform_bufs, self._transform_mems, self._aabb_bufs, self._aabb_mems, strict=True))
self._transform_bufs, self._transform_mems = [], []
self._aabb_bufs, self._aabb_mems = [], []
for i in range(FRAMES_IN_FLIGHT):
tbuf, tmem = create_buffer(device, phys, transform_size, _SSBO_USAGE, _HOST_FLAGS)
abuf, amem = create_buffer(device, phys, aabb_size, _SSBO_USAGE, _HOST_FLAGS)
self._transform_bufs.append(tbuf)
self._transform_mems.append(tmem)
self._aabb_bufs.append(abuf)
self._aabb_mems.append(amem)
write_ssbo_descriptor(device, self.ssbo_sets[i], 0, tbuf, transform_size)
self.max_objects = new_capacity
# The occlusion pass cached sets pointing at the freed transform/aabb.
occ = getattr(self._engine, "_renderer", None)
occ = getattr(occ, "_occlusion_pass", None) if occ is not None else None
if occ is not None and hasattr(occ, "invalidate_cache"):
occ.invalidate_cache()
for tbuf, tmem, abuf, amem in old:
for buf in (tbuf, abuf):
if buf:
vk.vkDestroyBuffer(device, buf, None)
for mem in (tmem, amem):
if mem:
vk.vkFreeMemory(device, mem, None)
self._pending_grow = 0
self._overflow_warned = False
log.info(
"Transform SSBO grown: %d -> %d slots (high-water %d)", old_cap, new_capacity, self.transform_high_water
)
def _create_cubearray_placeholder(self, device: Any, phys: Any) -> None:
"""Create a 1-probe (6-layer) black cube-array as the binding 10/11 fallback."""
from ..assets.cubemap_loader import load_cubemap
# load_cubemap builds a 6-layer CUBE image + sampler. Re-view it as a
# CUBE_ARRAY (layerCount=6 = one probe) so it satisfies samplerCubeArray.
(
_cube_view,
self.placeholder_cubearray_sampler,
self.placeholder_cubearray_img,
self.placeholder_cubearray_mem,
) = load_cubemap(
device, phys, self._engine.ctx.graphics_queue, self._engine.ctx.command_pool, colour=(0.0, 0.0, 0.0)
)
vk.vkDestroyImageView(device, _cube_view, None)
self.placeholder_cubearray_view = vk.vkCreateImageView(
device,
vk.VkImageViewCreateInfo(
image=self.placeholder_cubearray_img,
viewType=vk.VK_IMAGE_VIEW_TYPE_CUBE_ARRAY,
format=vk.VK_FORMAT_R32G32B32A32_SFLOAT,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=6,
),
),
None,
)
def _write_probe_array_fallback(self) -> None:
"""Bind the placeholder cube-array to the probe slots (10=irradiance, 11=prefilter)."""
s = self.placeholder_cubearray_sampler
self._write_image_all(10, self.placeholder_cubearray_view, s)
self._write_image_all(11, self.placeholder_cubearray_view, s)
[docs]
def write_probe_descriptors(self, irradiance_array_view: Any, prefilter_array_view: Any, sampler: Any) -> None:
"""Bind the ReflectionProbePass's cubemap arrays to the forward set (bindings 10/11)."""
self._write_image_all(10, irradiance_array_view, sampler)
self._write_image_all(11, prefilter_array_view, sampler)
[docs]
def write_probe_buffer(self, data: np.ndarray) -> None:
"""Upload the probe box SSBO bytes (count header + Probe array, binding 12)."""
upload_numpy(self._engine.ctx.device, self.probe_mem, data)
[docs]
def write_decals(self, data: np.ndarray) -> None:
"""Upload the decal SSBO bytes (count header + Decal array + tile masks, b19).
``data`` is the block from ``decal_pack.build_decal_buffer``. The base
(header + record) region is a fixed :data:`DECAL_BUFFER_SIZE`; when the
screen-tile cull is active the per-tile bitmask array is appended,
so the block can exceed the initial allocation. It is grown (and its b19
descriptor repointed) on demand at the frame boundary; the growth only ever
adds tile-mask capacity, so it never perturbs the record region. An all-zero
block (no decals) leaves ``decal_count`` at 0, a byte-identical no-op.
"""
device = self._engine.ctx.device
if data.nbytes > self.decal_buf_size:
phys = self._engine.ctx.physical_device
new_size = int(data.nbytes)
new_buf, new_mem = create_buffer(device, phys, new_size, _SSBO_USAGE, _HOST_FLAGS)
self._write_ssbo_all(19, new_buf, new_size)
old_buf, old_mem = self.decal_buf, self.decal_mem
self.decal_buf, self.decal_mem, self.decal_buf_size = new_buf, new_mem, new_size
if old_buf:
vk.vkDestroyBuffer(device, old_buf, None)
if old_mem:
vk.vkFreeMemory(device, old_mem, None)
upload_numpy(device, self.decal_mem, data)
def _write_ibl_fallback(self) -> None:
"""Bind placeholder textures to the IBL slots (7=irradiance, 8=prefilter,
9=BRDF). Used at init and whenever the skybox is cleared."""
s = self.placeholder_cubemap_sampler
self._write_image_all(7, self.placeholder_cubemap_view, s)
self._write_image_all(8, self.placeholder_cubemap_view, s)
self._write_image_all(9, self.placeholder_lut_view, s)
[docs]
def write_ibl_descriptors(self, irradiance_view: Any, prefilter_view: Any, brdf_view: Any, sampler: Any) -> None:
"""Bind an IBLPass's precomputed maps to the forward set (bindings 7/8/9)."""
self._write_image_all(7, irradiance_view, sampler)
self._write_image_all(8, prefilter_view, sampler)
self._write_image_all(9, brdf_view, sampler)
# ---------------------------------------------------------------- uploads
[docs]
def set_materials(self, materials: np.ndarray) -> np.ndarray:
"""Upload material array. Returns the (possibly clamped) array stored."""
if len(materials) > self.max_materials:
log.warning("Material count (%d) exceeds max (%d), clamping", len(materials), self.max_materials)
materials = materials[: self.max_materials]
if self.material_mem:
h = hash(materials.tobytes())
if h != self._materials_hash:
self._materials_hash = h
upload_numpy(self._engine.ctx.device, self.material_mem, materials)
return materials
[docs]
def set_lights(self, lights: np.ndarray) -> None:
"""Upload light array prefixed with the uint32 count (GLSL LightBuffer layout)."""
if not self.light_mem:
return
h = hash(lights.tobytes())
if h == self._lights_hash:
return
self._lights_hash = h
count = np.array([len(lights)], dtype=np.uint32)
padding = np.zeros(3, dtype=np.uint32)
header = np.concatenate([count, padding])
buf = np.concatenate([header.view(np.uint8), lights.view(np.uint8)])
upload_numpy(self._engine.ctx.device, self.light_mem, buf)
[docs]
def write_shadow_data(self, shadow_data: np.ndarray) -> None:
"""Upload raw shadow SSBO bytes (used by shadow passes + IBL-only fallback)."""
upload_numpy(self._engine.ctx.device, self.shadow_mem, shadow_data)
[docs]
def write_cubemap_descriptor(self, view: Any, sampler: Any) -> None:
"""Bind a cubemap view+sampler to the IBL slot (binding 4)."""
self._write_image_all(4, view, sampler)
[docs]
def update_frame_globals(
self,
now: float,
env: fg.FrameGlobalsEnv,
internal_extent: tuple[int, int],
output_extent: tuple[int, int],
render_scale: float = 1.0,
) -> None:
"""Repack + upload the per-frame FrameGlobals UBO (binding 13).
``delta`` is derived from the scene clock (``now`` minus the previous
call's ``now``) and ``frame_index`` from an internal counter, so the block
needs no renderer-wide frame clock. ``internal_extent`` is the HDR-chain
render extent, ``output_extent`` the swapchain extent; they
are equal (and ``render_scale`` 1.0) on the unscaled path, packing the
exact pre-scale block.
"""
delta = now - self._frame_globals_prev_now
if delta < 0.0:
delta = 0.0 # scene clock reset (fresh SceneTree)
self._frame_globals_prev_now = now
block = fg.pack(
time=now,
delta=delta,
frame_index=self._frame_globals_index,
wind_direction=env.wind_direction,
wind_strength=env.wind_strength,
wind_gustiness=env.wind_gustiness,
wetness=env.wetness,
rain_intensity=env.rain_intensity,
ripple_strength=env.ripple_strength,
internal_w=float(internal_extent[0]),
internal_h=float(internal_extent[1]),
render_scale=float(render_scale),
output_w=float(output_extent[0]),
output_h=float(output_extent[1]),
)
upload_numpy(self._engine.ctx.device, self.frame_globals_mem, block)
self._frame_globals_index += 1
# ---------------------------------------------------------------- cleanup
[docs]
def cleanup(self) -> None:
"""Destroy all buffers, descriptor pools/layouts, and placeholder cubemap."""
device = self._engine.ctx.device
# Ringed per-frame buffers: destroy every slot (the *_buf properties only
# expose the current one).
ringed = (
(self._transform_bufs, self._transform_mems),
(self._aabb_bufs, self._aabb_mems),
(self._shadow_bufs, self._shadow_mems),
(self._frame_globals_bufs, self._frame_globals_mems),
(self._joint_bufs, self._joint_mems),
)
for bufs, mems in ringed:
for buf, mem in zip(bufs, mems, strict=True):
if buf:
vk.vkDestroyBuffer(device, buf, None)
if mem:
vk.vkFreeMemory(device, mem, None)
for buf, mem in (
(self.material_buf, self.material_mem),
(self.light_buf, self.light_mem),
(self.tile_light_idx_buf, self.tile_light_idx_mem),
(self.tile_info_buf, self.tile_info_mem),
(self.probe_buf, self.probe_mem),
(self.irradiance_volume_buf, self.irradiance_volume_mem),
(self.decal_buf, self.decal_mem),
):
if buf:
vk.vkDestroyBuffer(device, buf, None)
if mem:
vk.vkFreeMemory(device, mem, None)
for layout in (self.joint_layout, self.ssbo_layout):
if layout:
vk.vkDestroyDescriptorSetLayout(device, layout, None)
for pool in (self.joint_pool, self.ssbo_pool):
if pool:
vk.vkDestroyDescriptorPool(device, pool, None)
if self.placeholder_cubemap_sampler:
vk.vkDestroySampler(device, self.placeholder_cubemap_sampler, None)
if self.placeholder_cubemap_view:
vk.vkDestroyImageView(device, self.placeholder_cubemap_view, None)
if self.placeholder_cubemap_img:
vk.vkDestroyImage(device, self.placeholder_cubemap_img, None)
if self.placeholder_cubemap_mem:
vk.vkFreeMemory(device, self.placeholder_cubemap_mem, None)
if self.placeholder_lut_view:
vk.vkDestroyImageView(device, self.placeholder_lut_view, None)
if self.placeholder_lut_img:
vk.vkDestroyImage(device, self.placeholder_lut_img, None)
if self.placeholder_lut_mem:
vk.vkFreeMemory(device, self.placeholder_lut_mem, None)
if self.placeholder_cubearray_sampler:
vk.vkDestroySampler(device, self.placeholder_cubearray_sampler, None)
if self.placeholder_cubearray_view:
vk.vkDestroyImageView(device, self.placeholder_cubearray_view, None)
if self.placeholder_cubearray_img:
vk.vkDestroyImage(device, self.placeholder_cubearray_img, None)
if self.placeholder_cubearray_mem:
vk.vkFreeMemory(device, self.placeholder_cubearray_mem, None)