concinnity-device 0.19.119

//! DirectX's implementation of the shared fullscreen post-pass seam
//! (`render::post::device::PostPassDevice`).
//!
//! The root signature is derived from what the single source declares rather
//! than hand-written per pass: N single-SRV descriptor tables at t0..tN-1, a
//! 32-bit constant block at b0 when the program declares constants, and N
//! static samplers at s0..sN-1. A probe-reading program adds the
//! reflection-probe cube array at tN with its one static sampler at sN, the
//! ProbeSet at b1, the probe records at tN+1, and the main camera's cluster
//! grid binning them: its params at b2 and its lists at tN+2. The tables are
//! single-SRV and separate because each source's texture and sampler sit at
//! their own register, and because a pass's sources come from unrelated owners,
//! so nothing makes them contiguous in the heap.
//!
//! Targets take their SRV and RTV from the shared post descriptor block
//! (post/descriptors.rs) instead of slots reserved for the effect by name in
//! `init/heap_layout.rs`.

use concinnity_core::render::error::{RenderError, RenderResult};
use concinnity_core::render::post::device::{
    PostBlend, PostDraw, PostExtent, PostLoadOp, PostPassDevice, PostTargetState, check_level,
    level_extent, resolved_texture,
};
use concinnity_core::render::post::program::{PostProgram, PostProgramBindings};
use concinnity_core::render::render_graph::{PixelFormat, TextureDesc};
use windows::Win32::Graphics::Direct3D::D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST;
use windows::Win32::Graphics::Direct3D12::*;

use crate::directx::builtin_shaders::{self, CompileProgram};
use crate::directx::com;
use crate::directx::context::dump_on_err;
use crate::directx::descriptor_slot::DescriptorTables;
use crate::directx::descriptor_slot::SrvSlot;
use crate::directx::post::descriptors::{PostDescriptors, PostTargetDescriptors};
use crate::directx::pso::GraphicsPso;
use crate::directx::root_constants::RootConstants;
use crate::directx::root_sig::{RootSig, SamplerState, Visibility};
use crate::directx::texture::{
    create_rt_chain, subresource_transition_barrier, write_level_rtv, write_levels_srv,
};
use crate::directx::transient_pool::dxgi_format;

// Pixel dimensions of a draw's target, which are its viewport and scissor.
#[derive(Clone, Copy)]
pub(in crate::directx) struct FullscreenExtent {
    pub width: u32,
    pub height: u32,
}

// The registers a probe-reading program declares its ProbeSet and the cluster
// params at, after the constants at b0.
const PROBE_SET_REGISTER: u32 = 1;
const CLUSTER_PARAMS_REGISTER: u32 = 2;

// A built fullscreen post pipeline: the PSO, the root signature it was derived
// against, and the declaration that derivation used, so a draw can check what
// it was handed and find its root parameters.
pub(in crate::directx) struct PostPipeline {
    pub(in crate::directx) pso: ID3D12PipelineState,
    pub(in crate::directx) root_sig: ID3D12RootSignature,
    bindings: PostProgramBindings,
}

impl PostPipeline {
    // Root parameter of the constant block, which follows the source tables.
    fn constants_parameter(&self) -> u32 {
        self.bindings.textures as u32
    }

    // Root parameter of the probe cube table; the ProbeSet, the probe records,
    // the cluster params and the cluster lists are the four after it.
    fn probes_parameter(&self) -> u32 {
        self.bindings.textures as u32 + u32::from(self.bindings.constants > 0)
    }
}

// A persistent post target: the resource, its descriptors, and the extent of its
// top level. `descriptors` samples every level and renders to the top one; a
// target with several levels also holds a set per level, sampling and
// rendering to that level alone.
pub(in crate::directx) struct PostTarget {
    pub(in crate::directx) resource: ID3D12Resource,
    pub(in crate::directx) descriptors: PostTargetDescriptors,
    label: &'static str,
    levels: Vec<PostTargetDescriptors>,
    level_count: u32,
    extent: FullscreenExtent,
}

