rlx-coreml 0.2.14

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0
// Quantized-weight ops, verified on-device. GGUF weights are stored
// `[N, K]` (B-transposed); the backend host-dequantizes them to f32 and
// matmuls with transpose_y. We quantize a known f32 weight, run through
// CoreML, and compare to the full-precision matmul within Q8_0 tolerance.
#![cfg(any(target_os = "macos", target_os = "ios"))]

use rlx_coreml::CoremlExecutable;
use rlx_ir::quant::QuantScheme;
use rlx_ir::{DType, Graph, Op, Shape};

fn approx(a: &[f32], b: &[f32], tol: f32) {
    assert_eq!(a.len(), b.len(), "len {} vs {}", a.len(), b.len());
    let mx = a
        .iter()
        .zip(b)
        .map(|(x, y)| (x - y).abs())
        .fold(0.0f32, f32::max);
    assert!(
        mx <= tol,
        "max abs diff {mx} > {tol}\n got {a:?}\n ref {b:?}"
    );
}

// row-major [M,K] @ [K,N]
fn matmul(x: &[f32], w: &[f32], m: usize, k: usize, n: usize) -> Vec<f32> {
    let mut o = vec![0.0f32; m * n];
    for i in 0..m {
        for j in 0..n {
            let mut acc = 0.0;
            for kk in 0..k {
                acc += x[i * k + kk] * w[kk * n + j];
            }
            o[i * n + j] = acc;
        }
    }
    o
}

#[test]
fn dequant_matmul_q8_0() {
    // M=2, K=64 (Q8_0 block = 32, so K divisible), N=3.
    let (m, k, n) = (2usize, 64usize, 3usize);

    // Reference weight in [K,N] f32; GGUF stores it transposed [N,K].
    let w_kn: Vec<f32> = (0..k * n)
        .map(|i| ((i as f32) * 0.013).sin() * 0.5)
        .collect();
    let mut w_nk = vec![0.0f32; n * k]; // [N,K] row-major
    for kk in 0..k {
        for j in 0..n {
            w_nk[j * k + kk] = w_kn[kk * n + j];
        }
    }
    // Quantize the [N,K] weight (each of the N rows is a length-K vector).
    let packed = rlx_gguf::quantize::quantize_q8_0(&w_nk).expect("quantize");

    let x: Vec<f32> = (0..m * k).map(|i| ((i as f32) * 0.02).cos()).collect();

    let mut g = Graph::new("dqmm");
    let xi = g.input("x", Shape::new(&[m, k], DType::F32));
    let w = g.param("W", Shape::new(&[n, k], DType::F32)); // logical [N,K]
    let y = g.append_node(
        Op::DequantMatMul {
            scheme: QuantScheme::GgufQ8_0,
        },
        vec![xi, w],
        Shape::new(&[m, n], DType::F32),
        None,
    );
    g.set_outputs(vec![y]);

    let mut e = CoremlExecutable::compile(g);
    e.set_param_typed("W", &packed, DType::U8);
    let out = e.run(&[("x", &x)]).expect("run").remove(0);

    // Reference: full-precision x @ w_kn.
    let want = matmul(&x, &w_kn, m, k, n);
    // Q8_0 ~ 8-bit; allow for quant error scaled by K.
    approx(&out, &want, 5e-2);
}

