"""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
"""
from __future__ import annotations
import logging
import math
from typing import Any
import numpy as np
import vulkan as vk
from ..gpu.memory import _find_memory_type
from ..gpu.pipeline import create_shader_module
from ..materials.shader_compiler import compile_shader
__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."""
def __init__(self, engine: Any):
self._engine = engine
# Compute pipelines
self._irradiance_pipeline: Any = None
self._irradiance_layout: Any = None
self._prefilter_pipeline: Any = None
self._prefilter_layout: Any = None
self._brdf_pipeline: Any = None
self._brdf_layout: Any = None
# Descriptor resources
self._irradiance_desc_layout: Any = None
self._irradiance_desc_pool: Any = None
self._irradiance_desc_set: Any = None
self._prefilter_desc_layout: Any = None
self._prefilter_desc_pool: Any = None
self._prefilter_desc_sets: list[Any] = [] # one per mip level
self._brdf_desc_layout: Any = None
self._brdf_desc_pool: Any = None
self._brdf_desc_set: Any = None
# Shader modules
self._irradiance_module: Any = None
self._prefilter_module: Any = None
self._brdf_module: Any = None
# Output images
self._irradiance_image: Any = None
self._irradiance_memory: Any = None
self._irradiance_view: Any = None
self._prefilter_image: Any = None
self._prefilter_memory: Any = None
self._prefilter_view: Any = None
self._prefilter_mip_views: list[Any] = []
self._brdf_image: Any = None
self._brdf_memory: Any = None
self._brdf_view: Any = None
# Sampler for IBL output textures
self._sampler: Any = None
self._ready = False
[docs]
def setup(self) -> None:
"""Create compute pipelines for all three IBL processing stages."""
e = self._engine
device = e.ctx.device
shader_dir = e.shader_dir
# Compile compute shaders
irr_spv = compile_shader(shader_dir / "ibl_irradiance.comp")
pre_spv = compile_shader(shader_dir / "ibl_prefilter.comp")
brdf_spv = compile_shader(shader_dir / "ibl_brdf.comp")
self._irradiance_module = create_shader_module(device, irr_spv)
self._prefilter_module = create_shader_module(device, pre_spv)
self._brdf_module = create_shader_module(device, brdf_spv)
# Create output images
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
# --- 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
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
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 compute pipeline for irradiance convolution."""
# Layout: binding 0 = input cubemap sampler, binding 1 = output storage image
bindings = [
vk.VkDescriptorSetLayoutBinding(
binding=0,
descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_COMPUTE_BIT,
),
vk.VkDescriptorSetLayoutBinding(
binding=1,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_COMPUTE_BIT,
),
]
self._irradiance_desc_layout = vk.vkCreateDescriptorSetLayout(
device, vk.VkDescriptorSetLayoutCreateInfo(bindingCount=2, pBindings=bindings), None
)
# Pool
pool_sizes = [
vk.VkDescriptorPoolSize(type=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, descriptorCount=1),
vk.VkDescriptorPoolSize(type=vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, descriptorCount=1),
]
self._irradiance_desc_pool = vk.vkCreateDescriptorPool(
device, vk.VkDescriptorPoolCreateInfo(maxSets=1, poolSizeCount=2, pPoolSizes=pool_sizes), None
)
# Allocate descriptor set
sets = vk.vkAllocateDescriptorSets(
device,
vk.VkDescriptorSetAllocateInfo(
descriptorPool=self._irradiance_desc_pool,
descriptorSetCount=1,
pSetLayouts=[self._irradiance_desc_layout],
),
)
self._irradiance_desc_set = sets[0]
# Pipeline layout (no push constants)
self._irradiance_layout = vk.vkCreatePipelineLayout(
device,
vk.VkPipelineLayoutCreateInfo(
setLayoutCount=1,
pSetLayouts=[self._irradiance_desc_layout],
),
None,
)
# Pipeline
self._irradiance_pipeline = self._create_compute_pipeline(
device, self._irradiance_module, self._irradiance_layout
)
def _create_prefilter_pipeline(self, device: Any) -> None:
"""Create compute pipeline for prefiltered specular map with push constants."""
bindings = [
vk.VkDescriptorSetLayoutBinding(
binding=0,
descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_COMPUTE_BIT,
),
vk.VkDescriptorSetLayoutBinding(
binding=1,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_COMPUTE_BIT,
),
]
self._prefilter_desc_layout = vk.vkCreateDescriptorSetLayout(
device, vk.VkDescriptorSetLayoutCreateInfo(bindingCount=2, pBindings=bindings), None
)
# Pool — one set per mip level
pool_sizes = [
vk.VkDescriptorPoolSize(
type=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, descriptorCount=PREFILTER_MIP_LEVELS
),
vk.VkDescriptorPoolSize(type=vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, descriptorCount=PREFILTER_MIP_LEVELS),
]
self._prefilter_desc_pool = vk.vkCreateDescriptorPool(
device,
vk.VkDescriptorPoolCreateInfo(
maxSets=PREFILTER_MIP_LEVELS,
poolSizeCount=2,
pPoolSizes=pool_sizes,
),
None,
)
# Allocate one descriptor set per mip
self._prefilter_desc_sets = []
for _ in range(PREFILTER_MIP_LEVELS):
sets = vk.vkAllocateDescriptorSets(
device,
vk.VkDescriptorSetAllocateInfo(
descriptorPool=self._prefilter_desc_pool,
descriptorSetCount=1,
pSetLayouts=[self._prefilter_desc_layout],
),
)
self._prefilter_desc_sets.append(sets[0])
# Push constant range: roughness (float) + mip_size (uint) = 8 bytes
push_range = vk.VkPushConstantRange(
stageFlags=vk.VK_SHADER_STAGE_COMPUTE_BIT,
offset=0,
size=8,
)
self._prefilter_layout = vk.vkCreatePipelineLayout(
device,
vk.VkPipelineLayoutCreateInfo(
setLayoutCount=1,
pSetLayouts=[self._prefilter_desc_layout],
pushConstantRangeCount=1,
pPushConstantRanges=[push_range],
),
None,
)
self._prefilter_pipeline = self._create_compute_pipeline(device, self._prefilter_module, self._prefilter_layout)
def _create_brdf_pipeline(self, device: Any) -> None:
"""Create compute pipeline for BRDF LUT generation."""
binding = vk.VkDescriptorSetLayoutBinding(
binding=0,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_COMPUTE_BIT,
)
self._brdf_desc_layout = vk.vkCreateDescriptorSetLayout(
device, vk.VkDescriptorSetLayoutCreateInfo(bindingCount=1, pBindings=[binding]), None
)
pool_size = vk.VkDescriptorPoolSize(type=vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, descriptorCount=1)
self._brdf_desc_pool = vk.vkCreateDescriptorPool(
device, vk.VkDescriptorPoolCreateInfo(maxSets=1, poolSizeCount=1, pPoolSizes=[pool_size]), None
)
sets = vk.vkAllocateDescriptorSets(
device,
vk.VkDescriptorSetAllocateInfo(
descriptorPool=self._brdf_desc_pool,
descriptorSetCount=1,
pSetLayouts=[self._brdf_desc_layout],
),
)
self._brdf_desc_set = sets[0]
# Write BRDF LUT descriptor
brdf_info = vk.VkDescriptorImageInfo(
imageView=self._brdf_view,
imageLayout=vk.VK_IMAGE_LAYOUT_GENERAL,
)
vk.vkUpdateDescriptorSets(
device,
1,
[
vk.VkWriteDescriptorSet(
dstSet=self._brdf_desc_set,
dstBinding=0,
dstArrayElement=0,
descriptorCount=1,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
pImageInfo=[brdf_info],
)
],
0,
None,
)
self._brdf_layout = vk.vkCreatePipelineLayout(
device,
vk.VkPipelineLayoutCreateInfo(
setLayoutCount=1,
pSetLayouts=[self._brdf_desc_layout],
),
None,
)
self._brdf_pipeline = self._create_compute_pipeline(device, self._brdf_module, self._brdf_layout)
def _create_compute_pipeline(self, device: Any, module: Any, layout: Any) -> Any:
"""Create a compute pipeline from a shader module and pipeline layout."""
ffi = vk.ffi
main_name = ffi.new("char[]", b"main")
stage = ffi.new("VkPipelineShaderStageCreateInfo*")
stage.sType = vk.VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO
stage.stage = vk.VK_SHADER_STAGE_COMPUTE_BIT
stage.module = module
stage.pName = main_name
ci = ffi.new("VkComputePipelineCreateInfo*")
ci.sType = vk.VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO
ci.stage = stage[0]
ci.layout = layout
pipeline_out = ffi.new("VkPipeline*")
result = vk._vulkan._callApi(
vk._vulkan.lib.vkCreateComputePipelines,
device,
ffi.NULL,
1,
ci,
ffi.NULL,
pipeline_out,
)
if result != vk.VK_SUCCESS:
raise RuntimeError(f"vkCreateComputePipelines failed: {result}")
return pipeline_out[0]
# --- Processing ---
[docs]
def process_cubemap(self, cubemap_view: Any, cubemap_sampler: Any) -> None:
"""Run all IBL processing on the given environment cubemap.
This is a one-shot operation executed via one-time command buffer.
Call after loading the skybox cubemap.
"""
if not self._ready:
raise RuntimeError("IBL pass not set up — call setup() first")
e = self._engine
device = e.ctx.device
# Update descriptors with the source cubemap
self._write_source_cubemap(device, cubemap_view, cubemap_sampler)
# Record and execute all IBL compute dispatches
from ..gpu.memory import begin_single_time_commands, end_single_time_commands
cmd = begin_single_time_commands(device, e.ctx.command_pool)
# Transition output images to GENERAL for compute writes
self._transition_outputs_to_general(cmd)
# 1. Irradiance convolution
self._dispatch_irradiance(cmd)
# 2. Prefiltered specular (one dispatch per mip level)
self._dispatch_prefilter(cmd)
# 3. BRDF LUT
self._dispatch_brdf(cmd)
# Transition outputs to SHADER_READ_ONLY for sampling
self._transition_outputs_to_shader_read(cmd)
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
)
def _write_source_cubemap(self, device: Any, cubemap_view: Any, cubemap_sampler: Any) -> None:
"""Write the source cubemap to irradiance and prefilter descriptor 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=self._irradiance_desc_set,
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=self._irradiance_desc_set,
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=self._prefilter_desc_sets[mip],
dstBinding=0,
dstArrayElement=0,
descriptorCount=1,
descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
pImageInfo=[src_info],
)
)
writes.append(
vk.VkWriteDescriptorSet(
dstSet=self._prefilter_desc_sets[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) -> None:
"""Transition all output images from UNDEFINED to GENERAL for compute shader writes."""
barriers = [
# Irradiance cubemap
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._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=0,
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,
),
),
# BRDF LUT
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,
vk.VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
0,
0,
None,
0,
None,
len(barriers),
barriers,
)
def _dispatch_irradiance(self, cmd: 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,
[self._irradiance_desc_set],
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) -> 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,
[self._prefilter_desc_sets[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) -> None:
"""Transition all output images from GENERAL to SHADER_READ_ONLY_OPTIMAL."""
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,
),
),
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,
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 sets)
for pool in (self._irradiance_desc_pool, self._prefilter_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