cortiq-engine 0.8.41

//! Resident Flux-VAE decoder on Metal (plan WP3, M6).
//!
//! NHWC activations: the residual stream in f32, every conv input in half.
//! Each 3×3 / 1×1 conv is an implicit GEMM (`zv_conv`, the DiT GEMM's
//! 64×64×32 tile, half weights `[oc][tap][ic]`, f32 accumulation, bias and
//! residual in the epilogue); the nearest-2× upsample is folded into the
//! gather. GroupNorm is two exact passes (sum, then centred sum of squares)
//! over per-block partials, fused with the affine, SiLU and the half cast.
//! The mid attention (one head, c = 512) runs as GEMMs: q (pre-scaled by
//! 1/√c), k, vᵀ, then per query chunk S = q·kᵀ (f32), row softmax → half P,
//! O = P·v, out projection + residual. One latent upload, one RGB readback.

use super::{Pipes, buf_from, buf_zeroed, pipes, set_p};
use crate::vae::{VaeChainArgs, VaeConvRef, VaeNormRef, VaeResnetRef};
use metal::{Buffer, ComputeCommandEncoderRef, MTLResourceOptions, MTLSize};
use std::ffi::c_void;
use std::sync::Mutex;

use super::super::Ctx;

#[repr(C)]
#[derive(Clone, Copy, Default)]
struct PVc {
    n: u32,
    rows: u32,
    k: u32,
    ic: u32,
    h: u32,
    w: u32,
    taps: u32,
    up: u32,
    epi: u32,
    ldy: u32,
    has_bias: u32,
    has_res: u32,
    mul: f32,
    ldx: u32,
    pad0: u32,
    pad1: u32,
}

#[repr(C)]
#[derive(Clone, Copy, Default)]
struct PGn {
    n: u32,
    c: u32,
    groups: u32,
    nblk: u32,
    pass: u32,
    eps: f32,
    silu: u32,
    half_out: u32,
}

/// One conv (or linear, taps = 0) with half weights [ocp][taps·icp] and
/// f32 bias [ocp] (zero-padded channels).
struct VConv {
    w: Buffer,
    b: Buffer,
    ocp: usize,
    icp: usize,
    taps: usize,
}

struct VNorm {
    w: Buffer,
    b: Buffer,
    groups: usize,
}

struct VRes {
    n1: VNorm,
    c1: VConv,
    n2: VNorm,
    c2: VConv,
    sc: Option<VConv>,
    cin: usize,
    cout: usize,
}

struct VAttn {
    norm: VNorm,
    q: VConv,
    k: VConv,
    v: VConv,
    o: VConv,
    c: usize,
}

struct VUp {
    res: Vec<VRes>,
    up: Option<VConv>,
}

struct VaeDev {
    key: u64,
    conv_in: VConv,
    mid1: VRes,
    attn: VAttn,
    mid2: VRes,
    ups: Vec<VUp>,
    norm_out: VNorm,
    conv_out: VConv,
    latent_c: usize,
}

struct VState(Option<VaeDev>);
unsafe impl Send for VState {}

static VSTATE: Mutex<VState> = Mutex::new(VState(None));

fn half_buf(c: &Ctx, v: &[f32]) -> Buffer {
    let h: Vec<u16> = v
        .iter()
        .map(|&x| cortiq_core::quant::f32_to_f16(x))
        .collect();
    c._device.new_buffer_with_data(
        h.as_ptr() as *const c_void,
        (h.len().max(8) * 2) as u64,
        MTLResourceOptions::StorageModeShared,
    )
}

/// conv weight [oc][ic][k][k] → half [ocp][k·k][icp]; `scale` multiplies
/// weight and bias (the attention's 1/√c for q).
fn conv_of(c: &Ctx, r: &VaeConvRef, scale: f32) -> Result<VConv, String> {
    let (oc, ic, k) = (r.oc, r.ic, r.k);
    if r.w.len() != oc * ic * k * k || r.b.len() != oc || (k != 1 && k != 3) {
        return Err(format!("vae conv shape oc {oc} ic {ic} k {k}"));
    }
    let taps = k * k;
    let ocp = oc.div_ceil(64) * 64;
    let icp = ic.div_ceil(32) * 32;
    let mut w = vec![0f32; ocp * taps * icp];
    for o in 0..oc {
        for i in 0..ic {
            for t in 0..taps {
                w[(o * taps + t) * icp + i] = r.w[(o * ic + i) * taps + t] * scale;
            }
        }
    }
    let mut b = vec![0f32; ocp];
    for o in 0..oc {
        b[o] = r.b[o] * scale;
    }
    Ok(VConv {
        w: half_buf(c, &w),
        b: buf_from(c, &b),
        ocp,
        icp,
        taps,
    })
}

