inferencelayer 0.2.1

//! GLiNER's span head on **wgpu** — the piece that makes the engine faster than torch at every
//! size, not just on short text.
//!
//! The head's cost grows with the SPAN COUNT (`n_spans ≈ max_width · W`), so on a 176-word input it
//! is ~2,300 spans and ~3.6 GMAC — and on the CPU the `gemm` crate delivers ~51 GFLOP/s against
//! torch dispatching to Apple's AMX-backed BLAS. That is a fight the CPU cannot win. On the GPU the
//! same arithmetic is a handful of dispatches.
//!
//! Everything here rides the encoder's existing WGSL GEMM (`enc_gemm_src`) — `y = act(x·Wᵀ + b)` —
//! except one new kernel, [`SPAN_ASSEMBLE`], which materializes the per-span rows. So this is
//! wgpu, i.e. Metal / Vulkan / DX12 / GL, not a Mac-specific path.
//!
//! ```text
//!   words [W, d]                                   (uploaded once per call)
//!     ├─ GEMM ─ GEMM(relu) ──────────► relu(start) [W, d]
//!     ├─ GEMM ─ GEMM(relu) ──────────► relu(end)   [W, d]
//!     ├─ GEMM  A = W_left ·relu(start) [W, 4d]      ┐ the concat-seam split: each half is
//!     ├─ GEMM  B = W_right·relu(end)   [W, 4d]      ┘ projected PER WORD, not per span
//!     ├─ SPAN_ASSEMBLE  h[s] = relu(A[l] + B[l+k])  [n_spans, 4d]
//!     ├─ GEMM  span_reps = h · W_downᵀ + b          [n_spans, d]
//!     └─ GEMM  scores    = span_reps · promptᵀ      [n_spans, C]  ← the only readback (tiny)
//! ```

use anyhow::{Context, Result};

use crate::GpuCtx;
use crate::encoder::{
    GEMM2_TILES, GEMM3_TILES, act_code, enc_gemm2_src, enc_gemm3_src, gemm2_tier, gemm2_tile,
    gemm3_tier, gemm3_tile,
};
use crate::encoder_weights::Act;
use crate::forward::{make_bg, pipeline, uni};

/// `h[s] = relu(A[l] + B[min(l+k, W-1)])` for span `s = l·max_width + k`.
///
/// The span index is derived IN-SHADER from the workgroup id — no `[n_spans, 2]` index buffer to
/// build and upload. Spans that run past the text (`l + k >= W`) are computed against a clamped
/// row and simply ignored by the caller, which is cheaper than branching around them and matches
/// GLiNER's own `span_mask`.
const SPAN_ASSEMBLE: &str = r#"
struct Meta { n_spans: u32, up_w: u32, max_width: u32, n_words: 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> h: array<f32>;
@group(0) @binding(3) var<uniform> mt: Meta;
@compute @workgroup_size(64)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) t: u32) {
    let s = wg.x;
    if (s >= mt.n_spans) { return; }
    let l = s / mt.max_width;
    let k = s % mt.max_width;
    let e = min(l + k, mt.n_words - 1u);
    let ao = l * mt.up_w;
    let bo = e * mt.up_w;
    let ho = s * mt.up_w;
    for (var j = t; j < mt.up_w; j += 64u) {
        h[ho + j] = max(a[ao + j] + b[bo + j], 0.0);
    }
}
"#;

/// One `nn.Linear` resident on the GPU (`W` as f32, plus a bias — the head is small enough that
/// f16 buys nothing and costs accuracy).
struct GpuLinear {
    w: wgpu::Buffer,
    b: wgpu::Buffer,
    n: u32,
    k: u32,
    /// Uploaded `[k, n]`-transposed for the v3 kernels (per-weight choice — see `new_v3`).
    v3: bool,
}

impl GpuLinear {
    fn new(ctx: &GpuCtx, w: &[f32], b: &[f32], n: usize, k: usize) -> Self {
        Self {
            w: ctx.storage(w),
            b: ctx.storage(b),
            n: n as u32,
            k: k as u32,
            v3: false,
        }
    }

    /// v3 upload: `[k, n]`-transposed at load, exactly like the encoder's `upload_linear`. The
    /// choice is PER WEIGHT because the head's `m` differs per GEMM: the n=2048 up-projections
    /// sit in v3's proven wide band at every m, and the span `down` runs at m = n_spans ≥ 132
    /// where plain v3 wins the mid band (the deberta-shape measurements) — while `ps_down`/
    /// `pe_down` run at m = n_words (as low as 11) and the runtime-weight score GEMM cannot
    /// pre-transpose, so those stay v2.
    fn new_v3(ctx: &GpuCtx, w: &[f32], b: &[f32], n: usize, k: usize) -> Self {
        let mut wt = vec![0f32; n * k];
        for nn in 0..n {
            for kk in 0..k {
                wt[kk * n + nn] = w[nn * k + kk];
            }
        }
        Self {
            w: ctx.storage(&wt),
            b: ctx.storage(b),
            n: n as u32,
            k: k as u32,
            v3: true,
        }
    }
}

