concinnity-device 0.19.119

//! The convolution half of a runtime reflection-probe bake on DirectX: the three
//! compute PSOs built from `probe_prefilter.hlsl`, their root signatures, the
//! capture cube one bake works from, and the dispatches that turn six captured
//! faces into the prefiltered radiance cube the specular term samples.
//! Mirrors `metal::probe_prefilter` and `vulkan::probe_prefilter`.
//!
//! The capture cube collects the six rendered faces (one array slice each) and
//! carries a mip chain the `probe_downsample` kernel fills; the probe cube is
//! the result, one cube of the probe cube array: mip 0 a firefly-clamped copy
//! of the capture and every mip after it a GGX convolution at that mip's
//! roughness. Both are R16G16B16A16_FLOAT: the faces are rendered as halfs, the
//! clamp caps luminance well inside the format's range, and it halves what a
//! probe costs against an R32G32B32A32 cube.
//!
//! Nothing reads back. The whole convolution stays on the direct queue, so the
//! frames that sample the finished cube are ordered after the dispatches that
//! wrote it by submission order alone.
//!
//! Resource states, which the barriers below are the whole of: the capture arrives
//! in COPY_DEST (the per-face copies write it), moves to UNORDERED_ACCESS for the
//! pyramid build, then to NON_PIXEL_SHADER_RESOURCE for the GGX dispatches that
//! sample it. The probe's cube of the array moves from PIXEL_SHADER_RESOURCE to
//! UNORDERED_ACCESS for the dispatches that write it and back after the last one,
//! one subresource at a time, so the cubes other frames sample never move.

use concinnity_core::render::error::{RenderError, RenderResult};
use concinnity_core::render::reflection_probe::PrefilterPlan;
use concinnity_core::render::uniforms::ProbePrefilterParams;
use windows::Win32::Graphics::Direct3D12::*;
use windows::Win32::Graphics::Dxgi::Common::*;

use super::builtin_shaders::CompileProgram;
use super::com;
use super::context::DxContext;
use super::probe_set::{CubeResource, create_cube_resource, write_cube_mip_uav};
use super::pso::compute_pso;
use crate::directx::descriptor_slot::DescriptorTables;
use crate::directx::descriptor_slot::SrvSlot;
use crate::directx::root_constants::RootConstants;
use crate::directx::root_sig::{RootSig, SamplerState, Visibility};

/// Color format of the capture and the probe cube array.
pub(in crate::directx) const PROBE_CUBE_FORMAT: DXGI_FORMAT = DXGI_FORMAT_R16G16B16A16_FLOAT;

/// Upper bound on a probe cube's mip count, sizing the descriptor block the SRV
/// heap reserves for one bake. `PrefilterPlan` stops at 4x4 faces, so 12 covers a
/// capture up to 8192 px on an edge.
pub(in crate::directx) const PROBE_MAX_MIPS: usize = 12;

// Threadgroup tile, matching the kernels' `[numthreads(8, 8, 1)]`. The third
// dispatch dimension is the six cube faces, one invocation deep.
const PREFILTER_TILE: u32 = 8;

/// The convolution PSOs and the two root signatures they bind. Built once at init
/// and reused by every bake.
pub(in crate::directx) struct ProbePrefilterPipelines {
    // The mirror-mip copy and the pyramid reduction bind the same shape (root
    // constants plus a two-descriptor UAV table), so they share a root signature;
    // the GGX kernel adds an SRV table and a static sampler.
    mip_root: ID3D12RootSignature,
    ggx_root: ID3D12RootSignature,
    mip0: ID3D12PipelineState,
    downsample: ID3D12PipelineState,
    ggx: ID3D12PipelineState,
}

/// Whether the device can load `R16G16B16A16_FLOAT` through an unordered-access
/// view. `probe_mip0` and `probe_downsample` both read their source mip as a
/// `RWTexture2DArray<float4>`, and D3D12 guarantees typed UAV loads for only four
/// formats without this optional capability. A device without it cannot run the
/// convolution at all, so the caller leaves the pipelines unbuilt and probes fall
/// back to the sky prefilter.
pub(in crate::directx) fn typed_uav_load_supported(device: &ID3D12Device) -> bool {
    let mut options = D3D12_FEATURE_DATA_D3D12_OPTIONS::default();
    // SAFETY: a query on a live COM object; the descriptor it reads and the out-parameters it fills
    // are live locals that outlive the call.
    let ok = unsafe {
        device.CheckFeatureSupport(
            D3D12_FEATURE_D3D12_OPTIONS,
            &mut options as *mut _ as *mut std::ffi::c_void,
            size_of::<D3D12_FEATURE_DATA_D3D12_OPTIONS>() as u32,
        )
    };
    ok.is_ok() && options.TypedUAVLoadAdditionalFormats.as_bool()
}