/// A [c_out, c_in] linear as a plain-mode "conv" (taps 0).
fn lin_of(c: &Ctx, w: &[f32], b: &[f32], ch: usize, scale: f32) -> Result<VConv, String> {
    if w.len() != ch * ch || b.len() != ch || ch % 64 != 0 {
        return Err(format!("vae attention linear {ch}"));
    }
    let ws: Vec<f32> = w.iter().map(|v| v * scale).collect();
    let bs: Vec<f32> = b.iter().map(|v| v * scale).collect();
    Ok(VConv {
        w: half_buf(c, &ws),
        b: buf_from(c, &bs),
        ocp: ch,
        icp: ch,
        taps: 0,
    })
}

fn norm_of(c: &Ctx, r: &VaeNormRef) -> VNorm {
    VNorm {
        w: buf_from(c, r.w),
        b: buf_from(c, r.b),
        groups: r.groups,
    }
}

/// The GroupNorm shader writes float4s, but its reduction has a scalar-tail
/// path for compact groups.  Keep admission deliberately narrow: every
/// channel vector must remain float4 aligned and the fixed scratch planes can
/// hold at most 64 groups.
fn norm_supported(r: &VaeNormRef, ch: usize) -> bool {
    ch % 4 == 0
        && r.groups != 0
        && r.groups <= 64
        && ch % r.groups == 0
        && r.w.len() == ch
        && r.b.len() == ch
}

fn res_of(c: &Ctx, r: &VaeResnetRef) -> Result<VRes, String> {
    let (cin, cout) = (r.conv1.ic, r.conv1.oc);
    // Conv inputs are loaded in 32-wide half lanes.  Outputs may be only
    // 32-wide (Flux2's final 192→96 stage); `zv_conv` guards their padded
    // 64-row tile before it stores into the logical NHWC stride.
    if cin % 32 != 0
        || cout % 32 != 0
        || !norm_supported(&r.norm1, cin)
        || !norm_supported(&r.norm2, cout)
    {
        return Err(format!("vae resnet {cin}->{cout}"));
    }
    Ok(VRes {
        n1: norm_of(c, &r.norm1),
        c1: conv_of(c, &r.conv1, 1.0)?,
        n2: norm_of(c, &r.norm2),
        c2: conv_of(c, &r.conv2, 1.0)?,
        sc: match &r.shortcut {
            Some(s) => Some(conv_of(c, s, 1.0)?),
            None => None,
        },
        cin,
        cout,
    })
}

fn build(c: &Ctx, a: &VaeChainArgs) -> Result<VaeDev, String> {
    let at = &a.mid_attn;
    let ch = at.c;
    if !norm_supported(&at.norm, ch) {
        return Err(format!("vae attention norm {ch}"));
    }
    if a.conv_out.ic % 32 != 0 || !norm_supported(&a.norm_out, a.conv_out.ic) {
        return Err(format!("vae output norm {}", a.conv_out.ic));
    }
    let scale = 1.0 / (ch as f32).sqrt();
    Ok(VaeDev {
        key: a.key,
        conv_in: conv_of(c, &a.conv_in, 1.0)?,
        mid1: res_of(c, &a.mid_res1)?,
        attn: VAttn {
            norm: norm_of(c, &at.norm),
            q: lin_of(c, at.q.0, at.q.1, ch, scale)?,
            k: lin_of(c, at.k.0, at.k.1, ch, 1.0)?,
            v: lin_of(c, at.v.0, at.v.1, ch, 1.0)?,
            o: lin_of(c, at.out.0, at.out.1, ch, 1.0)?,
            c: ch,
        },
        mid2: res_of(c, &a.mid_res2)?,
        ups: a
            .ups
            .iter()
            .map(|u| {
                Ok(VUp {
                    res: u
                        .resnets
                        .iter()
                        .map(|r| res_of(c, r))
                        .collect::<Result<Vec<_>, String>>()?,
                    up: match &u.upsample {
                        Some(s) => Some(conv_of(c, s, 1.0)?),
                        None => None,
                    },
                })
            })
            .collect::<Result<Vec<_>, String>>()?,
        norm_out: norm_of(c, &a.norm_out),
        conv_out: conv_of(c, &a.conv_out, 1.0)?,
        latent_c: a.latent_channels,
    })
}