impl PostTarget {
    // Mip `level`'s own descriptors, or the target's when it has one level.
    fn level(&self, level: u32) -> RenderResult<&PostTargetDescriptors> {
        check_level(self.label, level, self.level_count)?;
        Ok(self.levels.get(level as usize).unwrap_or(&self.descriptors))
    }
}

impl PostTarget {
    // The shader-visible handle a consumer samples this target through.
    pub(in crate::directx) fn srv_gpu(&self) -> SrvSlot {
        self.descriptors.srv_gpu
    }

    // The resource, for a consumer that copies from it or barriers it itself.
    pub(in crate::directx) fn resource(&self) -> &ID3D12Resource {
        &self.resource
    }

    // The render-target view of its top level.
    pub(in crate::directx) fn rtv(&self) -> D3D12_CPU_DESCRIPTOR_HANDLE {
        self.descriptors.rtv
    }
}

// A transient-pool resource a shared post pass writes, viewed through
// descriptors from the post block. The pool owns the memory and repacks it on
// a rebuild, so the view is recreated after every rebuild rather than kept.
pub(in crate::directx) struct PooledTarget {
    resource: ID3D12Resource,
    descriptors: PostTargetDescriptors,
    extent: FullscreenExtent,
}

impl PooledTarget {
    // The shader-visible handle a consumer samples it through.
    pub(in crate::directx) fn srv_gpu(&self) -> SrvSlot {
        self.descriptors.srv_gpu
    }

    // The resource, for a consumer that barriers it itself.
    pub(in crate::directx) fn resource(&self) -> &ID3D12Resource {
        &self.resource
    }

    // As a draw's color target.
    pub(in crate::directx) fn attachment(&self) -> DxAttachment<'_> {
        DxAttachment {
            resource: &self.resource,
            subresource: 0,
            rtv: self.descriptors.rtv,
            extent: self.extent,
        }
    }
}

// A draw's color target: the resource and the subresource it writes (for the
// transitions of a target its pass owns), the RTV it is written through, and its
// extent, which is the viewport.
#[derive(Clone, Copy)]
pub(in crate::directx) struct DxAttachment<'a> {
    pub resource: &'a ID3D12Resource,
    pub subresource: u32,
    pub rtv: D3D12_CPU_DESCRIPTOR_HANDLE,
    pub extent: FullscreenExtent,
}

// The world's reflection-probe set, as a probe-reading program binds it.
#[derive(Clone, Copy)]
pub(in crate::directx) struct DxPostProbes {
    // The cube array's SRV.
    pub cube_table: SrvSlot,
    // This frame's ProbeSet constant buffer.
    pub set_cbv: u64,
    // This frame's probe records buffer.
    pub records: u64,
    // This frame's `ClusterParams` constant buffer and the per-cluster lists,
    // which bin the probes a fragment blends.
    pub cluster_cbv: u64,
    pub cluster_list: u64,
}

impl DxPostProbes {
    // Bind the set at five consecutive root parameters from `first`: the cube
    // array table, the ProbeSet CBV, the records, the cluster params and the
    // cluster lists.
    pub(in crate::directx) fn bind(&self, cmd: &ID3D12GraphicsCommandList, first: u32) {
        // SAFETY: the command list is in the recording state with the SRV heap
        // bound, and the table slot and every buffer are live for the frame the
        // list records.
        unsafe {
            cmd.set_graphics_srv_table(first, self.cube_table);
            cmd.SetGraphicsRootConstantBufferView(first + 1, self.set_cbv);
            cmd.SetGraphicsRootShaderResourceView(first + 2, self.records);
            cmd.SetGraphicsRootConstantBufferView(first + 3, self.cluster_cbv);
            cmd.SetGraphicsRootShaderResourceView(first + 4, self.cluster_list);
        }
    }
}