impl ProbePrefilterPipelines {
    pub(in crate::directx) fn new(device: &ID3D12Device, hot_reload: bool) -> RenderResult<Self> {
        use super::builtin_shaders;
        let mip_root = create_mip_root_signature(device)?;
        let ggx_root = create_ggx_root_signature(device)?;
        let mip0 = compute_pso(
            device,
            &mip_root,
            &builtin_shaders::PROBE_MIP0.compile(hot_reload)?,
            "probe_mip0",
        )?;
        let downsample = compute_pso(
            device,
            &mip_root,
            &builtin_shaders::PROBE_DOWNSAMPLE.compile(hot_reload)?,
            "probe_downsample",
        )?;
        let ggx = compute_pso(
            device,
            &ggx_root,
            &builtin_shaders::PROBE_GGX.compile(hot_reload)?,
            "probe_ggx",
        )?;
        Ok(Self {
            mip_root,
            ggx_root,
            mip0,
            downsample,
            ggx,
        })
    }
}

/// One bake's capture cube and the probe-array cube it convolves into. Their
/// descriptors live in the SRV heap's reserved probe-prefilter block, which [`PrefilterGpu::new`]
/// rewrites for each bake. One block for every bake is what serializes them on
/// this backend: a capture starts only once the prefiltering slot is empty, and
/// the install that empties it is gated on the fence covering the prior bake's
/// last dispatch, so nothing in flight still binds the block when it is
/// rewritten.
pub(in crate::directx) struct PrefilterGpu {
    capture: ID3D12Resource,
    // The cube of the probe cube array the convolution writes.
    cube: usize,
    mips: u32,
}

impl PrefilterGpu {
    /// Allocate the capture and write every descriptor the bake's dispatches
    /// bind, writing into cube `cube` of `cubes`. The capture starts in COPY_DEST
    /// so the per-face copies can write it.
    pub(in crate::directx) fn new(
        ctx: &DxContext,
        plan: &PrefilterPlan,
        cubes: &super::probe_set::ProbeCubeArray,
        cube: usize,
    ) -> RenderResult<PrefilterGpu> {
        let mips = plan.mips();
        if mips as usize > PROBE_MAX_MIPS {
            return Err(RenderError::Other(format!(
                "probe: {mips} mips exceeds the {PROBE_MAX_MIPS} descriptors reserved for one bake"
            )));
        }
        if cube >= cubes.capacity() || cubes.mips() != mips {
            return Err(RenderError::Other(format!(
                "probe: the cube array has no cube {cube} at {mips} mips"
            )));
        }
        let capture = create_cube_resource(
            &ctx.hw.device,
            CubeResource {
                face_size: plan.face_size(),
                mips,
                cubes: 1,
                state: D3D12_RESOURCE_STATE_COPY_DEST,
                label: "probe capture cube",
            },
        )?;

        let device = &ctx.hw.device;
        let d = &ctx.descriptors;
        write_cube_srv(device, &capture, mips, ctx.probe_capture_srv_cpu());
        for mip in 0..mips {
            write_cube_mip_uav(
                device,
                &capture,
                0,
                mip,
                cpu_slot(ctx, d.layout.probe_capture_uav_base_slot + mip as usize),
            );
            cubes.write_mip_uav(
                device,
                cube,
                mip,
                cpu_slot(ctx, d.layout.probe_cube_uav_base_slot + mip as usize),
            );
        }
        // The mirror-mip copy binds capture mip 0 and probe mip 0 as one contiguous
        // pair, which neither per-cube block can supply, so it gets its own.
        write_cube_mip_uav(
            device,
            &capture,
            0,
            0,
            cpu_slot(ctx, d.layout.probe_mip0_pair_slot),
        );
        cubes.write_mip_uav(
            device,
            cube,
            0,
            cpu_slot(ctx, d.layout.probe_mip0_pair_slot + 1),
        );

        Ok(PrefilterGpu {
            capture,
            cube,
            mips,
        })
    }

    /// The capture resource the six face copies write into.
    pub(in crate::directx) fn capture(&self) -> &ID3D12Resource {
        &self.capture
    }

    /// The cube of the probe cube array being written.
    pub(in crate::directx) fn cube(&self) -> usize {
        self.cube
    }

    /// Mip levels the capture carries, which the face copies address it by.
    pub(in crate::directx) fn mips(&self) -> u32 {
        self.mips
    }
}

impl DxContext {
    /// CPU handle of the capture cube's all-mips SRV slot.
    pub(in crate::directx) fn probe_capture_srv_cpu(&self) -> D3D12_CPU_DESCRIPTOR_HANDLE {
        cpu_slot(self, self.descriptors.layout.probe_capture_srv_slot)
    }