fn ensure<'s>(st: &'s mut VState, c: &Ctx, a: &VaeChainArgs) -> Result<&'s VaeDev, String> {
    if st.0.as_ref().is_none_or(|d| d.key != a.key) {
        st.0 = None;
        st.0 = Some(build(c, a)?);
    }
    Ok(st.0.as_ref().unwrap())
}

pub(crate) fn prewarm(a: &VaeChainArgs) -> bool {
    let Some(c) = super::super::ctx() else {
        return false;
    };
    if pipes(c).is_none() {
        return false;
    }
    let mut st = VSTATE.lock().unwrap();
    match ensure(&mut st, c, a) {
        Ok(_) => true,
        Err(e) => super::decline(&e),
    }
}

pub(crate) fn release() {
    if let Ok(mut st) = VSTATE.lock() {
        st.0 = None;
    }
}

/// Encoder helper: one encoder per command buffer, a new buffer at every
/// `cut` (committed without a wait).
struct VRec<'a> {
    c: &'a Ctx,
    p: &'static Pipes,
    done: Vec<metal::CommandBuffer>,
    cur: Option<(metal::CommandBuffer, metal::ComputeCommandEncoder)>,
    part: Buffer,
    stat: Buffer,
}

impl VRec<'_> {
    fn enc(&mut self) -> &ComputeCommandEncoderRef {
        if self.cur.is_none() {
            let cmd = self.c.queue.new_command_buffer().to_owned();
            let enc = cmd.new_compute_command_encoder().to_owned();
            self.cur = Some((cmd, enc));
        }
        &self.cur.as_ref().unwrap().1
    }
    fn cut(&mut self) {
        if let Some((cmd, enc)) = self.cur.take() {
            enc.end_encoding();
            cmd.commit();
            self.done.push(cmd);
        }
    }
    fn finish(&mut self) -> bool {
        self.cut();
        let mut ok = true;
        for cmd in &self.done {
            cmd.wait_until_completed();
            ok &= cmd.status() == metal::MTLCommandBufferStatus::Completed;
        }
        ok
    }

    /// Y = conv(X) (+bias) (+R). `hw` = output (H, W); `up` = the source is
    /// (H/2, W/2). Plain mode (taps 0): X [n][ldx], `n` rows.
    #[allow(clippy::too_many_arguments)]
    fn conv(
        &mut self,
        cv: &VConv,
        x: (&Buffer, u64),
        y: (&Buffer, u64),
        res: Option<(&Buffer, u64)>,
        n: usize,
        hw: (usize, usize),
        up: bool,
        epi: u32,
        ldy: usize,
        ldx: usize,
        k_plain: usize,
    ) {
        let k = if cv.taps == 0 {
            k_plain
        } else {
            cv.taps * cv.icp
        };
        let pv = PVc {
            n: n as u32,
            // Normal NHWC outputs use `ldy` as their logical channel
            // stride.  The MMA tile still computes padded rows, and the
            // shader drops those rows before they can spill into the next
            // pixel.  EPI=2 is the attention V transpose [channel][token],
            // whose `ldy` is a token stride rather than a channel count.
            rows: (if epi == 2 { cv.ocp } else { ldy }) as u32,
            k: k as u32,
            ic: cv.icp as u32,
            h: hw.0 as u32,
            w: hw.1 as u32,
            taps: cv.taps as u32,
            up: up as u32,
            epi,
            ldy: ldy as u32,
            has_bias: 1,
            has_res: res.is_some() as u32,
            mul: 1.0,
            ldx: ldx as u32,
            ..Default::default()
        };
        let p = self.p;
        let enc = self.enc();
        enc.set_compute_pipeline_state(&p.vconv);
        enc.set_buffer(0, Some(&cv.w), 0);
        enc.set_buffer(1, Some(&cv.b), 0);
        enc.set_buffer(2, Some(x.0), x.1);
        enc.set_buffer(3, Some(y.0), y.1);
        let r = res.unwrap_or(y);
        enc.set_buffer(4, Some(r.0), r.1);
        set_p(enc, 5, &pv);
        enc.dispatch_thread_groups(
            MTLSize::new(n.div_ceil(64) as u64, (cv.ocp / 64) as u64, 1),
            MTLSize::new(128, 1, 1),
        );
    }

    /// Plain GEMM with an explicit weight buffer (the attention's k / vᵀ
    /// as "weights"): Y[n][rows] = X[n][K] · W[rows][K]ᵀ, no bias.
    #[allow(clippy::too_many_arguments)]
    fn gemm(
        &mut self,
        w: (&Buffer, u64),
        rows: usize,
        k: usize,
        x: (&Buffer, u64),
        ldx: usize,
        y: (&Buffer, u64),
        ldy: usize,
        n: usize,
        epi: u32,
    ) {
        let pv = PVc {
            n: n as u32,
            rows: rows as u32,
            k: k as u32,
            ic: k as u32,
            epi,
            ldy: ldy as u32,
            mul: 1.0,
            ldx: ldx as u32,
            ..Default::default()
        };
        let p = self.p;
        let enc = self.enc();
        enc.set_compute_pipeline_state(&p.vconv);
        enc.set_buffer(0, Some(w.0), w.1);
        enc.set_buffer(1, Some(y.0), y.1);
        enc.set_buffer(2, Some(x.0), x.1);
        enc.set_buffer(3, Some(y.0), y.1);
        enc.set_buffer(4, Some(y.0), y.1);
        set_p(enc, 5, &pv);
        enc.dispatch_thread_groups(
            MTLSize::new(n.div_ceil(64) as u64, (rows / 64) as u64, 1),
            MTLSize::new(128, 1, 1),
        );
    }

    /// GroupNorm (+SiLU) of x [n][ch] f32 → y (half or f32).
    fn gn(
        &mut self,
        nm: &VNorm,
        x: &Buffer,
        y: &Buffer,
        n: usize,
        ch: usize,
        silu: bool,
        half_out: bool,
    ) {
        let nblk = n.div_ceil(1024).clamp(1, 256);
        let mut pg = PGn {
            n: n as u32,
            c: ch as u32,
            groups: nm.groups as u32,
            nblk: nblk as u32,
            pass: 0,
            eps: 1e-6,
            silu: silu as u32,
            half_out: half_out as u32,
        };
        let p = self.p;
        let (part, stat) = (self.part.clone(), self.stat.clone());
        let enc = self.enc();
        for pass in 0..2u32 {
            pg.pass = pass;
            enc.set_compute_pipeline_state(&p.vgnpart);
            enc.set_buffer(0, Some(x), 0);
            enc.set_buffer(1, Some(&part), 0);
            enc.set_buffer(2, Some(&stat), 0);
            set_p(enc, 3, &pg);
            enc.dispatch_thread_groups(
                MTLSize::new(nm.groups as u64, nblk as u64, 1),
                MTLSize::new(256, 1, 1),
            );
            enc.set_compute_pipeline_state(&p.vgnfin);
            enc.set_buffer(0, Some(&part), 0);
            enc.set_buffer(1, Some(&stat), 0);
            set_p(enc, 2, &pg);
            enc.dispatch_thread_groups(
                MTLSize::new(nm.groups as u64, 1, 1),
                MTLSize::new(32, 1, 1),
            );
        }
        enc.set_compute_pipeline_state(&p.vgnapply);
        enc.set_buffer(0, Some(x), 0);
        enc.set_buffer(1, Some(y), 0);
        enc.set_buffer(2, Some(&stat), 0);
        enc.set_buffer(3, Some(&nm.w), 0);
        enc.set_buffer(4, Some(&nm.b), 0);
        set_p(enc, 5, &pg);
        let n4 = (n * ch / 4) as u64;
        enc.dispatch_thread_groups(
            MTLSize::new(n4.div_ceil(256), 1, 1),
            MTLSize::new(256, 1, 1),
        );
    }

    fn cvt(&mut self, x: &Buffer, y: &Buffer, count: usize) {
        let p = self.p;
        let n4 = (count / 4) as u32;
        let enc = self.enc();
        enc.set_compute_pipeline_state(&p.vcvt);
        enc.set_buffer(0, Some(x), 0);
        enc.set_buffer(1, Some(y), 0);
        set_p(enc, 2, &n4);
        enc.dispatch_thread_groups(
            MTLSize::new((n4 as u64).div_ceil(256), 1, 1),
            MTLSize::new(256, 1, 1),
        );
    }
}

