"""Custom shader material system: user-facing API for custom GLSL shaders.
Provides ShaderMaterial for per-object custom shaders, UniformBuffer for GPU-side
uniform data, and ShaderMaterialManager for pipeline caching and hot-reload.
"""
import logging
import struct
import tempfile
from pathlib import Path
from typing import Any
import numpy as np
import vulkan as vk
from ..gpu.pipeline import _make_vertex_input
from ..gpu.pipeline_cache import pipeline_cache_for
from .shader_compiler import compile_shader, resolve_includes
__all__ = ["ShaderMaterial", "UniformBuffer", "ShaderMaterialManager"]
log = logging.getLogger(__name__)
# Mapping from Python/numpy types to GLSL uniform metadata
_UNIFORM_FORMATS: dict[str, tuple[str, int]] = {
"float": ("f", 4),
"int": ("i", 4),
"uint": ("I", 4),
"vec2": ("2f", 8),
"vec3": ("3f", 12),
"vec4": ("4f", 16),
"ivec2": ("2i", 8),
"ivec3": ("3i", 12),
"ivec4": ("4i", 16),
"mat4": ("16f", 64),
}
def _infer_uniform_type(value: Any) -> str:
"""Infer the GLSL uniform type from a Python value."""
if isinstance(value, int | np.integer):
return "int"
if isinstance(value, float | np.floating):
return "float"
if isinstance(value, np.ndarray):
if value.shape == (4, 4):
return "mat4"
size = value.size
return {2: "vec2", 3: "vec3", 4: "vec4"}.get(size, "float")
if isinstance(value, tuple | list):
n = len(value)
return {2: "vec2", 3: "vec3", 4: "vec4"}.get(n, "float")
return "float"
def _pack_uniform(value: Any, utype: str) -> bytes:
"""Pack a uniform value into bytes matching the GLSL layout."""
fmt, expected_size = _UNIFORM_FORMATS.get(utype, ("f", 4))
if isinstance(value, np.ndarray):
flat = value.astype(np.float32).ravel()
return flat.tobytes()[:expected_size]
if isinstance(value, tuple | list):
return struct.pack(fmt, *value)
if isinstance(value, int | np.integer):
return struct.pack("i", int(value))
if isinstance(value, float | np.floating):
return struct.pack("f", float(value))
return struct.pack("f", float(value))
def _align_to(offset: int, alignment: int) -> int:
"""Round offset up to the next multiple of alignment (std140 rules)."""
return (offset + alignment - 1) & ~(alignment - 1)
def _std140_alignment(utype: str) -> int:
"""Return std140 base alignment for a given GLSL type."""
if utype in ("float", "int", "uint"):
return 4
if utype in ("vec2", "ivec2"):
return 8
if utype in ("vec3", "ivec3", "vec4", "ivec4"):
return 16
if utype == "mat4":
return 16
return 4
[docs]
class ShaderMaterial:
"""User-facing custom shader material.
Allows using custom GLSL vertex/fragment shaders with user-defined uniforms.
Works alongside the engine's existing uber-shader pipeline by creating its own
separate Vulkan pipeline.
Example::
mat = ShaderMaterial(
vertex_path="shaders/wave.vert",
fragment_path="shaders/gradient.frag",
)
mat.set_uniform("time", 0.0)
mat.set_uniform("colour", (1.0, 0.5, 0.2, 1.0))
Or with inline source::
mat = ShaderMaterial(
vertex_source=\"\"\"
#version 450
layout(location=0) in vec3 pos;
void main() { gl_Position = vec4(pos, 1.0); }
\"\"\",
fragment_source=\"\"\"
#version 450
layout(location=0) out vec4 out_color;
void main() { out_color = vec4(1.0, 0.0, 0.0, 1.0); }
\"\"\",
)
"""
def __init__(
self,
vertex_path: str | Path | None = None,
fragment_path: str | Path | None = None,
*,
vertex_source: str | None = None,
fragment_source: str | None = None,
language: str = "glsl",
wgsl_vertex: str | None = None,
wgsl_fragment: str | None = None,
wgsl_vertex_path: str | Path | None = None,
wgsl_fragment_path: str | Path | None = None,
) -> None:
self._vertex_path = Path(vertex_path) if vertex_path else None
self._fragment_path = Path(fragment_path) if fragment_path else None
self.vertex_source = vertex_source
self.fragment_source = fragment_source
self.language = language
# Web escape hatch: hand-written WGSL used verbatim by the web exporter,
# bypassing naga GLSL->WGSL transpilation. Inert on desktop (GLSL stays
# the desktop source of truth). The hand WGSL must still conform to the
# unified ABI (group0 camera UBO, group1 transforms SSBO, group2 per-
# material) and is naga-validated at export time.
self.wgsl_vertex = wgsl_vertex
self.wgsl_fragment = wgsl_fragment
self._wgsl_vertex_path = Path(wgsl_vertex_path) if wgsl_vertex_path else None
self._wgsl_fragment_path = Path(wgsl_fragment_path) if wgsl_fragment_path else None
self._uniforms: dict[str, Any] = {}
self._uniform_types: dict[str, str] = {}
self._vert_module: Any = None
self._frag_module: Any = None
self._is_compiled = False
self._vert_mtime: float = 0.0
self._frag_mtime: float = 0.0
[docs]
@property
def is_compiled(self) -> bool:
"""Whether shaders have been compiled to SPIR-V and loaded."""
return self._is_compiled
[docs]
def compile(self, device: Any, shader_dir: Path | None = None) -> None:
"""Compile shaders to SPIR-V and create Vulkan shader modules.
Uses file paths if provided, otherwise writes inline source to temp files
for compilation via glslc.
Args:
device: Vulkan logical device handle.
shader_dir: Base directory for resolving relative shader paths and includes.
"""
from ..gpu.pipeline import create_shader_module
base_dir = shader_dir or Path.cwd()
# Compile vertex shader
vert_spv = self._compile_stage("vertex", base_dir)
self._vert_module = create_shader_module(device, vert_spv)
# Compile fragment shader
frag_spv = self._compile_stage("fragment", base_dir)
self._frag_module = create_shader_module(device, frag_spv)
self._is_compiled = True
log.debug("ShaderMaterial compiled successfully")
def _compile_stage(self, stage: str, base_dir: Path) -> Path:
"""Compile a single shader stage, handling paths vs inline source."""
is_vertex = stage == "vertex"
path = self._vertex_path if is_vertex else self._fragment_path
source = self.vertex_source if is_vertex else self.fragment_source
ext = ".vert" if is_vertex else ".frag"
if path is not None:
resolved = path if path.is_absolute() else base_dir / path
if not resolved.exists():
raise FileNotFoundError(f"Shader file not found: {resolved}")
# Process includes
raw_source = resolved.read_text()
processed = resolve_includes(raw_source, resolved.parent)
# Write processed source to temp file for glslc
tmp = Path(tempfile.mktemp(suffix=ext))
tmp.write_text(processed)
try:
spv = compile_shader(tmp)
finally:
tmp.unlink(missing_ok=True)
# Track mtime for hot-reload
if is_vertex:
self._vert_mtime = resolved.stat().st_mtime
else:
self._frag_mtime = resolved.stat().st_mtime
return spv
if source is not None:
processed = resolve_includes(source, base_dir)
tmp = Path(tempfile.mktemp(suffix=ext))
tmp.write_text(processed)
try:
spv = compile_shader(tmp)
finally:
tmp.unlink(missing_ok=True)
return spv
raise ValueError(f"No {stage} shader source or path provided")
[docs]
def resolve_wgsl(self, base_dir: Path | None = None) -> tuple[str | None, str | None]:
"""Return (vertex_wgsl, fragment_wgsl) for the web escape hatch, or (None, None).
Inline ``wgsl_vertex``/``wgsl_fragment`` take precedence over their
``*_path`` counterparts. Used by the web exporter to skip naga
transpilation when hand-written WGSL is supplied. A partial escape hatch
(only one stage) is an error: both stages must be provided together.
"""
base = base_dir or Path.cwd()
def _load(src: str | None, path: Path | None) -> str | None:
if src is not None:
return src
if path is not None:
resolved = path if path.is_absolute() else base / path
if not resolved.exists():
raise FileNotFoundError(f"WGSL escape-hatch file not found: {resolved}")
return resolved.read_text()
return None
vert = _load(self.wgsl_vertex, self._wgsl_vertex_path)
frag = _load(self.wgsl_fragment, self._wgsl_fragment_path)
if (vert is None) != (frag is None):
raise ValueError(
"ShaderMaterial WGSL escape hatch requires BOTH wgsl_vertex and wgsl_fragment "
"(a partial escape hatch is not supported)."
)
return vert, frag
[docs]
@property
def has_wgsl_escape_hatch(self) -> bool:
"""Whether hand-written WGSL is supplied for both stages (web only)."""
v, f = self.resolve_wgsl()
return v is not None and f is not None
[docs]
def get_pipeline_key(self) -> tuple:
"""Return a hashable key unique to this shader combination.
Used for pipeline caching in ShaderMaterialManager.
"""
vert_key: str | None = None
frag_key: str | None = None
if self._vertex_path:
vert_key = str(self._vertex_path.resolve())
elif self.vertex_source:
vert_key = self.vertex_source
if self._fragment_path:
frag_key = str(self._fragment_path.resolve())
elif self.fragment_source:
frag_key = self.fragment_source
return (vert_key, frag_key)
[docs]
def has_source_changed(self) -> bool:
"""Check if shader source files have been modified since last compile."""
if self._vertex_path and self._vertex_path.exists():
if self._vertex_path.stat().st_mtime > self._vert_mtime:
return True
if self._fragment_path and self._fragment_path.exists():
if self._fragment_path.stat().st_mtime > self._frag_mtime:
return True
return False
[docs]
def cleanup(self, device: Any) -> None:
"""Destroy Vulkan shader modules."""
if self._vert_module:
vk.vkDestroyShaderModule(device, self._vert_module, None)
self._vert_module = None
if self._frag_module:
vk.vkDestroyShaderModule(device, self._frag_module, None)
self._frag_module = None
self._is_compiled = False
def _create_custom_pipeline(
device: Any,
vert_module: Any,
frag_module: Any,
render_pass: Any,
extent: tuple[int, int],
camera_layout: Any,
transforms_layout: Any,
material_layout: Any | None = None,
gbuffer: bool = False,
) -> tuple[Any, Any]:
"""Create a Vulkan graphics pipeline for a custom shader (unified ABI).
Vertex format: position(vec3) + normal(vec3) + uv(vec2) = 32 bytes stride.
Unified ABI (matches the web WebGPU custom-shader ABI; NO push constants):
- set 0 (group0): camera UBO ``{ mat4 view; mat4 proj; }``.
- set 1 (group1): transforms SSBO ``array<mat4>`` (read-only storage),
addressed by ``gl_InstanceIndex`` -> ``builtin(instance_index)``.
- set 2 (group2): per-material resources -- binding 0 custom-uniform UBO,
then binding 1.. for SEPARATED textures + samplers.
The custom pipeline gets its own pipeline layout, distinct from the forward3d
descriptor groups, so there is no collision with the main pass. Returns
``(pipeline, pipeline_layout)``.
"""
ffi = vk.ffi
# Collect descriptor set layouts in ABI order: camera (0), transforms (1),
# per-material (2). NO push constants in the custom-shader path.
layouts = [camera_layout, transforms_layout]
if material_layout:
layouts.append(material_layout)
set_layouts = ffi.new(f"VkDescriptorSetLayout[{len(layouts)}]", layouts)
layout_ci = ffi.new("VkPipelineLayoutCreateInfo*")
layout_ci.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO
layout_ci.setLayoutCount = len(layouts)
layout_ci.pSetLayouts = set_layouts
layout_ci.pushConstantRangeCount = 0
layout_ci.pPushConstantRanges = ffi.NULL
layout_out = ffi.new("VkPipelineLayout*")
result = vk._vulkan._callApi(
vk._vulkan.lib.vkCreatePipelineLayout,
device,
layout_ci,
ffi.NULL,
layout_out,
)
if result != vk.VK_SUCCESS:
raise RuntimeError(f"vkCreatePipelineLayout failed: {result}")
pipeline_layout = layout_out[0]
# Build pipeline create info
pi = ffi.new("VkGraphicsPipelineCreateInfo*")
pi.sType = vk.VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO
# Shader stages
stages = ffi.new("VkPipelineShaderStageCreateInfo[2]")
main_name = ffi.new("char[]", b"main")
stages[0].sType = vk.VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO
stages[0].stage = vk.VK_SHADER_STAGE_VERTEX_BIT
stages[0].module = vert_module
stages[0].pName = main_name
stages[1].sType = vk.VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO
stages[1].stage = vk.VK_SHADER_STAGE_FRAGMENT_BIT
stages[1].module = frag_module
stages[1].pName = main_name
pi.stageCount = 2
pi.pStages = stages
# Vertex input: position + shading streams (vertex_layouts.MESH_BINDINGS).
# The user-facing ShaderMaterial ABI is unchanged: GLSL inputs stay at
# locations 0 (position), 1 (normal), 2 (uv); only the buffer layout moved
# from one interleaved binding to the split streams.
from ..renderer.vertex_layouts import MESH_BINDINGS
vi, _vi_bindings, _vi_attrs = _make_vertex_input(ffi, MESH_BINDINGS) # keep all three alive until creation
pi.pVertexInputState = vi
# Input assembly
ia = ffi.new("VkPipelineInputAssemblyStateCreateInfo*")
ia.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO
ia.topology = vk.VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST
pi.pInputAssemblyState = ia
# Viewport state
vps = ffi.new("VkPipelineViewportStateCreateInfo*")
vps.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO
vps.viewportCount = 1
viewport = ffi.new("VkViewport*")
viewport.width = float(extent[0])
viewport.height = float(extent[1])
viewport.maxDepth = 1.0
vps.pViewports = viewport
scissor = ffi.new("VkRect2D*")
scissor.extent.width = extent[0]
scissor.extent.height = extent[1]
vps.scissorCount = 1
vps.pScissors = scissor
pi.pViewportState = vps
# Rasterization
rs = ffi.new("VkPipelineRasterizationStateCreateInfo*")
rs.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO
rs.polygonMode = vk.VK_POLYGON_MODE_FILL
rs.lineWidth = 1.0
rs.cullMode = vk.VK_CULL_MODE_BACK_BIT
rs.frontFace = vk.VK_FRONT_FACE_CLOCKWISE
pi.pRasterizationState = rs
# Multisample
ms = ffi.new("VkPipelineMultisampleStateCreateInfo*")
ms.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO
ms.rasterizationSamples = vk.VK_SAMPLE_COUNT_1_BIT
pi.pMultisampleState = ms
# Depth stencil
dss = ffi.new("VkPipelineDepthStencilStateCreateInfo*")
dss.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO
dss.depthTestEnable = 1
dss.depthWriteEnable = 1
dss.depthCompareOp = vk.VK_COMPARE_OP_LESS
pi.pDepthStencilState = dss
# Colour blend. Thin G-buffer: a ShaderMaterial is a user shader with no
# normal output, so in the 2-attachment HDR pass it declares the second blend
# attachment but masks its writes (leaves the normal target untouched).
n_att = 2 if gbuffer else 1
full_mask = (
vk.VK_COLOR_COMPONENT_R_BIT
| vk.VK_COLOR_COMPONENT_G_BIT
| vk.VK_COLOR_COMPONENT_B_BIT
| vk.VK_COLOR_COMPONENT_A_BIT
)
cba = ffi.new(f"VkPipelineColorBlendAttachmentState[{n_att}]")
cba[0].colorWriteMask = full_mask
for i in range(1, n_att):
cba[i].colorWriteMask = 0
cb = ffi.new("VkPipelineColorBlendStateCreateInfo*")
cb.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO
cb.attachmentCount = n_att
cb.pAttachments = cba
pi.pColorBlendState = cb
# Dynamic state
dyn_states = ffi.new("VkDynamicState[2]", [vk.VK_DYNAMIC_STATE_VIEWPORT, vk.VK_DYNAMIC_STATE_SCISSOR])
ds = ffi.new("VkPipelineDynamicStateCreateInfo*")
ds.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO
ds.dynamicStateCount = 2
ds.pDynamicStates = dyn_states
pi.pDynamicState = ds
pi.layout = pipeline_layout
pi.renderPass = render_pass
pipeline_out = ffi.new("VkPipeline*")
result = vk._vulkan._callApi(
vk._vulkan.lib.vkCreateGraphicsPipelines,
device,
pipeline_cache_for(device),
1,
pi,
ffi.NULL,
pipeline_out,
)
if result != vk.VK_SUCCESS:
raise RuntimeError(f"vkCreateGraphicsPipelines failed: {result}")
pipeline = pipeline_out[0]
log.debug("Custom shader pipeline created")
return pipeline, pipeline_layout
def _create_camera_ubo_layout(device: Any) -> Any:
"""Create the group0 descriptor set layout: camera UBO (binding 0) + the
optional FrameGlobals UBO (binding 1).
Binding 1 is always present in the layout so a user shader MAY declare the
FrameGlobals block; shaders that only declare binding 0 compile and render
unchanged (an unused layout binding is legal). Mirror of the web
group0-binding1 ABI.
"""
bindings = [
vk.VkDescriptorSetLayoutBinding(
binding=0,
descriptorType=vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_VERTEX_BIT | vk.VK_SHADER_STAGE_FRAGMENT_BIT,
),
vk.VkDescriptorSetLayoutBinding(
binding=1,
descriptorType=vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_VERTEX_BIT | vk.VK_SHADER_STAGE_FRAGMENT_BIT,
),
]
return vk.vkCreateDescriptorSetLayout(
device, vk.VkDescriptorSetLayoutCreateInfo(bindingCount=len(bindings), pBindings=bindings), None
)
def _create_transforms_ssbo_layout(device: Any) -> Any:
"""Create the group1 transforms SSBO descriptor set layout ``array<mat4>``."""
binding = vk.VkDescriptorSetLayoutBinding(
binding=0,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_VERTEX_BIT | vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
return vk.vkCreateDescriptorSetLayout(
device, vk.VkDescriptorSetLayoutCreateInfo(bindingCount=1, pBindings=[binding]), None
)
[docs]
class ShaderMaterialManager:
"""Caches custom-shader pipelines and drives the desktop custom-shader pass.
Tracks all registered ShaderMaterial instances and their compiled pipelines.
Pipelines are cached by the shader source/path combination so that multiple
objects sharing the same shaders reuse one pipeline.
The manager also owns the shared per-frame **camera UBO** (group0) and a
per-frame **transforms SSBO** (group1) used by every custom-shader draw.
Per-material uniforms live in a group2 UBO created lazily per material.
The descriptor layout ABI is identical to the web WebGPU custom-shader ABI:
group0 camera UBO, group1 transforms SSBO addressed by ``gl_InstanceIndex``,
group2 per-material UBO (+ separated textures/samplers at binding 1.., not
yet wired for the canary).
"""
def __init__(self) -> None:
self._pipeline_cache: dict[tuple, tuple[Any, Any]] = {} # key -> (pipeline, layout)
self._materials: list[ShaderMaterial] = []
self._uniform_buffers: dict[int, UniformBuffer] = {} # id(material) -> UniformBuffer (group2)
self._device: Any = None
self._physical_device: Any = None
# Shared ABI descriptor layouts (created lazily on first use).
self._camera_layout: Any = None
self._transforms_layout: Any = None
# Per-frame camera UBO (group0): host-visible, holds {mat4 view; mat4 proj}.
self._camera_buf: Any = None
self._camera_mem: Any = None
self._camera_set: Any = None
# FrameGlobals UBO (group0 binding 1): the renderer's shared
# per-frame block, bound here so user shaders MAY read it. The buffer is
# owned by the BufferManager (single source of truth); we only hold the
# handle and write the descriptor once, when it first becomes available.
self._frame_globals_buf: Any = None
self._fg_bound: bool = False
# Per-frame transforms SSBO (group1): host-visible array<mat4>.
self._transforms_buf: Any = None
self._transforms_mem: Any = None
self._transforms_set: Any = None
self._transforms_capacity: int = 0 # in number of mat4 entries
self._shared_pool: Any = None
# The render pass the cached pipelines were compiled against. Pipelines
# must be render-pass-compatible with the framebuffer they draw into, and
# the active pass flips when post-processing (HDR offscreen) toggles, so
# we rebuild the cache when it changes.
self._render_pass: Any = None
# Thin G-buffer active state for the current cache. Set each draw
# from the caller; the pipeline cache already rebuilds when the render
# pass handle flips (which it does when the G-buffer toggles).
self._gbuffer: bool = False
# -- Shared ABI resources (camera UBO group0 + transforms SSBO group1) --
def _ensure_shared_resources(self, device: Any, physical_device: Any, capacity: int) -> None:
"""Create/grow the shared camera UBO + transforms SSBO and their descriptors."""
from ..gpu.descriptors import allocate_descriptor_set
from ..gpu.memory import create_buffer
self._device = device
self._physical_device = physical_device
if self._camera_layout is None:
self._camera_layout = _create_camera_ubo_layout(device)
if self._transforms_layout is None:
self._transforms_layout = _create_transforms_ssbo_layout(device)
if self._shared_pool is None:
# Two UBO descriptors in the camera set (binding 0 camera + binding 1
# FrameGlobals), plus one SSBO for the transforms set.
pool_sizes = [
vk.VkDescriptorPoolSize(type=vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, descriptorCount=2),
vk.VkDescriptorPoolSize(type=vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, descriptorCount=1),
]
self._shared_pool = vk.vkCreateDescriptorPool(
device,
vk.VkDescriptorPoolCreateInfo(maxSets=2, poolSizeCount=2, pPoolSizes=pool_sizes),
None,
)
# Camera UBO: 2x mat4 = 128 bytes, created once.
if self._camera_buf is None:
self._camera_buf, self._camera_mem = create_buffer(
device, physical_device, 128,
vk.VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
)
self._camera_set = allocate_descriptor_set(device, self._shared_pool, self._camera_layout)
buf_info = vk.VkDescriptorBufferInfo(buffer=self._camera_buf, offset=0, range=128)
vk.vkUpdateDescriptorSets(device, 1, [vk.VkWriteDescriptorSet(
dstSet=self._camera_set, dstBinding=0, dstArrayElement=0, descriptorCount=1,
descriptorType=vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, pBufferInfo=[buf_info],
)], 0, None)
# FrameGlobals UBO (binding 1): write once when the renderer's shared block
# is available. The handle never changes for the manager's lifetime, so
# this one-time write precedes any bind of the set and needs no rewrite.
if self._camera_set is not None and self._frame_globals_buf is not None and not self._fg_bound:
from ..frame_globals import FRAME_GLOBALS_SIZE
fg_info = vk.VkDescriptorBufferInfo(buffer=self._frame_globals_buf, offset=0, range=FRAME_GLOBALS_SIZE)
vk.vkUpdateDescriptorSets(device, 1, [vk.VkWriteDescriptorSet(
dstSet=self._camera_set, dstBinding=1, dstArrayElement=0, descriptorCount=1,
descriptorType=vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, pBufferInfo=[fg_info],
)], 0, None)
self._fg_bound = True
# Transforms SSBO: grow to hold at least *capacity* mat4 entries.
need = max(capacity, 1)
if self._transforms_set is None or need > self._transforms_capacity:
if self._transforms_buf is not None:
vk.vkDestroyBuffer(device, self._transforms_buf, None)
vk.vkFreeMemory(device, self._transforms_mem, None)
self._transforms_capacity = need
size = need * 64
self._transforms_buf, self._transforms_mem = create_buffer(
device, physical_device, size,
vk.VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
)
if self._transforms_set is None:
self._transforms_set = allocate_descriptor_set(
device, self._shared_pool, self._transforms_layout
)
buf_info = vk.VkDescriptorBufferInfo(buffer=self._transforms_buf, offset=0, range=size)
vk.vkUpdateDescriptorSets(device, 1, [vk.VkWriteDescriptorSet(
dstSet=self._transforms_set, dstBinding=0, dstArrayElement=0, descriptorCount=1,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, pBufferInfo=[buf_info],
)], 0, None)
def _upload_camera(self, device: Any, view: np.ndarray, proj: np.ndarray) -> None:
"""Write column-major view+proj into the shared camera UBO."""
data = (
np.ascontiguousarray(view.T, dtype=np.float32).tobytes()
+ np.ascontiguousarray(proj.T, dtype=np.float32).tobytes()
)
src = vk.ffi.from_buffer(data)
dst = vk.vkMapMemory(device, self._camera_mem, 0, len(data), 0)
vk.ffi.memmove(dst, src, len(data))
vk.vkUnmapMemory(device, self._camera_mem)
def _upload_transforms(self, device: Any, transforms: list[np.ndarray]) -> None:
"""Write a contiguous array of column-major mat4 into the transforms SSBO."""
if not transforms:
return
flat = np.concatenate(
[np.ascontiguousarray(t.T, dtype=np.float32).ravel() for t in transforms]
)
data = flat.tobytes()
src = vk.ffi.from_buffer(data)
dst = vk.vkMapMemory(device, self._transforms_mem, 0, len(data), 0)
vk.ffi.memmove(dst, src, len(data))
vk.vkUnmapMemory(device, self._transforms_mem)
[docs]
def register_material(self, material: ShaderMaterial) -> None:
"""Track a ShaderMaterial for hot-reload monitoring."""
if material not in self._materials:
self._materials.append(material)
[docs]
def get_or_create_pipeline(
self,
material: ShaderMaterial,
device: Any,
physical_device: Any,
render_pass: Any,
extent: tuple[int, int],
shader_dir: Path | None = None,
) -> tuple[Any, Any]:
"""Get a cached pipeline for this material, or compile and create one.
Uses the unified ABI layouts: group0 camera UBO, group1 transforms SSBO,
group2 per-material UBO. Returns ``(VkPipeline, VkPipelineLayout)``.
"""
self._ensure_shared_resources(device, physical_device, capacity=1)
key = material.get_pipeline_key()
if key in self._pipeline_cache:
return self._pipeline_cache[key]
if not material.is_compiled:
material.compile(device, shader_dir)
self.register_material(material)
# group2 per-material UBO (custom uniforms at binding 0).
material_layout = None
if material.uniforms:
ubo = UniformBuffer()
ubo.create(device, physical_device)
self._uniform_buffers[id(material)] = ubo
material_layout = ubo.get_descriptor_layout()
pipeline, layout = _create_custom_pipeline(
device,
material._vert_module,
material._frag_module,
render_pass,
extent,
self._camera_layout,
self._transforms_layout,
material_layout=material_layout,
gbuffer=self._gbuffer,
)
self._pipeline_cache[key] = (pipeline, layout)
log.debug("Cached custom pipeline for key=%s", key)
return pipeline, layout
[docs]
def draw(
self,
cmd: Any,
submissions: list,
device: Any,
physical_device: Any,
render_pass: Any,
extent: tuple[int, int],
view: np.ndarray,
proj: np.ndarray,
registry: Any,
shader_dir: Path | None = None,
frame_globals_buf: Any = None,
gbuffer: bool = False,
) -> None:
"""Draw all ShaderMaterial-backed submissions inside the main render pass.
``submissions`` is the renderer's per-frame
``[(mesh_handle, transform, material_id, shader_material)]`` list. Each
submission is drawn individually (one transforms-SSBO slot per draw,
addressed via ``gl_InstanceIndex``), binding the cached custom pipeline,
the shared camera UBO (group0), the transforms SSBO (group1), and the
per-material UBO (group2). Per-material pipeline switching is the cost of
the custom-shader path.
"""
if not submissions:
return
# Render-pass change (post-process toggled): drop pipelines so they are
# recompiled against the now-active pass (layouts/buffers are reused).
if render_pass is not self._render_pass and self._render_pass is not None:
for pipeline, layout in self._pipeline_cache.values():
vk.vkDestroyPipeline(device, pipeline, None)
vk.vkDestroyPipelineLayout(device, layout, None)
self._pipeline_cache.clear()
self._render_pass = render_pass
self._gbuffer = gbuffer
# Publish the renderer's shared FrameGlobals buffer so binding 1 is wired
# on the camera set (one-time write inside _ensure_shared_resources).
if frame_globals_buf is not None:
self._frame_globals_buf = frame_globals_buf
self._ensure_shared_resources(device, physical_device, capacity=len(submissions))
self._upload_camera(device, view, proj)
self._upload_transforms(device, [s[1] for s in submissions])
for slot, (mesh_handle, _transform, _material_id, material) in enumerate(submissions):
pipeline, layout = self.get_or_create_pipeline(
material, device, physical_device, render_pass, extent, shader_dir,
)
self.update_uniforms(material, device)
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline)
# Dynamic viewport/scissor (pipeline declares them dynamic).
vp = vk.VkViewport(x=0.0, y=0.0, width=float(extent[0]), height=float(extent[1]),
minDepth=0.0, maxDepth=1.0)
vk.vkCmdSetViewport(cmd, 0, 1, [vp])
sc = vk.VkRect2D(offset=vk.VkOffset2D(x=0, y=0),
extent=vk.VkExtent2D(width=extent[0], height=extent[1]))
vk.vkCmdSetScissor(cmd, 0, 1, [sc])
# group0 camera UBO, group1 transforms SSBO.
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, layout, 0, 1, [self._camera_set], 0, None,
)
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, layout, 1, 1, [self._transforms_set], 0, None,
)
# group2 per-material UBO, if any.
ubo = self._uniform_buffers.get(id(material))
if ubo is not None:
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, layout, 2, 1,
[ubo.get_descriptor_set()], 0, None,
)
bufs = registry.get_buffers(mesh_handle)
vk.vkCmdBindVertexBuffers(cmd, 0, 2, [bufs.position, bufs.shading], [0, 0])
vk.vkCmdBindIndexBuffer(cmd, bufs.index, 0, vk.VK_INDEX_TYPE_UINT32)
# gl_InstanceIndex == firstInstance == slot, indexing the SSBO.
vk.vkCmdDrawIndexed(cmd, mesh_handle.index_count, 1, 0, 0, slot)
[docs]
def check_hot_reload(
self,
device: Any,
physical_device: Any,
render_pass: Any,
extent: tuple[int, int],
shader_dir: Path | None = None,
) -> list[ShaderMaterial]:
"""Check all registered materials for source file changes and recompile.
Returns a list of materials that were recompiled.
"""
recompiled = []
for material in self._materials:
if not material.has_source_changed():
continue
key = material.get_pipeline_key()
log.info("Hot-reloading shader: %s", key)
# Destroy old pipeline
old = self._pipeline_cache.pop(key, None)
if old:
vk.vkDestroyPipeline(device, old[0], None)
vk.vkDestroyPipelineLayout(device, old[1], None)
# Destroy old shader modules and recompile
material.cleanup(device)
try:
material.compile(device, shader_dir)
except Exception:
log.exception("Hot-reload compilation failed for %s", key)
continue
# Recreate pipeline
material_layout = None
ubo = self._uniform_buffers.get(id(material))
if ubo:
material_layout = ubo.get_descriptor_layout()
pipeline, layout = _create_custom_pipeline(
device,
material._vert_module,
material._frag_module,
render_pass,
extent,
self._camera_layout,
self._transforms_layout,
material_layout=material_layout,
gbuffer=self._gbuffer,
)
self._pipeline_cache[key] = (pipeline, layout)
recompiled.append(material)
return recompiled
[docs]
def cleanup(self, device: Any) -> None:
"""Destroy all cached pipelines, uniform buffers, shared resources, and modules."""
for pipeline, layout in self._pipeline_cache.values():
vk.vkDestroyPipeline(device, pipeline, None)
vk.vkDestroyPipelineLayout(device, layout, None)
self._pipeline_cache.clear()
for ubo in self._uniform_buffers.values():
ubo.cleanup(device)
self._uniform_buffers.clear()
for material in self._materials:
material.cleanup(device)
self._materials.clear()
if self._camera_buf is not None:
vk.vkDestroyBuffer(device, self._camera_buf, None)
vk.vkFreeMemory(device, self._camera_mem, None)
self._camera_buf = self._camera_mem = self._camera_set = None
# The FrameGlobals buffer is owned by the BufferManager, not freed here;
# just drop our handle so a rebuilt manager re-binds it.
self._frame_globals_buf = None
self._fg_bound = False
if self._transforms_buf is not None:
vk.vkDestroyBuffer(device, self._transforms_buf, None)
vk.vkFreeMemory(device, self._transforms_mem, None)
self._transforms_buf = self._transforms_mem = self._transforms_set = None
self._transforms_capacity = 0
if self._shared_pool is not None:
vk.vkDestroyDescriptorPool(device, self._shared_pool, None)
self._shared_pool = None
if self._camera_layout is not None:
vk.vkDestroyDescriptorSetLayout(device, self._camera_layout, None)
self._camera_layout = None
if self._transforms_layout is not None:
vk.vkDestroyDescriptorSetLayout(device, self._transforms_layout, None)
self._transforms_layout = None