nordisk/nornir

tests/embed_rocm_buckets.rs

//! MIGraphX/ROCm fixed-shape bucketing PARITY test.
//!
//! ONNX Runtime 1.24's AMD GPU path is MIGraphX-only, and MIGraphX graph-compiles
//! for a CONSTANT input shape (full dynamic shapes are still in design — AMD
//! Discussion #2046). So the ort embedder pads variable-length inputs up to a
//! fixed bucket and mask-pools the real tokens. This test proves that path is
//! CORRECT: across inputs that land in DIFFERENT buckets, the GPU (MIGraphX)
//! embedding matches the dynamic CPU embedding of the same text.
//!
//! It only does anything on a box where the ROCm GPU + the ROCm onnxruntime are
//! actually present (run with ORT_DYLIB_PATH + LD_LIBRARY_PATH set, an
//! `embed-ort-rocm` build); everywhere else it SKIPS, so it never fails CI.

#![cfg(feature = "embed-ort-rocm")]

use nornir::vector::embed_ort::OrtEmbedder;
use nornir::vector::store::Embedder;

fn cosine(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
    let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    if na == 0.0 || nb == 0.0 {
        0.0
    } else {
        dot / (na * nb)
    }
}

#[test]
fn rocm_bucketed_embeddings_match_cpu_across_buckets() {
    // Load the GPU embedder; if ROCm/MIGraphX didn't actually engage (no GPU, no
    // ORT_DYLIB_PATH, non-AMD box), there is nothing to validate — SKIP.
    let gpu = match OrtEmbedder::load() {
        Ok(e) => e,
        Err(e) => {
            eprintln!("SKIP: ort embedder did not load ({e:#})");
            return;
        }
    };
    if gpu.gpu_vendor() != "rocm" {
        eprintln!(
            "SKIP: ROCm/MIGraphX EP not active (gpu_vendor={})",
            gpu.gpu_vendor()
        );
        return;
    }
    let cpu = OrtEmbedder::load_cpu().expect("load CPU reference embedder");

    // Texts chosen to land in DIFFERENT shape buckets (≈6, ≈90, ≈700 tokens) so a
    // single session exercises multiple compiled MIGraphX programs — the case
    // that shape-locking would break.
    let long = "fn scheduler() { let g = lock(); ".repeat(80) + "}";
    let texts: Vec<String> = vec![
        "fn main() {}".to_string(),
        "Mutex ordering avoids deadlock when every task acquires locks in a single \
         global order; the warehouse writer takes the redb lock before the iceberg \
         catalog so readers never observe a torn snapshot."
            .to_string(),
        long,
    ];

    // `Embedder::embed` takes `&[&str]`; these are owned `String`s because two of
    // them are built at runtime (`repeat`), so borrow them once for both legs
    // rather than cloning per call.
    let refs: Vec<&str> = texts.iter().map(String::as_str).collect();
    let g = gpu.embed(&refs).expect("gpu embed");
    let c = cpu.embed(&refs).expect("cpu embed");
    assert_eq!(g.len(), texts.len());
    assert_eq!(c.len(), texts.len());

    for (i, (gv, cv)) in g.iter().zip(&c).enumerate() {
        assert_eq!(gv.len(), cv.len(), "dim mismatch for text {i}");
        // Each vector is L2-normalized (the pooler normalizes).
        let norm: f32 = gv.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(
            (norm - 1.0).abs() < 1e-3,
            "gpu vec {i} not normalized: {norm}"
        );
        let cos = cosine(gv, cv);
        // Bucketing + mask-pooling must reproduce the dynamic CPU vector. Allow a
        // small slack for GPU fp arithmetic / fp16 kernels, but a shape-locked
        // (truncated) embedding would score far below this.
        assert!(
            cos > 0.99,
            "text {i}: GPU (MIGraphX, bucketed) vs CPU cosine {cos:.4} < 0.99 — \
             the fixed-shape path is not reproducing the dynamic embedding"
        );
    }
    eprintln!(
        "OK: {} bucketed GPU embeddings match CPU (cos > 0.99)",
        texts.len()
    );
}