inferencelayer 0.2.5

//! RT-DETRv2 backbone + CCFM on wgpu — the measured answer to "win the layout benchmark".
//!
//! The CPU profile (task #9) put the backbone at 52% and the hybrid encoder at 28% of the layout
//! forward, and the per-shape A/B proved torch-CPU's edge on this box is Apple AMX (implied
//! 2.5-2.7 TFLOP/s on the 1×1 GEMM rows — beyond NEON's ceiling by construction). The GPU is the
//! only unit here that outruns AMX, so exactly those two stages move: every conv becomes
//! `im2col_pack → enc_gemm3` (bias + relu/silu FUSED in the existing GEMM epilogue), pools /
//! upsample / concat / residual-add are five small WGSL kernels, and the whole backbone+CCFM runs
//! GPU-resident in ONE submit per section. AIFI (400 tokens) stays on the proven CPU path — a
//! 400 KB readback beats porting attention for a layer that costs ~2 ms. Query selection and the
//! deformable decoder stay CPU per the plan (300 queries).
//!
//! Weights upload FROM the already-built CPU model ([`RtDetr`]) — `PackedWeight::unpack` +
//! transpose-to-[k,n], the same from-cpu construction whisper_gpu uses. The CPU path is the
//! parity oracle: `tests/rtdetr_gpu_parity.rs` gates every map (cos) and the end-to-end
//! detections against it.

use anyhow::Result;

use crate::GpuCtx;
use crate::encoder::{GEMM3_TILES, act_code, enc_gemm3_src, gemm3_tier, gemm3_tile};
use crate::encoder_weights::Act as EncAct;
use crate::forward::{make_bg, pipeline, uni};
use crate::rtdetr::{Act, Bottleneck, ConvNorm, RtDetr, Shortcut, Stages, sincos_pos_embed};

/// One conv's GPU residency: `[k, n]`-transposed weight + bias + geometry + fused activation.
struct GpuConv {
    w: wgpu::Buffer,
    b: wgpu::Buffer,
    n: u32,
    k: u32,
    ic: u32,
    kh: u32,
    kw: u32,
    stride: u32,
    pad: u32,
    act: u32,
}

impl GpuConv {
    fn new(ctx: &GpuCtx, cn: &ConvNorm) -> Self {
        let (pw, bias, [ic, kh, kw, stride, pad]) = cn.conv.gpu_parts();
        let (n, k) = (pw.n(), pw.k());
        let flat = pw.unpack(); // [n, k]
        let mut wt = vec![0f32; n * k];
        for nn in 0..n {
            for kk in 0..k {
                wt[kk * n + nn] = flat[nn * k + kk];
            }
        }
        let act = match cn.act {
            Act::None => act_code(None),
            Act::Relu => act_code(Some(EncAct::Relu)),
            Act::Silu => act_code(Some(EncAct::Silu)),
        };
        Self {
            w: ctx.storage(&wt),
            b: ctx.storage(bias),
            n: n as u32,
            k: k as u32,
            ic: ic as u32,
            kh: kh as u32,
            kw: kw as u32,
            stride: stride as u32,
            pad: pad as u32,
            act,
        }
    }

    fn from_conv(ctx: &GpuCtx, conv: &crate::conv2d::Conv2d, act: u32) -> Self {
        let (pw, bias, [ic, kh, kw, stride, pad]) = conv.gpu_parts();
        let (n, k) = (pw.n(), pw.k());
        let flat = pw.unpack();
        let mut wt = vec![0f32; n * k];
        for nn in 0..n {
            for kk in 0..k {
                wt[kk * n + nn] = flat[nn * k + kk];
            }
        }
        Self {
            w: ctx.storage(&wt),
            b: ctx.storage(bias),
            n: n as u32,
            k: k as u32,
            ic: ic as u32,
            kh: kh as u32,
            kw: kw as u32,
            stride: stride as u32,
            pad: pad as u32,
            act,
        }
    }
}

