concinnity-device 0.19.119

//! Auto-exposure (EV adaptation) on D3D12: a per-frame CPU readback of the
//! previous frame's average log-luminance, an EMA step that updates the adapted
//! EV, and the histogram build + average compute dispatches that produce next
//! frame's average. The compute passes are encoded after the main HDR resolve
//! (where `hdr_srv_gpu` carries this frame's scene color) and read CPU-side at
//! the top of a later frame, so there is `FRAMES - 1` frames of latency between
//! the scene's actual luminance and the exposure applied to it, invisible at
//! human-scale eye-adaptation rates. Mirrors `metal/auto_exposure.rs`.

use concinnity_core::gfx::auto_exposure;
use concinnity_core::gfx::auto_exposure::HISTOGRAM_BINS;
use concinnity_core::render::error::RenderResult;
use concinnity_core::render::uniforms::AutoExposureParams;
use windows::Win32::Graphics::Direct3D12::*;

use super::allocator::{DeviceAllocator, PooledBuffer};
use super::com;
use crate::directx::builtin_shaders;
use crate::directx::builtin_shaders::CompileProgram;
use crate::directx::context::{DxContext, FRAMES};
use crate::directx::descriptor_slot::DescriptorTables;
use crate::directx::error::map_hresult;
use crate::directx::pso::compute_pso;
use crate::directx::root_constants::RootConstants;
use crate::directx::root_sig::{RootSig, Visibility};
use crate::directx::texture::{create_uav_buffer, transition_barrier, uav_barrier};

// Auto-exposure (EV adaptation) state. `resources` is `Some` only when the
// world's `PostProcessConfig` opts in; it holds the histogram + average compute
// PSOs, the histogram UAV, the output UAV, and the per-frame readback buffers.
// `adaptation` carries the clamped tunables, the authored EV bias and the EMA
// target; `last_elapsed` is the previous frame's elapsed time used to derive
// `dt` for the EMA.
pub(in crate::directx) struct AutoExposureState {
    pub resources: Option<AutoExposureResources>,
    pub adaptation: Option<auto_exposure::ExposureAdaptation>,
    pub last_elapsed: f32,
}

// Compile the auto-exposure `build` + `average` compute kernels. Used at
// init and by shader hot-reload to rebuild the two compute PSOs.
pub(in crate::directx) fn compile_auto_exposure_shaders(
    hot_reload: bool,
) -> RenderResult<(Vec<u8>, Vec<u8>)> {
    let build_cs = builtin_shaders::AUTO_EXPOSURE_BUILD.compile(hot_reload)?;
    let average_cs = builtin_shaders::AUTO_EXPOSURE_AVERAGE.compile(hot_reload)?;
    Ok((build_cs, average_cs))
}

// Pair of compute pipelines + GPU buffers + per-frame readback driving the
// auto-exposure histogram path. Built only when the world's
// `PostProcessConfig` opts in; the encoder is a no-op otherwise.
pub(super) struct AutoExposureResources {
    // Build kernel: one thread per HDR-resolve pixel; produces the 256-bin
    // log-luminance histogram in the global UAV.
    build_pso: ID3D12PipelineState,
    build_root_sig: ID3D12RootSignature,
    // Average kernel: one threadgroup of `HISTOGRAM_BINS` threads that reduces
    // the histogram, clears it, and writes the weighted-average log-luminance
    // to the output UAV.
    average_pso: ID3D12PipelineState,
    average_root_sig: ID3D12RootSignature,

    // 256-bin u32 histogram (UNORDERED_ACCESS, DEFAULT heap). The build kernel
    // `InterlockedAdd`s into it; the average kernel reads and clears each bin.
    // Held only to keep the resource alive; the GVA is read directly through
    // `histogram.GetGPUVirtualAddress()` at encode time.
    histogram: ID3D12Resource,
    // Single f32 (UNORDERED_ACCESS, DEFAULT heap) the average kernel writes
    // the weighted-average log-luminance into.
    output_buf: ID3D12Resource,
    // Per-frame-slot READBACK buffer (4 bytes). Persistently mapped (READBACK
    // resources allow leaving Map active across submissions). The end of each
    // frame's command list copies `output_buf` into the matching slot; at the
    // top of a later frame (after the fence wait gates this slot's previous
    // use) the CPU reads its pointer for the EMA update.
    readback_bufs: Vec<PooledBuffer>,
    readback_ptrs: Vec<*const f32>,
}