    /// Record the cheap half of the convolution: the capture moves from the
    /// per-face copies' COPY_DEST into UNORDERED_ACCESS and the probe's cube of
    /// the array into UNORDERED_ACCESS too, the mirror mip is copied through with
    /// the firefly clamp, the source pyramid is reduced level by level, and the
    /// capture ends in NON_PIXEL_SHADER_RESOURCE for the GGX dispatches that
    /// follow.
    ///
    /// All of it goes in one command list: the reductions are a few taps per texel
    /// and each depends on the one before, so spreading them over frames would only
    /// lengthen the bake.
    pub(in crate::directx) fn encode_probe_pyramid(
        &self,
        cmd: &ID3D12GraphicsCommandList,
        gpu: &PrefilterGpu,
        plan: &PrefilterPlan,
    ) -> RenderResult<()> {
        let pipelines =
            self.probe.prefilter.as_ref().ok_or_else(|| {
                RenderError::Other("probe: prefilter pipelines missing".to_string())
            })?;
        let mut barriers = vec![super::texture::transition_barrier(
            gpu.capture(),
            D3D12_RESOURCE_STATE_COPY_DEST,
            D3D12_RESOURCE_STATE_UNORDERED_ACCESS,
        )];
        barriers.extend(self.probe.gpu.bound_cubes().cube_barriers(
            gpu.cube(),
            super::probe_set::PROBE_CUBES_STATE,
            D3D12_RESOURCE_STATE_UNORDERED_ACCESS,
        ));
        // SAFETY: the command list is in the recording state, and every resource, descriptor and
        // slice these commands name is live for the call.
        unsafe {
            cmd.ResourceBarrier(&barriers);
            cmd.SetComputeRootSignature(&pipelines.mip_root);
        }
        let d = &self.descriptors;
        self.dispatch_prefilter(
            cmd,
            &pipelines.mip0,
            gpu_slot(self, d.layout.probe_mip0_pair_slot),
            None,
            &plan.mip0_params(),
            plan.face_size(),
        );
        for mip in 1..plan.mips() {
            // Each level reads the one the previous dispatch wrote.
            uav_barrier(cmd, gpu.capture());
            self.dispatch_prefilter(
                cmd,
                &pipelines.downsample,
                gpu_slot(
                    self,
                    d.layout.probe_capture_uav_base_slot + (mip - 1) as usize,
                ),
                None,
                &plan.downsample_params(mip),
                plan.mip_face_size(mip),
            );
        }
        // SAFETY: as above.
        unsafe {
            cmd.ResourceBarrier(&[super::texture::transition_barrier(
                gpu.capture(),
                D3D12_RESOURCE_STATE_UNORDERED_ACCESS,
                D3D12_RESOURCE_STATE_NON_PIXEL_SHADER_RESOURCE,
            )]);
        }
        Ok(())
    }

    /// Record the GGX convolution producing probe-cube mip `dst_mip`, sampling the
    /// finished pyramid. Nothing else writes that mip, and the capture is read-only
    /// from here on, so consecutive mips need no barrier between them.
    pub(in crate::directx) fn encode_probe_ggx_mip(
        &self,
        cmd: &ID3D12GraphicsCommandList,
        plan: &PrefilterPlan,
        dst_mip: u32,
    ) -> RenderResult<()> {
        let pipelines =
            self.probe.prefilter.as_ref().ok_or_else(|| {
                RenderError::Other("probe: prefilter pipelines missing".to_string())
            })?;
        // SAFETY: the command list is in the recording state and the root signature is live.
        unsafe { cmd.SetComputeRootSignature(&pipelines.ggx_root) };
        let d = &self.descriptors;
        self.dispatch_prefilter(
            cmd,
            &pipelines.ggx,
            gpu_slot(self, d.layout.probe_cube_uav_base_slot + dst_mip as usize),
            Some(gpu_slot(self, d.layout.probe_capture_srv_slot)),
            &plan.ggx_params(dst_mip),
            plan.mip_face_size(dst_mip),
        );
        Ok(())
    }