// The D3D12 handles a shared post pass builds and encodes through.
pub(in crate::directx) struct DxPostDevice<'a> {
    pub device: &'a ID3D12Device,
    // Where a target's SRV and RTV come from.
    pub descriptors: &'a PostDescriptors,
    // The shader-visible heap a draw's root tables index, bound per pass.
    pub srv_heap: &'a ID3D12DescriptorHeap,
    // The debug-layer message queue, so a root-signature or PSO failure prints
    // what the layer said rather than just an HRESULT.
    pub info_queue: Option<&'a ID3D12InfoQueue>,
    // The probe set a probe-reading program binds. Absent at init, where no
    // draw is encoded.
    pub probes: Option<DxPostProbes>,
    pub hot_reload: bool,
}

// The DXIL 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>)> {
    Ok((
        builtin_shaders::FULLSCREEN_VERT.compile(hot_reload)?,
        program.program().compile(hot_reload)?,
    ))
}

// A root signature for `bindings`: one single-SRV table per source, a 32-bit
// constant block at b0 when the program declares constants, and for a
// probe-reading program the cube array table, a root CBV for the ProbeSet, a
// root SRV for the probe records at the register after the cube array, a root
// CBV for the cluster params and a root SRV for the cluster lists after that.
// Static samplers at s0..sN-1, plus sN for the cube array: every sampler kind
// the seam names resolves to the same trilinear clamp-to-edge state here, so the
// signature does not depend on which kind a draw picks for a slot.
fn create_root_signature(
    device: &ID3D12Device,
    bindings: PostProgramBindings,
) -> RenderResult<ID3D12RootSignature> {
    let textures = bindings.textures as u32;
    let mut sig = (0..textures).fold(RootSig::new(), |sig, reg| {
        sig.srv_table(reg, 1, Visibility::Pixel)
    });
    if bindings.constants > 0 {
        let dwords = bindings.constants.div_ceil(4) as u32;
        sig = sig.constant_dwords(0, dwords, Visibility::Pixel);
    }
    if bindings.probes {
        // The cube array takes the next texture register after the declared sources.
        sig = sig
            .srv_table(textures, 1, Visibility::Pixel)
            .cbv(PROBE_SET_REGISTER, Visibility::Pixel)
            .srv(textures + 1, Visibility::Pixel)
            .cbv(CLUSTER_PARAMS_REGISTER, Visibility::Pixel)
            .srv(textures + 2, Visibility::Pixel);
    }
    (0..textures + u32::from(bindings.probes))
        .fold(sig, |sig, reg| {
            sig.static_sampler(SamplerState::LinearClamp, reg, Visibility::Pixel)
        })
        .build(device, "post root sig")
}