/// The span head, resident on the GPU. Weights upload once at load; per call only `words` goes up
/// and `scores` comes back.
pub(crate) struct GlinerGpuHead {
    /// GEMM v2 [wide, skinny] — selected per call by the row count.
    gemm: Vec<wgpu::ComputePipeline>,
    /// GEMM v3 (vec4, `[k,n]`-transposed weights) — the up-projections and the span down ride
    /// these; tile picked per call like the encoder's.
    gemm3: Vec<wgpu::ComputePipeline>,
    assemble: wgpu::ComputePipeline,

    ps_up: GpuLinear,
    ps_down: GpuLinear,
    pe_up: GpuLinear,
    pe_down: GpuLinear,
    up_left: GpuLinear,
    up_right: GpuLinear,
    down: GpuLinear,

    /// Scratch, sized for `max_words` so a call allocates nothing.
    words: wgpu::Buffer,
    mid_a: wgpu::Buffer,
    mid_b: wgpu::Buffer,
    start: wgpu::Buffer,
    end: wgpu::Buffer,
    a: wgpu::Buffer,
    b: wgpu::Buffer,
    h: wgpu::Buffer,
    span_reps: wgpu::Buffer,
    scores: wgpu::Buffer,
    /// The entity-type reps, uploaded per call (they depend on the prompt) into a fixed buffer.
    prompt: wgpu::Buffer,
    /// The score GEMM has no bias; a zero vector bound in its place.
    no_bias: wgpu::Buffer,

    d: usize,
    up_w: usize,
    max_width: usize,
    max_words: usize,
    max_types: usize,
}

