"""GPU particle rendering — camera-facing billboards via SSBO."""
from __future__ import annotations
import logging
from typing import Any
import numpy as np
import vulkan as vk
from ..gpu.memory import create_buffer, upload_numpy
from .pass_helpers import load_shader_modules
__all__ = ["ParticlePass"]
log = logging.getLogger(__name__)
# Must match PARTICLE_DTYPE from core
_PARTICLE_GPU_STRIDE = 16 * 4 # 16 floats × 4 bytes = 64 bytes
MAX_PARTICLES = 10_000
[docs]
class ParticlePass:
"""Renders particles as camera-facing billboards.
Each particle is a 6-vertex quad (2 triangles) expanded in the vertex shader.
Particle data is uploaded to an SSBO each frame.
"""
def __init__(self, engine: Any):
self._engine = engine
self._pipeline: Any = None
self._pipeline_layout: Any = None
self._vert_module: Any = None
self._frag_module: Any = None
self._ssbo_layout: Any = None
self._ssbo_pool: Any = None
self._ssbo_set: Any = None
self._particle_buf: Any = None
self._particle_mem: Any = None
self._ready = False
[docs]
def setup(self) -> None:
e = self._engine
device = e.ctx.device
phys = e.ctx.physical_device
# Particle SSBO
buf_size = MAX_PARTICLES * _PARTICLE_GPU_STRIDE
self._particle_buf, self._particle_mem = create_buffer(
device, phys, buf_size,
vk.VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
)
# Descriptor set for particle SSBO
from ..gpu.descriptors import (
allocate_descriptor_set,
create_descriptor_pool,
create_ssbo_layout,
write_ssbo_descriptor,
)
self._ssbo_layout = create_ssbo_layout(device, binding_count=1)
self._ssbo_pool = create_descriptor_pool(device, max_sets=1)
self._ssbo_set = allocate_descriptor_set(device, self._ssbo_pool, self._ssbo_layout)
write_ssbo_descriptor(device, self._ssbo_set, 0, self._particle_buf, buf_size)
# Shaders
self._vert_module, self._frag_module = load_shader_modules(
device, e.shader_dir, "particle.vert", "particle.frag",
)
# Pipeline
self._pipeline, self._pipeline_layout = _create_particle_pipeline(
device, self._vert_module, self._frag_module,
e.render_pass, self._ssbo_layout,
)
self._ready = True
log.debug("Particle pass initialized (max %d particles)", MAX_PARTICLES)
[docs]
def rebuild_pipeline(self, render_pass: Any) -> None:
"""Recreate the particle pipeline against a different render pass (e.g. HDR)."""
if not self._ready:
return
device = self._engine.ctx.device
if self._pipeline:
vk.vkDestroyPipeline(device, self._pipeline, None)
if self._pipeline_layout:
vk.vkDestroyPipelineLayout(device, self._pipeline_layout, None)
self._pipeline, self._pipeline_layout = _create_particle_pipeline(
device, self._vert_module, self._frag_module,
render_pass, self._ssbo_layout,
)
[docs]
def render(
self,
cmd: Any,
particle_data: np.ndarray,
view_proj: np.ndarray,
camera_right: np.ndarray,
camera_up: np.ndarray,
extent: tuple[int, int],
) -> None:
"""Record particle draw commands."""
if not self._ready or len(particle_data) == 0:
return
count = min(len(particle_data), MAX_PARTICLES)
upload_numpy(self._engine.ctx.device, self._particle_mem, particle_data[:count])
# Push constants: mat4 view_proj (64) + vec3 camera_right + pad (16) + vec3 camera_up + pad (16) = 96
pc = np.zeros(24, dtype=np.float32)
pc[:16] = view_proj.T.ravel() # Transpose for GLSL column-major
pc[16:19] = camera_right
pc[20:23] = camera_up
pc_bytes = pc.tobytes()
ffi = vk.ffi
cbuf = ffi.new("char[]", pc_bytes)
# Viewport + scissor
vk_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, [vk_vp])
scissor = vk.VkRect2D(
offset=vk.VkOffset2D(x=0, y=0),
extent=vk.VkExtent2D(width=extent[0], height=extent[1]),
)
vk.vkCmdSetScissor(cmd, 0, 1, [scissor])
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, self._pipeline)
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, self._pipeline_layout,
0, 1, [self._ssbo_set], 0, None,
)
vk._vulkan.lib.vkCmdPushConstants(
cmd, self._pipeline_layout,
vk.VK_SHADER_STAGE_VERTEX_BIT,
0, len(pc_bytes), cbuf,
)
# 6 vertices per particle (billboard quad)
vk.vkCmdDraw(cmd, count * 6, 1, 0, 0)
[docs]
def cleanup(self) -> None:
if not self._ready:
return
device = self._engine.ctx.device
for obj, fn in [
(self._pipeline, vk.vkDestroyPipeline),
(self._pipeline_layout, vk.vkDestroyPipelineLayout),
(self._vert_module, vk.vkDestroyShaderModule),
(self._frag_module, vk.vkDestroyShaderModule),
(self._ssbo_layout, vk.vkDestroyDescriptorSetLayout),
(self._ssbo_pool, vk.vkDestroyDescriptorPool),
]:
if obj:
fn(device, obj, None)
if self._particle_buf:
vk.vkDestroyBuffer(device, self._particle_buf, None)
if self._particle_mem:
vk.vkFreeMemory(device, self._particle_mem, None)
self._ready = False
def _create_particle_pipeline(
device: Any, vert_module: Any, frag_module: Any,
render_pass: Any, ssbo_layout: Any,
) -> tuple[Any, Any]:
"""Create particle pipeline: alpha blend, no depth write, no vertex input."""
ffi = vk.ffi
# Push constant: mat4(64) + vec3+pad(16) + vec3+pad(16) = 96 bytes
push_range = ffi.new("VkPushConstantRange*")
push_range.stageFlags = vk.VK_SHADER_STAGE_VERTEX_BIT
push_range.offset = 0
push_range.size = 96
set_layouts = ffi.new("VkDescriptorSetLayout[1]", [ssbo_layout])
layout_ci = ffi.new("VkPipelineLayoutCreateInfo*")
layout_ci.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO
layout_ci.setLayoutCount = 1
layout_ci.pSetLayouts = set_layouts
layout_ci.pushConstantRangeCount = 1
layout_ci.pPushConstantRanges = push_range
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]
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
# No vertex input (vertices generated in shader)
vi = ffi.new("VkPipelineVertexInputStateCreateInfo*")
vi.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO
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 (dynamic)
vps = ffi.new("VkPipelineViewportStateCreateInfo*")
vps.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO
vps.viewportCount = 1
vps.pViewports = ffi.new("VkViewport*")
vps.scissorCount = 1
vps.pScissors = ffi.new("VkRect2D*")
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_NONE # Billboards face camera
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 — test but no write (particles blend over geometry)
dss = ffi.new("VkPipelineDepthStencilStateCreateInfo*")
dss.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO
dss.depthTestEnable = 1
dss.depthWriteEnable = 0
dss.depthCompareOp = vk.VK_COMPARE_OP_LESS_OR_EQUAL
pi.pDepthStencilState = dss
# Alpha blend
blend_att = ffi.new("VkPipelineColorBlendAttachmentState*")
blend_att.blendEnable = 1
blend_att.srcColorBlendFactor = vk.VK_BLEND_FACTOR_SRC_ALPHA
blend_att.dstColorBlendFactor = vk.VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA
blend_att.colorBlendOp = vk.VK_BLEND_OP_ADD
blend_att.srcAlphaBlendFactor = vk.VK_BLEND_FACTOR_ONE
blend_att.dstAlphaBlendFactor = vk.VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA
blend_att.alphaBlendOp = vk.VK_BLEND_OP_ADD
blend_att.colorWriteMask = 0xF
cb = ffi.new("VkPipelineColorBlendStateCreateInfo*")
cb.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO
cb.attachmentCount = 1
cb.pAttachments = blend_att
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, ffi.NULL, 1, pi, ffi.NULL, pipeline_out,
)
if result != vk.VK_SUCCESS:
raise RuntimeError(f"vkCreateGraphicsPipelines failed: {result}")
log.debug("Particle pipeline created")
return pipeline_out[0], pipeline_layout