concinnity-device 0.19.119

//! MetalFX temporal upscaling: the engine renders the 3D scene at a fraction
//! of drawable size and this pass reconstructs a drawable-resolution image
//! the bloom + composite stack consumes. Lives in `metal/post/` next to TAA +
//! bloom + SSAO + SSR; the scaler descriptor, the output texture, and the
//! per-frame encoder all live together so the effect is a single unit.
//!
//! The scaler does temporal accumulation itself, so the existing TAA pass is
//! bypassed while upscaling is on (`PostProcessConfig.aa_mode` is ignored). The
//! existing velocity pre-pass still runs: the scaler consumes its motion
//! vectors. The existing projection jitter still runs: the scaler consumes
//! its sub-pixel offset.
#![deny(unsafe_op_in_unsafe_fn)]

use concinnity_core::render::depth::CAMERA_DEPTH;
use concinnity_core::render::error::{RenderError, RenderResult};
use concinnity_core::render::history_reset::UpscalerResetLatch;
use objc2::rc::Retained;
use objc2::runtime::ProtocolObject;
use objc2_metal::{MTLDevice as _, MTLPixelFormat, MTLTexture, MTLTextureUsage};
use objc2_metal_fx::{MTLFXTemporalScaler, MTLFXTemporalScalerBase, MTLFXTemporalScalerDescriptor};

use crate::metal::context::MtlContext;
use crate::metal::descriptors::TextureDesc;
use crate::metal::error::allocation_failed;
use crate::metal::texture::{REACTIVE_MASK_FORMAT, REACTIVE_MASK_USAGE};

// All MetalFX-temporal-upscaling state grouped into one feature unit: the
// scaler instance, the input/output scale ratio, the per-frame projection
// jitter, and the history-reset flag. `scaler` is `Some` only when the world
// requested temporal upscaling AND the GPU supports it; otherwise `scale` is
// `1.0` (render-res == output-res) and the rest are inert.
pub(crate) struct UpscaleState {
    // The MetalFX scaler. `Some` only when upscaling is active.
    pub scaler: Option<MetalFXUpscaler>,
    // The per-axis input-to-output ratio the world asked for, `1.0` when the
    // scaler is absent. Kept so a resize can rebuild the scaler from the same
    // request. Deliberately not the realized `input / output`: that is rounded
    // to whole pixels, and reusing it as the next request would shrink the input
    // further on every resize. The realized size lives on the scaler itself.
    pub scale: f32,
    // Pixel-space jitter offset the projection applied this frame. Written on
    // the main thread before fan-out, read by `encode_upscale` on a worker;
    // packed atomically so the worker snapshots it without a mutex.
    pub jitter: UpscaleJitter,
    // Whether the scaler discards its temporal history on the next encode:
    // pending after a scaler rebuild (resize / startup) and on a frame that
    // resets the temporal history; consumed by `encode_upscale`.
    pub reset: UpscalerResetLatch,
}

// MetalFX temporal upscaler. Owns the descriptor-bound scaler instance plus
// the output texture the bloom + composite stack consumes (it must be at
// output resolution, GPU-private, with `ShaderRead | RenderTarget` usage
// per Apple's compatibility requirements).
//
// The scaler is sized once at construction. A window resize or a runtime
// scale change rebuilds the whole struct via [`MetalFXUpscaler::new`]; the
// old one drops and Metal retires its GPU storage on the next idle.
pub(crate) struct MetalFXUpscaler {
    // The MetalFX scaler instance. Its `colorTexture` / `depthTexture` /
    // `motionTexture` / `outputTexture` / `jitterOffset` / `reset` properties
    // are set per frame by `MtlContext::encode_upscale` before
    // `encodeToCommandBuffer:`.
    pub(crate) scaler: Retained<ProtocolObject<dyn MTLFXTemporalScaler>>,
    // The output texture the scaler writes into and the post stack reads
    // from. Always sized at `(output_width, output_height)`. Bloom +
    // Composite see this as `scene_color` when upscaling is on.
    pub(crate) output: Retained<ProtocolObject<dyn MTLTexture>>,
    // Render-resolution width (= input width fed to the scaler).
    pub(crate) input_width: u32,
    // Render-resolution height.
    pub(crate) input_height: u32,
    // Drawable-resolution width (= scaler output width).
    pub(crate) output_width: u32,
    // Drawable-resolution height.
    pub(crate) output_height: u32,
    // Whether the scaler takes a reactive mask: the OS has the input
    // (macOS 14.4) and the engine's mask carries the usage the scaler needs.
    pub(crate) reactive: bool,
}

