"""Scene colour/depth copy targets for the desktop pass split (design D1/A7).
When a material (or pass) this frame declares ``needs_scene_colour`` /
``needs_scene_depth``, the forward renderer ends the HDR pass after the opaque
family draws, copies the HDR colour + depth into the sampleable textures owned
here, then re-begins a LOAD-op pass and draws the transparent phase. The
transparent uber shader samples set0 b14 (scene colour) / b15 (scene depth) for
refraction.
Zero-cost when unused: the full-size textures are allocated lazily on the first
frame that actually needs them. Until then set0 b14/b15 point at a shared 1x1
black fallback (written once at setup) so the descriptors are never unbound and
feature-off frames record no copies and allocate nothing.
"""
import logging
from typing import Any
import numpy as np
import vulkan as vk
from ..gpu.descriptors import DescriptorWriteBatch
from ..gpu.memory import create_image, upload_image_data
log = logging.getLogger(__name__)
__all__ = ["SceneCopyTargets", "SCENE_COLOUR_BINDING", "SCENE_DEPTH_BINDING"]
# Forward descriptor set (set 0) bindings, design D14.
SCENE_COLOUR_BINDING = 14
SCENE_DEPTH_BINDING = 15
_HDR_COLOUR_FORMAT = vk.VK_FORMAT_R16G16B16A16_SFLOAT
_DEPTH_FORMAT = vk.VK_FORMAT_D32_SFLOAT
def _barrier(image, aspect, old, new, src_access, dst_access):
return vk.VkImageMemoryBarrier(
srcAccessMask=src_access,
dstAccessMask=dst_access,
oldLayout=old,
newLayout=new,
srcQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
dstQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
image=image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=aspect, baseMipLevel=0, levelCount=1, baseArrayLayer=0, layerCount=1,
),
)
[docs]
class SceneCopyTargets:
"""Owns the scene colour + depth copy textures and their descriptor writes."""
def __init__(self, engine: Any) -> None:
self._engine = engine
self._width = 0
self._height = 0
# Full-size copy targets (lazy). None until the first split.
self.colour_image: Any = None
self.colour_mem: Any = None
self.colour_view: Any = None
self.depth_image: Any = None
self.depth_mem: Any = None
self.depth_view: Any = None
self.sampler: Any = None
# 1x1 black fallback bound to b14/b15 whenever no full-size texture exists.
self._fallback_img: Any = None
self._fallback_mem: Any = None
self._fallback_view: Any = None
self._fallback_sampler: Any = None
# Which view is currently written into the descriptors (avoids redundant
# UPDATE_AFTER_BIND writes): "fallback" or "scene".
self._bound: str = "fallback"
# -------------------------------------------------------------- setup
[docs]
def setup_fallback(self, ssbo_set: Any) -> None:
"""Create the shared 1x1 fallback and write it to b14/b15 (never-unbound)."""
e = self._engine
device = e.ctx.device
self._fallback_img, self._fallback_mem = upload_image_data(
device, e.ctx.physical_device, e.ctx.graphics_queue, e.ctx.command_pool,
np.zeros((1, 1, 4), dtype=np.uint8), 1, 1,
)
self._fallback_view = vk.vkCreateImageView(device, vk.VkImageViewCreateInfo(
image=self._fallback_img,
viewType=vk.VK_IMAGE_VIEW_TYPE_2D,
format=vk.VK_FORMAT_R8G8B8A8_UNORM,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0, levelCount=1, baseArrayLayer=0, layerCount=1,
),
), None)
self._fallback_sampler = self._make_sampler(device)
with DescriptorWriteBatch(device) as b:
b.image(ssbo_set, SCENE_COLOUR_BINDING, self._fallback_view, self._fallback_sampler)
b.image(ssbo_set, SCENE_DEPTH_BINDING, self._fallback_view, self._fallback_sampler)
@staticmethod
def _make_sampler(device: Any) -> Any:
"""Linear, clamp-to-edge sampler for screen-space scene reads."""
return vk.vkCreateSampler(device, vk.VkSamplerCreateInfo(
magFilter=vk.VK_FILTER_LINEAR,
minFilter=vk.VK_FILTER_LINEAR,
mipmapMode=vk.VK_SAMPLER_MIPMAP_MODE_NEAREST,
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=0.0,
), None)
def _ensure(self, width: int, height: int, ssbo_set: Any) -> None:
"""Allocate (or resize) the full-size copy textures and bind them to b14/b15."""
if self.colour_image is not None and self._width == width and self._height == height:
if self._bound != "scene":
self._write_scene_descriptors(ssbo_set)
return
self._destroy_targets()
e = self._engine
device = e.ctx.device
phys = e.ctx.physical_device
self._width, self._height = width, height
self.colour_image, self.colour_mem = create_image(
device, phys, width, height, _HDR_COLOUR_FORMAT,
vk.VK_IMAGE_USAGE_SAMPLED_BIT | vk.VK_IMAGE_USAGE_TRANSFER_DST_BIT,
)
self.colour_view = vk.vkCreateImageView(device, vk.VkImageViewCreateInfo(
image=self.colour_image, viewType=vk.VK_IMAGE_VIEW_TYPE_2D, format=_HDR_COLOUR_FORMAT,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0, levelCount=1, baseArrayLayer=0, layerCount=1,
),
), None)
self.depth_image, self.depth_mem = create_image(
device, phys, width, height, _DEPTH_FORMAT,
vk.VK_IMAGE_USAGE_SAMPLED_BIT | vk.VK_IMAGE_USAGE_TRANSFER_DST_BIT,
)
self.depth_view = vk.vkCreateImageView(device, vk.VkImageViewCreateInfo(
image=self.depth_image, viewType=vk.VK_IMAGE_VIEW_TYPE_2D, format=_DEPTH_FORMAT,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_DEPTH_BIT,
baseMipLevel=0, levelCount=1, baseArrayLayer=0, layerCount=1,
),
), None)
self.sampler = self._make_sampler(device)
self._write_scene_descriptors(ssbo_set)
def _write_scene_descriptors(self, ssbo_set: Any) -> None:
device = self._engine.ctx.device
with DescriptorWriteBatch(device) as b:
b.image(ssbo_set, SCENE_COLOUR_BINDING, self.colour_view, self.sampler)
b.image(ssbo_set, SCENE_DEPTH_BINDING, self.depth_view, self.sampler)
self._bound = "scene"
# -------------------------------------------------------------- capture
[docs]
def capture(self, cmd: Any, hdr_target: Any, ssbo_set: Any) -> None:
"""Copy HDR colour + depth into the scene textures (records into ``cmd``).
Called after ``end_hdr_pass`` and before ``rebegin_hdr_pass``. On entry
the HDR colour is SHADER_READ_ONLY and the HDR depth is
DEPTH_STENCIL_READ_ONLY (the offscreen pass final layouts). Both HDR
images are restored to those layouts so the reload pass loads them and
the velocity / SSAO samplers still read the same depth. The scene
textures end in SHADER_READ_ONLY for the transparent shader sample.
"""
self._ensure(hdr_target.width, hdr_target.height, ssbo_set)
w, h = hdr_target.width, hdr_target.height
c_aspect = vk.VK_IMAGE_ASPECT_COLOR_BIT
d_aspect = vk.VK_IMAGE_ASPECT_DEPTH_BIT
RO = vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
DRO = vk.VK_IMAGE_LAYOUT_DEPTH_STENCIL_READ_ONLY_OPTIMAL
TSRC = vk.VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL
TDST = vk.VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
UNDEF = vk.VK_IMAGE_LAYOUT_UNDEFINED
SR = vk.VK_ACCESS_SHADER_READ_BIT
TR = vk.VK_ACCESS_TRANSFER_READ_BIT
TW = vk.VK_ACCESS_TRANSFER_WRITE_BIT
# To transfer layouts (sources: colour SHADER_READ_ONLY, depth DEPTH_RO;
# dsts fully overwritten so oldLayout=UNDEFINED).
vk.vkCmdPipelineBarrier(
cmd,
vk.VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT | vk.VK_PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, None, 0, None, 4, [
_barrier(hdr_target.colour_image, c_aspect, RO, TSRC, SR, TR),
_barrier(hdr_target.depth_image, d_aspect, DRO, TSRC, SR, TR),
_barrier(self.colour_image, c_aspect, UNDEF, TDST, 0, TW),
_barrier(self.depth_image, d_aspect, UNDEF, TDST, 0, TW),
])
def region(aspect):
layers = vk.VkImageSubresourceLayers(
aspectMask=aspect, mipLevel=0, baseArrayLayer=0, layerCount=1)
return vk.VkImageCopy(
srcSubresource=layers, srcOffset=vk.VkOffset3D(x=0, y=0, z=0),
dstSubresource=layers, dstOffset=vk.VkOffset3D(x=0, y=0, z=0),
extent=vk.VkExtent3D(width=w, height=h, depth=1),
)
vk.vkCmdCopyImage(cmd, hdr_target.colour_image, TSRC, self.colour_image, TDST, 1, [region(c_aspect)])
vk.vkCmdCopyImage(cmd, hdr_target.depth_image, TSRC, self.depth_image, TDST, 1, [region(d_aspect)])
# Restore HDR images to their pass-final layouts; scene textures to
# SHADER_READ_ONLY for the transparent shader sample.
vk.vkCmdPipelineBarrier(
cmd, vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
vk.VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT | vk.VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT,
0, 0, None, 0, None, 4, [
_barrier(hdr_target.colour_image, c_aspect, TSRC, RO, TR, SR),
_barrier(hdr_target.depth_image, d_aspect, TSRC, DRO, TR, SR),
_barrier(self.colour_image, c_aspect, TDST, RO, TW, SR),
_barrier(self.depth_image, d_aspect, TDST, RO, TW, SR),
])
# -------------------------------------------------------------- teardown
def _destroy_targets(self) -> None:
device = self._engine.ctx.device
if self.sampler is not None:
vk.vkDestroySampler(device, self.sampler, None)
self.sampler = None
for view, img, mem in (
(self.colour_view, self.colour_image, self.colour_mem),
(self.depth_view, self.depth_image, self.depth_mem),
):
if view is not None:
vk.vkDestroyImageView(device, view, None)
if img is not None:
vk.vkDestroyImage(device, img, None)
if mem is not None:
vk.vkFreeMemory(device, mem, None)
self.colour_view = self.colour_image = self.colour_mem = None
self.depth_view = self.depth_image = self.depth_mem = None
self._bound = "fallback"
[docs]
def destroy(self) -> None:
device = self._engine.ctx.device
self._destroy_targets()
if self._fallback_view is not None:
vk.vkDestroyImageView(device, self._fallback_view, None)
vk.vkDestroyImage(device, self._fallback_img, None)
vk.vkFreeMemory(device, self._fallback_mem, None)
vk.vkDestroySampler(device, self._fallback_sampler, None)
self._fallback_view = None