concinnity-device 0.19.53

//! Vulkan's implementation of the shared fullscreen post-pass seam
//! (`render::post::device::PostPassDevice`).
//!
//! Each pass's objects are derived here from what the single source declares:
//! the descriptor set layout is N sampled images and their N samplers, the
//! pipeline layout adds a fragment push-constant range of the declared size and,
//! for a probe-reading program, the forward global set as set 1, and the render
//! pass comes from the target's format and load action. The sets are allocated
//! per frame (post/set_arena.rs) rather than pre-wired per effect, so no pass
//! has to rewire an input another effect owns.

use ash::vk;
use concinnity_core::render::error::{RenderError, RenderResult};
use concinnity_core::render::post::device::{
    PostBlend, PostDraw, PostExtent, PostLoadOp, PostPassDevice, PostSampler, resolved_texture,
};
use concinnity_core::render::post::program::{PostProgram, PostProgramBindings};
use concinnity_core::render::render_graph::{PixelFormat, TextureDesc};

use crate::vulkan::allocator::DeviceAllocator;
use crate::vulkan::builtin_shaders::{self, CompileProgram};
use crate::vulkan::error::map_vk_result;
use crate::vulkan::owned::{OwnedPipeline, OwnedPipelineLayout, OwnedSetLayout, VkDevice};
use crate::vulkan::pipeline::GraphicsStages;
use crate::vulkan::post::pass_cache::PostPassCache;
use crate::vulkan::post::set_arena::PostSetArena;
use crate::vulkan::texture::{
    GpuImage, ImageSpec, create_image, create_image_view, one_shot_submit, transition_image_layout,
};
use crate::vulkan::transient_pool::{image_format, image_usage, sample_count};

// The set index a probe-reading program declares the forward global set at.
const PROBE_SET_INDEX: u32 = 1;

// A built fullscreen post pipeline plus the layouts a draw binds it through.
// The two layouts travel with the pipeline because they are derived from the
// same program declaration, so nothing else has to know the binding count.
pub(in crate::vulkan) struct PostPipeline {
    pipeline: OwnedPipeline,
    layout: OwnedPipelineLayout,
    set_layout: OwnedSetLayout,
    // What the program declares, so a draw can check what it was handed.
    bindings: PostProgramBindings,
}

// A persistent post target: the image, its view, and the extent its
// framebuffers were sized at.
pub(in crate::vulkan) struct PostTarget {
    image: GpuImage,
    extent: vk::Extent2D,
    format: PixelFormat,
}

impl PostTarget {
    /// The sampled view of this target, for a consumer that binds it directly.
    pub(in crate::vulkan) fn view(&self) -> vk::ImageView {
        self.image.view
    }
}

// A draw's color target: the view a framebuffer binds, the extent it is sized
// at, and the format its render pass is built for. A created target names
// itself this way through `target_attachment`; an image another subsystem owns
// (the HDR scene) is described directly.
#[derive(Clone, Copy)]
pub(in crate::vulkan) struct VkAttachment {
    pub view: vk::ImageView,
    pub extent: vk::Extent2D,
    pub format: PixelFormat,
}

// The forward global set, as a probe-reading program binds it for the probe
// count (binding 7), the cube array (binding 8) and the records (binding 17).
#[derive(Clone, Copy)]
pub(in crate::vulkan) struct VkPostProbes<'a> {
    pub layout: vk::DescriptorSetLayout,
    // One set per frame in flight. Empty at init, where no draw is encoded.
    pub sets: &'a [vk::DescriptorSet],
}

// The one-shot submit a freshly created target's initial layout transition
// needs. Held by the device value because target creation is the only operation
// that touches a queue.
#[derive(Clone, Copy)]
pub(in crate::vulkan) struct PostQueue {
    pub command_pool: vk::CommandPool,
    pub queue: vk::Queue,
}