// Does the active device support MetalFX temporal scaling at all? Used at
// init to log a clear "MetalFX not supported on this GPU; rendering at
// native resolution" warning and silently fall back rather than fail.
pub(crate) fn temporal_scaler_supported(
    device: &ProtocolObject<dyn objc2_metal::MTLDevice>,
) -> bool {
    // SAFETY: a device-capability query taking a live MTLDevice; it only reads.
    unsafe { MTLFXTemporalScalerDescriptor::supportsDevice(device) }
}

// Input (render) dimensions for an output size and requested per-axis scale.
// `ratio_range` is the device's supported OUTPUT-over-INPUT range, which is the
// inverse of the user-facing scale, so the request is flipped into the device's
// units, clamped, and flipped back.
//
// The one place this arithmetic happens. The realized size does not round-trip:
// `input / output` is whole pixels over whole pixels, so re-deriving a render
// size from it lands up to a pixel below what the scaler was built for, and
// MetalFX asserts that the input content exceeds the input texture. Callers take
// the render resolution from the built scaler instead of recomputing it.
fn scaler_input_size(output: (u32, u32), scale: f32, ratio_range: (f32, f32)) -> (u32, u32) {
    let (min_ratio, max_ratio) = ratio_range;
    let requested_ratio = if scale > 0.0 { 1.0 / scale } else { 1.0 };
    let lo = min_ratio.max(1.0);
    let clamped_ratio = requested_ratio.clamp(lo, max_ratio.max(lo));
    let scale = 1.0 / clamped_ratio;
    (
        ((output.0 as f32) * scale).max(1.0) as u32,
        ((output.1 as f32) * scale).max(1.0) as u32,
    )
}

