concinnity-device 0.19.119

//! The convolution half of a runtime reflection-probe bake: the compute
//! pipelines built from `probe_prefilter.hlsl`, the two cubes one bake works
//! between, and the dispatches that turn six captured faces into the prefiltered
//! radiance mip chain the specular term samples.
//!
//! The capture cube is the render target the six faces resolve into, one cube
//! slice each, with 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, every mip after it a GGX convolution at
//! that mip's roughness. Both are RGBA16Float --
//! the faces are captured as halfs, the clamp caps luminance well inside the
//! format's range, and it halves what a probe costs in memory against the
//! RGBA32Float cube the CPU convolution used to upload.
//!
//! Nothing here reads back. A dispatch per destination mip, one mip per frame,
//! is what replaced the readback plus the off-thread CPU convolution; the whole
//! bake now stays on the GPU timeline.
#![deny(unsafe_op_in_unsafe_fn)]

use concinnity_core::render::error::{RenderError, RenderResult};
use concinnity_core::render::reflection_probe::PrefilterPlan;
use objc2::rc::Retained;
use objc2::runtime::ProtocolObject;
use objc2_foundation::NSRange;
use objc2_foundation::ns_string;
use objc2_metal::{
    MTLCommandBuffer as _, MTLComputeCommandEncoder as _, MTLComputePipelineState, MTLDevice,
    MTLPixelFormat, MTLSize, MTLTexture, MTLTextureType, MTLTextureUsage,
};

use super::builtin_shaders::compute_pipeline;
use super::descriptors::TextureDesc;
use super::encode::ComputeEncode;
use super::error::allocation_failed;

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

// Color format of the capture and the probe cube array. RGBA16Float is what the
// faces resolve as, and what the read_write views the kernels bind require (an
// Apple7 device and later reads and writes it; the engine's Metal floor is
// Apple7).
pub(in crate::metal) const PROBE_CUBE_FORMAT: MTLPixelFormat = MTLPixelFormat::RGBA16Float;

/// The three compute pipelines a probe bake convolves with, built once from the
/// precompiled `probe_prefilter.hlsl` variants.
pub(in crate::metal) struct ProbePrefilterPipelines {
    mip0: Retained<ProtocolObject<dyn MTLComputePipelineState>>,
    downsample: Retained<ProtocolObject<dyn MTLComputePipelineState>>,
    ggx: Retained<ProtocolObject<dyn MTLComputePipelineState>>,
}

impl ProbePrefilterPipelines {
    // Build all three kernels. Called from init under the same gate the bake
    // itself needs (the bindless cull pipeline), so a probe never discovers a
    // missing pipeline mid-capture.
    pub(in crate::metal) fn new(
        device: &ProtocolObject<dyn MTLDevice>,
        hot_reload: bool,
    ) -> RenderResult<ProbePrefilterPipelines> {
        Ok(ProbePrefilterPipelines {
            mip0: compute_pipeline(device, &super::builtin_shaders::PROBE_MIP0, hot_reload)?,
            downsample: compute_pipeline(
                device,
                &super::builtin_shaders::PROBE_DOWNSAMPLE,
                hot_reload,
            )?,
            ggx: compute_pipeline(device, &super::builtin_shaders::PROBE_GGX, hot_reload)?,
        })
    }
}

/// The capture cube six faces render into: RGBA16Float, one slice per face, with
/// the mip chain the convolution's source pyramid occupies.
///
/// Created unpooled, like the bake's other render targets: it is transient, and
/// Metal keeps a resource alive while a command buffer references it, so it can
/// be dropped the moment the bake ends without waiting on a fence.
pub(in crate::metal) fn create_capture_cube(
    device: &ProtocolObject<dyn MTLDevice>,
    plan: &PrefilterPlan,
) -> RenderResult<Retained<ProtocolObject<dyn MTLTexture>>> {
    let desc = TextureDesc {
        kind: MTLTextureType::TypeCube,
        format: PROBE_CUBE_FORMAT,
        width: plan.face_size() as usize,
        height: plan.face_size() as usize,
        mip_count: plan.mips() as usize,
        usage: MTLTextureUsage(
            MTLTextureUsage::RenderTarget.0
                | MTLTextureUsage::ShaderRead.0
                | MTLTextureUsage::ShaderWrite.0,
        ),
        ..Default::default()
    }
    .build();
    device
        .newTextureWithDescriptor(&desc)
        .ok_or_else(|| allocation_failed("probe capture cube"))
}