/// im2col: one invocation per (output pixel, kernel tap) copies an `ic`-run. NHWC in, `[m, k]`
/// out with k = kh·kw·ic in (dy, dx, c) order — matching the CPU packing and the weight permute.
const IM2COL_SRC: &str = r#"
struct Meta { h: u32, w: u32, ic: u32, kh: u32, kw: u32, stride: u32, pad: u32, ow: u32 }
@group(0) @binding(0) var<storage, read>       x: array<f32>;
@group(0) @binding(1) var<storage, read_write> y: array<f32>;
@group(0) @binding(2) var<uniform>             mt: Meta;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    let taps = mt.kh * mt.kw;
    let total = mt.ow * taps;          // one output ROW handled per workgroup-y
    if (g.x >= total) { return; }
    let oy = g.y;
    let ox = g.x / taps;
    let tap = g.x % taps;
    let dy = tap / mt.kw;
    let dx = tap % mt.kw;
    let k = taps * mt.ic;
    let dst_base = (oy * mt.ow + ox) * k + tap * mt.ic;
    let iy = i32(oy * mt.stride + dy) - i32(mt.pad);
    let ix = i32(ox * mt.stride + dx) - i32(mt.pad);
    if (iy < 0 || iy >= i32(mt.h) || ix < 0 || ix >= i32(mt.w)) {
        for (var c = 0u; c < mt.ic; c++) { y[dst_base + c] = 0.0; }
        return;
    }
    let src_base = (u32(iy) * mt.w + u32(ix)) * mt.ic;
    for (var c = 0u; c < mt.ic; c++) { y[dst_base + c] = x[src_base + c]; }
}
"#;

/// 3×3/s2/p1 max-pool (the one geometry the model uses), NHWC, −∞ padding.
const MAXPOOL_SRC: &str = r#"
struct Meta { h: u32, w: u32, c: u32, oh: u32, ow: u32, pad0: u32, pad1: u32, pad2: u32 }
@group(0) @binding(0) var<storage, read>       x: array<f32>;
@group(0) @binding(1) var<storage, read_write> y: array<f32>;
@group(0) @binding(2) var<uniform>             mt: Meta;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    if (g.x >= mt.ow * mt.c) { return; }
    let ox = g.x / mt.c;
    let ch = g.x % mt.c;
    let oy = g.y;
    var best = -3.4e38;
    for (var dy = 0u; dy < 3u; dy++) {
        let iy = i32(oy * 2u + dy) - 1;
        if (iy < 0 || iy >= i32(mt.h)) { continue; }
        for (var dx = 0u; dx < 3u; dx++) {
            let ix = i32(ox * 2u + dx) - 1;
            if (ix < 0 || ix >= i32(mt.w)) { continue; }
            let v = x[(u32(iy) * mt.w + u32(ix)) * mt.c + ch];
            if (v > best) { best = v; }
        }
    }
    y[(oy * mt.ow + ox) * mt.c + ch] = best;
}
"#;

/// 2×2/s2 average pool (ResNet-D shortcut; even inputs only, asserted on the CPU twin).
const AVGPOOL_SRC: &str = r#"
struct Meta { h: u32, w: u32, c: u32, ow: u32 }
@group(0) @binding(0) var<storage, read>       x: array<f32>;
@group(0) @binding(1) var<storage, read_write> y: array<f32>;
@group(0) @binding(2) var<uniform>             mt: Meta;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    if (g.x >= mt.ow * mt.c) { return; }
    let ox = g.x / mt.c;
    let ch = g.x % mt.c;
    let oy = g.y;
    let base = (oy * 2u * mt.w + ox * 2u) * mt.c + ch;
    let s = x[base] + x[base + mt.c] + x[base + mt.w * mt.c] + x[base + mt.w * mt.c + mt.c];
    y[(oy * mt.ow + ox) * mt.c + ch] = s * 0.25;
}
"#;