// The Vulkan handles a shared post pass builds and encodes through. Borrowed,
// so the same value shape serves init (where no context exists yet) and each
// frame's encode.
pub(in crate::vulkan) struct VkPostDevice<'a> {
    pub device: &'a VkDevice,
    pub alloc: &'a DeviceAllocator,
    pub queue: PostQueue,
    // Render passes + framebuffers, cached across frames and passes.
    pub cache: &'a PostPassCache,
    // Per-frame descriptor sets.
    pub arena: &'a PostSetArena,
    // The linear clamp-to-edge state every screen-space source is sampled
    // through.
    pub sampler: vk::Sampler,
    // The trilinear clamp-to-edge state an environment cube is sampled through.
    pub cube_sampler: vk::Sampler,
    // The global set a probe-reading program binds.
    pub probes: VkPostProbes<'a>,
    // Which frame in flight is recording.
    pub frame: usize,
    pub hot_reload: bool,
}

// The SPIR-V a post program's two stages compile to: the one shared fullscreen
// triangle vertex plus the program's own fragment.
fn compile(program: PostProgram, hot_reload: bool) -> RenderResult<(Vec<u8>, Vec<u8>)> {
    let vert = builtin_shaders::FULLSCREEN_VERT.compile(hot_reload)?;
    Ok((vert, program.program().compile(hot_reload)?))
}

fn blend_attachment(blend: PostBlend) -> vk::PipelineColorBlendAttachmentState {
    let base = vk::PipelineColorBlendAttachmentState::default()
        .color_write_mask(vk::ColorComponentFlags::RGBA);
    match blend {
        PostBlend::Replace => base.blend_enable(false),
        PostBlend::Additive => base
            .blend_enable(true)
            .src_color_blend_factor(vk::BlendFactor::ONE)
            .dst_color_blend_factor(vk::BlendFactor::ONE)
            .src_alpha_blend_factor(vk::BlendFactor::ONE)
            .dst_alpha_blend_factor(vk::BlendFactor::ONE),
        PostBlend::PremultipliedOver => base
            .blend_enable(true)
            .src_color_blend_factor(vk::BlendFactor::ONE)
            .dst_color_blend_factor(vk::BlendFactor::ONE_MINUS_SRC_ALPHA)
            .src_alpha_blend_factor(vk::BlendFactor::ONE)
            .dst_alpha_blend_factor(vk::BlendFactor::ONE_MINUS_SRC_ALPHA),
    }
}