/// The capture and the per-mip views one convolution works between, held by
/// the prefiltering bake slot until the probe's cube is installed.
pub(in crate::metal) struct PrefilterGpu {
    // The capture, sampled whole (all mips) by the GGX kernel.
    capture: Retained<ProtocolObject<dyn MTLTexture>>,
    // One single-level 2D-array view of the capture per mip: mip M is the
    // downsample's source and mip M+1 its destination, so no dispatch reads the
    // texels it writes.
    capture_mip_views: Vec<Retained<ProtocolObject<dyn MTLTexture>>>,
    // One single-level 2D-array view per mip of this probe's six slices of the
    // probe cube array, the destination of the mip-0 copy and of each GGX
    // dispatch. A view keeps its parent alive, so a bake parked behind the
    // fence keeps a replaced array alive with it.
    probe_mip_views: Vec<Retained<ProtocolObject<dyn MTLTexture>>>,
}

impl PrefilterGpu {
    /// Take ownership of a finished `capture` and make the per-mip write views
    /// of it and of cube `slot` of the probe cube array `cubes`.
    pub(in crate::metal) fn new(
        capture: Retained<ProtocolObject<dyn MTLTexture>>,
        cubes: &ProtocolObject<dyn MTLTexture>,
        slot: usize,
        plan: &PrefilterPlan,
    ) -> RenderResult<PrefilterGpu> {
        let capture_mip_views = mip_array_views(&capture, 0, plan.mips(), "capture")?;
        let probe_mip_views = mip_array_views(cubes, slot, plan.mips(), "probe")?;
        Ok(PrefilterGpu {
            capture,
            capture_mip_views,
            probe_mip_views,
        })
    }
}

// One single-level 2D-array view per mip of cube `cube` of a cube or cube-array
// texture. A cube is six slices, so the view is what lets a kernel address
// (x, y, face) directly; the format is the parent's, so no reinterpretation
// occurs.
fn mip_array_views(
    texture: &ProtocolObject<dyn MTLTexture>,
    cube: usize,
    mips: u32,
    label: &str,
) -> RenderResult<Vec<Retained<ProtocolObject<dyn MTLTexture>>>> {
    let slices = texture.arrayLength() * 6;
    if (cube + 1) * 6 > slices || mips as usize > texture.mipmapLevelCount() {
        return Err(RenderError::Other(format!(
            "probe: {label} has no cube {cube} at {mips} mips"
        )));
    }
    (0..mips)
        .map(|mip| {
            // SAFETY: `mip` is below the texture's level count and cube `cube`'s
            // six slices are within it, both checked above; the view shares the
            // parent's pixel format, so it reinterprets nothing.
            unsafe {
                texture.newTextureViewWithPixelFormat_textureType_levels_slices(
                    PROBE_CUBE_FORMAT,
                    MTLTextureType::Type2DArray,
                    NSRange::new(mip as usize, 1),
                    NSRange::new(cube * 6, 6),
                )
            }
            .ok_or_else(|| {
                RenderError::Other(format!("probe: failed to create {label} mip {mip} view"))
            })
        })
        .collect()
}

