"""Image-Based Lighting pass: compute shader pipeline for IBL map generation.
Generates three IBL textures from an environment cubemap:
1. Irradiance cubemap (32x32): diffuse hemisphere convolution
2. Prefiltered specular cubemap (128x128 with mip chain): GGX importance sampling
3. BRDF integration LUT (512x512, RG16F): split-sum approximation
"""
import logging
import math
from typing import Any
import numpy as np
import vulkan as vk
from ..gpu.descriptors import (
DescriptorWriteBatch,
allocate_descriptor_set,
create_descriptor_set_layout,
create_pool_for_types,
)
from ..gpu.memory import _find_memory_type
from ..gpu.pipeline_compute import create_compute_pipeline
__all__ = ["IBLPass"]
log = logging.getLogger(__name__)
# IBL texture dimensions
IRRADIANCE_SIZE = 32
PREFILTER_SIZE = 128
PREFILTER_MIP_LEVELS = 5 # log2(128) - log2(8) + 1 = 5 mip levels
BRDF_LUT_SIZE = 512
[docs]
class IBLPass:
"""Compute-shader IBL processing: irradiance, prefiltered specular, and BRDF LUT."""
# Bundles = (default skybox) + one per (probe, frame-in-flight). Each bundle
# owns its own irradiance + per-mip prefilter descriptor sets so:
# 1. two probes convolved in the same command buffer read from their OWN
# source cubes and never cross-contaminate (vkUpdateDescriptorSets
# mutates a set at record time, so a shared set's last write would win);
# 2. a probe re-convolved on consecutive frames updates a DIFFERENT set
# than the one the previous (still in-flight) frame bound, so there is no
# update-while-in-use hazard across frames-in-flight.
# 1 skybox + MAX_PROBES probes * FRAMES_IN_FLIGHT.
MAX_BUNDLES = 1 + 8 * 2
__slots__ = (
"_engine",
"_irradiance_pipeline", "_irradiance_layout",
"_prefilter_pipeline", "_prefilter_layout",
"_brdf_pipeline", "_brdf_layout",
"_irradiance_desc_layout", "_prefilter_desc_layout",
"_brdf_desc_layout", "_brdf_desc_pool", "_brdf_desc_set",
"_desc_pool",
"_default_bundle", "_bundles",
"_irradiance_module", "_prefilter_module", "_brdf_module",
"_irradiance_image", "_irradiance_memory", "_irradiance_view",
"_prefilter_image", "_prefilter_memory", "_prefilter_view", "_prefilter_mip_views",
"_brdf_image", "_brdf_memory", "_brdf_view",
"_sampler",
"_ready",
)
def __init__(self, engine: Any):
for slot in self.__slots__:
object.__setattr__(self, slot, None)
self._engine = engine
self._prefilter_mip_views = []
# owner-id -> {"irradiance": set, "prefilter": [set, ...]}
self._bundles = {}
self._ready = False
[docs]
def setup(self) -> None:
"""Create compute pipelines for all three IBL processing stages."""
e = self._engine
device = e.ctx.device
# Create output images first (needed by descriptor writes during pipeline setup)
self._create_irradiance_image(device, e.ctx.physical_device)
self._create_prefilter_image(device, e.ctx.physical_device)
self._create_brdf_image(device, e.ctx.physical_device)
# Create sampler for IBL textures (clamp-to-edge, linear mip)
self._sampler = vk.vkCreateSampler(
device,
vk.VkSamplerCreateInfo(
magFilter=vk.VK_FILTER_LINEAR,
minFilter=vk.VK_FILTER_LINEAR,
mipmapMode=vk.VK_SAMPLER_MIPMAP_MODE_LINEAR,
addressModeU=vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
addressModeV=vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
addressModeW=vk.VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
minLod=0.0,
maxLod=float(PREFILTER_MIP_LEVELS),
),
None,
)
# Create pipelines
self._create_irradiance_pipeline(device)
self._create_prefilter_pipeline(device)
self._create_brdf_pipeline(device)
self._ready = True
log.debug("IBL pass initialized")
# --- Output accessors ---
[docs]
def get_irradiance_view(self) -> Any:
"""Return the irradiance cubemap image view."""
return self._irradiance_view
[docs]
def get_prefiltered_view(self) -> Any:
"""Return the prefiltered specular cubemap image view."""
return self._prefilter_view
[docs]
def get_brdf_lut_view(self) -> Any:
"""Return the BRDF LUT image view."""
return self._brdf_view
[docs]
def get_sampler(self) -> Any:
"""Return the IBL sampler."""
return self._sampler
[docs]
def get_irradiance_image(self) -> Any:
"""Return the irradiance cubemap image (for copy into a probe array)."""
return self._irradiance_image
[docs]
def get_prefiltered_image(self) -> Any:
"""Return the prefiltered specular cubemap image (for copy into a probe array)."""
return self._prefilter_image
# --- Image creation ---
def _create_cubemap_image(
self,
device: Any,
phys: Any,
size: int,
mip_levels: int,
usage: int,
) -> tuple[Any, Any]:
"""Create a cubemap image with given size and mip levels. Returns (image, memory)."""
ffi = vk.ffi
ci = ffi.new("VkImageCreateInfo*")
ci.sType = vk.VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO
ci.imageType = vk.VK_IMAGE_TYPE_2D
ci.format = vk.VK_FORMAT_R16G16B16A16_SFLOAT
ci.extent.width = size
ci.extent.height = size
ci.extent.depth = 1
ci.mipLevels = mip_levels
ci.arrayLayers = 6
ci.samples = vk.VK_SAMPLE_COUNT_1_BIT
ci.tiling = vk.VK_IMAGE_TILING_OPTIMAL
ci.usage = usage
ci.sharingMode = vk.VK_SHARING_MODE_EXCLUSIVE
ci.initialLayout = vk.VK_IMAGE_LAYOUT_UNDEFINED
ci.flags = vk.VK_IMAGE_CREATE_CUBE_COMPATIBLE_BIT
img_out = ffi.new("VkImage*")
result = vk._vulkan._callApi(vk._vulkan.lib.vkCreateImage, device, ci, ffi.NULL, img_out)
if result != vk.VK_SUCCESS:
raise RuntimeError(f"vkCreateImage failed: {result}")
image = img_out[0]
mem_req = vk.vkGetImageMemoryRequirements(device, image)
mem_type = _find_memory_type(phys, mem_req.memoryTypeBits, vk.VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)
memory = vk.vkAllocateMemory(
device,
vk.VkMemoryAllocateInfo(
allocationSize=mem_req.size,
memoryTypeIndex=mem_type,
),
None,
)
vk.vkBindImageMemory(device, image, memory, 0)
return image, memory
def _create_irradiance_image(self, device: Any, phys: Any) -> None:
"""Create the 32x32 irradiance cubemap."""
usage = (vk.VK_IMAGE_USAGE_STORAGE_BIT | vk.VK_IMAGE_USAGE_SAMPLED_BIT
| vk.VK_IMAGE_USAGE_TRANSFER_SRC_BIT)
self._irradiance_image, self._irradiance_memory = self._create_cubemap_image(
device,
phys,
IRRADIANCE_SIZE,
1,
usage,
)
self._irradiance_view = vk.vkCreateImageView(
device,
vk.VkImageViewCreateInfo(
image=self._irradiance_image,
viewType=vk.VK_IMAGE_VIEW_TYPE_CUBE,
format=vk.VK_FORMAT_R16G16B16A16_SFLOAT,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=6,
),
),
None,
)
def _create_prefilter_image(self, device: Any, phys: Any) -> None:
"""Create the prefiltered specular cubemap with mip chain."""
usage = (vk.VK_IMAGE_USAGE_STORAGE_BIT | vk.VK_IMAGE_USAGE_SAMPLED_BIT
| vk.VK_IMAGE_USAGE_TRANSFER_SRC_BIT)
self._prefilter_image, self._prefilter_memory = self._create_cubemap_image(
device,
phys,
PREFILTER_SIZE,
PREFILTER_MIP_LEVELS,
usage,
)
# Full view (all mips) for sampling
self._prefilter_view = vk.vkCreateImageView(
device,
vk.VkImageViewCreateInfo(
image=self._prefilter_image,
viewType=vk.VK_IMAGE_VIEW_TYPE_CUBE,
format=vk.VK_FORMAT_R16G16B16A16_SFLOAT,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=PREFILTER_MIP_LEVELS,
baseArrayLayer=0,
layerCount=6,
),
),
None,
)
# Per-mip views for compute shader writes
self._prefilter_mip_views = []
for mip in range(PREFILTER_MIP_LEVELS):
view = vk.vkCreateImageView(
device,
vk.VkImageViewCreateInfo(
image=self._prefilter_image,
viewType=vk.VK_IMAGE_VIEW_TYPE_CUBE,
format=vk.VK_FORMAT_R16G16B16A16_SFLOAT,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=mip,
levelCount=1,
baseArrayLayer=0,
layerCount=6,
),
),
None,
)
self._prefilter_mip_views.append(view)
def _create_brdf_image(self, device: Any, phys: Any) -> None:
"""Create the 512x512 BRDF LUT (RG16F)."""
from ..gpu.memory import create_image
self._brdf_image, self._brdf_memory = create_image(
device,
phys,
BRDF_LUT_SIZE,
BRDF_LUT_SIZE,
vk.VK_FORMAT_R16G16_SFLOAT,
vk.VK_IMAGE_USAGE_STORAGE_BIT | vk.VK_IMAGE_USAGE_SAMPLED_BIT,
)
self._brdf_view = vk.vkCreateImageView(
device,
vk.VkImageViewCreateInfo(
image=self._brdf_image,
viewType=vk.VK_IMAGE_VIEW_TYPE_2D,
format=vk.VK_FORMAT_R16G16_SFLOAT,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=1,
),
),
None,
)
# --- Pipeline creation ---
def _create_irradiance_pipeline(self, device: Any) -> None:
"""Create the irradiance + prefilter layouts, the shared bundle pool, the
default (skybox) bundle, and the irradiance compute pipeline."""
cs = vk.VK_SHADER_STAGE_COMPUTE_BIT
self._irradiance_desc_layout = create_descriptor_set_layout(device, [
(0, vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, cs, 1),
(1, vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, cs, 1),
])
self._prefilter_desc_layout = create_descriptor_set_layout(device, [
(0, vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, cs, 1),
(1, vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, cs, 1),
])
# One shared pool for every bundle's sets: each bundle is 1 irradiance set
# + PREFILTER_MIP_LEVELS prefilter sets. Size for MAX_BUNDLES of them.
sets_per_bundle = 1 + PREFILTER_MIP_LEVELS
self._desc_pool = create_pool_for_types(device, {
vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER: self.MAX_BUNDLES * sets_per_bundle,
vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE: self.MAX_BUNDLES * sets_per_bundle,
}, max_sets=self.MAX_BUNDLES * sets_per_bundle)
# Default bundle (used by the skybox / one-shot process_cubemap path).
self._default_bundle = self._allocate_bundle(device)
self._irradiance_pipeline, self._irradiance_layout, self._irradiance_module = create_compute_pipeline(
device, self._engine.shader_dir / "ibl_irradiance.comp",
[self._irradiance_desc_layout], 0,
)
def _allocate_bundle(self, device: Any) -> dict[str, Any]:
"""Allocate one irradiance set + per-mip prefilter sets from the shared pool."""
return {
"irradiance": allocate_descriptor_set(device, self._desc_pool, self._irradiance_desc_layout),
"prefilter": [
allocate_descriptor_set(device, self._desc_pool, self._prefilter_desc_layout)
for _ in range(PREFILTER_MIP_LEVELS)
],
}
[docs]
def acquire_bundle(self, owner_id: int, frame_slot: int) -> dict[str, Any]:
"""Return a per-(owner, frame-in-flight) descriptor bundle, allocated on first use.
``owner_id`` is the caller's bounded slot (reflection probes pass their array
slot, 0..N-1), so the bundle set ``{(slot, frame_slot)}`` is bounded by
``MAX_BUNDLES`` regardless of probe churn: bundles are reused when a slot is
reassigned (``record_convolution`` rebinds the source cube each call), so the
descriptor pool never leaks or exhausts. A distinct bundle *per frame-in-flight*
ensures a probe re-convolved on consecutive frames never updates a set still in
use by the previous in-flight frame.
"""
key = (owner_id, frame_slot)
bundle = self._bundles.get(key)
if bundle is None:
bundle = self._allocate_bundle(self._engine.ctx.device)
self._bundles[key] = bundle
return bundle
def _create_prefilter_pipeline(self, device: Any) -> None:
"""Create compute pipeline for prefiltered specular map with push constants."""
# Push constants: roughness (float) + mip_size (uint) = 8 bytes
self._prefilter_pipeline, self._prefilter_layout, self._prefilter_module = create_compute_pipeline(
device, self._engine.shader_dir / "ibl_prefilter.comp",
[self._prefilter_desc_layout], 8,
)
def _create_brdf_pipeline(self, device: Any) -> None:
"""Create compute pipeline for BRDF LUT generation."""
cs = vk.VK_SHADER_STAGE_COMPUTE_BIT
self._brdf_desc_layout = create_descriptor_set_layout(device, [
(0, vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, cs, 1),
])
self._brdf_desc_pool = create_pool_for_types(
device, {vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE: 1},
)
self._brdf_desc_set = allocate_descriptor_set(
device, self._brdf_desc_pool, self._brdf_desc_layout,
)
with DescriptorWriteBatch(device) as batch:
batch.storage_image(self._brdf_desc_set, 0, self._brdf_view)
self._brdf_pipeline, self._brdf_layout, self._brdf_module = create_compute_pipeline(
device, self._engine.shader_dir / "ibl_brdf.comp",
[self._brdf_desc_layout], 0,
)
# --- Processing ---
[docs]
def process_cubemap(self, cubemap_view: Any, cubemap_sampler: Any) -> None:
"""Run all IBL processing on the given environment cubemap (one-shot).
Submits + waits on its own command buffer. Used by the skybox install
path at setup; not the steady-state hot path. Reflection probes use
:meth:`record_convolution` to record into the primary frame cmd instead.
"""
if not self._ready:
raise RuntimeError("IBL pass not set up: call setup() first")
e = self._engine
device = e.ctx.device
from ..gpu.memory import begin_single_time_commands, end_single_time_commands
cmd = begin_single_time_commands(device, e.ctx.command_pool)
self.record_convolution(cmd, cubemap_view, cubemap_sampler, self._default_bundle, with_brdf=True)
end_single_time_commands(device, e.ctx.graphics_queue, e.ctx.command_pool, cmd)
log.debug(
"IBL maps generated (irradiance=%d, prefilter=%d, brdf=%d)", IRRADIANCE_SIZE, PREFILTER_SIZE, BRDF_LUT_SIZE
)
[docs]
def record_convolution(
self,
cmd: Any,
cubemap_view: Any,
cubemap_sampler: Any,
bundle: dict[str, Any] | None = None,
*,
with_brdf: bool = False,
) -> None:
"""Record the irradiance + prefilter convolution of *cubemap_view* into *cmd*.
No submit, no wait: the caller owns the command buffer (the primary frame
cmd for probes). Reads from the given per-owner *bundle* so concurrent
probe convolutions in one cmd never share a descriptor set. Leaves the
irradiance + prefilter output images in SHADER_READ_ONLY_OPTIMAL ready for
the per-probe array copy. ``with_brdf`` also (re)generates the shared BRDF
LUT (only needed for the one-shot skybox path).
"""
if not self._ready:
raise RuntimeError("IBL pass not set up: call setup() first")
if bundle is None:
bundle = self._default_bundle
self._write_source_cubemap(self._engine.ctx.device, cubemap_view, cubemap_sampler, bundle)
self._transition_outputs_to_general(cmd, with_brdf=with_brdf)
self._dispatch_irradiance(cmd, bundle)
self._dispatch_prefilter(cmd, bundle)
if with_brdf:
self._dispatch_brdf(cmd)
self._transition_outputs_to_shader_read(cmd, with_brdf=with_brdf)
def _write_source_cubemap(
self, device: Any, cubemap_view: Any, cubemap_sampler: Any, bundle: dict[str, Any]
) -> None:
"""Write the source cubemap into the bundle's irradiance + prefilter sets."""
src_info = vk.VkDescriptorImageInfo(
sampler=cubemap_sampler,
imageView=cubemap_view,
imageLayout=vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
)
irr_out_info = vk.VkDescriptorImageInfo(
imageView=self._irradiance_view,
imageLayout=vk.VK_IMAGE_LAYOUT_GENERAL,
)
writes = [
# Irradiance: binding 0 = source cubemap
vk.VkWriteDescriptorSet(
dstSet=bundle["irradiance"],
dstBinding=0,
dstArrayElement=0,
descriptorCount=1,
descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
pImageInfo=[src_info],
),
# Irradiance: binding 1 = output irradiance map
vk.VkWriteDescriptorSet(
dstSet=bundle["irradiance"],
dstBinding=1,
dstArrayElement=0,
descriptorCount=1,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
pImageInfo=[irr_out_info],
),
]
# Prefilter: write source cubemap + per-mip output views
for mip in range(PREFILTER_MIP_LEVELS):
mip_out_info = vk.VkDescriptorImageInfo(
imageView=self._prefilter_mip_views[mip],
imageLayout=vk.VK_IMAGE_LAYOUT_GENERAL,
)
writes.append(
vk.VkWriteDescriptorSet(
dstSet=bundle["prefilter"][mip],
dstBinding=0,
dstArrayElement=0,
descriptorCount=1,
descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
pImageInfo=[src_info],
)
)
writes.append(
vk.VkWriteDescriptorSet(
dstSet=bundle["prefilter"][mip],
dstBinding=1,
dstArrayElement=0,
descriptorCount=1,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
pImageInfo=[mip_out_info],
)
)
vk.vkUpdateDescriptorSets(device, len(writes), writes, 0, None)
def _transition_outputs_to_general(self, cmd: Any, with_brdf: bool = False) -> None:
"""Transition the irradiance + prefilter outputs to GENERAL for compute writes.
``oldLayout=UNDEFINED`` discards prior contents (each convolution fully
rewrites them). The barrier still serialises against a previous probe's
array-copy (TRANSFER read) and convolution (COMPUTE write) recorded
earlier in the SAME cmd, so two probes sharing these output images in one
frame do not race: probe B's compute write waits for probe A's copy.
"""
src_stage = (vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT
| vk.VK_PIPELINE_STAGE_TRANSFER_BIT
| vk.VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT)
src_access = vk.VK_ACCESS_SHADER_WRITE_BIT | vk.VK_ACCESS_TRANSFER_READ_BIT
barriers = [
# Irradiance cubemap
vk.VkImageMemoryBarrier(
srcAccessMask=src_access,
dstAccessMask=vk.VK_ACCESS_SHADER_WRITE_BIT,
oldLayout=vk.VK_IMAGE_LAYOUT_UNDEFINED,
newLayout=vk.VK_IMAGE_LAYOUT_GENERAL,
image=self._irradiance_image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=6,
),
),
# Prefiltered cubemap (all mips)
vk.VkImageMemoryBarrier(
srcAccessMask=src_access,
dstAccessMask=vk.VK_ACCESS_SHADER_WRITE_BIT,
oldLayout=vk.VK_IMAGE_LAYOUT_UNDEFINED,
newLayout=vk.VK_IMAGE_LAYOUT_GENERAL,
image=self._prefilter_image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=PREFILTER_MIP_LEVELS,
baseArrayLayer=0,
layerCount=6,
),
),
]
if with_brdf:
# BRDF LUT
barriers.append(vk.VkImageMemoryBarrier(
srcAccessMask=0,
dstAccessMask=vk.VK_ACCESS_SHADER_WRITE_BIT,
oldLayout=vk.VK_IMAGE_LAYOUT_UNDEFINED,
newLayout=vk.VK_IMAGE_LAYOUT_GENERAL,
image=self._brdf_image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=1,
),
))
vk.vkCmdPipelineBarrier(
cmd,
src_stage,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
0,
0,
None,
0,
None,
len(barriers),
barriers,
)
def _dispatch_irradiance(self, cmd: Any, bundle: dict[str, Any]) -> None:
"""Dispatch irradiance convolution compute shader."""
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._irradiance_pipeline)
vk.vkCmdBindDescriptorSets(
cmd,
vk.VK_PIPELINE_BIND_POINT_COMPUTE,
self._irradiance_layout,
0,
1,
[bundle["irradiance"]],
0,
None,
)
# Dispatch: ceil(32/8) x ceil(32/8) x 6 faces
groups_xy = math.ceil(IRRADIANCE_SIZE / 8)
vk.vkCmdDispatch(cmd, groups_xy, groups_xy, 6)
# Barrier between irradiance and prefilter
vk.vkCmdPipelineBarrier(
cmd,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
0,
0,
None,
0,
None,
0,
None,
)
def _dispatch_prefilter(self, cmd: Any, bundle: dict[str, Any]) -> None:
"""Dispatch prefiltered specular map for each mip level."""
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._prefilter_pipeline)
ffi = vk.ffi
for mip in range(PREFILTER_MIP_LEVELS):
mip_size = PREFILTER_SIZE >> mip
roughness = mip / max(PREFILTER_MIP_LEVELS - 1, 1)
vk.vkCmdBindDescriptorSets(
cmd,
vk.VK_PIPELINE_BIND_POINT_COMPUTE,
self._prefilter_layout,
0,
1,
[bundle["prefilter"][mip]],
0,
None,
)
# Push constants: roughness (float) + mip_size (uint32)
pc_data = np.array([roughness], dtype=np.float32).tobytes()
pc_data += np.array([mip_size], dtype=np.uint32).tobytes()
cbuf = ffi.new("char[]", pc_data)
vk._vulkan.lib.vkCmdPushConstants(
cmd,
self._prefilter_layout,
vk.VK_SHADER_STAGE_COMPUTE_BIT,
0,
8,
cbuf,
)
groups_xy = max(1, math.ceil(mip_size / 8))
vk.vkCmdDispatch(cmd, groups_xy, groups_xy, 6)
# Barrier after prefilter
vk.vkCmdPipelineBarrier(
cmd,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
0,
0,
None,
0,
None,
0,
None,
)
def _dispatch_brdf(self, cmd: Any) -> None:
"""Dispatch BRDF LUT generation."""
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._brdf_pipeline)
vk.vkCmdBindDescriptorSets(
cmd,
vk.VK_PIPELINE_BIND_POINT_COMPUTE,
self._brdf_layout,
0,
1,
[self._brdf_desc_set],
0,
None,
)
groups = math.ceil(BRDF_LUT_SIZE / 8)
vk.vkCmdDispatch(cmd, groups, groups, 1)
def _transition_outputs_to_shader_read(self, cmd: Any, with_brdf: bool = False) -> None:
"""Transition the irradiance + prefilter outputs to SHADER_READ_ONLY_OPTIMAL.
For the probe path the next reader is the per-probe array copy (TRANSFER),
so the destination stage covers FRAGMENT (skybox sampling) and TRANSFER
(the copy). ``with_brdf`` also transitions the shared BRDF LUT.
"""
barriers = [
vk.VkImageMemoryBarrier(
srcAccessMask=vk.VK_ACCESS_SHADER_WRITE_BIT,
dstAccessMask=vk.VK_ACCESS_SHADER_READ_BIT,
oldLayout=vk.VK_IMAGE_LAYOUT_GENERAL,
newLayout=vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
image=self._irradiance_image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=6,
),
),
vk.VkImageMemoryBarrier(
srcAccessMask=vk.VK_ACCESS_SHADER_WRITE_BIT,
dstAccessMask=vk.VK_ACCESS_SHADER_READ_BIT,
oldLayout=vk.VK_IMAGE_LAYOUT_GENERAL,
newLayout=vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
image=self._prefilter_image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=PREFILTER_MIP_LEVELS,
baseArrayLayer=0,
layerCount=6,
),
),
]
if with_brdf:
barriers.append(vk.VkImageMemoryBarrier(
srcAccessMask=vk.VK_ACCESS_SHADER_WRITE_BIT,
dstAccessMask=vk.VK_ACCESS_SHADER_READ_BIT,
oldLayout=vk.VK_IMAGE_LAYOUT_GENERAL,
newLayout=vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
image=self._brdf_image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=1,
),
))
vk.vkCmdPipelineBarrier(
cmd,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
vk.VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT | vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
0,
0,
None,
0,
None,
len(barriers),
barriers,
)
# --- Cleanup ---
[docs]
def cleanup(self) -> None:
"""Destroy all GPU resources owned by the IBL pass."""
if not self._ready:
return
device = self._engine.ctx.device
# Pipelines
for pipeline in (self._irradiance_pipeline, self._prefilter_pipeline, self._brdf_pipeline):
if pipeline:
vk.vkDestroyPipeline(device, pipeline, None)
for layout in (self._irradiance_layout, self._prefilter_layout, self._brdf_layout):
if layout:
vk.vkDestroyPipelineLayout(device, layout, None)
# Shader modules
for module in (self._irradiance_module, self._prefilter_module, self._brdf_module):
if module:
vk.vkDestroyShaderModule(device, module, None)
# Descriptor pools (implicitly frees all bundle sets)
for pool in (self._desc_pool, self._brdf_desc_pool):
if pool:
vk.vkDestroyDescriptorPool(device, pool, None)
for layout in (self._irradiance_desc_layout, self._prefilter_desc_layout, self._brdf_desc_layout):
if layout:
vk.vkDestroyDescriptorSetLayout(device, layout, None)
# Sampler
if self._sampler:
vk.vkDestroySampler(device, self._sampler, None)
# Image views
if self._irradiance_view:
vk.vkDestroyImageView(device, self._irradiance_view, None)
if self._prefilter_view:
vk.vkDestroyImageView(device, self._prefilter_view, None)
for view in self._prefilter_mip_views:
vk.vkDestroyImageView(device, view, None)
if self._brdf_view:
vk.vkDestroyImageView(device, self._brdf_view, None)
# Images and memory
for img, mem in [
(self._irradiance_image, self._irradiance_memory),
(self._prefilter_image, self._prefilter_memory),
(self._brdf_image, self._brdf_memory),
]:
if img:
vk.vkDestroyImage(device, img, None)
if mem:
vk.vkFreeMemory(device, mem, None)
self._ready = False