impl VkPostDevice<'_> {
    // The descriptor set layout for `n` sources, fragment-visible: their
    // sampled images at bindings `0..n` and their samplers at `n..2n`. Derived
    // from the program's declared count rather than written per pass, which is
    // what keeps the single source the contract.
    fn set_layout(&self, n: usize) -> RenderResult<OwnedSetLayout> {
        crate::vulkan::resources::create_descriptor_set_layout(
            self.device,
            &crate::vulkan::resources::source_set_bindings(n as u32),
        )
    }

    // The probe set a program that declares one binds.
    fn probes_for(&self, bindings: PostProgramBindings) -> Option<VkPostProbes<'_>> {
        bindings.probes.then_some(self.probes)
    }

    fn sampler_for(&self, sampler: PostSampler) -> vk::Sampler {
        match sampler {
            PostSampler::LinearClamp => self.sampler,
            PostSampler::LinearCube => self.cube_sampler,
        }
    }

    // The graphics pipeline itself: a vertex-buffer-less fullscreen triangle
    // into one color attachment, no depth, dynamic viewport and scissor. Every
    // fullscreen post pass has this shape, which is why it is built once here
    // instead of once per effect.
    fn build_pipeline(
        &self,
        shaders: (Vec<u8>, Vec<u8>),
        render_pass: vk::RenderPass,
        layout: vk::PipelineLayout,
        blend: PostBlend,
    ) -> RenderResult<OwnedPipeline> {
        let (vert_spv, frag_spv) = shaders;
        let modules = GraphicsStages::new(self.device, &vert_spv, &frag_spv)?;
        let stages = modules.infos();
        let vert_input = vk::PipelineVertexInputStateCreateInfo::default();
        let input_assembly = vk::PipelineInputAssemblyStateCreateInfo::default()
            .topology(vk::PrimitiveTopology::TRIANGLE_LIST)
            .primitive_restart_enable(false);
        let viewport_state = vk::PipelineViewportStateCreateInfo::default()
            .viewport_count(1)
            .scissor_count(1);
        let raster = vk::PipelineRasterizationStateCreateInfo::default()
            .depth_clamp_enable(false)
            .rasterizer_discard_enable(false)
            .polygon_mode(vk::PolygonMode::FILL)
            .line_width(1.0)
            .cull_mode(vk::CullModeFlags::NONE)
            .front_face(vk::FrontFace::COUNTER_CLOCKWISE)
            .depth_bias_enable(false);
        let multisample = vk::PipelineMultisampleStateCreateInfo::default()
            .sample_shading_enable(false)
            .rasterization_samples(vk::SampleCountFlags::TYPE_1);
        let depth_stencil = vk::PipelineDepthStencilStateCreateInfo::default()
            .depth_test_enable(false)
            .depth_write_enable(false)
            .depth_compare_op(vk::CompareOp::ALWAYS);
        let attach = blend_attachment(blend);
        let color_blend = vk::PipelineColorBlendStateCreateInfo::default()
            .logic_op_enable(false)
            .attachments(std::slice::from_ref(&attach));
        let dynamic_states = [vk::DynamicState::VIEWPORT, vk::DynamicState::SCISSOR];
        let dynamic = vk::PipelineDynamicStateCreateInfo::default().dynamic_states(&dynamic_states);
        let info = vk::GraphicsPipelineCreateInfo::default()
            .stages(&stages)
            .vertex_input_state(&vert_input)
            .input_assembly_state(&input_assembly)
            .viewport_state(&viewport_state)
            .rasterization_state(&raster)
            .multisample_state(&multisample)
            .depth_stencil_state(&depth_stencil)
            .color_blend_state(&color_blend)
            .dynamic_state(&dynamic)
            .layout(layout)
            .render_pass(render_pass)
            .subpass(0);
        crate::vulkan::pipeline_cache::create_graphics_pipeline(self.device, &info)
            .map_err(|e| map_vk_result(e, "create post pipeline"))
    }
}