/// Nearest-neighbour 2× upsample, NHWC.
const UPSAMPLE_SRC: &str = r#"
struct Meta { h: u32, w: u32, c: u32, ow: u32 }
@group(0) @binding(0) var<storage, read>       x: array<f32>;
@group(0) @binding(1) var<storage, read_write> y: array<f32>;
@group(0) @binding(2) var<uniform>             mt: Meta;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    if (g.x >= mt.ow * mt.c) { return; }
    let ox = g.x / mt.c;
    let ch = g.x % mt.c;
    let oy = g.y;
    y[(oy * mt.ow + ox) * mt.c + ch] = x[((oy / 2u) * mt.w + (ox / 2u)) * mt.c + ch];
}
"#;

/// Per-pixel channel concat of two NHWC maps.
const CONCAT_SRC: &str = r#"
struct Meta { hw: u32, ca: u32, cb: u32, pad: u32 }
@group(0) @binding(0) var<storage, read>       a: array<f32>;
@group(0) @binding(1) var<storage, read>       b: array<f32>;
@group(0) @binding(2) var<storage, read_write> y: array<f32>;
@group(0) @binding(3) var<uniform>             mt: Meta;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    let cc = mt.ca + mt.cb;
    if (g.x >= mt.hw * cc) { return; }
    let p = g.x / cc;
    let ch = g.x % cc;
    y[g.x] = select(b[p * mt.cb + (ch - mt.ca)], a[p * mt.ca + ch], ch < mt.ca);
}
"#;

/// `y = act(a + b)` — residual adds (relu), RepVGG sums (silu), CSP merge (none).
/// act codes: 0 none, 3 silu, 5 relu (matching the GEMM epilogue's table).
const ADDACT_SRC: &str = r#"
struct Meta { len: u32, act: u32, pad0: u32, pad1: u32 }
@group(0) @binding(0) var<storage, read>       a: array<f32>;
@group(0) @binding(1) var<storage, read>       b: array<f32>;
@group(0) @binding(2) var<storage, read_write> y: array<f32>;
@group(0) @binding(3) var<uniform>             mt: Meta;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    if (g.x >= mt.len) { return; }
    var v = a[g.x] + b[g.x];
    if (mt.act == 5u) { v = max(v, 0.0); }
    else if (mt.act == 3u) { v = v / (1.0 + exp(-v)); }
    y[g.x] = v;
}
"#;

struct GpuBottleneck {
    layer: [GpuConv; 3],
    shortcut_conv: Option<GpuConv>,
    shortcut_pool: bool,
}

struct GpuCsp {
    conv1: GpuConv,
    conv2: GpuConv,
    bottlenecks: Vec<(GpuConv, GpuConv)>, // (3×3, 1×1), silu after sum
}

pub struct RtDetrGpu {
    cpu: RtDetr,
    dec_input_proj: Vec<GpuConv>,
    /// per decoder layer: value_proj as ([k,n]-transposed weight, bias) — precomputed on-device.
    value_proj: Vec<(wgpu::Buffer, wgpu::Buffer)>,
    gemm3: Vec<wgpu::ComputePipeline>,
    im2col: wgpu::ComputePipeline,
    maxpool: wgpu::ComputePipeline,
    avgpool: wgpu::ComputePipeline,
    upsample: wgpu::ComputePipeline,
    concat: wgpu::ComputePipeline,
    addact: wgpu::ComputePipeline,
    stem: [GpuConv; 3],
    stages: Vec<Vec<GpuBottleneck>>,
    enc_input_proj: Vec<GpuConv>,
    lateral: Vec<GpuConv>,
    fpn: Vec<GpuCsp>,
    downsample: Vec<GpuConv>,
    pan: Vec<GpuCsp>,
}

/// A recorded pass list executed in one submit — the whisper_gpu batching pattern.
struct Recorder<'c> {
    ctx: &'c GpuCtx,
    passes: Vec<(*const wgpu::ComputePipeline, wgpu::BindGroup, u32, u32)>,
    keep: Vec<wgpu::Buffer>,
}

impl<'c> Recorder<'c> {
    fn new(ctx: &'c GpuCtx) -> Self {
        Self {
            ctx,
            passes: Vec::new(),
            keep: Vec::new(),
        }
    }

    fn push(&mut self, pl: &wgpu::ComputePipeline, bg: wgpu::BindGroup, gx: u32, gy: u32) {
        self.passes.push((pl as *const _, bg, gx, gy));
    }

