"""Descriptor pool, layout, and set management."""
import logging
from typing import Any
import vulkan as vk
from ..types import MAX_TEXTURES
log = logging.getLogger(__name__)
__all__ = [
"DescriptorWriteBatch",
"allocate_descriptor_set",
"create_descriptor_pool",
"create_descriptor_set_layout",
"create_pool_for_types",
"create_ssbo_layout",
"create_texture_descriptor_layout",
"create_texture_descriptor_pool",
"write_image_descriptor",
"write_ssbo_descriptor",
"write_texture_descriptor",
"write_ubo_descriptor",
]
# This ``vulkan`` build omits the descriptor-indexing CREATE flags (their ``_EXT``
# aliases resolve to ``None``), so use the stable Vulkan ABI bit values directly.
# Both are bit 1 (0x2) in their respective flag enums and are frozen by the spec.
_POOL_CREATE_UPDATE_AFTER_BIND = 0x00000002 # VK_DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT
_LAYOUT_CREATE_UPDATE_AFTER_BIND = 0x00000002 # VK_DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT
[docs]
def create_pool_for_types(
device: Any,
sizes: dict[int, int],
max_sets: int = 1,
) -> Any:
"""Create a descriptor pool from a ``{descriptor_type: count}`` mapping.
General-purpose alternative to :func:`create_descriptor_pool` (which is
specialised for SSBO + sampler workloads). Use this for compute passes
that mix STORAGE_IMAGE, COMBINED_IMAGE_SAMPLER, UNIFORM_BUFFER, etc.
"""
pool_sizes = [vk.VkDescriptorPoolSize(type=t, descriptorCount=c) for t, c in sizes.items()]
return vk.vkCreateDescriptorPool(
device,
vk.VkDescriptorPoolCreateInfo(maxSets=max_sets, poolSizeCount=len(pool_sizes), pPoolSizes=pool_sizes),
None,
)
[docs]
def create_descriptor_set_layout(
device: Any,
bindings: list[tuple[int, int, int, int]],
) -> Any:
"""Create a descriptor set layout from a list of ``(binding, type, stage_flags, count)`` tuples."""
vk_bindings = [
vk.VkDescriptorSetLayoutBinding(binding=b, descriptorType=t, descriptorCount=c, stageFlags=s)
for b, t, s, c in bindings
]
return vk.vkCreateDescriptorSetLayout(
device,
vk.VkDescriptorSetLayoutCreateInfo(bindingCount=len(vk_bindings), pBindings=vk_bindings),
None,
)
[docs]
def create_descriptor_pool(
device: Any,
max_sets: int = 4,
extra_samplers: int = 0,
ssbo_count: int = 0,
ubo_count: int = 0,
update_after_bind: bool = False,
) -> Any:
"""Create a descriptor pool for SSBO descriptors (+ optional image samplers / UBOs).
If *ssbo_count* is given it overrides the default ``max_sets * 4`` SSBO descriptor count.
*ubo_count* reserves that many uniform-buffer descriptors (e.g. the FrameGlobals UBO).
Pass ``update_after_bind=True`` to allocate UPDATE_AFTER_BIND-capable sets from it.
"""
pool_sizes = [
vk.VkDescriptorPoolSize(
type=vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
descriptorCount=ssbo_count if ssbo_count > 0 else max_sets * 4,
),
]
if extra_samplers > 0:
pool_sizes.append(
vk.VkDescriptorPoolSize(
type=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
descriptorCount=extra_samplers,
)
)
if ubo_count > 0:
pool_sizes.append(
vk.VkDescriptorPoolSize(
type=vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
descriptorCount=ubo_count,
)
)
pool_info = vk.VkDescriptorPoolCreateInfo(
flags=_POOL_CREATE_UPDATE_AFTER_BIND if update_after_bind else 0,
maxSets=max_sets,
poolSizeCount=len(pool_sizes),
pPoolSizes=pool_sizes,
)
pool = vk.vkCreateDescriptorPool(device, pool_info, None)
log.debug("Descriptor pool created (max_sets=%d)", max_sets)
return pool
[docs]
def create_ssbo_layout(
device: Any,
binding_count: int = 3,
extra_samplers: int = 0,
trailing_ssbos: int = 0,
extra_samplers_tail: int = 0,
tail2_samplers: int = 0,
tail2_ssbos: int = 0,
tail_ubos: int = 0,
tail3_samplers: int = 0,
tail4_samplers: int = 0,
tail5_ssbos: int = 0,
tail6_ssbos: int = 0,
update_after_bind: bool = False,
) -> Any:
"""Create a descriptor set layout with N SSBO bindings + optional sampler bindings.
Binding order: ``binding_count`` SSBOs, then ``extra_samplers`` image samplers,
then ``trailing_ssbos`` additional SSBOs (fragment-only, for tile light data etc.),
then ``extra_samplers_tail`` more image samplers (fragment-only, e.g. the IBL
irradiance / prefilter / BRDF maps that must sit after the trailing SSBOs),
then ``tail2_samplers`` image samplers (e.g. reflection-probe cubemap arrays),
then ``tail2_ssbos`` SSBOs (e.g. the reflection-probe box buffer),
then ``tail_ubos`` uniform buffers (vertex + fragment, e.g. the FrameGlobals UBO),
then ``tail3_samplers`` fragment image samplers that must sit AFTER the UBO
tail (e.g. the scene colour / depth copy samplers, set0 b14/b15),
then ``tail4_samplers`` fragment image samplers after those (e.g. the pluggable
ambient indirect specular / diffuse hooks, set0 b16/b17),
then ``tail5_ssbos`` fragment SSBOs after those (e.g. the irradiance-volume
SH buffer, set0 b18),
then ``tail6_ssbos`` fragment SSBOs after those (e.g. the clustered-decal
buffer, set0 b19).
"""
bindings = []
for i in range(binding_count):
bindings.append(
vk.VkDescriptorSetLayoutBinding(
binding=i,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_VERTEX_BIT | vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
)
for i in range(extra_samplers):
bindings.append(
vk.VkDescriptorSetLayoutBinding(
binding=binding_count + i,
descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
)
trailing_start = binding_count + extra_samplers
for i in range(trailing_ssbos):
bindings.append(
vk.VkDescriptorSetLayoutBinding(
binding=trailing_start + i,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
)
tail_start = trailing_start + trailing_ssbos
for i in range(extra_samplers_tail):
bindings.append(
vk.VkDescriptorSetLayoutBinding(
binding=tail_start + i,
descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
)
tail2_start = tail_start + extra_samplers_tail
for i in range(tail2_samplers):
bindings.append(
vk.VkDescriptorSetLayoutBinding(
binding=tail2_start + i,
descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
)
tail2_ssbo_start = tail2_start + tail2_samplers
for i in range(tail2_ssbos):
bindings.append(
vk.VkDescriptorSetLayoutBinding(
binding=tail2_ssbo_start + i,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
)
tail_ubo_start = tail2_ssbo_start + tail2_ssbos
for i in range(tail_ubos):
bindings.append(
vk.VkDescriptorSetLayoutBinding(
binding=tail_ubo_start + i,
descriptorType=vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_VERTEX_BIT | vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
)
tail3_start = tail_ubo_start + tail_ubos
for i in range(tail3_samplers):
bindings.append(
vk.VkDescriptorSetLayoutBinding(
binding=tail3_start + i,
descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
)
tail4_start = tail3_start + tail3_samplers
for i in range(tail4_samplers):
bindings.append(
vk.VkDescriptorSetLayoutBinding(
binding=tail4_start + i,
descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
)
tail5_start = tail4_start + tail4_samplers
for i in range(tail5_ssbos):
bindings.append(
vk.VkDescriptorSetLayoutBinding(
binding=tail5_start + i,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
)
tail6_start = tail5_start + tail5_ssbos
for i in range(tail6_ssbos):
bindings.append(
vk.VkDescriptorSetLayoutBinding(
binding=tail6_start + i,
descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
descriptorCount=1,
stageFlags=vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
)
if update_after_bind:
# The image-sampler bindings in this set are written after the set is bound
# into a still-recording command buffer -- the IBL maps on skybox install,
# the reflection-probe cube arrays per capture -- so flag them
# UPDATE_AFTER_BIND + PARTIALLY_BOUND. The SSBO bindings stay normal: their
# contents change via mapped memory, not descriptor writes. This is the
# canonical replacement for the brittle "update every set before the first
# bind" sequencing the renderer used to rely on.
binding_flag_values = [
(vk.VK_DESCRIPTOR_BINDING_UPDATE_AFTER_BIND_BIT | vk.VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT)
if b.descriptorType == vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
else 0
for b in bindings
]
binding_flags = vk.VkDescriptorSetLayoutBindingFlagsCreateInfo(
bindingCount=len(binding_flag_values),
pBindingFlags=binding_flag_values,
)
layout_info = vk.VkDescriptorSetLayoutCreateInfo(
pNext=binding_flags,
flags=_LAYOUT_CREATE_UPDATE_AFTER_BIND,
bindingCount=len(bindings),
pBindings=bindings,
)
else:
layout_info = vk.VkDescriptorSetLayoutCreateInfo(
bindingCount=len(bindings),
pBindings=bindings,
)
layout = vk.vkCreateDescriptorSetLayout(device, layout_info, None)
log.debug("SSBO descriptor set layout created (%d bindings)", len(bindings))
return layout
[docs]
def allocate_descriptor_set(device: Any, pool: Any, layout: Any) -> Any:
"""Allocate a single descriptor set from the pool."""
alloc_info = vk.VkDescriptorSetAllocateInfo(
descriptorPool=pool,
descriptorSetCount=1,
pSetLayouts=[layout],
)
sets = vk.vkAllocateDescriptorSets(device, alloc_info)
return sets[0]
[docs]
def create_texture_descriptor_pool(
device: Any, max_textures: int = MAX_TEXTURES, update_after_bind: bool = False
) -> Any:
"""Create a descriptor pool for combined image samplers.
``update_after_bind`` must match the layout: pass ``True`` for the global
bindless texture array (whose slots are written while bound), ``False`` for
plain per-pass single-texture sets written once at setup.
"""
pool_size = vk.VkDescriptorPoolSize(
type=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
descriptorCount=max_textures,
)
pool_info = vk.VkDescriptorPoolCreateInfo(
flags=_POOL_CREATE_UPDATE_AFTER_BIND if update_after_bind else 0,
maxSets=1,
poolSizeCount=1,
pPoolSizes=[pool_size],
)
return vk.vkCreateDescriptorPool(device, pool_info, None)
[docs]
def create_texture_descriptor_layout(
device: Any, max_textures: int = MAX_TEXTURES, update_after_bind: bool = False
) -> Any:
"""Create a set layout for a combined-image-sampler array.
When ``update_after_bind`` is set (the global bindless texture array), the
binding is UPDATE_AFTER_BIND so individual slots can be (re)written while the
set is bound in a recording command buffer -- the canonical bindless pattern
-- and PARTIALLY_BOUND so unregistered slots need not be written. Plain
per-pass single-texture sets pass ``False`` (written once at setup, and they
allocate from a non-UAB pool).
"""
binding = vk.VkDescriptorSetLayoutBinding(
binding=0,
descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
descriptorCount=max_textures,
stageFlags=vk.VK_SHADER_STAGE_FRAGMENT_BIT,
)
if update_after_bind:
binding_flags = vk.VkDescriptorSetLayoutBindingFlagsCreateInfo(
bindingCount=1,
pBindingFlags=[
vk.VK_DESCRIPTOR_BINDING_UPDATE_AFTER_BIND_BIT
| vk.VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT
],
)
layout_info = vk.VkDescriptorSetLayoutCreateInfo(
pNext=binding_flags,
flags=_LAYOUT_CREATE_UPDATE_AFTER_BIND,
bindingCount=1,
pBindings=[binding],
)
else:
layout_info = vk.VkDescriptorSetLayoutCreateInfo(
bindingCount=1,
pBindings=[binding],
)
return vk.vkCreateDescriptorSetLayout(device, layout_info, None)
def _make_texture_write(
descriptor_set: Any, texture_index: int, image_view: Any, sampler: Any,
) -> vk.VkWriteDescriptorSet:
"""Build a VkWriteDescriptorSet for a texture array element without submitting it."""
image_info = vk.VkDescriptorImageInfo(
sampler=sampler, imageView=image_view,
imageLayout=vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
)
return vk.VkWriteDescriptorSet(
dstSet=descriptor_set, dstBinding=0, dstArrayElement=texture_index,
descriptorCount=1, descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
pImageInfo=[image_info],
)
def _make_image_write(
descriptor_set: Any, binding: int, image_view: Any, sampler: Any,
image_layout: int = vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
) -> vk.VkWriteDescriptorSet:
"""Build a VkWriteDescriptorSet for a combined image sampler without submitting it."""
image_info = vk.VkDescriptorImageInfo(sampler=sampler, imageView=image_view, imageLayout=image_layout)
return vk.VkWriteDescriptorSet(
dstSet=descriptor_set, dstBinding=binding, dstArrayElement=0,
descriptorCount=1, descriptorType=vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
pImageInfo=[image_info],
)
def _make_storage_image_write(
descriptor_set: Any, binding: int, image_view: Any,
image_layout: int = vk.VK_IMAGE_LAYOUT_GENERAL,
) -> vk.VkWriteDescriptorSet:
"""Build a VkWriteDescriptorSet for a storage image (compute write target)."""
image_info = vk.VkDescriptorImageInfo(imageView=image_view, imageLayout=image_layout)
return vk.VkWriteDescriptorSet(
dstSet=descriptor_set, dstBinding=binding, dstArrayElement=0,
descriptorCount=1, descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
pImageInfo=[image_info],
)
def _make_ssbo_write(
descriptor_set: Any, binding: int, buffer: Any, size: int,
) -> vk.VkWriteDescriptorSet:
"""Build a VkWriteDescriptorSet for an SSBO binding without submitting it."""
buf_info = vk.VkDescriptorBufferInfo(buffer=buffer, offset=0, range=size)
return vk.VkWriteDescriptorSet(
dstSet=descriptor_set, dstBinding=binding, dstArrayElement=0,
descriptorCount=1, descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
pBufferInfo=[buf_info],
)
def _make_uniform_buffer_write(
descriptor_set: Any, binding: int, buffer: Any, size: int,
) -> vk.VkWriteDescriptorSet:
"""Build a VkWriteDescriptorSet for a uniform buffer binding."""
buf_info = vk.VkDescriptorBufferInfo(buffer=buffer, offset=0, range=size)
return vk.VkWriteDescriptorSet(
dstSet=descriptor_set, dstBinding=binding, dstArrayElement=0,
descriptorCount=1, descriptorType=vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
pBufferInfo=[buf_info],
)
[docs]
def write_texture_descriptor(
device: Any, descriptor_set: Any, texture_index: int, image_view: Any, sampler: Any,
) -> None:
"""Write a single texture to the texture array at the given index."""
w = _make_texture_write(descriptor_set, texture_index, image_view, sampler)
vk.vkUpdateDescriptorSets(device, 1, [w], 0, None)
[docs]
def write_image_descriptor(
device: Any, descriptor_set: Any, binding: int, image_view: Any, sampler: Any,
) -> None:
"""Write a combined image sampler to a descriptor set at the given binding."""
w = _make_image_write(descriptor_set, binding, image_view, sampler)
vk.vkUpdateDescriptorSets(device, 1, [w], 0, None)
[docs]
def write_ssbo_descriptor(
device: Any, descriptor_set: Any, binding: int, buffer: Any, size: int,
) -> None:
"""Write a single SSBO buffer binding to a descriptor set."""
w = _make_ssbo_write(descriptor_set, binding, buffer, size)
vk.vkUpdateDescriptorSets(device, 1, [w], 0, None)
[docs]
def write_ubo_descriptor(
device: Any, descriptor_set: Any, binding: int, buffer: Any, size: int,
) -> None:
"""Write a single uniform-buffer binding to a descriptor set."""
w = _make_uniform_buffer_write(descriptor_set, binding, buffer, size)
vk.vkUpdateDescriptorSets(device, 1, [w], 0, None)
[docs]
class DescriptorWriteBatch:
"""Collects VkWriteDescriptorSet structs and flushes them in a single Vulkan call.
Usage::
batch = DescriptorWriteBatch(device)
batch.ssbo(ds, 0, buf_a, size_a)
batch.ssbo(ds, 1, buf_b, size_b)
batch.image(ds, 2, view, sampler)
batch.flush()
Can also be used as a context manager -- ``flush()`` is called on exit::
with DescriptorWriteBatch(device) as batch:
batch.ssbo(ds, 0, buf, size)
"""
__slots__ = ("_device", "_writes")
def __init__(self, device: Any) -> None:
self._device = device
self._writes: list[vk.VkWriteDescriptorSet] = []
# -- Accumulate writes --------------------------------------------------
[docs]
def ssbo(self, descriptor_set: Any, binding: int, buffer: Any, size: int) -> DescriptorWriteBatch:
"""Queue an SSBO descriptor write."""
self._writes.append(_make_ssbo_write(descriptor_set, binding, buffer, size))
return self
[docs]
def image(
self, descriptor_set: Any, binding: int, image_view: Any, sampler: Any,
image_layout: int = vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
) -> DescriptorWriteBatch:
"""Queue a combined image sampler descriptor write (default layout: SHADER_READ_ONLY_OPTIMAL)."""
self._writes.append(_make_image_write(descriptor_set, binding, image_view, sampler, image_layout))
return self
[docs]
def storage_image(
self, descriptor_set: Any, binding: int, image_view: Any,
image_layout: int = vk.VK_IMAGE_LAYOUT_GENERAL,
) -> DescriptorWriteBatch:
"""Queue a storage image descriptor write (compute writeable target)."""
self._writes.append(_make_storage_image_write(descriptor_set, binding, image_view, image_layout))
return self
[docs]
def texture(
self, descriptor_set: Any, texture_index: int, image_view: Any, sampler: Any,
) -> DescriptorWriteBatch:
"""Queue a texture array element descriptor write."""
self._writes.append(_make_texture_write(descriptor_set, texture_index, image_view, sampler))
return self
[docs]
def raw(self, write: vk.VkWriteDescriptorSet) -> DescriptorWriteBatch:
"""Queue a pre-built VkWriteDescriptorSet."""
self._writes.append(write)
return self
# -- Submit --------------------------------------------------------------
[docs]
def flush(self) -> int:
"""Submit all queued writes in a single ``vkUpdateDescriptorSets`` call.
Returns the number of writes submitted. The internal queue is cleared.
"""
count = len(self._writes)
if count:
vk.vkUpdateDescriptorSets(self._device, count, self._writes, 0, None)
self._writes.clear()
return count
# -- Context manager -----------------------------------------------------
[docs]
def __enter__(self) -> DescriptorWriteBatch:
return self
[docs]
def __exit__(self, *exc: object) -> None:
self.flush()
[docs]
def __len__(self) -> int:
return len(self._writes)