impl GlinerGpuHead {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        ctx: &GpuCtx,
        ps: (&[f32], &[f32], &[f32], &[f32]), // project_start: up.w, up.b, down.w, down.b
        pe: (&[f32], &[f32], &[f32], &[f32]), // project_end
        up_left: (&[f32], &[f32]),            // span_up.left  (w, b=zeros)
        up_right: (&[f32], &[f32]),           // span_up.right (w, b=b_up)
        down: (&[f32], &[f32]),               // span_down
        d: usize,
        up_w: usize,
        max_width: usize,
        max_words: usize,
        max_types: usize,
    ) -> Result<Self> {
        let n_spans = max_words * max_width;
        let zero = |n: usize| ctx.storage(&vec![0f32; n]);
        Ok(Self {
            // The head's own GEMMs ride the register-blocked v2 family, like the encoder's.
            // One pipeline per tile: the span GEMMs are M≈2300 and the per-word ones M=W, and the
            // right tile is a function of the GRID (M *and* the output width), not of M alone.
            gemm: GEMM2_TILES
                .iter()
                .map(|&(bm, bn)| {
                    pipeline(
                        ctx,
                        &format!("gliner_gemm2_{bm}x{bn}"),
                        &enc_gemm2_src(false, bm, bn),
                    )
                })
                .collect(),
            gemm3: GEMM3_TILES
                .iter()
                .map(|&(bm, bn, bk)| {
                    pipeline(
                        ctx,
                        &format!("gliner_gemm3_{bm}x{bn}x{bk}"),
                        &enc_gemm3_src(false, bm, bn, bk),
                    )
                })
                .collect(),
            assemble: pipeline(ctx, "gliner_span_assemble", SPAN_ASSEMBLE),

            ps_up: GpuLinear::new_v3(ctx, ps.0, ps.1, up_w, d),
            ps_down: GpuLinear::new(ctx, ps.2, ps.3, d, up_w),
            pe_up: GpuLinear::new_v3(ctx, pe.0, pe.1, up_w, d),
            pe_down: GpuLinear::new(ctx, pe.2, pe.3, d, up_w),
            up_left: GpuLinear::new_v3(ctx, up_left.0, up_left.1, up_w, d),
            up_right: GpuLinear::new_v3(ctx, up_right.0, up_right.1, up_w, d),
            down: GpuLinear::new_v3(ctx, down.0, down.1, d, up_w),

            words: zero(max_words * d),
            mid_a: zero(max_words * up_w),
            mid_b: zero(max_words * up_w),
            start: zero(max_words * d),
            end: zero(max_words * d),
            a: zero(max_words * up_w),
            b: zero(max_words * up_w),
            h: zero(n_spans * up_w),
            span_reps: zero(n_spans * d),
            scores: zero(n_spans * max_types),
            prompt: zero(max_types * d),
            no_bias: zero(max_types),

            d,
            up_w,
            max_width,
            max_words,
            max_types,
        })
    }

    /// `words` `[W, d]` (post-BiLSTM) and `prompt` `[C, d]` → raw scores `[n_spans, C]`, with
    /// `n_spans = W · max_width` and span `s = l·max_width + k` meaning the word range `[l, l+k]`.
    /// The whole head runs in ONE submit; only the (small) score matrix comes back.
    pub(crate) fn scores(
        &self,
        ctx: &GpuCtx,
        words: &[f32],
        prompt: &[f32],
        n_words: usize,
        n_types: usize,
    ) -> Result<Vec<f32>> {
        anyhow::ensure!(
            n_words <= self.max_words,
            "gliner gpu head: {n_words} words exceeds the {} it was sized for",
            self.max_words
        );
        anyhow::ensure!(
            n_types <= self.max_types,
            "gliner gpu head: {n_types} entity types exceeds the {} it was sized for",
            self.max_types
        );
        let (d, up_w) = (self.d, self.up_w);
        let n_spans = n_words * self.max_width;

        ctx.queue
            .write_buffer(&self.words, 0, bytemuck::cast_slice(words));
        // `prompt` is the RHS of the score GEMM, so it plays the role of a weight matrix [C, d].
        ctx.queue
            .write_buffer(&self.prompt, 0, bytemuck::cast_slice(prompt));

        // Bind groups + uniforms for the whole head, built up front, then dispatched in ONE compute
        // pass. Within a pass WebGPU orders dispatches and makes their writes visible to the next,
        // so the chain is safe — and a pass per GEMM (8 of them) is pure per-call overhead. This is
        // the same plan-replay shape `EncoderGpu::encode` uses.
        let mut passes: Vec<(&wgpu::ComputePipeline, wgpu::BindGroup, u32, u32)> = Vec::new();
        let mut keep: Vec<wgpu::Buffer> = Vec::new(); // uniforms must outlive the submit

        // `gemm!` rather than a closure: it needs &mut on both `passes` and `keep`, and a closure
        // capturing them would hold those borrows across every call.
        macro_rules! gemm {
            ($x:expr, $lw:expr, $y:expr, $m:expr, $act:expr, $bias:expr) => {{
                let lw: &GpuLinear = $lw;
                let m: usize = $m;
                let flags = u32::from($bias as bool) | (act_code($act) << 8);
                let meta = uni(ctx, bytemuck::cast_slice(&[m as u32, lw.n, lw.k, flags]));
                let (pl, bm, bn) = if lw.v3 {
                    let tile = gemm3_tile(m, lw.n as usize);
                    (&self.gemm3[gemm3_tier(tile)], tile.0, tile.1)
                } else {
                    let tile = gemm2_tile(m, lw.n as usize);
                    (&self.gemm[gemm2_tier(tile)], tile.0, tile.1)
                };
                let bg = make_bg(ctx, pl, &[$x, &lw.w, &lw.b, $y], &meta);
                passes.push((
                    pl,
                    bg,
                    lw.n.div_ceil(bn as u32),
                    (m as u32).div_ceil(bm as u32),
                ));
                keep.push(meta);
            }};
        }

        // project_start / project_end: Linear → ReLU → Linear (GLiNER's create_projection_layer),
        // activations folded into the GEMM epilogue.
        //
        // The SECOND ReLU is SpanMarkerV0's, applied to the concatenation `relu(start ‖ end)`.
        // ReLU is elementwise, so `relu(start ‖ end) == relu(start) ‖ relu(end)` — which is what
        // lets the concat seam be split, and it means `start`/`end` leave these GEMMs already
        // rectified. No extra pass.
        let r = Some(Act::Relu);
        gemm!(&self.words, &self.ps_up, &self.mid_a, n_words, r, true);
        gemm!(&self.mid_a, &self.ps_down, &self.start, n_words, r, true);
        gemm!(&self.words, &self.pe_up, &self.mid_b, n_words, r, true);
        gemm!(&self.mid_b, &self.pe_down, &self.end, n_words, r, true);

        // A = W_left·relu(start), B = W_right·relu(end) + b_up — the concat-seam split.
        gemm!(&self.start, &self.up_left, &self.a, n_words, None, false);
        gemm!(&self.end, &self.up_right, &self.b, n_words, None, true);

        // h[s] = relu(A[l] + B[l+k])
        let meta = uni(
            ctx,
            bytemuck::cast_slice(&[
                n_spans as u32,
                up_w as u32,
                self.max_width as u32,
                n_words as u32,
            ]),
        );
        let bg = make_bg(ctx, &self.assemble, &[&self.a, &self.b, &self.h], &meta);
        passes.push((&self.assemble, bg, n_spans as u32, 1));
        keep.push(meta);

        // span_reps = h · W_downᵀ + b   →   scores = span_reps · promptᵀ
        gemm!(&self.h, &self.down, &self.span_reps, n_spans, None, true);
        let score_w = GpuLinear {
            w: self.prompt.clone(),
            b: self.no_bias.clone(),
            n: n_types as u32,
            k: d as u32,
            v3: false, // per-call upload in [C, d] — no load-time transpose to ride
        };
        gemm!(
            &self.span_reps,
            &score_w,
            &self.scores,
            n_spans,
            None,
            false
        );

        let mut enc = ctx
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        {
            let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
            for (pl, bg, gx, gy) in &passes {
                pass.set_pipeline(pl);
                pass.set_bind_group(0, bg, &[]);
                pass.dispatch_workgroups(*gx, *gy, 1);
            }
        }
        ctx.queue.submit([enc.finish()]);
        let out = ctx
            .read(&self.scores, n_spans * n_types)
            .context("gliner gpu head readback")?;
        drop(keep);
        Ok(out)
    }
}