impl PostPassDevice for DxPostDevice<'_> {
    type Recorder = ID3D12GraphicsCommandList;
    type Pipeline = PostPipeline;
    type Target = PostTarget;
    type TextureRef<'a> = SrvSlot;
    type Attachment<'a> = DxAttachment<'a>;

    fn create_pipeline(
        &self,
        program: PostProgram,
        format: PixelFormat,
        blend: PostBlend,
    ) -> RenderResult<Self::Pipeline> {
        let bindings = program.bindings();
        let root_sig = dump_on_err(
            self.info_queue,
            create_root_signature(self.device, bindings),
        )?;
        let (vs, ps) = compile(program, self.hot_reload)?;
        let pso = dump_on_err(
            self.info_queue,
            GraphicsPso::fullscreen(&root_sig, &vs, &ps, dxgi_format(format), blend.into())
                .build(self.device, program.label()),
        )?;
        Ok(PostPipeline {
            pso,
            root_sig,
            bindings,
        })
    }

    fn create_target(
        &self,
        label: &'static str,
        desc: &TextureDesc,
        extent: PostExtent,
    ) -> RenderResult<Self::Target> {
        let spec = resolved_texture(label, desc, extent);
        let format = dxgi_format(spec.format);
        let level_count = spec.mip_levels.max(1);
        let resource = create_rt_chain(self.device, spec.width, spec.height, format, level_count)?;
        let descriptors = self.descriptors.allocate()?;
        write_levels_srv(
            self.device,
            &resource,
            descriptors.srv_cpu,
            format,
            (0, level_count),
        );
        write_level_rtv(self.device, &resource, descriptors.rtv, format, 0);
        let mut levels = Vec::new();
        if level_count > 1 {
            for level in 0..level_count {
                let d = self.descriptors.allocate()?;
                write_levels_srv(self.device, &resource, d.srv_cpu, format, (level, 1));
                write_level_rtv(self.device, &resource, d.rtv, format, level);
                levels.push(d);
            }
        }
        Ok(PostTarget {
            resource,
            descriptors,
            label,
            levels,
            level_count,
            extent: FullscreenExtent {
                width: spec.width,
                height: spec.height,
            },
        })
    }

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

    fn target_attachment<'a>(&self, target: &'a Self::Target) -> Self::Attachment<'a> {
        DxAttachment {
            resource: &target.resource,
            // The top level is subresource 0 of a single-slice texture.
            subresource: 0,
            rtv: target.descriptors.rtv,
            extent: target.extent,
        }
    }

    fn target_level_ref<'a>(
        &self,
        target: &'a Self::Target,
        level: u32,
    ) -> RenderResult<Self::TextureRef<'a>> {
        Ok(target.level(level)?.srv_gpu)
    }

    fn target_level_attachment<'a>(
        &self,
        target: &'a Self::Target,
        level: u32,
    ) -> RenderResult<Self::Attachment<'a>> {
        let extent = level_extent(
            PostExtent {
                width: target.extent.width,
                height: target.extent.height,
            },
            level,
        );
        Ok(DxAttachment {
            resource: &target.resource,
            // A single-slice texture's subresources are its mip levels.
            subresource: level,
            rtv: target.level(level)?.rtv,
            extent: FullscreenExtent {
                width: extent.width,
                height: extent.height,
            },
        })
    }

    fn encode(&self, cmd: &Self::Recorder, draw: &PostDraw<'_, '_, Self>) -> RenderResult<()> {
        let pipe = draw.pipeline;
        draw.check(pipe.bindings)?;
        let probes = match (pipe.bindings.probes, self.probes) {
            (false, _) => None,
            (true, Some(probes)) => Some(probes),
            (true, None) => {
                return Err(RenderError::Other(format!(
                    "{}: the program reads the reflection-probe set, but this device holds none",
                    draw.label
                )));
            }
        };
        let target = draw.target;
        // A target the graph declares is already in RENDER_TARGET and its next
        // consumer's barrier takes it back out. One private to its pass rests
        // readable, so the draw moves it itself. Either way there is nothing to
        // clear: the fullscreen triangle covers every pixel, and a load keeps
        // what the target holds.
        let owns_state = draw.state == PostTargetState::Pass;
        debug_assert!(matches!(draw.load, PostLoadOp::DontCare | PostLoadOp::Load));
        // SAFETY: the command list is in the recording state, and every resource, descriptor and
        // slice these commands name is live for the call.
        unsafe {
            if owns_state {
                cmd.ResourceBarrier(&[subresource_transition_barrier(
                    target.resource,
                    target.subresource,
                    D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
                    D3D12_RESOURCE_STATE_RENDER_TARGET,
                )]);
            }
            cmd.OMSetRenderTargets(1, Some(&target.rtv), false, None);
            let vp = D3D12_VIEWPORT {
                TopLeftX: 0.0,
                TopLeftY: 0.0,
                Width: target.extent.width as f32,
                Height: target.extent.height as f32,
                MinDepth: 0.0,
                MaxDepth: 1.0,
            };
            cmd.RSSetViewports(&[vp]);
            cmd.RSSetScissorRects(&[windows::Win32::Foundation::RECT {
                left: 0,
                top: 0,
                right: target.extent.width as i32,
                bottom: target.extent.height as i32,
            }]);
            cmd.SetDescriptorHeaps(&[Some(self.srv_heap.clone())]);
            cmd.SetPipelineState(&pipe.pso);
            cmd.SetGraphicsRootSignature(&pipe.root_sig);
            // Every source binds through a static sampler the root signature
            // declared at its own slot, so a bind is just its table.
            for (slot, bind) in draw.binds.iter().enumerate() {
                cmd.set_graphics_srv_table(slot as u32, bind.texture);
            }
            if !draw.constants.is_empty() {
                cmd.set_graphics_root_constant_bytes(pipe.constants_parameter(), draw.constants);
            }
            if let Some(probes) = probes {
                probes.bind(cmd, pipe.probes_parameter());
            }
            cmd.IASetPrimitiveTopology(D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
            // The vertex stage builds the fullscreen triangle from SV_VertexID.
            cmd.IASetVertexBuffers(0, None);
            cmd.IASetIndexBuffer(None);
            cmd.DrawInstanced(3, 1, 0, 0);
            if owns_state {
                cmd.ResourceBarrier(&[subresource_transition_barrier(
                    target.resource,
                    target.subresource,
                    D3D12_RESOURCE_STATE_RENDER_TARGET,
                    D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
                )]);
            }
        }
        Ok(())
    }
}