impl PostPassDevice for VkPostDevice<'_> {
    type Recorder = vk::CommandBuffer;
    type Pipeline = PostPipeline;
    type Target = PostTarget;
    type TextureRef<'a> = vk::ImageView;
    type Attachment<'a> = VkAttachment;

    fn create_pipeline(
        &self,
        program: PostProgram,
        format: PixelFormat,
        blend: PostBlend,
    ) -> RenderResult<Self::Pipeline> {
        let bindings = program.bindings();
        let probes = self.probes_for(bindings);
        let set_layout = self.set_layout(bindings.textures)?;
        let mut set_layouts = vec![set_layout.handle()];
        set_layouts.extend(probes.map(|p| p.layout));
        let push = vk::PushConstantRange::default()
            .stage_flags(vk::ShaderStageFlags::FRAGMENT)
            .offset(0)
            .size(bindings.constants as u32);
        let mut layout_info = vk::PipelineLayoutCreateInfo::default().set_layouts(&set_layouts);
        if bindings.constants > 0 {
            layout_info = layout_info.push_constant_ranges(std::slice::from_ref(&push));
        }
        let layout = self
            .device
            .create_pipeline_layout(&layout_info)
            .map_err(|e| map_vk_result(e, "post pipeline layout"))?;

        // A pipeline is created against a render pass but is compatible with any
        // pass of the same attachment shape, so the cached one for this format
        // serves both the build and every draw.
        let render_pass = self
            .cache
            .render_pass(self.device, format, PostLoadOp::DontCare)?;
        let shaders = compile(program, self.hot_reload)?;
        let pipeline = self.build_pipeline(shaders, render_pass, layout.handle(), blend)?;
        Ok(PostPipeline {
            pipeline,
            layout,
            set_layout,
            bindings,
        })
    }

    fn create_target(
        &self,
        label: &'static str,
        desc: &TextureDesc,
        extent: PostExtent,
    ) -> RenderResult<Self::Target> {
        let spec = resolved_texture(label, desc, extent);
        let pooled = create_image(
            self.alloc,
            &ImageSpec {
                width: spec.width,
                height: spec.height,
                format: image_format(spec.format),
                tiling: vk::ImageTiling::OPTIMAL,
                usage: image_usage(spec.usage),
                mem_props: vk::MemoryPropertyFlags::DEVICE_LOCAL,
                samples: sample_count(spec.sample_count),
            },
        )
        .map_err(|e| e.context(format_args!("{label} post target")))?;
        let image = pooled.image();
        // Pre-transitioned so the first frame can sample a slot before anything
        // has rendered into it: a temporal pass binds its history on the very
        // first draw, gated to ignore what it reads.
        one_shot_submit(
            self.device,
            self.queue.command_pool,
            self.queue.queue,
            |cmd| {
                transition_image_layout(
                    self.device,
                    cmd,
                    image,
                    vk::ImageLayout::UNDEFINED,
                    vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL,
                    vk::ImageAspectFlags::COLOR,
                );
            },
        )?;
        let view = create_image_view(
            self.device,
            image,
            image_format(spec.format),
            vk::ImageAspectFlags::COLOR,
        )?;
        Ok(PostTarget {
            image: GpuImage::from_pooled(pooled, view),
            extent: vk::Extent2D {
                width: spec.width,
                height: spec.height,
            },
            format: spec.format,
        })
    }

    fn target_ref<'a>(&self, target: &'a Self::Target) -> Self::TextureRef<'a> {
        target.image.view
    }

    fn target_attachment<'a>(&self, target: &'a Self::Target) -> Self::Attachment<'a> {
        VkAttachment {
            view: target.image.view,
            extent: target.extent,
            format: target.format,
        }
    }

    fn encode(&self, rec: &Self::Recorder, draw: &PostDraw<'_, '_, Self>) -> RenderResult<()> {
        let cmd = *rec;
        let pipe = draw.pipeline;
        draw.check(pipe.bindings)?;
        // A render pass writes its attachment's layout in and out itself, so
        // who owns the target's state between passes changes nothing here.
        let probe_set = match self.probes_for(pipe.bindings) {
            None => None,
            Some(probes) => Some(*probes.sets.get(self.frame).ok_or_else(|| {
                RenderError::Other(format!(
                    "{}: no global set for frame {}",
                    draw.label, self.frame
                ))
            })?),
        };
        let target = draw.target;
        let render_pass = self
            .cache
            .render_pass(self.device, target.format, draw.load)?;
        let framebuffer =
            self.cache
                .framebuffer(self.device, render_pass, target.view, target.extent)?;

        // One set per draw from this frame's arena, written from what the pass
        // holds now. Nothing here is cached across frames, so no input needs a
        // re-point when another effect takes over the scene image.
        let set = self
            .arena
            .alloc(self.device, self.frame, pipe.set_layout.handle())?;
        let mut sources =
            [(vk::ImageView::null(), vk::Sampler::null()); crate::vulkan::resources::MAX_SOURCES];
        for (source, b) in sources.iter_mut().zip(draw.binds) {
            *source = (b.texture, self.sampler_for(b.sampler));
        }
        crate::vulkan::resources::write_source_set(self.device, set, &sources[..draw.binds.len()]);

        let extent = target.extent;
        let rp_begin = vk::RenderPassBeginInfo::default()
            .render_pass(render_pass)
            .framebuffer(framebuffer)
            .render_area(vk::Rect2D::default().extent(extent));
        let vp = vk::Viewport {
            x: 0.0,
            y: 0.0,
            width: extent.width as f32,
            height: extent.height as f32,
            min_depth: 0.0,
            max_depth: 1.0,
        };
        let scissor = vk::Rect2D::default().extent(extent);
        let device = self.device;
        // SAFETY: `cmd` is a command buffer in the recording state, and every handle and slice
        // these commands name is live for the call.
        unsafe {
            device.cmd_begin_render_pass(cmd, &rp_begin, vk::SubpassContents::INLINE);
            device.cmd_set_viewport(cmd, 0, std::slice::from_ref(&vp));
            device.cmd_set_scissor(cmd, 0, std::slice::from_ref(&scissor));
            device.cmd_bind_pipeline(cmd, vk::PipelineBindPoint::GRAPHICS, pipe.pipeline.handle());
            device.cmd_bind_descriptor_sets(
                cmd,
                vk::PipelineBindPoint::GRAPHICS,
                pipe.layout.handle(),
                0,
                std::slice::from_ref(&set),
                &[],
            );
            if let Some(global) = probe_set {
                device.cmd_bind_descriptor_sets(
                    cmd,
                    vk::PipelineBindPoint::GRAPHICS,
                    pipe.layout.handle(),
                    PROBE_SET_INDEX,
                    std::slice::from_ref(&global),
                    &[],
                );
            }
            if !draw.constants.is_empty() {
                device.cmd_push_constants(
                    cmd,
                    pipe.layout.handle(),
                    vk::ShaderStageFlags::FRAGMENT,
                    0,
                    draw.constants,
                );
            }
            // The vertex stage builds the fullscreen triangle from the vertex id,
            // so the draw reads no vertex buffer.
            device.cmd_draw(cmd, 3, 1, 0, 0);
            device.cmd_end_render_pass(cmd);
        }
        Ok(())
    }
}