#[test]
fn dequant_matmul_q1_0_nonunfolding_variations() {
    // Bonsai-27B Q1_0 (1-bit ±d, 128-elem blocks). Validate NON-UNFOLDING CoreML
    // paths ON-DEVICE: F32 (legacy unfold) and Lut (1-bit `constexpr_lut_to_dense`,
    // iOS18 opset — packed UINT1 indices + per-block LUT). Both encode the SAME
    // quantized weight, so both must match `x @ dequant_q1_0(packed)` — proving
    // the 1-bit blob format + LSB-first packing + grouped LUT are correct on real
    // coremlc, not just structurally. The Lut blob is ~n·k/8 (≈3.4 GB for 27B).
    use rlx_coreml::mil::{LowerOptions, Q1Mode};
    let (m, k, n) = (2usize, 128usize, 3usize); // k multiple of 128
    let w_nk: Vec<f32> = (0..n * k)
        .map(|i| ((i as f32) * 0.017).sin() * 0.4)
        .collect();
    let packed = rlx_gguf::quantize(&w_nk, rlx_gguf::GgmlType::Q1_0).expect("quantize");
    let w_deq = rlx_gguf::q1_dequant::dequant_q1_0(&packed, n * k).expect("deq"); // [n,k] ±d
    let x: Vec<f32> = (0..m * k).map(|i| ((i as f32) * 0.02).cos()).collect();
    // reference: out[i,j] = sum_k x[i,k] * w_deq[j,k]
    let mut want = vec![0.0f32; m * n];
    for i in 0..m {
        for j in 0..n {
            let mut a = 0.0f32;
            for kk in 0..k {
                a += x[i * k + kk] * w_deq[j * k + kk];
            }
            want[i * n + j] = a;
        }
    }
    for mode in [Q1Mode::F32, Q1Mode::Lut] {
        let mut g = Graph::new("dq1");
        let xi = g.input("x", Shape::new(&[m, k], DType::F32));
        let w = g.param("W", Shape::new(&[n, k], DType::F32));
        let y = g.append_node(
            Op::DequantMatMul {
                scheme: QuantScheme::GgufQ1_0,
            },
            vec![xi, w],
            Shape::new(&[m, n], DType::F32),
            None,
        );
        g.set_outputs(vec![y]);
        let opts = LowerOptions {
            ondevice_dequant: true,
            q1_mode: Some(mode),
            ..Default::default()
        };
        let mut e = CoremlExecutable::compile_with_lower_opts(g, opts);
        e.set_param_typed("W", &packed, DType::U8);
        let out = e
            .run(&[("x", &x)])
            .unwrap_or_else(|e| panic!("Q1_0 mode {mode:?} run failed: {e:?}"))
            .remove(0);
        // Same ±d weight in every mode; f16 scale precision only.
        approx(&out, &want, 2e-2);
    }
}

/// Timing: Bonsai decode-shaped Q1_0 matmul (m=1, k=5120, n=6144) on-device.
/// Run with `--nocapture`. Measures compile (coremlc) + warm run latency for the
/// 1-bit LUT path vs the F32 unfold, to answer "how fast is the no-unfold path".
#[test]
fn bench_q1_0_lut_vs_unfold_decode_shape() {
    use rlx_coreml::mil::{LowerOptions, Q1Mode};
    use std::time::Instant;
    let (m, k, n) = (1usize, 5120usize, 6144usize); // Bonsai GDN-proj decode shape
    let w_nk: Vec<f32> = (0..n * k)
        .map(|i| ((i as f32) * 0.0007).sin() * 0.3)
        .collect();
    let packed = rlx_gguf::quantize(&w_nk, rlx_gguf::GgmlType::Q1_0).expect("quantize");
    let x: Vec<f32> = (0..m * k).map(|i| ((i as f32) * 0.003).cos()).collect();
    let bench = |mode: Option<Q1Mode>| {
        let mut g = Graph::new("dq1b");
        let xi = g.input("x", Shape::new(&[m, k], DType::F32));
        let w = g.param("W", Shape::new(&[n, k], DType::F32));
        let y = g.append_node(
            Op::DequantMatMul {
                scheme: QuantScheme::GgufQ1_0,
            },
            vec![xi, w],
            Shape::new(&[m, n], DType::F32),
            None,
        );
        g.set_outputs(vec![y]);
        let opts = LowerOptions {
            ondevice_dequant: true,
            q1_mode: mode,
            ..Default::default()
        };
        let t0 = Instant::now();
        let mut e = CoremlExecutable::compile_with_lower_opts(g, opts);
        e.set_param_typed("W", &packed, DType::U8);
        let _ = e.run(&[("x", &x)]).expect("warm"); // compile + first run (warm)
        let compile_ms = t0.elapsed().as_secs_f64() * 1e3;
        let iters = 30;
        let t1 = Instant::now();
        for _ in 0..iters {
            let _ = e.run(&[("x", &x)]).expect("run");
        }
        let run_ms = t1.elapsed().as_secs_f64() * 1e3 / iters as f64;
        (compile_ms, run_ms)
    };
    let (c_lut, r_lut) = bench(Some(Q1Mode::Lut));
    let (c_f32, r_f32) = bench(Some(Q1Mode::F32));
    let packed_mb = packed.len() as f64 / 1e6;
    let unfold_mb = (n * k * 4) as f64 / 1e6;
    eprintln!(
        "\n[bench Q1_0 m={m} k={k} n={n}] blob: LUT~{:.1}MB (indices n·k/8) vs unfold {:.1}MB\n\
         LUT : compile+warm {c_lut:.0}ms, steady {r_lut:.2}ms/run\n\
         F32 : compile+warm {c_f32:.0}ms, steady {r_f32:.2}ms/run",
        packed.len() as f64 / 8.0 / 1e6 + packed_mb * 0.0, // indices ~ n*k/8
        unfold_mb
    );
}