impl super::context::MtlContext {
    /// Encode the source pyramid: the firefly-clamped copy of the capture into
    /// probe-cube mip 0, then the box reduction of each capture mip into the
    /// next. Both run in one encoder because Metal orders successive dispatches
    /// in a serial compute encoder, which is exactly the chain's dependency.
    ///
    /// These are the cheap dispatches (a few taps per texel), so they all go in
    /// the frame that starts the convolution; the GGX mips that follow are the
    /// expensive ones and take a frame each.
    pub(in crate::metal) fn encode_probe_pyramid(
        &self,
        cmd_buf: &ProtocolObject<dyn objc2_metal::MTLCommandBuffer>,
        gpu: &PrefilterGpu,
        plan: &PrefilterPlan,
    ) -> RenderResult<()> {
        let pipelines = self
            .probe
            .prefilter
            .as_ref()
            .ok_or_else(|| RenderError::Other("probe: prefilter pipelines missing".into()))?;
        let enc = super::scoped_encoder::ScopedEncoder::new(
            cmd_buf.computeCommandEncoder().ok_or_else(|| {
                RenderError::Other("probe: failed to get prefilter compute encoder".into())
            })?,
            ns_string!("probe-pyramid"),
        );

        let params = plan.mip0_params();
        enc.set_pipeline(&pipelines.mip0);
        enc.set_value(&params, 0);
        enc.set_texture(gpu.capture_mip_views[0].as_ref(), 0);
        enc.set_texture(gpu.probe_mip_views[0].as_ref(), 1);
        dispatch_cube(&enc, plan.face_size());

        for mip in 1..plan.mips() {
            let params = plan.downsample_params(mip);
            enc.set_pipeline(&pipelines.downsample);
            enc.set_value(&params, 0);
            enc.set_texture(gpu.capture_mip_views[(mip - 1) as usize].as_ref(), 0);
            enc.set_texture(gpu.capture_mip_views[mip as usize].as_ref(), 1);
            dispatch_cube(&enc, plan.mip_face_size(mip));
        }
        Ok(())
    }

    /// Encode the GGX convolution producing probe-cube mip `dst_mip`, sampling
    /// the whole capture pyramid through the engine's cube sampler (linear,
    /// clamped, mipmapped -- the solid-angle lod the kernel picks needs the
    /// trilinear tap).
    pub(in crate::metal) fn encode_probe_ggx_mip(
        &self,
        cmd_buf: &ProtocolObject<dyn objc2_metal::MTLCommandBuffer>,
        gpu: &PrefilterGpu,
        plan: &PrefilterPlan,
        dst_mip: u32,
    ) -> RenderResult<()> {
        let pipelines = self
            .probe
            .prefilter
            .as_ref()
            .ok_or_else(|| RenderError::Other("probe: prefilter pipelines missing".into()))?;
        let enc = super::scoped_encoder::ScopedEncoder::new(
            cmd_buf.computeCommandEncoder().ok_or_else(|| {
                RenderError::Other("probe: failed to get prefilter compute encoder".into())
            })?,
            ns_string!("probe-ggx"),
        );
        let params = plan.ggx_params(dst_mip);
        enc.set_pipeline(&pipelines.ggx);
        enc.set_value(&params, 0);
        enc.set_texture(gpu.capture.as_ref(), 0);
        enc.set_sampler(&self.scene.cube_sampler, 0);
        enc.set_texture(gpu.probe_mip_views[dst_mip as usize].as_ref(), 1);
        dispatch_cube(&enc, plan.mip_face_size(dst_mip));
        Ok(())
    }
}

// Dispatch one thread per texel of a `size`-square cube face, six faces deep.
// The kernels bounds-guard against `dst_size`, so a non-uniform remainder
// returns early.
fn dispatch_cube(enc: &ProtocolObject<dyn objc2_metal::MTLComputeCommandEncoder>, size: u32) {
    let grid = MTLSize {
        width: size.max(1) as usize,
        height: size.max(1) as usize,
        depth: 6,
    };
    let tg = MTLSize {
        width: PREFILTER_TILE,
        height: PREFILTER_TILE,
        depth: 1,
    };
    enc.dispatchThreads_threadsPerThreadgroup(grid, tg);
}