impl DxPostDevice<'_> {
    // View `resource`, a single-level `format` texture of `extent` the transient
    // pool holds, through a slot of the post block.
    pub(in crate::directx) fn pooled_target(
        &self,
        resource: &ID3D12Resource,
        format: PixelFormat,
        extent: PostExtent,
    ) -> RenderResult<PooledTarget> {
        let format = dxgi_format(format);
        let descriptors = self.descriptors.allocate()?;
        write_levels_srv(self.device, resource, descriptors.srv_cpu, format, (0, 1));
        write_level_rtv(self.device, resource, descriptors.rtv, format, 0);
        Ok(PooledTarget {
            resource: resource.clone(),
            descriptors,
            extent: FullscreenExtent {
                width: extent.width,
                height: extent.height,
            },
        })
    }
}

impl crate::directx::context::DxContext {
    // The post-pass device over this context, binding frame slot `frame`'s
    // reflection-probe set.
    pub(in crate::directx) fn post_device(&self, frame: usize) -> DxPostDevice<'_> {
        DxPostDevice {
            device: &self.hw.device,
            descriptors: &self.post,
            srv_heap: &self.descriptors.srv_heap,
            info_queue: self.hw.info_queue.as_ref(),
            probes: Some(self.probe_bindings(frame)),
            hot_reload: self.hot_reload.enabled,
        }
    }

    // Frame slot `frame`'s live reflection-probe set and the main camera's
    // cluster grid binning it.
    pub(in crate::directx) fn probe_bindings(&self, frame: usize) -> DxPostProbes {
        DxPostProbes {
            cube_table: self.probe_cube_table_gpu(),
            set_cbv: com::gpu_va(&self.uniforms.probe_set_cbvs[frame]),
            records: self.probe.gpu.records[frame].gpu_va(),
            cluster_cbv: self.cluster_params_gva(frame, true),
            cluster_list: self.cluster_list_gva(),
        }
    }

    // The single-sample HDR scene the graph drives as `hdr_resolve`, as a post
    // draw's target.
    pub(in crate::directx) fn hdr_scene_attachment(&self) -> DxAttachment<'_> {
        DxAttachment {
            resource: self.hdr_scene_target(),
            subresource: D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES,
            rtv: self.hdr_scene_rtv(),
            extent: FullscreenExtent {
                width: self.targets.extent.render_width,
                height: self.targets.extent.render_height,
            },
        }
    }
}

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

    // Every post program's fragment, and the shared vertex, compile to DXIL, so a
    // source or register error fails a test instead of only surfacing as an
    // init failure on the GPU host. Skipped on a host without dxc.
    #[test]
    fn every_post_program_compiles() {
        concinnity_shader::require_dxc!();
        for program in PostProgram::ALL {
            compile(program, false).unwrap_or_else(|e| panic!("{program:?}: {e}"));
        }
    }
}