// A post-pass extent from a Vulkan one.
pub(in crate::vulkan) fn post_extent(extent: vk::Extent2D) -> PostExtent {
    PostExtent {
        width: extent.width,
        height: extent.height,
    }
}

impl crate::vulkan::context::VkContext {
    // The post-pass device over this context, recording into frame slot `frame`.
    pub(in crate::vulkan) fn post_device(&self, frame: usize) -> VkPostDevice<'_> {
        VkPostDevice {
            device: &self.hw.device,
            alloc: &self.hw.alloc,
            queue: PostQueue {
                command_pool: self.commands.command_pool,
                queue: self.hw.graphics_queue,
            },
            cache: &self.post.cache,
            arena: &self.post.arena,
            sampler: self.post.sampler.handle(),
            cube_sampler: self.scene.cube_sampler.handle(),
            probes: VkPostProbes {
                layout: self.descriptors.global_set_layout.handle(),
                sets: &self.descriptors.global_sets,
            },
            frame,
            hot_reload: self.hot_reload.enabled,
        }
    }

    // This frame slot's HDR scene, as a post draw's target.
    pub(in crate::vulkan) fn hdr_scene_attachment(&self, frame: usize) -> VkAttachment {
        VkAttachment {
            view: self.targets.hdr_resolve_images[frame % self.targets.hdr_resolve_images.len()]
                .view,
            extent: self.targets.render_extent,
            format: PixelFormat::Rgba16Float,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // Every post program's fragment, and the shared vertex, compile to SPIR-V.
    #[test]
    fn every_post_program_compiles() {
        if !concinnity_shader::dxc_available() {
            return;
        }
        for program in [
            PostProgram::TaaResolve,
            PostProgram::SsrResolve,
            PostProgram::SsgiGather,
            PostProgram::SsgiComposite,
        ] {
            let (vert, frag) =
                compile(program, false).unwrap_or_else(|e| panic!("{program:?}: {e}"));
            assert!(crate::vulkan::pipeline::is_spirv(&vert));
            assert!(crate::vulkan::pipeline::is_spirv(&frag));
        }
    }
}