diff --git a/.claude/board/entries/2026-11-08-dr-virtual-csr-bf16-gemm.md b/.claude/board/entries/2026-11-08-dr-virtual-csr-bf16-gemm.md new file mode 100644 index 000000000..bb8c7d1ec --- /dev/null +++ b/.claude/board/entries/2026-11-08-dr-virtual-csr-bf16-gemm.md @@ -0,0 +1,112 @@ +# 2026-11-08 — Virtual CSR over the DR cube; aggregation as BF16 GEMM (tile, ndarray CBLAS) + +**Status:** MEASURED (lab). This **fails the ms gate**: every aggregation row +is in milliseconds. Crate: `crates/dr-isa-stream-lab` (`vcsr.rs`, `spmm.rs`, +`bin/bench_vcsr.rs`). Stacked on the DR × ISA stream lab entry of the same +day. Depends on AdaWorldAPI/ndarray#414: the strided GEMM entry. + +## What it is + +**A virtual CSR.** Adjacency lives in the lanes' own 256 × 256 cube +coordinates. + +- Each row tile lists its nonempty column tiles. +- Each (row tile, column tile) pair is one 256 × 256 bitmask block. +- There is no `row_ptr` or `col_idx` array. Row `i`'s neighbours are the set + bits of block row `slot(i)` across the blocks of `tile(i)`. +- `ExplicitCsr` is the materialised form of the same edges. It is only an + oracle, and the row-by-row test shows the two agree. + +**Neighbour aggregation (plus-times).** `Y[p][i] = Σ_{j ∈ row(i)} lane_p[j]` +for 16 bound lanes, computed per block as BF16 GEMMs: + +- **ndarray tile GEMM** (AMX TDPBF16PS on this host). A is read in place from + the slab with stride = pitch. B is staged per 16-row group in VNNI layout, + and empty groups are skipped. +- **ndarray's own** `backend::cblas_gemm_bf16bf16f32`: the CBLAS-compatible + entry (AMX, then AVX-512 BF16, then scalar), with no external library. It + takes dense `&[u16]` and no `lda`, so A is staged; one call per block. (An + earlier revision of this lab used upstream `mkl_rt` instead. That arm was + removed; the comparator is ndarray's own GEMM.) +- **Staged A**, for comparison. + +All three are exact against the CSR oracle at n = 256, 513 (a partial last +tile) and 4096 for the tile paths. The tile paths are also exact at 65,536, +and ndarray's CBLAS entry is exact at 513 and 65,536. + +**Not claimed:** blasgraph's `CsrStorage` holds `BitVec` values under HDR +semirings (XorBundle, HammingMin, …). A floating-point GEMM does not realise +those, and they were not wired: depending on `lance-graph` pulls in DataFusion +and Lance unconditionally. + +## Measured + +See `crates/dr-isa-stream-lab/BENCH.md`, which has every arm of both labs +in one table, from one run. For the graph rows, ndarray's CBLAS entry costs +581 ms, 4.7 s and 51 ms on the three graphs, because it repacks both +operands on every call (see BENCH.md, Reading). + +## What the numbers say + +1. **The cost of a block GEMM is set by block area, not by edges.** Each block + costs 16 × 256 × 256 MACs regardless of content. + - At ~101 edges per block (0.15% fill), the GEMM loses 2× to a scalar CSR + loop over the same lanes. + - At 25% fill (community structure), it wins 6.4×. + - This is the same crossover the WHERE-mask GEMM showed from the other + side: there B was ~50% dense, and the GEMM was the right tool. +2. **The block layout is only as sparse as the local order is local.** The + same degree with one long-range edge per identity grows from 5,172 to + 41,916 blocks (8.1×) and 330 MB of bitmask. A virtual CSR therefore wants a + locality-preserving G2L order for graph workloads. Nothing in DR provides + that order yet; it is a property of how local ordinals are assigned. +3. **Reading A in place does not matter here.** Direct and staged A are within + 5–10%: per call, B dominates, not A. +4. **Per call, the community case costs ~3 µs** against the ~0.45 µs measured + for a bare K = 256 GEMM. The likely cause is restaging B for every call + (1,024 VNNI writes plus clears), though the adjacency is constant per graph + and could be packed once. **Not measured separately.** +5. **ndarray's CBLAS entry is 3–45× slower than the strided tile path.** It + repacks A and B on every call: a prepacked or strided entry is the gap. + +## Production crate: not now (DECISION) + +**DECISION:** no production crate for the virtual CSR or the graph +aggregation. + +**BASIS (measured, BENCH.md):** + +- Every graph-aggregation arm misses the gate, and the best one loses to a + scalar CSR loop on the sparse graph. +- Even the dense community case is bounded below by its own work: + 16 × 65,536 × 256 MACs ≈ 268M. At the measured AMX rate (~43 ns per + 8,192-MAC block) that is ≈ 1.4 ms before any staging. So the batch cannot + reach µs at this size, whichever arm runs it. +- The one arm that meets the gate is `SUM(lane) WHERE mask` over the whole + universe: 172 µs per 16 lanes × 16 masks with A in place. That is an + aggregation terminal over masks. If it is promoted, its home is the + Quack / mask-risc terminal set, behind DT64 binding, not a new crate. + +**SCOPE:** `crates/dr-isa-stream-lab` stays a lab. + +**REVISIT WHEN** one of these is measured: + +- each block's B packed once per graph version; +- a per-block dense/sparse choice; +- a locality-preserving G2L order; +- a prepacked or strided entry in ndarray's CBLAS GEMM; + +and a graph arm then beats the scalar CSR loop on a sparse graph, or the +WHERE terminal is proposed for Quack. + +## Open + +- Pack each block's B once per graph version and re-measure; this would + confirm or refute finding 4. +- A locality-preserving G2L order (for example a bandwidth-reducing order) and + its effect on the block count. +- A sparse-block realization (CSR-in-block or mask-fold) for low-fill blocks, + with the crossover chosen per block. +- The blasgraph HDR-semiring adapter: a named seam, not wired. +- Every number above misses the ms gate. The work is real dense block math, + but the gate is not met and this is not a pass. diff --git a/.claude/board/entries/README.md b/.claude/board/entries/README.md index 180b28810..8c220eb0d 100644 --- a/.claude/board/entries/README.md +++ b/.claude/board/entries/README.md @@ -25,10 +25,11 @@ index row, (3) no duplicate entry id. Checks 1 and 2 are deliberately opposite directions; the stranding this convention prevents shows up in exactly one of them, never both. -587 entries, 2026-08-06 .. 2026-11-08. +588 entries, 2026-08-06 .. 2026-11-08. | date | entry id | finding | file | |---|---|---|---| +| 2026-11-08 | `dr-virtual-csr-bf16-gemm` | | [2026-11-08-dr-virtual-csr-bf16-gemm.md](2026-11-08-dr-virtual-csr-bf16-gemm.md) | | 2026-11-08 | `dr-isa-stream-lab` | | [2026-11-08-dr-isa-stream-lab.md](2026-11-08-dr-isa-stream-lab.md) | | 2026-11-07 | `D-DT64-7-L5b` | | [2026-11-07-dt64-l5b-historical-basis.md](2026-11-07-dt64-l5b-historical-basis.md) | | 2026-11-07 | `D-DT64-7-L5a` | | [2026-11-07-dt64-l5a-exact-basis-witness.md](2026-11-07-dt64-l5a-exact-basis-witness.md) | diff --git a/.github/workflows/dr-isa-stream-lab.yml b/.github/workflows/dr-isa-stream-lab.yml index 0f50b837c..2a19f2f65 100644 --- a/.github/workflows/dr-isa-stream-lab.yml +++ b/.github/workflows/dr-isa-stream-lab.yml @@ -69,5 +69,5 @@ jobs: run: cargo clippy --manifest-path crates/dr-isa-stream-lab/Cargo.toml --all-targets -- -D warnings - name: tests (prints the GEMM tier) run: cargo test --release --manifest-path crates/dr-isa-stream-lab/Cargo.toml -- --nocapture - - name: stage costs (every timed result is checked against the oracle) - run: cargo run --release --manifest-path crates/dr-isa-stream-lab/Cargo.toml --bin bench + - name: benchmark of every arm (each result checked against its oracle) + run: cargo run --release --manifest-path crates/dr-isa-stream-lab/Cargo.toml --bin bench_all diff --git a/crates/dr-isa-stream-lab/BENCH.md b/crates/dr-isa-stream-lab/BENCH.md new file mode 100644 index 000000000..5e48c5dfb --- /dev/null +++ b/crates/dr-isa-stream-lab/BENCH.md @@ -0,0 +1,45 @@ +# dr-isa-stream-lab: benchmark of every arm + +Generated by `cargo run --release --bin bench_all`. Every result is checked against its oracle before timing and on every timed run. + +Host: Granite Rapids VM, 4 vCPU, rustc 1.99.0. lance-graph `7c301f5` + this branch; ndarray `f6c1a07` (AdaWorldAPI/ndarray#414). Single-threaded throughout. This is not certification on dedicated silicon: AMX tile state can corrupt under host contention, which is why every timed result is checked. + +The lance-graph gate is ns per operation and µs per batch. Rows in milliseconds fail it and are recorded as findings, not passes. + +Host tile tier: `AMX TDPBF16PS`. n = 65536, 16 lanes, 21 warm runs (median). + +| Workload | Arm | Time | Notes | +|---|---|---|---| +| WHERE 16×16 | bind (cold) | 993.0 ns | DT64 lane binding, 16 fields | +| WHERE 16×16 | 16 WHERE counts (popcount) | 2163.0 ns | the cube's fold terminal | +| WHERE 16×16 | stage selection B (cold) | 713.8 µs | once per mask set, 2 MiB | +| WHERE 16×16 | tile GEMM, whole lane, A in place | 172.5 µs | 1 call, A staged 0 B | +| WHERE 16×16 | tile GEMM, whole lane, A staged (slab) | 394.7 µs | A copied once | +| WHERE 16×16 | tile GEMM, whole lane, A staged (scattered) | 389.0 µs | lanes in separate allocations | +| WHERE 16×16 | ndarray cblas_gemm_bf16bf16f32, whole 64K | 3397.3 µs | 1 call; A is the slab (pitch = K); B row-major | +| WHERE 16×16 | reference: tile GEMM per 32-identity block | 2222.5 µs | 2,048 calls, B restaged per block | +| WHERE 16×16 | reference: cube per-bit sum fold | 13.5 ms | scalar gather | +| graph deg 8, 1/16 far | build virtual CSR (cold) | 35.1 ms | nnz 522603, 5194 blocks (7.9% of tile pairs) | +| graph deg 8, 1/16 far | tile GEMM, A in place | 32.3 ms | 24280 calls | +| graph deg 8, 1/16 far | tile GEMM, A staged | 33.6 ms | | +| graph deg 8, 1/16 far | ndarray cblas_gemm_bf16bf16f32 | 580.9 ms | 5194 calls, A staged (no lda) | +| graph deg 8, 1/16 far | scalar CSR over the same lanes | 11.6 ms | non-GEMM baseline | +| graph deg 8, all far | build virtual CSR (cold) | 182.3 ms | nnz 522981, 41756 blocks (63.7% of tile pairs) | +| graph deg 8, all far | tile GEMM, A in place | 105.9 ms | 82647 calls | +| graph deg 8, all far | tile GEMM, A staged | 120.2 ms | | +| graph deg 8, all far | ndarray cblas_gemm_bf16bf16f32 | 4730.9 ms | 41756 calls, A staged (no lda) | +| graph deg 8, all far | scalar CSR over the same lanes | 17.4 ms | non-GEMM baseline | +| graph deg 64, own tile | build virtual CSR (cold) | 108.2 ms | nnz 4194304, 256 blocks (0.4% of tile pairs) | +| graph deg 64, own tile | tile GEMM, A in place | 15.3 ms | 4096 calls | +| graph deg 64, own tile | tile GEMM, A staged | 12.5 ms | | +| graph deg 64, own tile | ndarray cblas_gemm_bf16bf16f32 | 51.4 ms | 256 calls, A staged (no lda) | +| graph deg 64, own tile | scalar CSR over the same lanes | 75.0 ms | non-GEMM baseline | + +## Reading + +- **The 256 × 256 cube is the whole universe.** One lane or one WHERE mask over 65,536 identities is a single 256 × 256 matrix, and the WHERE rows run over all of it in one call per 16 lanes. The graph's edge relation is identity × identity: 256 × 256 tile pairs, each a 256 × 256 block. The virtual CSR addresses that second level by `(tile_r, tile_c, slot_r, slot_c)`. +- **WHERE (selection about 50% dense).** The strided tile GEMM with A in place is the fastest arm: 172 µs for 16.7M MACs. Copying A first doubles the cost. +- **ndarray `cblas_gemm_bf16bf16f32` is about 20× slower than the strided tile path on the same problem**, although A is the slab itself at pitch = K. `amx_matmul::matmul_bf16_to_f32` copies both operands (`pack_contig`) and calls `bf16_tile_gemm_16x16` per 16×16 tile, which VNNI-packs B and allocates on every call. Before this entry can serve a stream it needs a prepacked or strided-B variant. +- **Graph aggregation.** A block GEMM's cost is set by block area, not by edges. It loses to a scalar CSR loop at about 0.15% fill (degree 8) and wins about 5–6× at 25% fill (degree 64 within the tile). +- **Long-range edges fill the block layout.** One long-range edge per identity grows the block count about 8×. A virtual CSR needs a locality-preserving G2L order. +- **Every graph row is in milliseconds.** None passes the gate. diff --git a/crates/dr-isa-stream-lab/src/bin/bench.rs b/crates/dr-isa-stream-lab/src/bin/bench.rs deleted file mode 100644 index 45435e26c..000000000 --- a/crates/dr-isa-stream-lab/src/bin/bench.rs +++ /dev/null @@ -1,148 +0,0 @@ -//! Stage costs for one binding feeding both arms. Every timed result is -//! checked against the fixture oracle; a mismatch aborts the run, because a -//! fast wrong answer is not a measurement (AMX tile state can corrupt under -//! host contention, ndarray `.claude/AMX_GOTCHAS.md` Gotcha 14). -//! -//! Units: ns per operation, us per batch. Rows marked "reference" are -//! oracle-grade paths kept for comparison; they are not gates. - -use std::hint::black_box; -use std::time::Instant; - -use dr_isa_stream_lab::bind::{World, SPACE_A}; -use dr_isa_stream_lab::cube::{fold, WhereMasks}; -use dr_isa_stream_lab::fixture::{Fixture, MASKS}; -use dr_isa_stream_lab::gemm::{sums, sums_whole, Mode, SelectionB}; -use dr_isa_stream_lab::oracle; -use dr_isa_stream_lab::source::{Bf16Lanes, Socket}; - -const RUNS: usize = 41; - -fn median_ns(mut f: impl FnMut()) -> u128 { - f(); - let mut t: Vec = (0..RUNS) - .map(|_| { - let s = Instant::now(); - f(); - s.elapsed().as_nanos() - }) - .collect(); - t.sort_unstable(); - t[RUNS / 2] -} - -fn us(ns: u128) -> f64 { - ns as f64 / 1000.0 -} - -fn main() { - println!("tier: {}", ndarray::simd::bf16_tile_gemm_tier()); - for (n, props) in [(65_536usize, 16usize), (65_536, 32)] { - let fx = Fixture { - n, - props, - seed: 0xBE7C, - }; - let fields: Vec = (0..props as u8).collect(); - let want = oracle::sums(&fx, &(0..props).collect::>()); - let want_f: Vec<[f32; MASKS]> = want.iter().map(|r| r.map(|v| v as f32)).collect(); - - let slab = Bf16Lanes::slab_in_order(&fx).unwrap(); - let scattered = Bf16Lanes::scattered(&fx).unwrap(); - let world = World::new(); - let s_slab = Socket::in_order(&slab, props); - let s_scat = Socket::in_order(&scattered, props); - - let bind_ns = median_ns(|| { - black_box(world.bind(&s_slab, &fields, SPACE_A).unwrap()); - }); - let binding = world.bind(&s_slab, &fields, SPACE_A).unwrap(); - let b_scat = world.bind(&s_scat, &fields, SPACE_A).unwrap(); - let masks = WhereMasks::from_fixture(&fx, SPACE_A); - let counts_ns = median_ns(|| { - black_box(masks.counts()); - }); - let stage_ns = median_ns(|| { - black_box(SelectionB::stage(&masks)); - }); - let sel = SelectionB::stage(&masks); - - let check = |label: &str, run: &dyn Fn() -> (Vec<[f32; MASKS]>, usize, usize)| { - let (got, gemms, staged) = run(); - assert_eq!(got, want_f, "{label}: wrong answer"); - let ns = median_ns(|| { - let (got, _, _) = run(); - assert_eq!(got, want_f, "{label}: wrong answer while timing"); - black_box(got); - }); - (ns, gemms, staged) - }; - let whole = |b: &_, s: &Socket<'_>, mode| { - let (got, r) = sums_whole(&world, b, s, &sel, &fields, mode).unwrap(); - (got, r.gemms, r.a_staged_bytes) - }; - let blocked = |mode| { - let (got, r) = sums(&world, &binding, &s_slab, &masks, &fields, mode).unwrap(); - (got, r.gemms, r.a_staged_bytes) - }; - let wd = check("whole direct", &|| whole(&binding, &s_slab, Mode::Direct)); - let ws = check("whole staged slab", &|| { - whole(&binding, &s_slab, Mode::Staged) - }); - let wc = check("whole staged scattered", &|| { - whole(&b_scat, &s_scat, Mode::Staged) - }); - let bd = check("blocked direct", &|| blocked(Mode::Direct)); - - assert_eq!( - fold(&world, &binding, &s_slab, &masks, &fields).unwrap(), - want - ); - let fold_ns = median_ns(|| { - black_box(fold(&world, &binding, &s_slab, &masks, &fields).unwrap()); - }); - - println!( - "\nn={n} props={props} masks={MASKS} selected={} lanes={} B selection={} B", - masks.counts().iter().sum::(), - slab.resident_bytes(), - sel.bytes() - ); - println!(" bind, cold, once {bind_ns:>9} ns"); - println!( - " ARM A: 16 WHERE counts (popcount) {:>9.1} us", - us(counts_ns) - ); - println!( - " stage selection B, cold, per masks {:>9.1} us", - us(stage_ns) - ); - println!( - " ARM B whole lane, A in place {:>9.1} us gemms={} A-staged={} B", - us(wd.0), - wd.1, - wd.2 - ); - println!( - " ARM B whole lane, A staged (slab) {:>9.1} us gemms={} A-staged={} B", - us(ws.0), - ws.1, - ws.2 - ); - println!( - " ARM B whole lane, A staged (scatter) {:>9.1} us gemms={} A-staged={} B", - us(wc.0), - wc.1, - wc.2 - ); - println!( - " reference: ARM B per-32-block calls {:>9.1} us gemms={} (B restaged per block)", - us(bd.0), - bd.1 - ); - println!( - " reference: ARM A per-bit sum fold {:>9.1} us", - us(fold_ns) - ); - } -} diff --git a/crates/dr-isa-stream-lab/src/bin/bench_all.rs b/crates/dr-isa-stream-lab/src/bin/bench_all.rs new file mode 100644 index 000000000..8aab40338 --- /dev/null +++ b/crates/dr-isa-stream-lab/src/bin/bench_all.rs @@ -0,0 +1,267 @@ +//! Every arm this lab tests, on one host, in one table. Each result is +//! checked against its oracle before timing and on every timed run; a +//! mismatch aborts the run. Prints markdown (the committed `BENCH.md`). +//! +//! Warm medians of `RUNS` runs. "cold" rows are one-off preparation. +//! "reference" rows are kept to show where cost goes; they are not gates. + +use std::hint::black_box; +use std::time::Instant; + +use dr_isa_stream_lab::bind::{World, SPACE_A}; +use dr_isa_stream_lab::cube::{fold, WhereMasks}; +use dr_isa_stream_lab::fixture::{bf16_value, Fixture, FAR_EVERY, MASKS}; +use dr_isa_stream_lab::gemm::{sums, sums_whole, Mode, SelectionB}; +use dr_isa_stream_lab::source::{Bf16Lanes, Socket}; +use dr_isa_stream_lab::spmm::{aggregate, aggregate_cblas}; +use dr_isa_stream_lab::vcsr::{ExplicitCsr, VirtualCsr}; +use dr_isa_stream_lab::{oracle, spmm}; + +const RUNS: usize = 21; +const N: usize = 65_536; + +fn median_ns(mut f: impl FnMut()) -> u128 { + f(); + let mut t: Vec = (0..RUNS) + .map(|_| { + let s = Instant::now(); + f(); + s.elapsed().as_nanos() + }) + .collect(); + t.sort_unstable(); + t[RUNS / 2] +} + +fn row(section: &str, arm: &str, ns: u128, note: &str) { + let (v, unit) = if ns < 10_000 { + (ns as f64, "ns") + } else if ns < 10_000_000 { + (ns as f64 / 1e3, "µs") + } else { + (ns as f64 / 1e6, "ms") + }; + println!("| {section} | {arm} | {v:.1} {unit} | {note} |"); +} + +fn main() { + println!( + "Host tile tier: `{}`. n = {N}, 16 lanes, {RUNS} warm runs (median).\n", + ndarray::simd::bf16_tile_gemm_tier() + ); + println!("| Workload | Arm | Time | Notes |"); + println!("|---|---|---|---|"); + + let fx = Fixture { + n: N, + props: 16, + seed: 0xA11, + }; + let fields: Vec = (0..16).collect(); + let slab = Bf16Lanes::slab_in_order(&fx).unwrap(); + let scattered = Bf16Lanes::scattered(&fx).unwrap(); + let world = World::new(); + let s_slab = Socket::in_order(&slab, 16); + let s_scat = Socket::in_order(&scattered, 16); + let binding = world.bind(&s_slab, &fields, SPACE_A).unwrap(); + let b_scat = world.bind(&s_scat, &fields, SPACE_A).unwrap(); + + // ── WHERE: SUM(lane p) WHERE mask m, 16 x 16 ────────────────────────── + let w = "WHERE 16×16"; + let bind_ns = median_ns(|| { + black_box(world.bind(&s_slab, &fields, SPACE_A).unwrap()); + }); + row(w, "bind (cold)", bind_ns, "DT64 lane binding, 16 fields"); + let masks = WhereMasks::from_fixture(&fx, SPACE_A); + let want = oracle::sums(&fx, &(0..16).collect::>()); + let want_f: Vec<[f32; MASKS]> = want.iter().map(|r| r.map(|v| v as f32)).collect(); + let counts_ns = median_ns(|| { + black_box(masks.counts()); + }); + row( + w, + "16 WHERE counts (popcount)", + counts_ns, + "the cube's fold terminal", + ); + let stage_ns = median_ns(|| { + black_box(SelectionB::stage(&masks)); + }); + row( + w, + "stage selection B (cold)", + stage_ns, + "once per mask set, 2 MiB", + ); + let sel = SelectionB::stage(&masks); + let whole = |b, s: &Socket<'_>, mode| { + let (y, r) = sums_whole(&world, b, s, &sel, &fields, mode).unwrap(); + assert_eq!(y, want_f); + r + }; + let r = whole(&binding, &s_slab, Mode::Direct); + let ns = median_ns(|| { + black_box(whole(&binding, &s_slab, Mode::Direct)); + }); + row( + w, + "tile GEMM, whole lane, A in place", + ns, + &format!("{} call, A staged {} B", r.gemms, r.a_staged_bytes), + ); + let ns = median_ns(|| { + black_box(whole(&binding, &s_slab, Mode::Staged)); + }); + row( + w, + "tile GEMM, whole lane, A staged (slab)", + ns, + "A copied once", + ); + let ns = median_ns(|| { + black_box(whole(&b_scat, &s_scat, Mode::Staged)); + }); + row( + w, + "tile GEMM, whole lane, A staged (scattered)", + ns, + "lanes in separate allocations", + ); + + // ndarray's CBLAS-compatible GEMM over the whole universe in one call. + // At n = 65,536 the padded pitch equals K, so the slab is already a dense + // 16 x 65,536 row-major matrix and is passed in place. + let a_rows = slab.row(0).unwrap(); + assert_eq!(a_rows.pitch, N, "pitch == K: the slab is a dense A"); + let a_dense = &a_rows.slab[..16 * N]; + let mut b_rm = vec![0u16; N * MASKS]; + for m in 0..MASKS { + for (wi, &word) in masks.words(m).iter().enumerate() { + let mut bits = word; + while bits != 0 { + b_rm[(wi * 64 + bits.trailing_zeros() as usize) * MASKS + m] = 0x3F80; + bits &= bits - 1; + } + } + } + let cblas = || { + let mut c = vec![0f32; 16 * MASKS]; + ndarray::backend::cblas_gemm_bf16bf16f32(a_dense, &b_rm, &mut c, 16, MASKS, N); + let y: Vec<[f32; MASKS]> = c.chunks(MASKS).map(|r| r.try_into().unwrap()).collect(); + assert_eq!(y, want_f); + y + }; + let ns = median_ns(|| { + black_box(cblas()); + }); + row( + w, + "ndarray cblas_gemm_bf16bf16f32, whole 64K", + ns, + "1 call; A is the slab (pitch = K); B row-major", + ); + let ns = median_ns(|| { + let (y, _) = sums(&world, &binding, &s_slab, &masks, &fields, Mode::Direct).unwrap(); + assert_eq!(y, want_f); + }); + row( + w, + "reference: tile GEMM per 32-identity block", + ns, + "2,048 calls, B restaged per block", + ); + let ns = median_ns(|| { + assert_eq!( + fold(&world, &binding, &s_slab, &masks, &fields).unwrap(), + want + ); + }); + row(w, "reference: cube per-bit sum fold", ns, "scalar gather"); + + // ── Graph: Y[p][i] = sum over row(i) of lane_p[j] ───────────────────── + let local = |far_every: usize| -> Vec<(usize, usize)> { + (0..N) + .flat_map(|i| fx.neighbors_with(i, far_every).map(move |j| (i, j))) + .collect() + }; + let community: Vec<(usize, usize)> = (0..N) + .flat_map(|i| { + let base = i - i % 256; + (0..64usize).map(move |k| (i, base + (i * 7 + k * 4 + k / 64) % 256)) + }) + .collect(); + let lane_vals: Vec<&[u16]> = (0..16).map(|p| slab.row(p).unwrap().values()).collect(); + for (label, edges) in [ + (format!("graph deg 8, 1/{FAR_EVERY} far"), local(FAR_EVERY)), + ("graph deg 8, all far".to_string(), local(1)), + ("graph deg 64, own tile".to_string(), community), + ] { + let g = VirtualCsr::from_edges(SPACE_A, N, edges.iter().copied()); + let csr = ExplicitCsr::from_edges(N, edges.iter().copied()); + let want: Vec> = spmm::oracle(&fx, &csr, &(0..16).collect::>()) + .into_iter() + .map(|r| r.into_iter().map(|v| v as f32).collect()) + .collect(); + let shape = format!( + "nnz {}, {} blocks ({:.1}% of tile pairs)", + g.nnz(), + g.blocks(), + 100.0 * g.blocks() as f64 / 65_536.0 + ); + let ns = median_ns(|| { + black_box(VirtualCsr::from_edges(SPACE_A, N, edges.iter().copied())); + }); + row(&label, "build virtual CSR (cold)", ns, &shape); + let tile = |mode| { + let (y, r) = aggregate(&world, &binding, &s_slab, &g, &fields, mode).unwrap(); + assert_eq!(y, want); + r + }; + let r = tile(Mode::Direct); + let ns = median_ns(|| { + black_box(tile(Mode::Direct)); + }); + row( + &label, + "tile GEMM, A in place", + ns, + &format!("{} calls", r.gemms), + ); + let ns = median_ns(|| { + black_box(tile(Mode::Staged)); + }); + row(&label, "tile GEMM, A staged", ns, ""); + let nat = || { + let (y, r) = aggregate_cblas(&world, &binding, &s_slab, &g, &fields).unwrap(); + assert_eq!(y, want); + r + }; + let r = nat(); + let ns = median_ns(|| { + black_box(nat()); + }); + row( + &label, + "ndarray cblas_gemm_bf16bf16f32", + ns, + &format!("{} calls, A staged (no lda)", r.gemms), + ); + let ns = median_ns(|| { + let mut y = vec![vec![0f32; N]; 16]; + for i in 0..N { + for &j in csr.row(i) { + for (p, yp) in y.iter_mut().enumerate() { + yp[i] += bf16_value(lane_vals[p][j]); + } + } + } + assert_eq!(y, want); + }); + row( + &label, + "scalar CSR over the same lanes", + ns, + "non-GEMM baseline", + ); + } +} diff --git a/crates/dr-isa-stream-lab/src/cube.rs b/crates/dr-isa-stream-lab/src/cube.rs index 625d4a438..b3e771e2a 100644 --- a/crates/dr-isa-stream-lab/src/cube.rs +++ b/crates/dr-isa-stream-lab/src/cube.rs @@ -78,6 +78,12 @@ pub enum ArmRefusal { masks: CoordSpace, selection: CoordSpace, }, + /// The graph's rows and columns are in another row order than the + /// binding's selection. + GraphSpace { + graph: CoordSpace, + selection: CoordSpace, + }, } impl From for ArmRefusal { diff --git a/crates/dr-isa-stream-lab/src/fixture.rs b/crates/dr-isa-stream-lab/src/fixture.rs index 7b115bf11..30343d6e4 100644 --- a/crates/dr-isa-stream-lab/src/fixture.rs +++ b/crates/dr-isa-stream-lab/src/fixture.rs @@ -62,3 +62,31 @@ pub fn bf16_round(v: f32) -> u16 { pub fn bf16_value(h: u16) -> f32 { f32::from_bits(u32::from(h) << 16) } + +/// Out-degree drawn per identity by [`Fixture::neighbors`]. +pub const DEGREE: usize = 8; +/// One identity in this many gets a long-range edge in [`Fixture::neighbors`]. +pub const FAR_EVERY: usize = 16; + +impl Fixture { + /// Identity `i`'s neighbors with the default long-range rate. + pub fn neighbors(&self, i: usize) -> impl Iterator + '_ { + self.neighbors_with(i, FAR_EVERY) + } + + /// Identity `i`'s `DEGREE` neighbors, within +-512 of `i` (its own and + /// nearby cube tiles), except that one identity in `far_every` has its + /// last edge go anywhere. Long-range edges are what make a block layout + /// fill up, so the rate is a parameter. May repeat a neighbor. + pub fn neighbors_with(&self, i: usize, far_every: usize) -> impl Iterator + '_ { + (0..DEGREE).map(move |k| { + let h = mix(self.seed.rotate_left(29) ^ ((k as u64) << 56) ^ i as u64); + if k + 1 == DEGREE && i.is_multiple_of(far_every) { + (h % self.n as u64) as usize + } else { + let off = (h % 1025) as i64 - 512; + (i as i64 + off).rem_euclid(self.n as i64) as usize + } + }) + } +} diff --git a/crates/dr-isa-stream-lab/src/gemm.rs b/crates/dr-isa-stream-lab/src/gemm.rs index 7f55e344c..4f4b139ae 100644 --- a/crates/dr-isa-stream-lab/src/gemm.rs +++ b/crates/dr-isa-stream-lab/src/gemm.rs @@ -23,7 +23,7 @@ use crate::fixture::{BLOCK, MASKS, ROWS}; use crate::source::{SlabLane, Socket}; /// bf16 1.0. -const ONE: u16 = 0x3F80; +pub(crate) const ONE: u16 = 0x3F80; /// How A reaches the tile. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -88,13 +88,16 @@ pub struct Report { /// In-place geometry of one row tile: A's row `r` is /// `slab[(first_row + r) * pitch..]`. -struct Geometry<'a> { - slab: &'a [u16], - first_row: usize, - pitch: usize, +pub(crate) struct Geometry<'a> { + pub(crate) slab: &'a [u16], + pub(crate) first_row: usize, + pub(crate) pitch: usize, } -fn geometry<'a>(fields: &[u8], lanes: &[SlabLane<'a>]) -> Result, DirectRefusal> { +pub(crate) fn geometry<'a>( + fields: &[u8], + lanes: &[SlabLane<'a>], +) -> Result, DirectRefusal> { if lanes.len() != ROWS { return Err(DirectRefusal::PartialRowTile { lanes: lanes.len() }); } diff --git a/crates/dr-isa-stream-lab/src/lib.rs b/crates/dr-isa-stream-lab/src/lib.rs index db745c729..9150523a9 100644 --- a/crates/dr-isa-stream-lab/src/lib.rs +++ b/crates/dr-isa-stream-lab/src/lib.rs @@ -13,6 +13,12 @@ //! `C[p][m] = sum_i lane_p[i] * mask_m(i)`, with A read in place from the //! slab when its geometry allows and staged otherwise. //! * [`oracle`] computes the answer from the fixture alone. +//! * [`vcsr`] is a virtual CSR: graph adjacency in the same 256 x 256 cube, +//! as a tile directory of 256 x 256 bitmask blocks, with no row-pointer or +//! column-index arrays. +//! * [`spmm`] aggregates property lanes over that adjacency (plus-times) as +//! BF16 GEMMs: the ndarray strided tile GEMM (AMX where present), reading A +//! in place, and ndarray's CBLAS-compatible `cblas_gemm_bf16bf16f32`. //! //! [`LaneSource`]: lance_graph_dt64::LaneSource @@ -22,3 +28,5 @@ pub mod fixture; pub mod gemm; pub mod oracle; pub mod source; +pub mod spmm; +pub mod vcsr; diff --git a/crates/dr-isa-stream-lab/src/spmm.rs b/crates/dr-isa-stream-lab/src/spmm.rs new file mode 100644 index 000000000..39c356d85 --- /dev/null +++ b/crates/dr-isa-stream-lab/src/spmm.rs @@ -0,0 +1,259 @@ +//! Neighbor aggregation over the virtual CSR, as BF16 GEMMs. +//! +//! `Y[p][i] = sum over j in row(i) of lane_p[j]` (plus-times). Per nonempty +//! block (row tile `tr`, column tile `tc`): +//! +//! * A is the 16 demanded lanes over column tile `tc`. In [`Mode::Direct`] +//! it is read in place from the slab at the slab pitch, exactly as in +//! [`crate::gemm`]; in [`Mode::Staged`] it is copied into a dense tile. +//! * B is the block's adjacency, `B[c][r] = edge(row r, col c)`, staged from +//! the bitmask: in VNNI layout 16 rows at a time for the tile GEMM (row +//! groups with no edges are skipped), or as a row-major 256-column tile for +//! ndarray's CBLAS-compatible GEMM ([`aggregate_cblas`]). +//! +//! This is the plus-times semiring. blasgraph's `CsrStorage` carries `BitVec` +//! values under HDR semirings (XorBundle, HammingMin, ...); those are not +//! realized by a floating-point GEMM, and this module does not claim them. + +use lance_graph_dt64::{LaneBinding, LaneRefusal}; +use ndarray::simd::{ + bf16_tile_gemm_16x16_packed, bf16_tile_gemm_16x16_packed_strided, bf16_tile_gemm_tier, + PackedBf16B, +}; + +use crate::bind::World; +use crate::cube::{ArmRefusal, TILE}; +use crate::fixture::{Fixture, BLOCK, ROWS}; +use crate::gemm::{geometry, GemmRefusal, Geometry, Mode, ONE}; +use crate::source::{SlabLane, Socket}; +use crate::vcsr::{Block, ExplicitCsr, VirtualCsr}; + +/// What one aggregation did. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct SpmmReport { + pub tier: &'static str, + /// Nonempty blocks visited. + pub blocks: usize, + /// GEMM calls. + pub gemms: usize, + /// Bytes of A copied into staging tiles. + pub a_staged_bytes: usize, + /// Nonzero B entries written. + pub b_entries: usize, +} + +/// `Y[p][i]` for the listed properties, from the fixture and the explicit +/// CSR alone. +pub fn oracle(fx: &Fixture, csr: &ExplicitCsr, props: &[usize]) -> Vec> { + props + .iter() + .map(|&p| { + (0..fx.n) + .map(|i| csr.row(i).iter().map(|&j| i64::from(fx.value(j, p))).sum()) + .collect() + }) + .collect() +} + +fn check_graph(binding: &LaneBinding, g: &VirtualCsr) -> Result<(), GemmRefusal> { + if g.space == binding.selection() { + Ok(()) + } else { + Err(GemmRefusal::Arm(ArmRefusal::GraphSpace { + graph: g.space, + selection: binding.selection(), + })) + } +} + +fn borrow<'s>( + world: &World, + binding: &LaneBinding, + socket: &'s Socket<'_>, + fields: &[u8], +) -> Result>, LaneRefusal> { + fields + .iter() + .map(|&f| binding.lane(&world.store, &world.view, socket.lanes.frames(), socket, f)) + .collect() +} + +/// K for column tile `tc`: 256, or the padded remainder of the last tile. +fn k_of(n: usize, tc: usize) -> usize { + (n.div_ceil(BLOCK) * BLOCK - tc * TILE).min(TILE) +} + +/// Stage block rows `g*16..g*16+16` into `b` as `B[c][r]`; returns the +/// number of entries written. +fn stage_rows(block: &Block, g: usize, k: usize, b: &mut PackedBf16B, value: u16) -> usize { + let mut written = 0; + for r in 0..ROWS { + for (w, &word) in block.row_words(g * ROWS + r).iter().enumerate() { + let mut bits = word; + while bits != 0 { + let c = w * 64 + bits.trailing_zeros() as usize; + bits &= bits - 1; + if c < k { + b.data_mut()[PackedBf16B::vnni_index(c, r)] = value; + written += 1; + } + } + } + } + written +} + +/// Aggregate 16 demanded lanes (or fewer, staged) over the virtual CSR with +/// the ndarray BF16 tile GEMM. +pub fn aggregate( + world: &World, + binding: &LaneBinding, + socket: &Socket<'_>, + g: &VirtualCsr, + fields: &[u8], + mode: Mode, +) -> Result<(Vec>, SpmmReport), GemmRefusal> { + check_graph(binding, g)?; + let lanes = borrow(world, binding, socket, fields)?; + let geo: Option> = match mode { + Mode::Direct => Some(geometry(fields, &lanes).map_err(GemmRefusal::Direct)?), + Mode::Staged => None, + }; + let mut report = SpmmReport { + tier: bf16_tile_gemm_tier(), + ..SpmmReport::default() + }; + let mut y = vec![vec![0f32; g.n]; fields.len()]; + let mut b_full = PackedBf16B::zeroed(TILE); + let mut b_tail = PackedBf16B::zeroed(k_of(g.n, g.row_tiles().saturating_sub(1)).max(BLOCK)); + let mut a = vec![0u16; ROWS * TILE]; + let mut c = vec![[0f32; ROWS * ROWS]; TILE / ROWS]; + for tr in 0..g.row_tiles() { + c.iter_mut().for_each(|t| t.fill(0.0)); + for block in g.row_tile(tr) { + report.blocks += 1; + let k = k_of(g.n, block.tc); + let b = if k == TILE { &mut b_full } else { &mut b_tail }; + if geo.is_none() { + a[..ROWS * k].fill(0); + for (r, lane) in lanes.iter().enumerate() { + let v = lane.values(); + let lo = block.tc * TILE; + let hi = (lo + k).min(v.len()); + a[r * k..r * k + (hi - lo)].copy_from_slice(&v[lo..hi]); + report.a_staged_bytes += (hi - lo) * 2; + } + } + for (grp, acc) in c.iter_mut().enumerate() { + if !block.any_in_rows(grp * ROWS..(grp + 1) * ROWS) { + continue; + } + report.b_entries += stage_rows(block, grp, k, b, ONE); + match &geo { + Some(geo) => { + let start = geo.first_row * geo.pitch + block.tc * TILE; + bf16_tile_gemm_16x16_packed_strided(&geo.slab[start..], geo.pitch, b, acc); + } + None => bf16_tile_gemm_16x16_packed(&a[..ROWS * k], b, acc), + } + report.gemms += 1; + stage_rows(block, grp, k, b, 0); + } + } + for (grp, acc) in c.iter().enumerate() { + for r in 0..ROWS { + let i = tr * TILE + grp * ROWS + r; + if i < g.n { + for (p, yp) in y.iter_mut().enumerate() { + yp[i] = acc[p * ROWS + r]; + } + } + } + } + } + Ok((y, report)) +} + +/// The same aggregation through ndarray's own CBLAS-compatible +/// `backend::cblas_gemm_bf16bf16f32` (AMX TDPBF16PS, then AVX-512 +/// VDPBF16PS, then scalar; no external library). One call per nonempty +/// block: M = 16 lanes, N = 256 block rows, K = the block's columns. +/// +/// That entry takes dense row-major `&[u16]` and no leading dimension, so A +/// cannot be passed in place: each block's 16 x K slice of the lanes is +/// copied into a dense buffer, and [`SpmmReport::a_staged_bytes`] counts it. +/// B is the block expanded to row-major `B[c][r]`. C is overwritten by the +/// call and added into the row tile's accumulator. +pub fn aggregate_cblas( + world: &World, + binding: &LaneBinding, + socket: &Socket<'_>, + g: &VirtualCsr, + fields: &[u8], +) -> Result<(Vec>, SpmmReport), GemmRefusal> { + check_graph(binding, g)?; + let lanes = borrow(world, binding, socket, fields)?; + let rows = fields.len(); + let mut report = SpmmReport { + tier: "ndarray backend::cblas_gemm_bf16bf16f32", + ..SpmmReport::default() + }; + let mut y = vec![vec![0f32; g.n]; rows]; + let mut a = vec![0u16; rows * TILE]; + let mut b = vec![0u16; TILE * TILE]; + let mut c = vec![0f32; rows * TILE]; + let mut acc = vec![0f32; rows * TILE]; + for tr in 0..g.row_tiles() { + acc.fill(0.0); + for block in g.row_tile(tr) { + report.blocks += 1; + let k = k_of(g.n, block.tc); + a[..rows * k].fill(0); + for (r, lane) in lanes.iter().enumerate() { + let v = lane.values(); + let lo = block.tc * TILE; + let hi = (lo + k).min(v.len()); + a[r * k..r * k + (hi - lo)].copy_from_slice(&v[lo..hi]); + report.a_staged_bytes += (hi - lo) * 2; + } + let set = |b: &mut [u16], value: u16| { + let mut written = 0; + for r in 0..TILE { + for (w, &word) in block.row_words(r).iter().enumerate() { + let mut bits = word; + while bits != 0 { + let col = w * 64 + bits.trailing_zeros() as usize; + bits &= bits - 1; + b[col * TILE + r] = value; + written += 1; + } + } + } + written + }; + report.b_entries += set(&mut b, ONE); + ndarray::backend::cblas_gemm_bf16bf16f32( + &a[..rows * k], + &b[..k * TILE], + &mut c, + rows, + TILE, + k, + ); + report.gemms += 1; + set(&mut b, 0); + for (x, &d) in acc.iter_mut().zip(&c) { + *x += d; + } + } + for r in 0..TILE { + let i = tr * TILE + r; + if i < g.n { + for (p, yp) in y.iter_mut().enumerate() { + yp[i] = acc[p * TILE + r]; + } + } + } + } + Ok((y, report)) +} diff --git a/crates/dr-isa-stream-lab/src/vcsr.rs b/crates/dr-isa-stream-lab/src/vcsr.rs new file mode 100644 index 000000000..b5544b875 --- /dev/null +++ b/crates/dr-isa-stream-lab/src/vcsr.rs @@ -0,0 +1,154 @@ +//! A virtual CSR over the 256 x 256 identity cube. +//! +//! The adjacency lives in the same local coordinates as the lanes: identity +//! `i` is row `slot(i)` of row tile `tile(i)`. Each row tile lists its +//! nonempty column tiles, and each (row tile, column tile) pair is one +//! 256 x 256 bitmask block, four words per row. There is no row-pointer or +//! column-index array: row `i`'s neighbors are the set bits of its row in the +//! blocks of its tile, read in place. [`ExplicitCsr`] builds the materialized +//! form from the same edges; it exists only as the differential oracle. + +use std::collections::BTreeMap; + +use lance_graph_dt64::CoordSpace; + +use crate::cube::{slot, tile, TILE}; + +/// Words per block row. +const ROW_WORDS: usize = TILE / 64; + +/// One 256 x 256 adjacency block: column tile `tc` of its row tile. +pub struct Block { + pub tc: usize, + words: Box<[u64]>, +} + +impl Block { + /// The four words of block row `r`; bit `c` is the edge to column `c`. + pub fn row_words(&self, r: usize) -> &[u64] { + &self.words[r * ROW_WORDS..(r + 1) * ROW_WORDS] + } + + /// Whether any of block rows `rows` has an edge. + pub fn any_in_rows(&self, rows: std::ops::Range) -> bool { + self.words[rows.start * ROW_WORDS..rows.end * ROW_WORDS] + .iter() + .any(|&w| w != 0) + } +} + +/// The virtual CSR: a tile directory of bitmask blocks, in one row order. +pub struct VirtualCsr { + pub space: CoordSpace, + pub n: usize, + dir: Vec>, +} + +impl VirtualCsr { + /// Build from `(row, col)` edges over `n` identities in `space`. + pub fn from_edges( + space: CoordSpace, + n: usize, + edges: impl IntoIterator, + ) -> VirtualCsr { + let tiles = n.div_ceil(TILE); + let mut blocks: BTreeMap<(usize, usize), Box<[u64]>> = BTreeMap::new(); + for (i, j) in edges { + assert!(i < n && j < n, "edge ({i}, {j}) outside n = {n}"); + let words = blocks + .entry((tile(i), tile(j))) + .or_insert_with(|| vec![0u64; TILE * ROW_WORDS].into_boxed_slice()); + let c = slot(j); + words[slot(i) * ROW_WORDS + c / 64] |= 1 << (c % 64); + } + let mut dir: Vec> = (0..tiles).map(|_| Vec::new()).collect(); + for ((tr, tc), words) in blocks { + dir[tr].push(Block { tc, words }); + } + VirtualCsr { space, n, dir } + } + + /// The blocks of row tile `tr`, by ascending column tile. + pub fn row_tile(&self, tr: usize) -> &[Block] { + &self.dir[tr] + } + + /// Row tiles. + pub fn row_tiles(&self) -> usize { + self.dir.len() + } + + /// Row `i`'s neighbors, ascending, read from the blocks in place. + pub fn row(&self, i: usize) -> impl Iterator + '_ { + let r = slot(i); + self.dir[tile(i)].iter().flat_map(move |b| { + b.row_words(r) + .iter() + .enumerate() + .flat_map(move |(w, &word)| { + let mut bits = word; + std::iter::from_fn(move || { + (bits != 0).then(|| { + let c = bits.trailing_zeros() as usize; + bits &= bits - 1; + b.tc * TILE + w * 64 + c + }) + }) + }) + }) + } + + /// Stored edges. + pub fn nnz(&self) -> usize { + self.dir + .iter() + .flatten() + .map(|b| { + b.words + .iter() + .map(|w| w.count_ones() as usize) + .sum::() + }) + .sum() + } + + /// Nonempty blocks. + pub fn blocks(&self) -> usize { + self.dir.iter().map(Vec::len).sum() + } + + /// Bytes of block storage. + pub fn bytes(&self) -> usize { + self.blocks() * TILE * ROW_WORDS * 8 + } +} + +/// The materialized CSR of the same edges: the differential oracle. +pub struct ExplicitCsr { + pub row_ptr: Vec, + pub col_idx: Vec, +} + +impl ExplicitCsr { + pub fn from_edges(n: usize, edges: impl IntoIterator) -> ExplicitCsr { + let mut rows: Vec> = vec![Vec::new(); n]; + for (i, j) in edges { + rows[i].push(j); + } + let mut row_ptr = Vec::with_capacity(n + 1); + let mut col_idx = Vec::new(); + row_ptr.push(0); + for mut r in rows { + r.sort_unstable(); + r.dedup(); + col_idx.extend(r); + row_ptr.push(col_idx.len()); + } + ExplicitCsr { row_ptr, col_idx } + } + + /// Row `i`'s neighbors, ascending. + pub fn row(&self, i: usize) -> &[usize] { + &self.col_idx[self.row_ptr[i]..self.row_ptr[i + 1]] + } +} diff --git a/crates/dr-isa-stream-lab/tests/integration/main.rs b/crates/dr-isa-stream-lab/tests/integration/main.rs index 91492ceb1..fa07b19aa 100644 --- a/crates/dr-isa-stream-lab/tests/integration/main.rs +++ b/crates/dr-isa-stream-lab/tests/integration/main.rs @@ -2,3 +2,4 @@ pub mod geometry; pub mod numeric; pub mod parity; pub mod refusals; +pub mod vcsr; diff --git a/crates/dr-isa-stream-lab/tests/integration/vcsr.rs b/crates/dr-isa-stream-lab/tests/integration/vcsr.rs new file mode 100644 index 000000000..de69e1b1c --- /dev/null +++ b/crates/dr-isa-stream-lab/tests/integration/vcsr.rs @@ -0,0 +1,260 @@ +//! The virtual CSR and neighbor aggregation over it. + +use dr_isa_stream_lab::bind::{World, SPACE_A, SPACE_OTHER}; +use dr_isa_stream_lab::cube::{ArmRefusal, TILE}; +use dr_isa_stream_lab::fixture::Fixture; +use dr_isa_stream_lab::gemm::{DirectRefusal, GemmRefusal, Mode}; +use dr_isa_stream_lab::source::{Bf16Lanes, Socket}; +use dr_isa_stream_lab::spmm::{aggregate, oracle}; +use dr_isa_stream_lab::vcsr::{ExplicitCsr, VirtualCsr}; +use lance_graph_dt64::dr::{FrameRefusal, FrameRevisions}; +use lance_graph_dt64::LaneRefusal; + +fn edges(fx: &Fixture) -> Vec<(usize, usize)> { + (0..fx.n) + .flat_map(|i| fx.neighbors(i).map(move |j| (i, j))) + .collect() +} + +fn graph(fx: &Fixture, space: lance_graph_dt64::CoordSpace) -> (VirtualCsr, ExplicitCsr) { + let e = edges(fx); + ( + VirtualCsr::from_edges(space, fx.n, e.iter().copied()), + ExplicitCsr::from_edges(fx.n, e), + ) +} + +fn as_f32(rows: Vec>) -> Vec> { + rows.into_iter() + .map(|r| r.into_iter().map(|v| v as f32).collect()) + .collect() +} + +#[test] +fn the_virtual_csr_has_the_explicit_csr_rows() { + for n in [300, 4096] { + let fx = Fixture { + n, + props: 16, + seed: 31, + }; + let (v, e) = graph(&fx, SPACE_A); + for i in 0..n { + assert_eq!(v.row(i).collect::>(), e.row(i), "row {i}, n={n}"); + } + assert_eq!(v.nnz(), e.col_idx.len()); + } +} + +/// The fixture has local and long-range blocks, and is far from dense: a +/// block-sparse layout is doing real work. +#[test] +fn the_graph_is_block_sparse_with_long_range_blocks() { + let fx = Fixture { + n: 65_536, + props: 16, + seed: 32, + }; + let (v, _) = graph(&fx, SPACE_A); + let tiles = v.row_tiles(); + assert_eq!(tiles, 256); + // Measured 5,186 of 65,536 tile pairs at the default rate. + assert!( + v.blocks() * 10 < tiles * tiles, + "{} of {} blocks present", + v.blocks(), + tiles * tiles + ); + let far = (0..tiles) + .flat_map(|tr| v.row_tile(tr).iter().map(move |b| (tr, b.tc))) + .filter(|&(tr, tc)| tr.abs_diff(tc) > 2 && tr.abs_diff(tc) < tiles - 2) + .count(); + assert!(far > tiles, "only {far} long-range blocks"); +} + +/// The block count is set by locality: the same degree with a long-range +/// edge on every identity fills most tile pairs. +#[test] +fn long_range_edges_fill_the_block_layout() { + let fx = Fixture { + n: 65_536, + props: 16, + seed: 32, + }; + let local = VirtualCsr::from_edges(SPACE_A, fx.n, edges(&fx)).blocks(); + let everywhere = VirtualCsr::from_edges( + SPACE_A, + fx.n, + (0..fx.n).flat_map(|i| fx.neighbors_with(i, 1).map(move |j| (i, j))), + ) + .blocks(); + // Measured 41,997 vs 5,186 blocks. + assert!(everywhere > 5 * local, "{everywhere} vs {local} blocks"); +} + +#[test] +fn aggregation_is_exact_direct_and_staged() { + for n in [256, 513, 4096, 65_536] { + let fx = Fixture { + n, + props: 16, + seed: 33 + n as u64, + }; + let (v, e) = graph(&fx, SPACE_A); + let want = as_f32(oracle(&fx, &e, &(0..16).collect::>())); + let lanes = Bf16Lanes::slab_in_order(&fx).unwrap(); + let socket = Socket::in_order(&lanes, 16); + let world = World::new(); + let fields: Vec = (0..16).collect(); + let binding = world.bind(&socket, &fields, SPACE_A).unwrap(); + let (direct, r) = aggregate(&world, &binding, &socket, &v, &fields, Mode::Direct).unwrap(); + assert_eq!(direct, want, "direct, n={n}, tier {}", r.tier); + assert_eq!(r.a_staged_bytes, 0, "direct copied A, n={n}"); + assert_eq!(r.blocks, v.blocks()); + let (staged, r) = aggregate(&world, &binding, &socket, &v, &fields, Mode::Staged).unwrap(); + assert_eq!(staged, want, "staged, n={n}"); + assert!(r.a_staged_bytes > 0); + } +} + +#[test] +fn the_cblas_aggregation_matches() { + use dr_isa_stream_lab::spmm::aggregate_cblas; + for n in [513, 65_536] { + let fx = Fixture { + n, + props: 16, + seed: 34, + }; + let (v, e) = graph(&fx, SPACE_A); + let want = as_f32(oracle(&fx, &e, &(0..16).collect::>())); + let lanes = Bf16Lanes::slab_in_order(&fx).unwrap(); + let socket = Socket::in_order(&lanes, 16); + let world = World::new(); + let fields: Vec = (0..16).collect(); + let binding = world.bind(&socket, &fields, SPACE_A).unwrap(); + let (got, r) = aggregate_cblas(&world, &binding, &socket, &v, &fields).unwrap(); + assert_eq!(got, want, "cblas, n={n}"); + assert_eq!(r.gemms, v.blocks(), "one call per block"); + } +} + +#[test] +fn lanes_are_borrowed_once_per_aggregation() { + for n in [256, 65_536] { + let fx = Fixture { + n, + props: 16, + seed: 35, + }; + let (v, _) = graph(&fx, SPACE_A); + let lanes = Bf16Lanes::slab_in_order(&fx).unwrap(); + let socket = Socket::in_order(&lanes, 16); + let world = World::new(); + let fields: Vec = (0..16).collect(); + let binding = world.bind(&socket, &fields, SPACE_A).unwrap(); + let before = socket.borrows.get(); + aggregate(&world, &binding, &socket, &v, &fields, Mode::Direct).unwrap(); + assert_eq!(socket.borrows.get() - before, 16, "n={n}"); + } +} + +#[test] +fn a_graph_in_another_row_order_is_refused() { + let fx = Fixture { + n: 1024, + props: 16, + seed: 36, + }; + let (v, _) = graph(&fx, SPACE_OTHER); + let lanes = Bf16Lanes::slab_in_order(&fx).unwrap(); + let socket = Socket::in_order(&lanes, 16); + let world = World::new(); + let fields: Vec = (0..16).collect(); + let binding = world.bind(&socket, &fields, SPACE_A).unwrap(); + assert_eq!( + aggregate(&world, &binding, &socket, &v, &fields, Mode::Direct).unwrap_err(), + GemmRefusal::Arm(ArmRefusal::GraphSpace { + graph: SPACE_OTHER, + selection: SPACE_A + }) + ); +} + +#[test] +fn a_stale_frame_refuses_aggregation() { + let fx = Fixture { + n: 1024, + props: 16, + seed: 37, + }; + let (v, _) = graph(&fx, SPACE_A); + let mut lanes = Bf16Lanes::slab_in_order(&fx).unwrap(); + let world = World::new(); + let fields: Vec = (0..16).collect(); + let binding = world + .bind(&Socket::in_order(&lanes, 16), &fields, SPACE_A) + .unwrap(); + lanes.frames_mut().advance(FrameRevisions { + source: 2, + relation_map: 1, + coordinates: 1, + }); + let socket = Socket::in_order(&lanes, 16); + assert_eq!( + aggregate(&world, &binding, &socket, &v, &fields, Mode::Direct).unwrap_err(), + GemmRefusal::Arm(ArmRefusal::Lane(LaneRefusal::Frame( + FrameRefusal::SourceChanged + ))) + ); +} + +#[test] +fn fewer_than_sixteen_lanes_refuse_direct_and_stage_correctly() { + let fx = Fixture { + n: 700, + props: 5, + seed: 38, + }; + let (v, e) = graph(&fx, SPACE_A); + let lanes = Bf16Lanes::slab_in_order(&fx).unwrap(); + let socket = Socket::in_order(&lanes, 5); + let world = World::new(); + let fields: Vec = (0..5).collect(); + let binding = world.bind(&socket, &fields, SPACE_A).unwrap(); + assert_eq!( + aggregate(&world, &binding, &socket, &v, &fields, Mode::Direct).unwrap_err(), + GemmRefusal::Direct(DirectRefusal::PartialRowTile { lanes: 5 }) + ); + let (got, _) = aggregate(&world, &binding, &socket, &v, &fields, Mode::Staged).unwrap(); + assert_eq!(got, as_f32(oracle(&fx, &e, &(0..5).collect::>()))); +} + +/// The one-edge case pins orientation: row i gets lane_p[j], not lane_p[i]. +#[test] +fn a_single_edge_aggregates_the_neighbor_not_the_row() { + let fx = Fixture { + n: TILE * 2, + props: 16, + seed: 39, + }; + let (i, j) = (3usize, TILE + 77); + let v = VirtualCsr::from_edges(SPACE_A, fx.n, [(i, j)]); + let lanes = Bf16Lanes::slab_in_order(&fx).unwrap(); + let socket = Socket::in_order(&lanes, 16); + let world = World::new(); + let fields: Vec = (0..16).collect(); + let binding = world.bind(&socket, &fields, SPACE_A).unwrap(); + let (y, _) = aggregate(&world, &binding, &socket, &v, &fields, Mode::Direct).unwrap(); + for (p, yp) in y.iter().enumerate() { + assert_eq!(yp[i], fx.value(j, p) as f32, "prop {p}"); + assert_eq!( + yp.iter().filter(|&&x| x != 0.0).count(), + usize::from(fx.value(j, p) != 0) + ); + } + assert!( + (0..16).any(|p| fx.value(i, p) != fx.value(j, p)), + "fixture cannot tell i from j" + ); +}