impl MetalFXUpscaler {
    // Build a fresh upscaler at the given output size and per-axis scale.
    // `scale` is clamped to the device's supported range; the resolved
    // input dimensions are stored on the struct so the caller knows what
    // resolution to actually render at.
    pub(crate) fn new(
        device: &ProtocolObject<dyn objc2_metal::MTLDevice>,
        output_width: u32,
        output_height: u32,
        scale: f32,
    ) -> RenderResult<Self> {
        // Apple's `inputContentMin/MaxScale` reports the OUTPUT-over-INPUT
        // ratio: a 1.0 minimum means "the smallest upscale is no upscale";
        // a 3.0 maximum means "the largest upscale is 3× per axis". Our
        // user-facing `scale` is the inverse (input/output, ≤ 1.0), so flip
        // it to the device's units, clamp, and flip back.
        // SAFETY: a device-capability query taking a live MTLDevice; it only reads.
        let min_ratio = unsafe {
            MTLFXTemporalScalerDescriptor::supportedInputContentMinScaleForDevice(device)
        };
        // SAFETY: a device-capability query taking a live MTLDevice; it only reads.
        let max_ratio = unsafe {
            MTLFXTemporalScalerDescriptor::supportedInputContentMaxScaleForDevice(device)
        };
        let (input_width, input_height) =
            scaler_input_size((output_width, output_height), scale, (min_ratio, max_ratio));

        // SAFETY: the designated initializer for a fresh MTLFXTemporalScalerDescriptor, which takes
        // no arguments.
        let descriptor = unsafe { MTLFXTemporalScalerDescriptor::new() };
        // SAFETY: plain descriptor property setters, all values in range.
        unsafe {
            // Color: matches the engine's pre-TAA scene format (RGBA16Float
            // single-sample, written by the SSR resolve or the HDR resolve
            // depending on configuration).
            descriptor.setColorTextureFormat(MTLPixelFormat::RGBA16Float);
            // Depth: the velocity pre-pass writes a single-sample Depth32Float
            // depth buffer at render-res that doubles as the scaler depth
            // input. (The main pass's MSAA depth is not directly usable.)
            descriptor.setDepthTextureFormat(MTLPixelFormat::Depth32Float);
            // Motion: the velocity pre-pass writes RG16Float UV-space motion
            // vectors at render-res. The scaler interprets them in input
            // pixel coords once `motionVectorScale` rescales the UV delta.
            descriptor.setMotionTextureFormat(MTLPixelFormat::RG16Float);
            // Output: drawable-res RGBA16Float so the existing post stack
            // (bloom, composite) can read it without an additional format
            // conversion. Composite handles ACES + LUT + FXAA when SDR;
            // skips them on the HDR path.
            descriptor.setOutputTextureFormat(MTLPixelFormat::RGBA16Float);
            descriptor.setInputWidth(input_width as usize);
            descriptor.setInputHeight(input_height as usize);
            descriptor.setOutputWidth(output_width as usize);
            descriptor.setOutputHeight(output_height as usize);
            // We do not author a pre-exposed color buffer (the input is
            // linear HDR fresh from the SSR / HDR resolve), so leave
            // auto-exposure off and let MetalFX use its built-in heuristic.
            descriptor.setAutoExposureEnabled(false);
        }
        // The reactive mask input arrived in macOS 14.4; on an older OS the
        // descriptor does not know the selector, and the scaler runs without it.
        let reactive_available = objc2::runtime::NSObjectProtocol::respondsToSelector(
            &*descriptor,
            objc2::sel!(setReactiveMaskTextureEnabled:),
        );
        if reactive_available {
            // SAFETY: plain descriptor property setters, guarded by the
            // selector check above; R8Unorm is a single-channel color format.
            unsafe {
                descriptor.setReactiveMaskTextureEnabled(true);
                descriptor.setReactiveMaskTextureFormat(REACTIVE_MASK_FORMAT);
            }
        }

        // SAFETY: `descriptor` is fully configured above and `device` is live; MetalFX returns None
        // rather than faulting when the configuration is unsupported.
        let scaler =
            unsafe { descriptor.newTemporalScalerWithDevice(device) }.ok_or_else(|| {
                RenderError::Other("MetalFX: failed to create temporal scaler".to_string())
            })?;

        // The scaler enforces a minimum texture-usage set on the output
        // texture; query it and union with the bloom + composite read
        // requirement so the post stack can sample the result.
        // SAFETY: a property read on the live scaler just created above.
        let required_output_usage = unsafe { scaler.outputTextureUsage() };
        let output_desc = TextureDesc {
            format: MTLPixelFormat::RGBA16Float,
            width: output_width.max(1) as usize,
            height: output_height.max(1) as usize,
            usage: MTLTextureUsage(
                required_output_usage.0
                    | MTLTextureUsage::ShaderRead.0
                    | MTLTextureUsage::RenderTarget.0,
            ),
            ..Default::default()
        }
        .build();
        let output = device
            .newTextureWithDescriptor(&output_desc)
            .ok_or_else(|| allocation_failed("MetalFX upscaler output texture"))?;
        let reactive = reactive_available && {
            // SAFETY: a property read on the live scaler, which a descriptor
            // with the reactive input enabled created.
            let needed = unsafe { scaler.reactiveTextureUsage() };
            needed.0 & !REACTIVE_MASK_USAGE.0 == 0
        };
        if reactive_available && !reactive {
            tracing::warn!(
                "MetalFX: the reactive mask lacks a usage the scaler needs; running without it"
            );
        }

        Ok(MetalFXUpscaler {
            scaler,
            output,
            input_width,
            input_height,
            output_width,
            output_height,
            reactive,
        })
    }
}