struct Bufs {
    xa: Buffer,
    t: Buffer,
    hh: Buffer,
}

/// One resnet on xa [n][cin] (f32) → xa [n][cout].
fn resnet(r: &mut VRec, b: &Bufs, rs: &VRes, n: usize, hw: (usize, usize)) {
    r.gn(&rs.n1, &b.xa, &b.hh, n, rs.cin, true, true);
    r.conv(
        &rs.c1,
        (&b.hh, 0),
        (&b.t, 0),
        None,
        n,
        hw,
        false,
        0,
        rs.cout,
        0,
        0,
    );
    if let Some(sc) = &rs.sc {
        r.cvt(&b.xa, &b.hh, n * rs.cin);
        r.conv(
            sc,
            (&b.hh, 0),
            (&b.xa, 0),
            None,
            n,
            hw,
            false,
            0,
            rs.cout,
            0,
            0,
        );
    }
    r.gn(&rs.n2, &b.t, &b.hh, n, rs.cout, true, true);
    r.conv(
        &rs.c2,
        (&b.hh, 0),
        (&b.xa, 0),
        Some((&b.xa, 0)),
        n,
        hw,
        false,
        0,
        rs.cout,
        0,
        0,
    );
}

/// The f16-plane arm of the codec A/B (plan M5c): the same 64×64×32 tile
/// with half weights (`zv_conv`, plain mode) on y[n, rows] = x[n, k]·Wᵀ.
/// Returns (min ms, median ms) of GPU time.
pub(crate) fn bench_plane(rows: usize, k: usize, n: usize, reps: usize) -> Option<(f64, f64)> {
    let c = super::super::ctx()?;
    let p = pipes(c)?;
    let w: Vec<u16> = (0..rows * k)
        .map(|i| cortiq_core::quant::f32_to_f16(((i * 7) % 255) as f32 - 127.0))
        .collect();
    let x: Vec<u16> = (0..(n.div_ceil(64) * 64 + 64) * k)
        .map(|i| cortiq_core::quant::f32_to_f16(((i * 5) % 17) as f32 * 0.01))
        .collect();
    let wb = c._device.new_buffer_with_data(
        w.as_ptr() as *const c_void,
        (w.len() * 2) as u64,
        MTLResourceOptions::StorageModeShared,
    );
    let xb = c._device.new_buffer_with_data(
        x.as_ptr() as *const c_void,
        (x.len() * 2) as u64,
        MTLResourceOptions::StorageModeShared,
    );
    let yb = buf_zeroed(c, n * rows * 4);
    let pv = PVc {
        n: n as u32,
        rows: rows as u32,
        k: k as u32,
        ic: k as u32,
        ldy: rows as u32,
        mul: 1.0,
        ldx: k as u32,
        ..Default::default()
    };
    let mut t = Vec::new();
    for _ in 0..reps + 1 {
        let cmd = c.queue.new_command_buffer();
        let enc = cmd.new_compute_command_encoder();
        enc.set_compute_pipeline_state(&p.vconv);
        enc.set_buffer(0, Some(&wb), 0);
        enc.set_buffer(1, Some(&yb), 0);
        enc.set_buffer(2, Some(&xb), 0);
        enc.set_buffer(3, Some(&yb), 0);
        enc.set_buffer(4, Some(&yb), 0);
        set_p(enc, 5, &pv);
        enc.dispatch_thread_groups(
            MTLSize::new(n.div_ceil(64) as u64, (rows / 64) as u64, 1),
            MTLSize::new(128, 1, 1),
        );
        enc.end_encoding();
        cmd.commit();
        cmd.wait_until_completed();
        t.push(super::super::cmd_gpu_ms(cmd));
    }
    t.remove(0);
    t.sort_by(|a, b| a.partial_cmp(b).unwrap());
    Some((t[0], t[t.len() / 2]))
}