impl AutoExposureResources {
    // Root signature for the build kernel. Exposed so the DirectX shader
    // hot-reload pass can rebuild the `build_pso` against the same root sig.
    pub(in crate::directx) fn build_root_sig(&self) -> &ID3D12RootSignature {
        &self.build_root_sig
    }
    // Root signature for the average kernel. Same purpose as
    // [`Self::build_root_sig`].
    pub(in crate::directx) fn average_root_sig(&self) -> &ID3D12RootSignature {
        &self.average_root_sig
    }
    // Swap the freshly-built build + average PSOs into the live resources.
    // Driven by the DirectX shader hot-reload pass after every replacement
    // successfully compiled.
    pub(in crate::directx) fn swap_pipelines(
        &mut self,
        build_pso: ID3D12PipelineState,
        average_pso: ID3D12PipelineState,
    ) {
        self.build_pso = build_pso;
        self.average_pso = average_pso;
    }

    // Build all auto-exposure resources. Called from `DxContext::new` only
    // when `PostProcessConfig.auto_exposure` is enabled.
    pub(super) fn new(alloc: &DeviceAllocator, hot_reload: bool) -> RenderResult<Self> {
        let device = alloc.device();
        let (build_cs, average_cs) = compile_auto_exposure_shaders(hot_reload)?;

        let build_root_sig = create_build_root_signature(device)?;
        let build_pso = compute_pso(device, &build_root_sig, &build_cs, "auto-exposure build")?;

        let average_root_sig = create_average_root_signature(device)?;
        let average_pso = compute_pso(
            device,
            &average_root_sig,
            &average_cs,
            "auto-exposure average",
        )?;

        // Histogram: 256 * 4 bytes UAV, cleared by the average kernel each
        // frame. Created in COMMON; a one-shot transition before the first
        // dispatch flips it into UNORDERED_ACCESS for steady state.
        let histogram = create_uav_buffer(
            device,
            (HISTOGRAM_BINS * std::mem::size_of::<u32>()) as u64,
            D3D12_RESOURCE_STATE_COMMON,
        )?;

        // Output: a single f32 the average kernel writes into.
        let output_buf = create_uav_buffer(
            device,
            std::mem::size_of::<f32>() as u64,
            D3D12_RESOURCE_STATE_COMMON,
        )?;

        // Per-frame readback buffers. READBACK heap resources start in
        // COPY_DEST and never need a barrier.
        let mut readback_bufs: Vec<PooledBuffer> = Vec::with_capacity(FRAMES);
        let mut readback_ptrs: Vec<*const f32> = Vec::with_capacity(FRAMES);
        for _ in 0..FRAMES {
            let buf = alloc.alloc_buffer(
                std::mem::size_of::<f32>() as u64,
                D3D12_HEAP_TYPE_READBACK,
                D3D12_RESOURCE_STATE_COPY_DEST,
            )?;
            let mut ptr = std::ptr::null_mut::<std::ffi::c_void>();
            // READBACK heaps allow leaving Map active across submissions; the
            // pointer stays valid for the resource's lifetime.
            // SAFETY: the resource is a live CPU-visible buffer, and the out-parameter is a live
            // local that receives the mapping.
            unsafe { buf.Map(0, None, Some(&mut ptr)) }
                .map_err(|e| map_hresult(e.code(), "auto-exposure readback map"))?;
            readback_ptrs.push(ptr as *const f32);
            readback_bufs.push(buf);
        }

        Ok(Self {
            build_pso,
            build_root_sig,
            average_pso,
            average_root_sig,
            histogram,
            output_buf,
            readback_bufs,
            readback_ptrs,
        })
    }
}