impl MtlContext {
    // Encode the MetalFX temporal upscale: feed it the pre-TAA scene
    // (post-SSR, post-Fog, post-particles), the velocity pre-pass's
    // motion vectors + depth, and the sub-pixel jitter offset applied to
    // this frame's projection. Outputs into `upscaler.output`, which the
    // bloom + composite passes then read as `scene_color`.
    //
    // `scene_pre_taa` is whatever pre-TAA texture the graph routed in for
    // this frame (SSR resolve output when SSR is on; otherwise the main
    // pass's `hdr_resolve`). The scaler needs its `depthTexture` to be
    // the depth that produced `scene_pre_taa`: we use the single-sample
    // depth the velocity pre-pass already writes at render resolution,
    // which is rasterized from the same geometry the main pass shaded.
    //
    // Reset is requested whenever the scaler was just rebuilt (resize or
    // first frame) or the frame resets the temporal history; the encode
    // consumes the latch `draw_frame` raises on the main thread between frames.
    pub(in crate::metal) fn encode_upscale(
        &self,
        cmd_buf: &ProtocolObject<dyn objc2_metal::MTLCommandBuffer>,
        scene_pre_taa: &Retained<ProtocolObject<dyn MTLTexture>>,
        reactive_written: bool,
    ) -> RenderResult<u32> {
        let upscaler = self.upscale.scaler.as_ref().ok_or_else(|| {
            RenderError::Other("Upscale enabled but upscaler missing".to_string())
        })?;
        // The pre-pass depth stays feature-owned; its motion channel is
        // pool-owned and fetched at encode time.
        let gbuf = self.gbuffer.targets.as_ref().ok_or_else(|| {
            RenderError::Other("Upscale enabled but G-buffer targets missing".to_string())
        })?;
        let velocity = self.gbuffer_velocity().ok_or_else(|| {
            RenderError::Other(
                "Upscale enabled but the pooled G-buffer velocity is missing".to_string(),
            )
        })?;

        // SAFETY: every texture set here is owned by `self` or the upscaler and outlives the encode
        // below, and each matches the format the descriptor declared for that slot.
        unsafe {
            upscaler
                .scaler
                .setColorTexture(Some(scene_pre_taa.as_ref()));
            upscaler.scaler.setDepthTexture(Some(gbuf.depth.as_ref()));
            upscaler.scaler.setMotionTexture(Some(velocity));
            upscaler
                .scaler
                .setOutputTexture(Some(upscaler.output.as_ref()));
            // A frame no particle or transparent pass wrote the mask for gives
            // the scaler none, which it reads as no pixel being reactive.
            if upscaler.reactive {
                let mask = reactive_written.then(|| self.targets.hdr.reactive_mask.as_ref());
                upscaler.scaler.setReactiveMaskTexture(mask);
            }

            // Motion vectors are stored as `prev_uv - cur_uv` in UV space
            // (RG16Float). The scaler expects motion in input-pixel coords,
            // so the per-axis scale is the input texture extent.
            upscaler
                .scaler
                .setMotionVectorScaleX(upscaler.input_width as f32);
            upscaler
                .scaler
                .setMotionVectorScaleY(upscaler.input_height as f32);

            // The jitter the projection matrix applied this frame in input
            // pixel space (each Halton sample is in [-0.5, 0.5] pixels).
            // `draw_frame` stashes it on the context before fan-out.
            let [jx, jy] = self
                .upscale
                .jitter
                .load(std::sync::atomic::Ordering::Relaxed);
            upscaler.scaler.setJitterOffsetX(jx);
            upscaler.scaler.setJitterOffsetY(jy);

            // `depthReversed` means near is device depth 1.
            upscaler.scaler.setDepthReversed(CAMERA_DEPTH.reversed);

            // The first frame after a rebuild or a history reset discards
            // history. The latch is raised on the main thread between frames.
            upscaler
                .scaler
                .setReset(crate::upscale_reset::consume(&self.upscale.reset));
        }

        // MetalFX's encode is a single dispatch that doesn't go through a render
        // or compute encoder we own, so the standard `attach_render` /
        // `attach_compute` plumbing doesn't apply. The pass-timing slot for
        // Upscale stays at zero until Apple exposes a sampling hook; the chip
        // simply omits the row.

        // SAFETY: the scaler's input, output, and parameter slots were all set above, and `cmd_buf`
        // is the live command buffer for this frame.
        unsafe {
            upscaler.scaler.encodeToCommandBuffer(cmd_buf);
        }
        Ok(0)
    }
}