/// Resident Flux-VAE decoder; `z` is already de-normalised, [lc, h, w];
/// `out` [3, 8h, 8w].
pub(crate) fn decode(a: &VaeChainArgs, z: &[f32], h: usize, w: usize, out: &mut [f32]) -> bool {
    if std::env::var("CMF_ZI_VAE").as_deref() == Ok("0") {
        return false;
    }
    let Some(c) = super::super::ctx() else {
        return false;
    };
    let Some(p) = pipes(c) else { return false };
    let mut st = VSTATE.lock().unwrap();
    let d = match ensure(&mut st, c, a) {
        Ok(d) => d,
        Err(e) => return super::decline(&e),
    };
    let lc = d.latent_c;
    let n0 = h * w;
    if z.len() != lc * n0 || out.len() != 3 * 64 * n0 || d.attn.c % 64 != 0 || n0 % 4 != 0 {
        return false;
    }
    // attention keys padded to the GEMM tile (zero keys, masked in the softmax)
    let n0p = n0.div_ceil(64) * 64;
    // largest n·C of any stage
    let mut cmax = 0usize;
    {
        let (mut hh, mut ww, mut ch) = (h, w, d.mid1.cout);
        cmax = cmax.max(hh * ww * ch.max(d.conv_in.icp));
        for u in &d.ups {
            for r in &u.res {
                cmax = cmax.max(hh * ww * r.cin.max(r.cout));
                ch = r.cout;
            }
            if u.up.is_some() {
                hh *= 2;
                ww *= 2;
                cmax = cmax.max(hh * ww * ch);
            }
        }
        cmax = cmax.max(hh * ww * 64);
    }
    // Metal returns nil past maxBufferLength (13.6 GB on a 24 GB M4) and the
    // uploads below would write through a null pointer: decline instead.
    if (cmax * 4) as u64 > c._device.max_buffer_length() {
        return super::decline(&format!(
            "the VAE activation buffer ({:.1} GB) exceeds the device's maxBufferLength",
            (cmax * 4) as f64 / 1e9
        ));
    }
    let bufs = Bufs {
        xa: buf_zeroed(c, cmax * 4),
        t: buf_zeroed(c, cmax * 4),
        hh: buf_zeroed(c, cmax * 2),
    };
    // latent → NHWC half, channels padded to the conv's icp
    {
        let icp = d.conv_in.icp;
        let mut zh = vec![0u16; n0 * icp];
        for ch in 0..lc {
            for px in 0..n0 {
                zh[px * icp + ch] = cortiq_core::quant::f32_to_f16(z[ch * n0 + px]);
            }
        }
        unsafe {
            std::ptr::copy_nonoverlapping(zh.as_ptr(), bufs.hh.contents() as *mut u16, zh.len());
        }
    }
    let ac = d.attn.c;
    let chunk = n0.min(env_chunk());
    let (qb, kb, vt, ob) = (
        buf_zeroed(c, n0p * ac * 2),
        buf_zeroed(c, n0p * ac * 2),
        buf_zeroed(c, n0p * ac * 2),
        buf_zeroed(c, n0p * ac * 2),
    );
    let sb = buf_zeroed(c, chunk * n0p * 4);
    let pb = buf_zeroed(c, chunk * n0p * 2);
    let mut r = VRec {
        c,
        p,
        done: Vec::new(),
        cur: None,
        part: buf_zeroed(c, 64 * 256 * 4),
        stat: buf_zeroed(c, 64 * 2 * 4),
    };
    let t0 = std::time::Instant::now();
    let b = &bufs;
    // conv_in
    r.conv(
        &d.conv_in,
        (&b.hh, 0),
        (&b.xa, 0),
        None,
        n0,
        (h, w),
        false,
        0,
        d.conv_in.ocp,
        0,
        0,
    );
    resnet(&mut r, b, &d.mid1, n0, (h, w));
    r.cut();
    // mid attention
    {
        let at = &d.attn;
        r.gn(&at.norm, &b.xa, &b.hh, n0, ac, false, true);
        r.conv(
            &at.q,
            (&b.hh, 0),
            (&qb, 0),
            None,
            n0,
            (0, 0),
            false,
            1,
            ac,
            ac,
            ac,
        );
        r.conv(
            &at.k,
            (&b.hh, 0),
            (&kb, 0),
            None,
            n0,
            (0, 0),
            false,
            1,
            ac,
            ac,
            ac,
        );
        r.conv(
            &at.v,
            (&b.hh, 0),
            (&vt, 0),
            None,
            n0,
            (0, 0),
            false,
            2,
            n0p,
            ac,
            ac,
        );
        let mut q0 = 0;
        while q0 < n0 {
            let m = chunk.min(n0 - q0);
            r.gemm(
                (&kb, 0),
                n0p,
                ac,
                (&qb, (q0 * ac * 2) as u64),
                ac,
                (&sb, 0),
                n0p,
                m,
                0,
            );
            {
                let pp = r.p;
                let (n0u, ldu) = (n0 as u32, n0p as u32);
                let enc = r.enc();
                enc.set_compute_pipeline_state(&pp.vsoftmax);
                enc.set_buffer(0, Some(&sb), 0);
                enc.set_buffer(1, Some(&pb), 0);
                set_p(enc, 2, &n0u);
                set_p(enc, 3, &ldu);
                enc.dispatch_thread_groups(MTLSize::new(m as u64, 1, 1), MTLSize::new(256, 1, 1));
            }
            r.gemm(
                (&vt, 0),
                ac,
                n0p,
                (&pb, 0),
                n0p,
                (&ob, (q0 * ac * 2) as u64),
                ac,
                m,
                1,
            );
            q0 += m;
        }
        r.conv(
            &at.o,
            (&ob, 0),
            (&b.xa, 0),
            Some((&b.xa, 0)),
            n0,
            (0, 0),
            false,
            0,
            ac,
            ac,
            ac,
        );
    }
    resnet(&mut r, b, &d.mid2, n0, (h, w));
    r.cut();
    let (mut hh, mut ww) = (h, w);
    let mut ch = d.mid2.cout;
    let mut bufs_swapped = false;
    let (mut xa, mut tt) = (b.xa.clone(), b.t.clone());
    for u in &d.ups {
        for rs in &u.res {
            let bb = Bufs {
                xa: xa.clone(),
                t: tt.clone(),
                hh: b.hh.clone(),
            };
            resnet(&mut r, &bb, rs, hh * ww, (hh, ww));
            ch = rs.cout;
            r.cut();
        }
        if let Some(upc) = &u.up {
            r.cvt(&xa, &b.hh, hh * ww * ch);
            hh *= 2;
            ww *= 2;
            r.conv(
                upc,
                (&b.hh, 0),
                (&tt, 0),
                None,
                hh * ww,
                (hh, ww),
                true,
                0,
                upc.ocp,
                0,
                0,
            );
            std::mem::swap(&mut xa, &mut tt);
            bufs_swapped = !bufs_swapped;
            r.cut();
        }
    }
    let n = hh * ww;
    r.gn(&d.norm_out, &xa, &b.hh, n, ch, true, true);
    r.conv(
        &d.conv_out,
        (&b.hh, 0),
        (&tt, 0),
        None,
        n,
        (hh, ww),
        false,
        0,
        64,
        0,
        0,
    );
    let rgb = buf_zeroed(c, 3 * n * 4);
    {
        let pp = r.p;
        let nu = n as u32;
        let enc = r.enc();
        enc.set_compute_pipeline_state(&pp.vrgb);
        enc.set_buffer(0, Some(&tt), 0);
        enc.set_buffer(1, Some(&rgb), 0);
        set_p(enc, 2, &nu);
        enc.dispatch_thread_groups(
            MTLSize::new((n as u64).div_ceil(256), 1, 1),
            MTLSize::new(256, 1, 1),
        );
    }
    let _ = bufs_swapped;
    if !r.finish() {
        return super::decline("a Metal command buffer failed (VAE)");
    }
    if std::env::var("CMF_ZIMAGE_PROF").is_ok_and(|v| v != "0") {
        let gpu: f64 = r.done.iter().map(|cmd| super::super::cmd_gpu_ms(cmd)).sum();
        eprintln!(
            "zimage metal vae: {:.3}s wall, {:.1} ms gpu ({}x{})",
            t0.elapsed().as_secs_f64(),
            gpu,
            hh,
            ww
        );
    }
    let v = unsafe { std::slice::from_raw_parts(rgb.contents() as *const f32, 3 * n) };
    if v.iter().any(|x| !x.is_finite()) {
        return super::decline("the VAE output is not finite");
    }
    out.copy_from_slice(v);
    true
}