    fn submit(self) {
        let mut enc = self
            .ctx
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        {
            let mut cpass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: None,
                timestamp_writes: None,
            });
            for (pl, bg, gx, gy) in &self.passes {
                // SAFETY: every pipeline outlives the recorder (owned by RtDetrGpu).
                cpass.set_pipeline(unsafe { &**pl });
                cpass.set_bind_group(0, bg, &[]);
                cpass.dispatch_workgroups(*gx, *gy, 1);
            }
        }
        self.ctx.queue.submit(Some(enc.finish()));
        drop(self.keep);
    }
}

impl RtDetrGpu {
    pub fn new(ctx: &GpuCtx, cpu: RtDetr) -> Result<Self> {
        let stem = [
            GpuConv::new(ctx, &cpu.backbone.stem[0]),
            GpuConv::new(ctx, &cpu.backbone.stem[1]),
            GpuConv::new(ctx, &cpu.backbone.stem[2]),
        ];
        let stages = cpu
            .backbone
            .stages
            .iter()
            .map(|blocks| {
                blocks
                    .iter()
                    .map(|b: &Bottleneck| {
                        let (sc, pool) = match &b.shortcut {
                            Shortcut::Identity => (None, false),
                            Shortcut::Conv(c) => (Some(GpuConv::new(ctx, c)), false),
                            Shortcut::PoolConv(c) => (Some(GpuConv::new(ctx, c)), true),
                        };
                        GpuBottleneck {
                            layer: [
                                GpuConv::new(ctx, &b.layer[0]),
                                GpuConv::new(ctx, &b.layer[1]),
                                GpuConv::new(ctx, &b.layer[2]),
                            ],
                            shortcut_conv: sc,
                            shortcut_pool: pool,
                        }
                    })
                    .collect()
            })
            .collect();
        let enc_input_proj = cpu
            .encoder
            .input_proj
            .iter()
            .map(|(c, ())| GpuConv::from_conv(ctx, c, act_code(None)))
            .collect();
        let dec_input_proj = cpu
            .dec_input_proj
            .iter()
            .map(|c| GpuConv::from_conv(ctx, c, act_code(None)))
            .collect();
        let value_proj = cpu
            .decoder
            .layers
            .iter()
            .map(|l| {
                let (n, k) = (l.value_proj.w.n(), l.value_proj.w.k());
                let flat = l.value_proj.w.unpack();
                let mut wt = vec![0f32; n * k];
                for nn in 0..n {
                    for kk in 0..k {
                        wt[kk * n + nn] = flat[nn * k + kk];
                    }
                }
                (ctx.storage(&wt), ctx.storage(&l.value_proj.b))
            })
            .collect();
        let csp = |c: &crate::rtdetr::Csp| GpuCsp {
            conv1: GpuConv::new(ctx, &c.conv1),
            conv2: GpuConv::new(ctx, &c.conv2),
            bottlenecks: c
                .bottlenecks
                .iter()
                .map(|r| (GpuConv::new(ctx, &r.conv1), GpuConv::new(ctx, &r.conv2)))
                .collect(),
        };
        Ok(Self {
            gemm3: GEMM3_TILES
                .iter()
                .map(|&(bm, bn, bk)| {
                    pipeline(ctx, "rtdetr_gemm3", &enc_gemm3_src(false, bm, bn, bk))
                })
                .collect(),
            im2col: pipeline(ctx, "rtdetr_im2col", IM2COL_SRC),
            maxpool: pipeline(ctx, "rtdetr_maxpool", MAXPOOL_SRC),
            avgpool: pipeline(ctx, "rtdetr_avgpool", AVGPOOL_SRC),
            upsample: pipeline(ctx, "rtdetr_upsample", UPSAMPLE_SRC),
            concat: pipeline(ctx, "rtdetr_concat", CONCAT_SRC),
            addact: pipeline(ctx, "rtdetr_addact", ADDACT_SRC),
            lateral: cpu
                .encoder
                .lateral
                .iter()
                .map(|c| GpuConv::new(ctx, c))
                .collect(),
            fpn: cpu.encoder.fpn.iter().map(csp).collect(),
            downsample: cpu
                .encoder
                .downsample
                .iter()
                .map(|c| GpuConv::new(ctx, c))
                .collect(),
            pan: cpu.encoder.pan.iter().map(csp).collect(),
            stem,
            stages,
            enc_input_proj,
            dec_input_proj,
            value_proj,
            cpu,
        })
    }