// Root signature for the build kernel: 4 root constants (b0, AutoExposureParams),
// a single-SRV descriptor table for the HDR texture (t0), and a root UAV for
// the histogram (u0). The HDR SRV needs a descriptor table because root SRVs
// are limited to raw / structured buffers, not Texture2D.
fn create_build_root_signature(device: &ID3D12Device) -> RenderResult<ID3D12RootSignature> {
    RootSig::new()
        // [0] Root constants b0: AutoExposureParams.
        .constants::<AutoExposureParams>(0, Visibility::All)
        // [1] Descriptor table SRV t0: HDR texture.
        .srv_table(0, 1, Visibility::All)
        // [2] Root UAV u0: histogram.
        .uav(0, Visibility::All)
        .build(device, "auto-exposure build root sig")
}

// Root signature for the average kernel: 4 root constants (b0), root UAV for
// the histogram (u0, read + clear), root UAV for the output (u1, write-once).
fn create_average_root_signature(device: &ID3D12Device) -> RenderResult<ID3D12RootSignature> {
    RootSig::new()
        .constants::<AutoExposureParams>(0, Visibility::All)
        .uav(0, Visibility::All)
        .uav(1, Visibility::All)
        .build(device, "auto-exposure average root sig")
}

impl DxContext {
    // Step the auto-exposure EMA from a previous frame's GPU measurement, then
    // push the new exposure multiplier into `self.post_process.exposure`.
    // A no-op when auto-exposure is disabled: the static authored EV then
    // drives `exposure` unchanged.
    //
    // Called at the top of `draw_frame` after the fence wait for this slot's
    // previous use completes, so the matching readback buffer holds a fully
    // committed GPU result (one or two frames stale, smoothed by the EMA).
    // `elapsed` is the total elapsed seconds since startup; the per-call diff
    // drives `dt` for the EMA.
    pub(super) fn update_auto_exposure(&mut self, elapsed: f32, frame_idx: usize) {
        let Some(adaptation) = self.auto_exposure.adaptation.as_mut() else {
            return;
        };
        let Some(resources) = self.auto_exposure.resources.as_ref() else {
            return;
        };
        let Some(&ptr) = resources.readback_ptrs.get(frame_idx) else {
            return;
        };

        // Read the previous frame's average log-luminance for this slot. The
        // fence wait above this call already gated the GPU work that wrote it,
        // so the READBACK-heap mapping reflects the committed value.
        // SAFETY: `ptr` is this frame slot's persistent mapping of the auto-exposure READBACK
        // buffer, which holds one `f32`, and the fence wait ahead of this call retired the compute
        // pass that wrote it.
        let avg_log_lum = unsafe { ptr.read() };

        let dt = (elapsed - self.auto_exposure.last_elapsed).max(0.0);
        self.auto_exposure.last_elapsed = elapsed;

        // `self.post_process.exposure` is the linear multiplier the bloom
        // prefilter and composite consume; it already folds in the authored
        // exposure_ev when auto-exposure is off, so we only overwrite it here
        // when the GPU path owns the value.
        self.post_process.exposure = adaptation.step(avg_log_lum, dt);
    }

    // Resource the histogram_build kernel samples through `hdr_srv_gpu`: the
    // resolved single-sample HDR scene with MSAA on, otherwise the raw
    // (single-sample) `hdr_color`. The auto-exposure measurement runs between
    // the main HDR resolve and any post pass that mutates the scene (decals,
    // fog, SSR, TAA, bloom, composite), so it samples the pre-post scene.
    fn auto_exposure_source(&self) -> &ID3D12Resource {
        self.targets
            .hdr
            .resolve
            .as_ref()
            .unwrap_or(&self.targets.hdr.color)
    }