    // Bind, push and dispatch one prefilter kernel over a `size`-square cube face,
    // six faces deep. The kernels bounds-guard against `dst_size`, so the
    // rounded-up remainder returns early. `srv_table` is the GGX kernel's sampled
    // capture; the mip kernels bind none.
    fn dispatch_prefilter(
        &self,
        cmd: &ID3D12GraphicsCommandList,
        pso: &ID3D12PipelineState,
        uav_table: SrvSlot,
        srv_table: Option<SrvSlot>,
        params: &ProbePrefilterParams,
        size: u32,
    ) {
        let groups = size.div_ceil(PREFILTER_TILE).max(1);
        // SAFETY: the command list is in the recording state, and every resource, descriptor and
        // slice these commands name is live for the call; the root-constant count matches the
        // signature's, both derived from `ProbePrefilterParams`.
        unsafe {
            cmd.SetPipelineState(pso);
            cmd.set_compute_root_constants(0, params);
            match srv_table {
                Some(srv) => {
                    cmd.set_compute_srv_table(1, srv);
                    cmd.set_compute_srv_table(2, uav_table);
                }
                None => cmd.set_compute_srv_table(1, uav_table),
            }
            cmd.Dispatch(groups, groups, 6);
        }
    }
}

// Root signature for the mirror-copy and downsample kernels: root constants at b0
// plus one two-descriptor UAV table (u0 the source mip, u1 the destination). The
// per-mip UAVs are contiguous in the heap, which is what lets one range cover the
// pair.
fn create_mip_root_signature(device: &ID3D12Device) -> RenderResult<ID3D12RootSignature> {
    RootSig::new()
        .constants::<ProbePrefilterParams>(0, Visibility::All)
        .uav_table(0, 2, Visibility::All) // u0..u1
        .build(device, "probe prefilter mip root sig")
}

// Root signature for the GGX kernel: root constants at b0, the sampled capture
// pyramid at t0, the destination mip at u0, and the linear-clamp mipmapped
// sampler at s0 as a static sampler (a shader sampler needs no heap of its own
// when it never varies). The solid-angle lod the kernel computes is fractional,
// so the trilinear filter is what makes the level selection continuous.
fn create_ggx_root_signature(device: &ID3D12Device) -> RenderResult<ID3D12RootSignature> {
    use Visibility::All;
    RootSig::new()
        .constants::<ProbePrefilterParams>(0, All)
        .srv_table(0, 1, All) // t0
        .uav_table(0, 1, All) // u0
        .static_sampler(SamplerState::LinearClamp, 0, All)
        .build(device, "probe prefilter ggx root sig")
}

// All-mips TEXTURECUBE SRV, the shape a sampler reads.
fn write_cube_srv(
    device: &ID3D12Device,
    resource: &ID3D12Resource,
    mips: u32,
    srv_cpu: D3D12_CPU_DESCRIPTOR_HANDLE,
) {
    let desc = D3D12_SHADER_RESOURCE_VIEW_DESC {
        Format: PROBE_CUBE_FORMAT,
        ViewDimension: D3D12_SRV_DIMENSION_TEXTURECUBE,
        Shader4ComponentMapping: D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING,
        Anonymous: D3D12_SHADER_RESOURCE_VIEW_DESC_0 {
            TextureCube: D3D12_TEXCUBE_SRV {
                MostDetailedMip: 0,
                MipLevels: mips,
                ResourceMinLODClamp: 0.0,
            },
        },
    };
    // SAFETY: the view descriptor and the resource it names are live for the call, and the
    // destination handle addresses a slot this context reserved for the view in a heap it owns.
    unsafe { device.CreateShaderResourceView(resource, Some(&desc), srv_cpu) };
}

// Order one dispatch's writes to `resource` before the next dispatch's reads.
fn uav_barrier(cmd: &ID3D12GraphicsCommandList, resource: &ID3D12Resource) {
    let barrier = D3D12_RESOURCE_BARRIER {
        Type: D3D12_RESOURCE_BARRIER_TYPE_UAV,
        Flags: D3D12_RESOURCE_BARRIER_FLAG_NONE,
        Anonymous: D3D12_RESOURCE_BARRIER_0 {
            UAV: std::mem::ManuallyDrop::new(D3D12_RESOURCE_UAV_BARRIER {
                pResource: com::borrowed(resource),
            }),
        },
    };
    // SAFETY: the command list is in the recording state, and the barrier and the resource it
    // borrows are live for the call.
    unsafe { cmd.ResourceBarrier(&[barrier]) };
}

fn cpu_slot(ctx: &DxContext, slot: usize) -> D3D12_CPU_DESCRIPTOR_HANDLE {
    // SAFETY: a property query on a live descriptor heap; it only reads.
    let base = unsafe {
        ctx.descriptors
            .srv_heap
            .GetCPUDescriptorHandleForHeapStart()
    };
    D3D12_CPU_DESCRIPTOR_HANDLE {
        ptr: base.ptr + slot * ctx.descriptors.srv_descriptor_size,
    }
}

fn gpu_slot(ctx: &DxContext, slot: usize) -> SrvSlot {
    SrvSlot::at(
        &ctx.descriptors.srv_heap,
        ctx.descriptors.srv_descriptor_size,
        slot,
    )
}