#[test]
fn dequant_matmul_q4_0() {
    // Q4_0 block = 32 → K multiple of 32. Exercises a second scheme.
    let (m, k, n) = (1usize, 64usize, 2usize);
    let w_kn: Vec<f32> = (0..k * n)
        .map(|i| ((i as f32) * 0.01).sin() * 0.3)
        .collect();
    let mut w_nk = vec![0.0f32; n * k];
    for kk in 0..k {
        for j in 0..n {
            w_nk[j * k + kk] = w_kn[kk * n + j];
        }
    }
    let packed = rlx_gguf::quantize::quantize_q4_0(&w_nk).expect("quantize");
    let x: Vec<f32> = (0..m * k).map(|i| ((i as f32) * 0.003).cos()).collect();

    let mut g = Graph::new("dqmm_q40");
    let xi = g.input("x", Shape::new(&[m, k], DType::F32));
    let w = g.param("W", Shape::new(&[n, k], DType::F32));
    let y = g.append_node(
        Op::DequantMatMul {
            scheme: QuantScheme::GgufQ4_0,
        },
        vec![xi, w],
        Shape::new(&[m, n], DType::F32),
        None,
    );
    g.set_outputs(vec![y]);

    let mut e = CoremlExecutable::compile(g);
    e.set_param_typed("W", &packed, DType::U8);
    let out = e.run(&[("x", &x)]).expect("run").remove(0);

    let want = matmul(&x, &w_kn, m, k, n);
    approx(&out, &want, 1e-1); // Q4_0 is ~4.5 bpw, looser tolerance
}

#[test]
fn dequant_grouped_matmul_q8_0() {
    // E=2 experts, M=3 tokens, K=64, N=2.
    let (e_n, m, k, n) = (2usize, 3usize, 64usize, 2usize);
    // per-expert weight stored [N,K]; reference keeps [K,N] for matmul.
    let mut packed = Vec::new();
    let mut w_kn_per = Vec::new();
    for e in 0..e_n {
        let w_kn: Vec<f32> = (0..k * n)
            .map(|i| (((e * 1000 + i) as f32) * 0.011).sin() * 0.4)
            .collect();
        let mut w_nk = vec![0.0f32; n * k];
        for kk in 0..k {
            for j in 0..n {
                w_nk[j * k + kk] = w_kn[kk * n + j];
            }
        }
        packed.extend(rlx_gguf::quantize::quantize_q8_0(&w_nk).expect("q"));
        w_kn_per.push(w_kn);
    }
    let x: Vec<f32> = (0..m * k).map(|i| ((i as f32) * 0.02).cos()).collect();
    let experts = [0.0f32, 1.0, 0.0];

    let mut g = Graph::new("dqgmm");
    let xi = g.input("x", Shape::new(&[m, k], DType::F32));
    let w = g.param("W", Shape::new(&[e_n, n, k], DType::F32));
    let ei = g.input("e", Shape::new(&[m], DType::F32));
    let y = g.append_node(
        Op::DequantGroupedMatMul {
            scheme: QuantScheme::GgufQ8_0,
        },
        vec![xi, w, ei],
        Shape::new(&[m, n], DType::F32),
        None,
    );
    g.set_outputs(vec![y]);

    let mut exe = CoremlExecutable::compile(g);
    exe.set_param_typed("W", &packed, DType::U8);
    let out = exe
        .run(&[("x", &x), ("e", &experts)])
        .expect("run")
        .remove(0);

    let mut want = vec![0.0f32; m * n];
    for t in 0..m {
        let e = experts[t] as usize;
        let part = matmul(&x[t * k..(t + 1) * k], &w_kn_per[e], 1, k, n);
        want[t * n..(t + 1) * n].copy_from_slice(&part);
    }
    approx(&out, &want, 5e-2);
}