fn env_chunk() -> usize {
    std::env::var("CMF_ZI_VAE_CHUNK")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(4096)
        .max(64)
        / 64
        * 64
}

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

    fn cpu_group_norm(
        x: &[f32],
        w: &[f32],
        b: &[f32],
        n: usize,
        ch: usize,
        groups: usize,
    ) -> Vec<f32> {
        let cg = ch / groups;
        let mut out = vec![0.0; x.len()];
        for g in 0..groups {
            let c0 = g * cg;
            let mut sum = 0.0;
            for px in 0..n {
                for c in c0..c0 + cg {
                    sum += x[px * ch + c];
                }
            }
            let mean = sum / (n * cg) as f32;
            let mut sumsq = 0.0;
            for px in 0..n {
                for c in c0..c0 + cg {
                    let d = x[px * ch + c] - mean;
                    sumsq += d * d;
                }
            }
            let inv = (sumsq / (n * cg) as f32 + 1e-6).sqrt().recip();
            for px in 0..n {
                for c in c0..c0 + cg {
                    out[px * ch + c] = (x[px * ch + c] - mean) * inv * w[c] + b[c];
                }
            }
        }
        out
    }

    /// Flux2's decoder uses 32 groups over 96 and 192 channels.  The
    /// compact group widths (3 and 6) are exactly the path that must not use
    /// a whole-float4 statistic.
    #[test]
    fn compact_group_norm_matches_cpu_reference() {
        let Some(c) = super::super::super::ctx() else {
            eprintln!("Metal unavailable; compact GroupNorm test skipped");
            return;
        };
        let Some(p) = pipes(c) else {
            eprintln!("Metal Z-Image pipelines unavailable; compact GroupNorm test skipped");
            return;
        };
        for &ch in &[96usize, 192] {
            let n = 19;
            let x: Vec<f32> = (0..n * ch)
                .map(|i| (((i * 17 % 97) as f32) - 48.0) / 17.0)
                .collect();
            let w: Vec<f32> = (0..ch).map(|i| 0.55 + (i % 13) as f32 * 0.03).collect();
            let b: Vec<f32> = (0..ch).map(|i| ((i % 11) as f32 - 5.0) * 0.02).collect();
            let xb = buf_from(c, &x);
            let yb = buf_zeroed(c, n * ch * 4);
            let norm = VNorm {
                w: buf_from(c, &w),
                b: buf_from(c, &b),
                groups: 32,
            };
            let mut rec = VRec {
                c,
                p,
                done: Vec::new(),
                cur: None,
                part: buf_zeroed(c, 64 * 256 * 4),
                stat: buf_zeroed(c, 64 * 2 * 4),
            };
            rec.gn(&norm, &xb, &yb, n, ch, false, false);
            assert!(
                rec.finish(),
                "Metal GroupNorm command failed for {ch} channels"
            );
            let got = unsafe { std::slice::from_raw_parts(yb.contents() as *const f32, n * ch) };
            let want = cpu_group_norm(&x, &w, &b, n, ch, 32);
            let max_abs = got
                .iter()
                .zip(&want)
                .map(|(a, b)| (a - b).abs())
                .fold(0.0f32, f32::max);
            assert!(
                max_abs < 3e-4,
                "{ch}-channel GroupNorm max abs error {max_abs}"
            );
        }
    }

    /// The 96-channel final Flux2 stage occupies two 64-row MMA tiles.  The
    /// second tile's padded rows must not write into the following NHWC
    /// pixel; compare every logical output against a half-input/half-weight
    /// CPU reference.
    #[test]
    fn conv_96_output_tail_matches_cpu_reference() {
        let Some(c) = super::super::super::ctx() else {
            eprintln!("Metal unavailable; 96-channel conv test skipped");
            return;
        };
        let Some(p) = pipes(c) else {
            eprintln!("Metal Z-Image pipelines unavailable; 96-channel conv test skipped");
            return;
        };
        let (n, ic, oc) = (5usize, 32usize, 96usize);
        let x: Vec<f32> = (0..n * ic)
            .map(|i| ((i * 11 % 31) as f32 - 15.0) * 0.02)
            .collect();
        let w: Vec<f32> = (0..oc * ic)
            .map(|i| ((i * 7 % 29) as f32 - 14.0) * 0.01)
            .collect();
        let b: Vec<f32> = (0..oc).map(|i| ((i % 9) as f32 - 4.0) * 0.015).collect();
        let cv = conv_of(
            c,
            &VaeConvRef {
                w: &w,
                b: &b,
                oc,
                ic,
                k: 1,
            },
            1.0,
        )
        .unwrap();
        let xb = half_buf(c, &x);
        let yb = buf_zeroed(c, n * oc * 4);
        let mut rec = VRec {
            c,
            p,
            done: Vec::new(),
            cur: None,
            part: buf_zeroed(c, 64 * 256 * 4),
            stat: buf_zeroed(c, 64 * 2 * 4),
        };
        rec.conv(&cv, (&xb, 0), (&yb, 0), None, n, (1, n), false, 0, oc, 0, 0);
        assert!(rec.finish(), "Metal 96-channel conv command failed");
        let got = unsafe { std::slice::from_raw_parts(yb.contents() as *const f32, n * oc) };
        let h = |v: f32| cortiq_core::quant::f16_to_f32(cortiq_core::quant::f32_to_f16(v));
        for px in 0..n {
            for out in 0..oc {
                let want = b[out]
                    + (0..ic)
                        .map(|inp| h(w[out * ic + inp]) * h(x[px * ic + inp]))
                        .sum::<f32>();
                let err = (got[px * oc + out] - want).abs();
                assert!(
                    err <= 2e-3 + want.abs() * 2e-3,
                    "pixel {px}, channel {out}: got {}, want {want}, abs error {err}",
                    got[px * oc + out],
                );
            }
        }
    }
}