"""Mesh registry: manages uploaded GPU mesh buffers and returns handles.
One GPU buffer per vertex stream (vertex stream split, D5): positions,
normal+uv shading, optional extras (tangent/colour/uv2), optional skin
(joints/weights), plus the index buffer. See ``vertex_layouts.py`` for the
matching pipeline binding layouts.
"""
import logging
from typing import Any, NamedTuple
import numpy as np
import vulkan as vk
from ..gpu.memory import create_buffer, upload_numpy
from ..types import MeshHandle, VertexStreams
log = logging.getLogger(__name__)
__all__ = ["MeshBuffers", "MeshRegistry"]
[docs]
class MeshBuffers(NamedTuple):
"""The GPU buffers of one registered mesh, per vertex stream.
``extras`` / ``skin`` are ``None`` when the mesh has no such stream.
Fields order-match the Vulkan binding numbers (0..3) plus the index buffer.
"""
position: Any # VkBuffer, binding 0
shading: Any # VkBuffer, binding 1
extras: Any | None # VkBuffer, binding 2
skin: Any | None # VkBuffer, binding 3
index: Any # VkBuffer
[docs]
class MeshRegistry:
"""Upload meshes to GPU, return handles for efficient referencing."""
def __init__(self, device: Any, physical_device: Any, *, retain_geometry: bool = False):
self.device = device
self.physical_device = physical_device
# id -> (MeshBuffers, memories tuple parallel to the buffers)
self._meshes: dict[int, tuple[MeshBuffers, tuple[Any, ...]]] = {}
self._next_id = 0
# Retain the source CPU streams/index arrays keyed by mesh id. OFF by default
# so the single-GPU path stores nothing extra (byte-identical memory). The
# D8 multi-GPU path opts in (engine sets it when multi-GPU is active) so an
# offloaded SubViewport SRU can mirror its geometry onto a secondary device's
# own registry: a VkBuffer cannot cross devices, only the CPU arrays can.
self._retain_geometry = retain_geometry
# mesh id -> (streams, indices) when retention is on; empty otherwise.
self._geometry: dict[int, tuple[VertexStreams, np.ndarray]] = {}
# Mesh ids whose extras stream carries real (non-zero) tangents: these
# draw with the tangent normal-mapping pipeline variant.
# Empty for tangent-free scenes, so the draw path's truthiness
# check keeps the default path zero-cost.
self.tangent_mesh_ids: set[int] = set()
def _upload(self, data: np.ndarray, usage: int) -> tuple[Any, Any]:
buf, mem = create_buffer(
self.device,
self.physical_device,
data.nbytes,
usage,
vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
)
upload_numpy(self.device, mem, data)
return buf, mem
[docs]
def register(self, streams: VertexStreams, indices: np.ndarray) -> MeshHandle:
"""Upload the mesh's vertex streams + indices to GPU, return a handle."""
vertex_usage = vk.VK_BUFFER_USAGE_VERTEX_BUFFER_BIT
buffers: list[Any] = []
memories: list[Any] = []
for stream in (streams.positions, streams.shading, streams.extras, streams.skin):
if stream is None:
buffers.append(None)
memories.append(None)
continue
buf, mem = self._upload(np.ascontiguousarray(stream), vertex_usage)
buffers.append(buf)
memories.append(mem)
ib, ib_mem = self._upload(indices, vk.VK_BUFFER_USAGE_INDEX_BUFFER_BIT)
buffers.append(ib)
memories.append(ib_mem)
# Compute bounding sphere from local origin.
# The frustum culler places the sphere at model_matrix[:3, 3] (the node's
# world position), which corresponds to local origin (0,0,0). The radius
# must therefore be the max distance from origin to any vertex: NOT from
# the mean vertex position, which would create a mismatched sphere center.
xyz = streams.positions
if len(xyz) > 0:
radius = float(np.linalg.norm(xyz, axis=1).max())
aabb_min = xyz.min(axis=0).astype(np.float32)
aabb_max = xyz.max(axis=0).astype(np.float32)
else:
radius = 1.0
aabb_min = np.zeros(3, dtype=np.float32)
aabb_max = np.zeros(3, dtype=np.float32)
mesh_id = self._next_id
self._next_id += 1
self._meshes[mesh_id] = (MeshBuffers(*buffers), tuple(memories))
if self._retain_geometry:
self._geometry[mesh_id] = (streams, indices)
# An extras stream may exist for colour/uv2 alone; only real tangent
# data (any non-zero component) selects the tangent pipeline. One-time
# vectorized check at upload; the draw path reads the set.
if streams.extras is not None and bool(streams.extras["tangent"].any()):
self.tangent_mesh_ids.add(mesh_id)
return MeshHandle(
id=mesh_id,
vertex_count=streams.vertex_count,
index_count=len(indices),
bounding_radius=radius,
aabb_min=aabb_min,
aabb_max=aabb_max,
)
[docs]
def get_buffers(self, handle: MeshHandle) -> MeshBuffers:
"""Get the per-stream :class:`MeshBuffers` for a mesh handle."""
return self._meshes[handle.id][0]
[docs]
def get_geometry(self, mesh_id: int) -> tuple[VertexStreams, np.ndarray] | None:
"""Return the retained source ``(streams, indices)`` for a mesh id, or ``None``.
Only populated when the registry was created with ``retain_geometry=True``
(the D8 multi-GPU path). ``None`` on the default single-GPU path (nothing is
retained) and for any mesh id never registered. The arrays are the same
device-independent CPU data the primary uploaded; the offload coordinator
re-uploads them into a secondary device's own ``MeshRegistry`` to make the
offloaded SRU's geometry resident there (a ``VkBuffer`` cannot cross devices).
"""
return self._geometry.get(mesh_id)
[docs]
def destroy(self) -> None:
"""Free all mesh buffers."""
for buffers, memories in self._meshes.values():
for buf, mem in zip(buffers, memories, strict=True):
if buf is not None:
vk.vkDestroyBuffer(self.device, buf, None)
vk.vkFreeMemory(self.device, mem, None)
self._meshes.clear()
self._geometry.clear()
self.tangent_mesh_ids.clear()