#[test]
fn dequant_matmul_through_session() {
    // Quantized weights via the public Session / set_param_typed path
    // (CompiledGraph forwarding), not CoremlExecutable directly.
    use rlx_runtime::{Device, Session};
    let (m, k, n) = (2usize, 64usize, 3usize);
    let w_kn: Vec<f32> = (0..k * n)
        .map(|i| ((i as f32) * 0.013).sin() * 0.5)
        .collect();
    let mut w_nk = vec![0.0f32; n * k];
    for kk in 0..k {
        for j in 0..n {
            w_nk[j * k + kk] = w_kn[kk * n + j];
        }
    }
    let packed = rlx_gguf::quantize::quantize_q8_0(&w_nk).expect("quantize");
    let x: Vec<f32> = (0..m * k).map(|i| ((i as f32) * 0.02).cos()).collect();

    let mut g = Graph::new("dqmm_session");
    let xi = g.input("x", Shape::new(&[m, k], DType::F32));
    let w = g.param("W", Shape::new(&[n, k], DType::F32));
    let y = g.append_node(
        Op::DequantMatMul {
            scheme: QuantScheme::GgufQ8_0,
        },
        vec![xi, w],
        Shape::new(&[m, n], DType::F32),
        None,
    );
    g.set_outputs(vec![y]);

    let mut compiled = Session::new(Device::Ane).compile(g);
    compiled.set_param_typed("W", &packed, DType::U8);
    let out = compiled.run(&[("x", &x)]).remove(0);

    approx(&out, &matmul(&x, &w_kn, m, k, n), 5e-2);
}

#[test]
fn quantize_dequantize_roundtrip() {
    // Per-tensor int8 fake-quant: x -> quantize -> dequantize -> ~x.
    let scale = 0.1f32;
    let zp = 0i32;
    let n = 6usize;
    let mut g = Graph::new("fakequant");
    let x = g.input("x", Shape::new(&[n], DType::F32));
    let q = g.append_node(
        Op::Quantize {
            axis: None,
            scales: vec![scale],
            zero_points: vec![zp],
        },
        vec![x],
        Shape::new(&[n], DType::I8),
        None,
    );
    let y = g.append_node(
        Op::Dequantize {
            axis: None,
            scales: vec![scale],
            zero_points: vec![zp],
        },
        vec![q],
        Shape::new(&[n], DType::F32),
        None,
    );
    g.set_outputs(vec![y]);

    // Values chosen off the .5 boundary so round-half convention is moot.
    let xs = [0.07f32, 0.23, -0.41, 1.04, -0.77, 0.34];
    let mut e = CoremlExecutable::compile(g);
    let out = e.run(&[("x", &xs)]).expect("run").remove(0);

    // Reference: round(x/scale)+zp clamped, then (q-zp)*scale.
    let want: Vec<f32> = xs
        .iter()
        .map(|&v| {
            let q = ((v / scale).round() + zp as f32).clamp(-128.0, 127.0);
            (q - zp as f32) * scale
        })
        .collect();
    approx(&out, &want, 1e-5);
}

#[test]
fn dequant_moe_weights_q8_0() {
    // Dequantize a packed weight back to f32 (no matmul).
    let n = 64usize; // one Q8_0-friendly block multiple
    let w: Vec<f32> = (0..n).map(|i| ((i as f32) * 0.05).sin()).collect();
    let packed = rlx_gguf::quantize::quantize_q8_0(&w).expect("quantize");

    let mut g = Graph::new("dqmoe");
    let wp = g.param("W", Shape::new(&[n], DType::F32));
    let dq = g.append_node(
        Op::DequantMoEWeights {
            scheme: QuantScheme::GgufQ8_0,
        },
        vec![wp],
        Shape::new(&[n], DType::F32),
        None,
    );
    // CoreML (spec < v9) requires ≥1 input; add a zero bias to satisfy it.
    let bias = g.input("bias", Shape::new(&[n], DType::F32));
    let y = g.binary(
        rlx_ir::op::BinaryOp::Add,
        dq,
        bias,
        Shape::new(&[n], DType::F32),
    );
    g.set_outputs(vec![y]);

    let mut e = CoremlExecutable::compile(g);
    e.set_param_typed("W", &packed, DType::U8);
    let out = e.run(&[("bias", &vec![0.0f32; n])]).expect("run").remove(0);

    let want = rlx_gguf::dequant_q8_0(&packed, n).unwrap();
    approx(&out, &want, 1e-5);
}