    pub fn cpu(&self) -> &RtDetr {
        &self.cpu
    }

    /// conv = (im2col if k>1 or strided) → gemm3 with fused bias+act. Returns (out, oh, ow).
    #[allow(clippy::too_many_arguments)]
    fn conv<'a>(
        &'a self,
        rec: &mut Recorder<'a>,
        gc: &GpuConv,
        x: &wgpu::Buffer,
        h: u32,
        w: u32,
    ) -> (wgpu::Buffer, u32, u32) {
        let ctx = rec.ctx;
        let oh = (h + 2 * gc.pad - gc.kh) / gc.stride + 1;
        let ow = (w + 2 * gc.pad - gc.kw) / gc.stride + 1;
        let m = (oh * ow) as usize;
        let y = ctx.empty(m * gc.n as usize);

        let gemm_in: wgpu::Buffer;
        if gc.kh == 1 && gc.kw == 1 && gc.stride == 1 && gc.pad == 0 {
            // 1×1 s1: the NHWC map IS the [m, k] GEMM input — no packing pass at all.
            gemm_in = x.clone();
        } else {
            let cols = ctx.empty(m * gc.k as usize);
            let meta = uni(
                ctx,
                bytemuck::cast_slice(&[h, w, gc.ic, gc.kh, gc.kw, gc.stride, gc.pad, ow]),
            );
            let bg = make_bg(ctx, &self.im2col, &[x, &cols], &meta);
            let taps = gc.kh * gc.kw;
            rec.push(&self.im2col, bg, (ow * taps).div_ceil(256), oh);
            rec.keep.push(meta);
            gemm_in = cols;
        }

        let flags = 1u32 | (gc.act << 8);
        let meta = uni(ctx, bytemuck::cast_slice(&[m as u32, gc.n, gc.k, flags]));
        let tile = gemm3_tile(m, gc.n as usize);
        let pl = &self.gemm3[gemm3_tier(tile)];
        let bg = make_bg(ctx, pl, &[&gemm_in, &gc.w, &gc.b, &y], &meta);
        rec.push(
            pl,
            bg,
            gc.n.div_ceil(tile.1 as u32),
            (m as u32).div_ceil(tile.0 as u32),
        );
        rec.keep.push(meta);
        rec.keep.push(gemm_in);
        (y, oh, ow)
    }

    fn addact<'a>(
        &'a self,
        rec: &mut Recorder<'a>,
        a: &wgpu::Buffer,
        b: &wgpu::Buffer,
        len: usize,
        act: u32,
    ) -> wgpu::Buffer {
        let y = rec.ctx.empty(len);
        let meta = uni(rec.ctx, bytemuck::cast_slice(&[len as u32, act, 0, 0]));
        let bg = make_bg(rec.ctx, &self.addact, &[a, b, &y], &meta);
        rec.push(&self.addact, bg, (len as u32).div_ceil(256), 1);
        rec.keep.push(meta);
        y
    }

    /// Full GPU backbone + CCFM; AIFI runs on the CPU between two GPU sections. Returns the three
    /// fused encoder maps on the CPU, ready for `RtDetr::stages_from_maps`.
    pub fn encoder_maps(
        &self,
        ctx: &GpuCtx,
        pixel_values_nhwc: &[f32],
    ) -> Result<Vec<(Vec<f32>, usize, usize)>> {
        // ── section 1: stem → stages → input_proj (one submit) ────────────────────────────────
        let mut rec = Recorder::new(ctx);
        let x0 = ctx.storage(pixel_values_nhwc);
        let (x1, h1, w1) = self.conv(&mut rec, &self.stem[0], &x0, 640, 640);
        let (x2, h2, w2) = self.conv(&mut rec, &self.stem[1], &x1, h1, w1);
        let (x3, h3, w3) = self.conv(&mut rec, &self.stem[2], &x2, h2, w2);
        // maxpool 3/2/1
        let (mh, mw) = ((h3 + 2 - 3) / 2 + 1, (w3 + 2 - 3) / 2 + 1);
        let pooled = ctx.empty((mh * mw * 64) as usize);
        {
            let meta = uni(ctx, bytemuck::cast_slice(&[h3, w3, 64u32, mh, mw, 0, 0, 0]));
            let bg = make_bg(ctx, &self.maxpool, &[&x3, &pooled], &meta);
            rec.push(&self.maxpool, bg, (mw * 64).div_ceil(256), mh);
            rec.keep.push(meta);
        }

        let (mut x, mut h, mut w) = (pooled, mh, mw);
        let mut backbone_maps: Vec<(wgpu::Buffer, u32, u32, u32)> = Vec::new(); // (buf, h, w, c)
        for (si, blocks) in self.stages.iter().enumerate() {
            for b in blocks {
                let (a1, ah, aw) = self.conv(&mut rec, &b.layer[0], &x, h, w);
                let (a2, bh, bw) = self.conv(&mut rec, &b.layer[1], &a1, ah, aw);
                let (a3, ch, cw) = self.conv(&mut rec, &b.layer[2], &a2, bh, bw);
                let res: wgpu::Buffer = match (&b.shortcut_conv, b.shortcut_pool) {
                    (None, _) => x.clone(),
                    (Some(sc), false) => self.conv(&mut rec, sc, &x, h, w).0,
                    (Some(sc), true) => {
                        let (ph, pw) = (h / 2, w / 2);
                        let ic = sc.ic;
                        let p = ctx.empty((ph * pw * ic) as usize);
                        let meta = uni(ctx, bytemuck::cast_slice(&[h, w, ic, pw]));
                        let bg = make_bg(ctx, &self.avgpool, &[&x, &p], &meta);
                        rec.push(&self.avgpool, bg, (pw * ic).div_ceil(256), ph);
                        rec.keep.push(meta);
                        self.conv(&mut rec, sc, &p, ph, pw).0
                    }
                };
                let len = (ch * cw) as usize * b.layer[2].n as usize;
                let summed = self.addact(&mut rec, &a3, &res, len, 5); // relu
                rec.keep.push(a1);
                rec.keep.push(a2);
                rec.keep.push(a3);
                rec.keep.push(res);
                rec.keep.push(x);
                x = summed;
                h = ch;
                w = cw;
            }
            if si >= 1 {
                backbone_maps.push((x.clone(), h, w, self.stages[si].last().unwrap().layer[2].n));
            }
        }

        // input_proj to 256 per map
        let mut proj: Vec<(wgpu::Buffer, u32, u32)> = Vec::new();
        for (i, (m, mh2, mw2, _)) in backbone_maps.iter().enumerate() {
            let (p, ..) = self.conv(&mut rec, &self.enc_input_proj[i], m, *mh2, *mw2);
            proj.push((p, *mh2, *mw2));
        }
        rec.submit();

        // ── AIFI on CPU (stride-32 map is 400 tokens — readback + proven path + upload) ───────
        let (m2, h2s, w2s) = &proj[2];
        let map2 = ctx.read(m2, (h2s * w2s * 256) as usize)?;
        let pos = sincos_pos_embed(*w2s as usize, *h2s as usize);
        let aifi_out = self
            .cpu
            .encoder
            .aifi
            .forward(&map2, &pos, (h2s * w2s) as usize);

        // ── section 2: CCFM (one submit) ───────────────────────────────────────────────────────
        let mut rec = Recorder::new(ctx);
        let maps: Vec<(wgpu::Buffer, u32, u32)> = vec![
            (proj[0].0.clone(), proj[0].1, proj[0].2),
            (proj[1].0.clone(), proj[1].1, proj[1].2),
            (ctx.storage(&aifi_out), *h2s, *w2s),
        ];

        // FPN top-down; fpn_maps[last] replaced by its lateral conv each iteration
        let mut fpn_maps: Vec<(wgpu::Buffer, u32, u32)> = vec![maps[2].clone()];
        for idx in 0..2 {
            let (bf, bh, bw) = &maps[1 - idx];
            let (top, th, tw) = fpn_maps.last().unwrap().clone();
            let (lat, ..) = self.conv(&mut rec, &self.lateral[idx], &top, th, tw);
            *fpn_maps.last_mut().unwrap() = (lat.clone(), th, tw);
            // upsample 2×
            let (uh, uw) = (th * 2, tw * 2);
            let up = ctx.empty((uh * uw * 256) as usize);
            {
                let meta = uni(ctx, bytemuck::cast_slice(&[th, tw, 256u32, uw]));
                let bg = make_bg(ctx, &self.upsample, &[&lat, &up], &meta);
                rec.push(&self.upsample, bg, (uw * 256).div_ceil(256), uh);
                rec.keep.push(meta);
            }
            // concat [up, backbone]
            let fused = ctx.empty((uh * uw * 512) as usize);
            {
                let meta = uni(ctx, bytemuck::cast_slice(&[uh * uw, 256u32, 256u32, 0]));
                let bg = make_bg(ctx, &self.concat, &[&up, bf, &fused], &meta);
                rec.push(&self.concat, bg, (uh * uw * 512).div_ceil(256), 1);
                rec.keep.push(meta);
            }
            let o = self.csp(&mut rec, &self.fpn[idx], &fused, uh, uw);
            rec.keep.push(up);
            rec.keep.push(fused);
            rec.keep.push(top);
            fpn_maps.push((o, uh, uw));
        }
        fpn_maps.reverse();

        // PAN bottom-up
        let mut pan_maps: Vec<(wgpu::Buffer, u32, u32)> = vec![fpn_maps[0].clone()];
        for idx in 0..2 {
            let (top, th, tw) = pan_maps.last().unwrap().clone();
            let (down, dh, dw) = self.conv(&mut rec, &self.downsample[idx], &top, th, tw);
            let (ff, ..) = &fpn_maps[idx + 1];
            let fused = ctx.empty((dh * dw * 512) as usize);
            {
                let meta = uni(ctx, bytemuck::cast_slice(&[dh * dw, 256u32, 256u32, 0]));
                let bg = make_bg(ctx, &self.concat, &[&down, ff, &fused], &meta);
                rec.push(&self.concat, bg, (dh * dw * 512).div_ceil(256), 1);
                rec.keep.push(meta);
            }
            let o = self.csp(&mut rec, &self.pan[idx], &fused, dh, dw);
            rec.keep.push(down);
            rec.keep.push(fused);
            pan_maps.push((o, dh, dw));
        }
        rec.submit();

        // readback the three fused maps
        let mut out = Vec::with_capacity(3);
        for (buf, mh3, mw3) in &pan_maps {
            let v = ctx.read(buf, (mh3 * mw3 * 256) as usize)?;
            out.push((v, *mh3 as usize, *mw3 as usize));
        }
        Ok(out)
    }

    fn csp<'a>(
        &'a self,
        rec: &mut Recorder<'a>,
        c: &'a GpuCsp,
        x: &wgpu::Buffer,
        h: u32,
        w: u32,
    ) -> wgpu::Buffer {
        let (mut a, ..) = self.conv(rec, &c.conv1, x, h, w);
        for (r3, r1) in &c.bottlenecks {
            let (p, ..) = self.conv(rec, r3, &a, h, w);
            let (q, ..) = self.conv(rec, r1, &a, h, w);
            let len = (h * w) as usize * r3.n as usize;
            let s = self.addact(rec, &p, &q, len, 3); // silu after the RepVGG sum
            rec.keep.push(p);
            rec.keep.push(q);
            rec.keep.push(a);
            a = s;
        }
        let (b, ..) = self.conv(rec, &c.conv2, x, h, w);
        let len = (h * w * 256) as usize;
        let out = self.addact(rec, &a, &b, len, 0); // conv3 = Identity: plain add
        rec.keep.push(a);
        rec.keep.push(b);
        out
    }

    /// GPU backbone/encoder + shared CPU tail. The GPU twin of `RtDetr::forward_stages`.
    pub fn forward_stages(&self, ctx: &GpuCtx, pixel_values: &[f32]) -> Result<Stages> {
        let maps = self.encoder_maps(ctx, pixel_values)?;

        // ── section 3: dec_input_proj + every decoder layer's value_proj, on-device ───────────
        // The CPU tail's dominant cost was these 6×[8400,256]·[256,256] GEMMs (measured — the
        // qs+decoder lap). Re-uploading the fused maps costs ~1 ms; the GEMMs leave with them.
        let mut rec = Recorder::new(ctx);
        let mut projs: Vec<(wgpu::Buffer, u32, u32)> = Vec::new();
        for (i, (m, mh, mw)) in maps.iter().enumerate() {
            let x = ctx.storage(m);
            let (p, ph, pw) = self.conv(
                &mut rec,
                &self.dec_input_proj[i],
                &x,
                *mh as u32,
                *mw as u32,
            );
            rec.keep.push(x);
            projs.push((p, ph, pw));
        }
        // per (layer, level) value GEMM — row-independent, so per-level GEMMs ≡ one concat GEMM
        let mut value_bufs: Vec<Vec<(wgpu::Buffer, usize)>> = Vec::new();
        for (wb, bb) in &self.value_proj {
            let mut per_level = Vec::new();
            for (p, ph, pw) in &projs {
                let m = (ph * pw) as usize;
                let y = ctx.empty(m * 256);
                let meta = uni(ctx, bytemuck::cast_slice(&[m as u32, 256u32, 256u32, 1u32]));
                let tile = gemm3_tile(m, 256);
                let pl = &self.gemm3[gemm3_tier(tile)];
                let bg = make_bg(ctx, pl, &[p, wb, bb, &y], &meta);
                rec.push(
                    pl,
                    bg,
                    256u32.div_ceil(tile.1 as u32),
                    (m as u32).div_ceil(tile.0 as u32),
                );
                rec.keep.push(meta);
                per_level.push((y, m));
            }
            value_bufs.push(per_level);
        }
        rec.submit();

        // readbacks: stitch flattened memory + per-layer values
        let mut memory: Vec<f32> = Vec::new();
        let mut shapes: Vec<(usize, usize)> = Vec::new();
        for (p, ph, pw) in &projs {
            memory.extend_from_slice(&ctx.read(p, (ph * pw * 256) as usize)?);
            shapes.push((*ph as usize, *pw as usize));
        }
        let mut values: Vec<Vec<f32>> = Vec::with_capacity(value_bufs.len());
        for per_level in &value_bufs {
            let mut v: Vec<f32> = Vec::with_capacity(memory.len());
            for (buf, m) in per_level {
                v.extend_from_slice(&ctx.read(buf, m * 256)?);
            }
            values.push(v);
        }

        // backbone maps aren't re-read from the GPU (nothing downstream consumes them on CPU);
        // hand empty placeholders so Stages keeps its shape for the parity test's encoder gates.
        Ok(self
            .cpu
            .stages_from_projected(Vec::new(), maps, memory, shapes, Some(values)))
    }

    /// End-to-end detections, GPU-accelerated — the drop-in for `RtDetr::detect`.
    pub fn detect(
        &self,
        ctx: &GpuCtx,
        rgb: &[u8],
        w: usize,
        h: usize,
    ) -> Result<Vec<crate::rtdetr::Detection>> {
        let resized = crate::vision::resize_rgb8_bilinear(rgb, w, h, 640, 640);
        let pixels: Vec<f32> = resized.iter().map(|&b| b as f32 / 255.0).collect();
        let stages = self.forward_stages(ctx, &pixels)?;
        Ok(crate::rtdetr::postprocess(
            stages.dec_logits.last().unwrap(),
            stages.dec_refs.last().unwrap(),
            w as f32,
            h as f32,
        ))
    }
}