    // Encode the auto-exposure histogram passes against the resolved HDR
    // scene. The build kernel runs one thread per HDR pixel; the average
    // kernel runs one threadgroup of 256 threads that reduces the histogram,
    // clears it for the next frame, and writes the average log-luminance to
    // the output UAV. The end of the encoder copies the output to this slot's
    // readback buffer for the CPU's EMA step at the top of a later frame.
    // A no-op when auto-exposure is disabled.
    pub(super) fn encode_auto_exposure(&self, cmd: &ID3D12GraphicsCommandList, frame_idx: usize) {
        let Some(resources) = self.auto_exposure.resources.as_ref() else {
            return;
        };

        let params = AutoExposureParams::HISTOGRAM;
        let source = self.auto_exposure_source();

        // The build kernel needs the HDR scene readable in a compute (i.e.
        // NON_PIXEL_SHADER_RESOURCE) stage. After encode_main_pass the
        // resolved scene rests in PIXEL_SHADER_RESOURCE; flip it to the
        // compute-readable state for the dispatch and back so the downstream
        // decal / fog / SSR / TAA / bloom / composite passes find it where
        // they expect.
        let to_compute = transition_barrier(
            source,
            D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
            D3D12_RESOURCE_STATE_NON_PIXEL_SHADER_RESOURCE,
        );
        // 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(&[to_compute]) };

        let histogram_gva = com::gpu_va(&resources.histogram);
        let output_gva = com::gpu_va(&resources.output_buf);

        // Build dispatch: 16×16 threadgroups, one thread per HDR pixel.
        // 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.SetComputeRootSignature(&resources.build_root_sig);
            cmd.SetPipelineState(&resources.build_pso);
            cmd.SetDescriptorHeaps(&[Some(self.descriptors.srv_heap.clone())]);
            cmd.set_compute_root_constants(0, &params);
            cmd.set_compute_srv_table(1, self.targets.hdr.srv_gpu);
            cmd.SetComputeRootUnorderedAccessView(2, histogram_gva);

            let groups_x = self.targets.extent.render_width.div_ceil(16);
            let groups_y = self.targets.extent.render_height.div_ceil(16);
            cmd.Dispatch(groups_x, groups_y, 1);
        }

        // UAV barrier so the average dispatch sees the build kernel's writes.
        let barrier = uav_barrier(&resources.histogram);
        // 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(&[barrier]) };

        // Average dispatch: one threadgroup of 256 threads.
        // 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.SetComputeRootSignature(&resources.average_root_sig);
            cmd.SetPipelineState(&resources.average_pso);
            cmd.set_compute_root_constants(0, &params);
            cmd.SetComputeRootUnorderedAccessView(1, histogram_gva);
            cmd.SetComputeRootUnorderedAccessView(2, output_gva);
            cmd.Dispatch(1, 1, 1);
        }

        // Copy the freshly-written average to this slot's readback buffer. A
        // later frame using the same slot reads it from the matching
        // `readback_ptrs[frame_idx]` after the fence wait gates the copy. The
        // transition out of UNORDERED_ACCESS is what makes the average kernel's
        // write visible to the copy, so no UAV barrier precedes it.
        let to_copy_src = transition_barrier(
            &resources.output_buf,
            D3D12_RESOURCE_STATE_UNORDERED_ACCESS,
            D3D12_RESOURCE_STATE_COPY_SOURCE,
        );
        // 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(&[to_copy_src]) };
        if let Some(readback) = resources.readback_bufs.get(frame_idx) {
            // 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.CopyBufferRegion(
                    &**readback,
                    0,
                    &resources.output_buf,
                    0,
                    std::mem::size_of::<f32>() as u64,
                );
            }
        }
        // Give the output buffer back to the next frame's kernels, and restore
        // the HDR source to PIXEL_SHADER_RESOURCE for the post stack. Different
        // resources with nothing between them, so they ride one call.
        let to_uav = transition_barrier(
            &resources.output_buf,
            D3D12_RESOURCE_STATE_COPY_SOURCE,
            D3D12_RESOURCE_STATE_UNORDERED_ACCESS,
        );
        let back_to_psr = transition_barrier(
            source,
            D3D12_RESOURCE_STATE_NON_PIXEL_SHADER_RESOURCE,
            D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
        );
        // 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(&[to_uav, back_to_psr]) };
    }
}