"""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 ..types import LIGHT_DTYPE, MATERIAL_DTYPE, TRANSFORM_DTYPE
from ..gpu.descriptors import (
allocate_descriptor_set,
create_descriptor_pool,
create_ssbo_layout,
write_image_descriptor,
write_ssbo_descriptor,
)
from ..gpu.memory import create_buffer, upload_image_data, upload_numpy
__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) = 352
SHADOW_DATA_SIZE = 352
# 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
# 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
[docs]
class BufferManager:
"""Owns the renderer's SSBOs and descriptor sets.
Main descriptor set (``ssbo_set``) exposes thirteen 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.
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
self.transform_buf: Any = None
self.transform_mem: Any = None
self.material_buf: Any = None
self.material_mem: Any = None
self.light_buf: Any = None
self.light_mem: Any = None
self.shadow_buf: Any = None
self.shadow_mem: Any = None
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
self.joint_buf: Any = None
self.joint_mem: Any = None
# Descriptors
self.ssbo_layout: Any = None
self.ssbo_pool: Any = None
self.ssbo_set: Any = None
self.joint_layout: Any = None
self.joint_pool: Any = None
self.joint_set: Any = None
# 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
# Dirty-tracking: skip redundant GPU uploads when data hasn't changed
self._materials_hash: int = 0
self._lights_hash: int = 0
# ------------------------------------------------------------------ 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
self.transform_buf, self.transform_mem = create_buffer(device, phys, transform_size, _SSBO_USAGE, _HOST_FLAGS)
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.shadow_buf, self.shadow_mem = create_buffer(device, phys, SHADOW_DATA_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)
self.joint_buf, self.joint_mem = create_buffer(device, phys, joint_buf_size, _SSBO_USAGE, _HOST_FLAGS)
# 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).
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,
)
self.ssbo_pool = create_descriptor_pool(device, max_sets=1, extra_samplers=6, ssbo_count=7)
self.ssbo_set = allocate_descriptor_set(device, self.ssbo_pool, self.ssbo_layout)
write_ssbo_descriptor(device, self.ssbo_set, 0, self.transform_buf, transform_size)
write_ssbo_descriptor(device, self.ssbo_set, 1, self.material_buf, material_size)
write_ssbo_descriptor(device, self.ssbo_set, 2, self.light_buf, light_size)
write_ssbo_descriptor(device, self.ssbo_set, 3, self.shadow_buf, SHADOW_DATA_SIZE)
write_ssbo_descriptor(device, self.ssbo_set, 5, self.tile_light_idx_buf, 16)
write_ssbo_descriptor(device, self.ssbo_set, 6, self.tile_info_buf, 16)
# Joint SSBO set (set 2, binding 0)
self.joint_layout = create_ssbo_layout(device, binding_count=1)
self.joint_pool = create_descriptor_pool(device, max_sets=2)
self.joint_set = allocate_descriptor_set(device, self.joint_pool, self.joint_layout)
write_ssbo_descriptor(device, self.joint_set, 0, self.joint_buf, 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)
upload_numpy(device, self.shadow_mem, 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))
write_image_descriptor(
device, self.ssbo_set, 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()
# 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))
write_ssbo_descriptor(device, self.ssbo_set, 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()
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)."""
device = self._engine.ctx.device
s = self.placeholder_cubearray_sampler
write_image_descriptor(device, self.ssbo_set, 10, self.placeholder_cubearray_view, s)
write_image_descriptor(device, self.ssbo_set, 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)."""
device = self._engine.ctx.device
write_image_descriptor(device, self.ssbo_set, 10, irradiance_array_view, sampler)
write_image_descriptor(device, self.ssbo_set, 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)
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."""
device = self._engine.ctx.device
s = self.placeholder_cubemap_sampler
write_image_descriptor(device, self.ssbo_set, 7, self.placeholder_cubemap_view, s)
write_image_descriptor(device, self.ssbo_set, 8, self.placeholder_cubemap_view, s)
write_image_descriptor(device, self.ssbo_set, 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)."""
device = self._engine.ctx.device
write_image_descriptor(device, self.ssbo_set, 7, irradiance_view, sampler)
write_image_descriptor(device, self.ssbo_set, 8, prefilter_view, sampler)
write_image_descriptor(device, self.ssbo_set, 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 set_hdr_flag(self, enabled: bool) -> None:
"""Toggle ``hdr_output`` (byte offset 216) in the shadow SSBO."""
flag = np.array([1 if enabled else 0], dtype=np.uint32)
shadow_data = np.zeros(SHADOW_DATA_SIZE, dtype=np.uint8)
ffi = vk.ffi
device = self._engine.ctx.device
src = vk.vkMapMemory(device, self.shadow_mem, 0, SHADOW_DATA_SIZE, 0)
ffi.memmove(ffi.cast("void*", shadow_data.ctypes.data), src, SHADOW_DATA_SIZE)
vk.vkUnmapMemory(device, self.shadow_mem)
shadow_data[216:220] = flag.view(np.uint8)
upload_numpy(device, self.shadow_mem, shadow_data)
[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)."""
write_image_descriptor(self._engine.ctx.device, self.ssbo_set, 4, view, sampler)
# ---------------------------------------------------------------- cleanup
[docs]
def cleanup(self) -> None:
"""Destroy all buffers, descriptor pools/layouts, and placeholder cubemap."""
device = self._engine.ctx.device
for buf, mem in (
(self.joint_buf, self.joint_mem),
(self.transform_buf, self.transform_mem),
(self.material_buf, self.material_mem),
(self.light_buf, self.light_mem),
(self.shadow_buf, self.shadow_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),
):
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)