// Per-axis upscale-jitter holder. Two `f32` slots packed atomically so the
// main thread (which computes Halton samples) and the worker thread (which
// reads them in `encode_upscale`) can synchronize without a mutex. Bit-cast
// through `u64` since `AtomicU64` is widely available on the targets
// MetalFX runs on (macOS 13+, aarch64 / x86_64).
#[derive(Default)]
pub(crate) struct UpscaleJitter(std::sync::atomic::AtomicU64);

impl UpscaleJitter {
    pub(crate) fn store(&self, jx: f32, jy: f32, ordering: std::sync::atomic::Ordering) {
        let packed = (jx.to_bits() as u64) | ((jy.to_bits() as u64) << 32);
        self.0.store(packed, ordering);
    }

    pub(crate) fn load(&self, ordering: std::sync::atomic::Ordering) -> [f32; 2] {
        let packed = self.0.load(ordering);
        let jx = f32::from_bits(packed as u32);
        let jy = f32::from_bits((packed >> 32) as u32);
        [jx, jy]
    }
}

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

    // Typical range: no upscale at all up to 3x per axis.
    const RANGE: (f32, f32) = (1.0, 3.0);

    #[test]
    fn the_realized_scale_does_not_round_trip() {
        // The bug this guards. A 2048x1536 drawable at the default 2/3 quality
        // scale gives a 1365x1024 input. Feeding the realized ratio
        // (1365/2048) back in as a request yields 1023 rows -- one short of what
        // the scaler was built for, which MetalFX rejects outright rather than
        // tolerating. So the render resolution is read off the scaler, and the
        // stored scale stays the original request.
        let out = (2048, 1536);
        let built = scaler_input_size(out, 2.0 / 3.0, RANGE);
        assert_eq!(built, (1365, 1024));

        let realized = built.0 as f32 / out.0 as f32;
        let rederived = scaler_input_size(out, realized, RANGE);
        assert_ne!(
            rederived, built,
            "if this ever round-trips the guard below is measuring nothing"
        );
        assert_eq!(rederived.1, 1023, "a row short of the built input");
    }

    #[test]
    fn a_request_outside_the_device_range_is_clamped() {
        // Asking for a 5x upscale on a 3x device resolves to 3x, and the
        // resulting input size is the one the scaler is really built for --
        // which is exactly why the caller must not assume its own request.
        let out = (1920, 1080);
        assert_eq!(scaler_input_size(out, 1.0 / 5.0, RANGE), (640, 360));
        // And below the minimum: a request to render larger than the output
        // clamps to no upscale at all.
        assert_eq!(scaler_input_size(out, 2.0, RANGE), out);
    }

    #[test]
    fn degenerate_inputs_stay_renderable() {
        // A zero or negative scale means "no upscale" rather than a zero-sized
        // target, and a tiny output never rounds to a zero dimension.
        assert_eq!(scaler_input_size((800, 600), 0.0, RANGE), (800, 600));
        assert_eq!(scaler_input_size((800, 600), -1.0, RANGE), (800, 600));
        assert_eq!(scaler_input_size((1, 1), 1.0 / 3.0, RANGE), (1, 1));
        // An inverted range still yields a usable size rather than a panic from
        // `clamp` (min > max).
        assert_eq!(scaler_input_size((1920, 1080), 0.5, (3.0, 1.0)), (640, 360));
    }
}