"""GPU-driven batch rendering with multi-draw indirect (MDI) and fallback."""
import logging
from typing import Any
import numpy as np
import vulkan as vk
from simvx.graphics.gpu.memory import create_indirect_buffer
from simvx.graphics.types import FRAMES_IN_FLIGHT, INDIRECT_DRAW_DTYPE
log = logging.getLogger(__name__)
__all__ = ["GPUBatch"]
[docs]
class GPUBatch:
"""Manages batched draw commands.
When ``use_mdi=True`` (default), uses ``vkCmdDrawIndexedIndirect`` to issue
all draw commands in a single GPU call. When ``use_mdi=False``, falls back
to a loop of ``vkCmdDrawIndexed`` calls: functionally identical but slower
on GPUs that support MDI.
Usage::
batch = GPUBatch(engine, device, physical_device, max_draws=100)
batch.add_draw(index_count=36, first_instance=0)
batch.upload()
batch.on_draw(cmd)
"""
def __init__(
self,
engine: Any,
device: Any,
physical_device: Any,
max_draws: int = 1000,
*,
use_mdi: bool = True,
concurrent_families: list[int] | None = None,
):
self._engine = engine
self.device = device
self.physical_device = physical_device
self.max_draws = max_draws
self.draw_count = 0
self._use_mdi = use_mdi
# Queue families that share the indirect buffer when the occlusion-cull
# compute that writes ``instance_count`` runs on a dedicated compute
# queue (async-compute path). ``None`` keeps the buffer EXCLUSIVE
# (single-queue path, byte-identical) and is retained for ``resize``.
self._concurrent_families = concurrent_families
# Buffers replaced by a grow are not freed immediately: a grow can happen
# mid-build, after this frame's earlier passes already recorded indirect
# draws against the old buffer. We retire them and free only after
# FRAMES_IN_FLIGHT resets, by which point the recording frame's fence has
# signalled. Each entry is ``[buffer, memory, resets_remaining]``.
self._retiring: list[list[Any]] = []
# The indirect command buffer is host-written every frame and read by the
# GPU's indirect draw (and patched in place by the occlusion cull); its
# contents change as frustum/occlusion culling shifts under camera motion.
# Ring it across FRAMES_IN_FLIGHT so a frame never overwrites the copy the
# previous in-flight frame is still drawing from. ``indirect_buffer`` /
# ``indirect_memory`` expose the current frame's slot.
self._indirect_bufs: list[Any] = []
self._indirect_mems: list[Any] = []
for _ in range(FRAMES_IN_FLIGHT):
buf, mem = create_indirect_buffer(
device, physical_device, max_draws, concurrent_families=concurrent_families
)
self._indirect_bufs.append(buf)
self._indirect_mems.append(mem)
self._commands = np.zeros(max_draws, dtype=INDIRECT_DRAW_DTYPE)
@property
def _frame(self) -> int:
return self._engine.current_frame
[docs]
@property
def indirect_buffer(self) -> Any:
return self._indirect_bufs[self._frame]
[docs]
@property
def indirect_memory(self) -> Any:
return self._indirect_mems[self._frame]
[docs]
def add_draw(
self,
index_count: int,
instance_count: int = 1,
first_index: int = 0,
vertex_offset: int = 0,
first_instance: int = 0,
) -> int:
"""Add a draw command. Returns the draw index."""
self._ensure_capacity(1)
idx = self.draw_count
cmd = self._commands[idx]
cmd["index_count"] = index_count
cmd["instance_count"] = instance_count
cmd["first_index"] = first_index
cmd["vertex_offset"] = vertex_offset
cmd["first_instance"] = first_instance
self.draw_count += 1
return idx
[docs]
def add_draws(
self,
index_count: int,
first_instances: np.ndarray | list[int],
) -> int:
"""Bulk-add draw commands sharing the same mesh: avoids per-instance Python loop.
Args:
index_count: Index count for the mesh (same for all draws).
first_instances: (N,) array of SSBO instance indices.
Returns:
Batch offset of the first added draw command.
"""
arr = np.asarray(first_instances, dtype=np.uint32)
n = len(arr)
self._ensure_capacity(n)
start = self.draw_count
sl = self._commands[start : start + n]
sl["index_count"] = index_count
sl["instance_count"] = 1
sl["first_index"] = 0
sl["vertex_offset"] = 0
sl["first_instance"] = arr
self.draw_count += n
return start
[docs]
def add_instanced_runs(
self,
index_count: int,
slots: np.ndarray | list[int],
) -> tuple[int, int]:
"""Add draws for SSBO ``slots``, coalescing contiguous slots into instanced draws.
Each maximal run of consecutive slots becomes ONE indirect command with
``instance_count = run length`` and ``first_instance = run start``. The
vertex shader reads ``transforms[gl_InstanceIndex]`` where
``gl_InstanceIndex = gl_InstanceID + first_instance``, so a run of
contiguous slots is drawn correctly by a single instanced command. A
MultiMesh (whose N instances occupy N consecutive slots) thus collapses
to a single draw.
Returns ``(batch_offset, command_count)`` where ``command_count`` is the
number of runs (== number of indirect commands added), for ``draw_range``.
Must NOT be used for the occlusion-cull batch: that path needs one
command per instance so the cull compute can zero ``instance_count``
per object. Use :meth:`add_draws` there.
"""
arr = np.unique(np.asarray(slots, dtype=np.uint32)) # sorted + de-duplicated
if arr.size == 0:
return self.draw_count, 0
# Run boundaries: a new run starts wherever the slot is not +1 of the prev.
boundaries = np.concatenate(([0], np.flatnonzero(np.diff(arr) != 1) + 1, [arr.size]))
run_first = arr[boundaries[:-1]]
run_len = np.diff(boundaries).astype(np.uint32)
nruns = int(run_first.size)
self._ensure_capacity(nruns)
start = self.draw_count
sl = self._commands[start : start + nruns]
sl["index_count"] = index_count
sl["instance_count"] = run_len
sl["first_index"] = 0
sl["vertex_offset"] = 0
sl["first_instance"] = run_first
self.draw_count += nruns
return start, nruns
def _ensure_capacity(self, extra: int) -> None:
"""Grow the batch in place so ``extra`` more commands fit, preserving existing ones.
The old GPU buffer is retired (freed after FRAMES_IN_FLIGHT resets), so a
mid-build grow is safe even though earlier passes this frame may already
have recorded indirect draws against it.
"""
if self.draw_count + extra <= self.max_draws:
return
new_cap = max(self.draw_count + extra, self.max_draws * 2)
self._realloc(new_cap, preserve=True)
[docs]
def upload(self) -> None:
"""Upload draw commands to GPU indirect buffer."""
if self.draw_count == 0:
return
data = self._commands[: self.draw_count]
ffi = vk.ffi
size = data.nbytes
dst = vk.vkMapMemory(self.device, self.indirect_memory, 0, size, 0)
ffi.memmove(dst, ffi.cast("void*", data.ctypes.data), size)
vk.vkUnmapMemory(self.device, self.indirect_memory)
[docs]
def draw(self, cmd: Any) -> None:
"""Record draw commands for the entire batch."""
if self.draw_count == 0:
return
if self._use_mdi:
vk.vkCmdDrawIndexedIndirect(cmd, self.indirect_buffer, 0, self.draw_count, INDIRECT_DRAW_DTYPE.itemsize)
else:
self._draw_individual(cmd, 0, self.draw_count)
[docs]
def draw_range(self, cmd: Any, offset: int, count: int) -> None:
"""Draw a sub-range of commands.
Args:
cmd: Vulkan command buffer
offset: First draw command index (not byte offset)
count: Number of draw commands to execute
"""
if count == 0:
return
if self._use_mdi:
byte_offset = offset * INDIRECT_DRAW_DTYPE.itemsize
vk.vkCmdDrawIndexedIndirect(cmd, self.indirect_buffer, byte_offset, count, INDIRECT_DRAW_DTYPE.itemsize)
else:
self._draw_individual(cmd, offset, count)
def _draw_individual(self, cmd: Any, offset: int, count: int) -> None:
"""Fallback: issue individual vkCmdDrawIndexed calls from the CPU-side command array."""
cmds = self._commands[offset : offset + count]
for c in cmds:
vk.vkCmdDrawIndexed(
cmd,
int(c["index_count"]),
int(c["instance_count"]),
int(c["first_index"]),
int(c["vertex_offset"]),
int(c["first_instance"]),
)
[docs]
def reset(self) -> None:
"""Clear batch for next frame and free buffers retired long enough ago."""
self.draw_count = 0
if self._retiring:
survivors: list[list[Any]] = []
for entry in self._retiring:
entry[2] -= 1
if entry[2] <= 0:
vk.vkDestroyBuffer(self.device, entry[0], None)
vk.vkFreeMemory(self.device, entry[1], None)
else:
survivors.append(entry)
self._retiring = survivors
[docs]
def grow(self, physical_device: Any, max_draws: int) -> None:
"""Reallocate the indirect buffer to hold at least ``max_draws`` commands.
A no-op when ``max_draws`` does not exceed the current capacity. Used when
the transform-SSBO arena grows so the indirect batch can keep pace. The
old buffer is retired (freed after FRAMES_IN_FLIGHT resets) rather than
freed immediately, and ``draw_count`` is reset (callers rebuild after an
arena grow).
"""
if max_draws <= self.max_draws:
return
self.physical_device = physical_device
self._realloc(max_draws, preserve=False)
self.draw_count = 0
def _realloc(self, new_cap: int, *, preserve: bool) -> None:
"""Grow the indirect buffer + command array to ``new_cap``, retiring the old buffer.
When ``preserve`` is true the existing commands are copied into the new
(larger) command array so a mid-build grow keeps the draws added so far.
"""
new_commands = np.zeros(new_cap, dtype=INDIRECT_DRAW_DTYPE)
if preserve and self.draw_count:
new_commands[: self.draw_count] = self._commands[: self.draw_count]
# Grow every slot in the ring; retire the old buffers (they may still be
# referenced by indirect draws recorded this frame or by in-flight frames).
old_bufs, old_mems = self._indirect_bufs, self._indirect_mems
self._indirect_bufs, self._indirect_mems = [], []
for _ in range(FRAMES_IN_FLIGHT):
buf, mem = create_indirect_buffer(
self.device, self.physical_device, new_cap, concurrent_families=self._concurrent_families
)
self._indirect_bufs.append(buf)
self._indirect_mems.append(mem)
for buf, mem in zip(old_bufs, old_mems, strict=True):
self._retiring.append([buf, mem, FRAMES_IN_FLIGHT])
self._commands = new_commands
self.max_draws = new_cap
[docs]
def destroy(self) -> None:
"""Free GPU resources."""
for entry in self._retiring:
vk.vkDestroyBuffer(self.device, entry[0], None)
vk.vkFreeMemory(self.device, entry[1], None)
self._retiring.clear()
for buf, mem in zip(self._indirect_bufs, self._indirect_mems, strict=True):
vk.vkDestroyBuffer(self.device, buf, None)
vk.vkFreeMemory(self.device, mem, None)
self._indirect_bufs.clear()
self._indirect_mems.clear()