diff --git a/ext/Cargo.toml b/ext/Cargo.toml index 8b3919aba8..29745ae957 100644 --- a/ext/Cargo.toml +++ b/ext/Cargo.toml @@ -59,9 +59,27 @@ concurrent = [ odd-primes = ["fp/odd-primes", "algebra/odd-primes", "sseq/odd-primes"] logging = [] nassau = [] +# Dispatch large p=2 matrix products to the Hopper GPU backend via `fp`. +gpu = ["fp/gpu"] [workspace] members = [ + "crates/algebra", + "crates/bivec", + "crates/fp", + "crates/fp-cuda", + "crates/maybe-rayon", + "crates/once", + "crates/query", + "crates/sseq", +] +# `fp-cuda` builds its CUDA kernel with nvcc (falling back to a stub PTX when +# nvcc is absent) and is opt-in. It stays out of the default member set so plain +# workspace commands (`cargo build`, `cargo test`, `nix run .#test`) don't pull +# it in; `fp`'s `gpu` feature depends on it explicitly when the GPU backend is +# wanted. +default-members = [ + ".", "crates/algebra", "crates/bivec", "crates/fp", diff --git a/ext/crates/fp-cuda/Cargo.toml b/ext/crates/fp-cuda/Cargo.toml new file mode 100644 index 0000000000..d6a2aed7f4 --- /dev/null +++ b/ext/crates/fp-cuda/Cargo.toml @@ -0,0 +1,50 @@ +[package] +name = "fp-cuda" +version = "0.1.0" +edition = "2024" +authors = ["Joey Beauvais-Feisthauer "] +description = "CUDA backend for fp's F_2 BLAS routines (Hopper wgmma.b1 via inline PTX)" +license = "MIT OR Apache-2.0" +publish = false + +# This crate ships a CUDA C++ kernel (cuda_kernels/matmul_b1.cu) that is +# compiled to PTX by nvcc at build time. The Rust side uses `cudarc` (stable +# Rust, runtime dynamic-loading of the CUDA driver) for the host driver-API +# surface: module load, device buffers, launch, and the raw `cuTensorMapEncodeTiled` +# call for the TMA descriptors. +# +# Build with: +# cargo build -p fp-cuda +# +# Prerequisites: nvcc (CUDA Toolkit 12.x+) on PATH, Hopper GPU at runtime. +# Plain `cargo build` from the workspace root skips this crate via the +# `default-members` entry in the workspace Cargo.toml. + +build = "build.rs" + +[dependencies] +# Driver API only; dynamic-loading means the crate builds with no CUDA present +# (we still need nvcc at build time to produce the PTX, and a driver at runtime). +# cuda-12080 only selects which pre-generated driver bindings to compile; with +# dynamic-loading the actual driver is dlopen'd at runtime, so no CUDA is needed +# to build and a newer driver (12.8+/13.x) still works. +cudarc = { version = "0.19", default-features = false, features = [ + "std", + "driver", + "nvrtc", + "cuda-12080", + "dynamic-loading", +] } + +[dev-dependencies] +# `fp` is a dev-dependency only: the library API is fp-agnostic (raw limbs), so +# the higher-level `fp` crate can depend on `fp-cuda` without a cycle. The +# examples/benches still use `fp::Matrix` for random generation and CPU +# cross-checks (dev-dependency cycles are permitted by Cargo). +fp = { path = "../fp", default-features = false } +criterion = { version = "0.5", features = ["html_reports"] } +rand = "0.9" + +[[bench]] +name = "matmul_b1" +harness = false diff --git a/ext/crates/fp-cuda/README.md b/ext/crates/fp-cuda/README.md new file mode 100644 index 0000000000..2fca6bd2c4 --- /dev/null +++ b/ext/crates/fp-cuda/README.md @@ -0,0 +1,212 @@ +# fp-cuda + +CUDA backend for the F₂ matrix multiplication implemented in `crates/fp/src/blas/`. +The Hopper memory pipeline is used end-to-end: kernel written in CUDA C++ with +inline PTX for **TMA bulk tensor loads** with **128B swizzle** +(`cp.async.bulk.tensor.2d`), **mbarrier**-based completion sync, and the binary +tensor cores +(`wgmma.mma_async.sync.aligned.m64n64k256.row.col.s32.b1.b1.and.popc`). +Both operands are pre-arranged into plain row-major K-major tiles on the host; +the TMA applies the swizzle that the wgmma matrix descriptors expect. +Rust-side glue uses [`cudarc`](https://crates.io/crates/cudarc) for the host +driver-API surface (module load, device buffers, typed launch) and its +`driver::sys` raw bindings for the `cuTensorMapEncodeTiled` call that builds the +TMA descriptors. `cudarc` is stable Rust and dynamically loads the CUDA driver +at runtime, so the Rust side builds with no CUDA present. + +This crate is **excluded from the workspace's `default-members`**, so plain +`cargo build` / `nix run .#test` ignore it. It is opt-in: building requires +nvcc on `PATH` and (at runtime) a Hopper-class GPU. + +## Prerequisites + +1. **nvcc** (CUDA Toolkit 12.x+, since TMA + wgmma require 12.0+) on `PATH`, + with Hopper (sm_90a) support. Override the binary location with the + `NVCC` env var if needed. +2. A **Hopper GPU** at runtime (sm_90a). The kernel is built for `sm_90a`, an + architecture-specific target that is **not** forward-compatible, so the PTX + runs only on Hopper — not on pre-Hopper devices (which lack the `wgmma.*` and + `cp.async.bulk.tensor.*` instructions it emits) nor on newer architectures + such as Blackwell (`sm_100`). + +Builds on **stable** Rust — no nightly toolchain required. (`nvcc` is still +needed at build time to compile the kernel to PTX, and a CUDA driver at runtime.) + +## Building + +```bash +# From the workspace root; the leading -p selects this crate explicitly. +cargo build -p fp-cuda +``` + +`build.rs` invokes nvcc on `cuda_kernels/matmul_b1.cu` and emits +`matmul_b1.ptx` into the cargo `OUT_DIR`. `src/lib.rs` embeds it via +`include_bytes!` and loads it at runtime through cudarc. + +**When nvcc is absent** (CI, or a contributor without the CUDA Toolkit) +`build.rs` emits a *stub* PTX and prints a warning instead of failing, so the +crate stays `cargo check`/`clippy`-able everywhere — including under +`cargo check --workspace --all-features`, which the workspace CI runs on a +machine with no CUDA. The stub carries no kernel, so at runtime +`GpuContext::new` returns an error and callers (e.g. `fp`'s `cuda` dispatch) +fall back to the CPU path. An nvcc that *is* present but fails to compile is +still a hard build error. + +## Running + +```bash +# Smoke test (multiplies a few small shapes, asserts CPU↔GPU equality): +cargo run -p fp-cuda --example matmul_b1_demo + +# Benchmark against the CPU AVX-512 path in fp::blas: +cargo bench -p fp-cuda +``` + +The bench compares each square size in `{128, 256, 512, 1024, 2048, 4096, 8192}` +against `fp::blas::fast_mul_concurrent`, asserts bit-equality of the outputs, +and prints binary TOPS for both backends. + +## Why excluded from `default-members`? + +Contributors without nvcc would otherwise see this crate's build fail every +time they run `cargo build`. Keeping it out of the default member set means: + +- `cargo build`, `cargo test`, `cargo fmt`, `nix run .#test` from the + workspace root behave exactly as before. +- Tooling that wants this crate explicitly opts in with `-p fp-cuda`. + +The crate is still a workspace **member**, so `cargo metadata` sees it, +`rust-analyzer` indexes it, and shared dependency resolution works. + +## Status + +The full pipeline (host row-major pre-arrangement → TMA 128B-swizzle loads → +mbarrier sync → persistent grid + clusters/multicast → register-blocked +`m64n128k256` wgmma.b1 strips → bit-pack → double-buffered TMA bulk output +store) is **validated on an H200 NVL (sm_90, CUDA 13.0 driver / 12.4 toolkit, +2026-07-07)** and earlier on an H100 NVL. The PTX JITs at module load, the +dynamic-SMEM opt-in (~166 KB) and all three TMA descriptors are accepted, and +outputs are **bit-exact** against the CPU `fp::blas` path across `matmul_b1_demo` +(64…8192) and the kernel-only bench (4096…32768, incl. a full 32768³ CPU +cross-check). + +Throughput, **kernel-only** (host setup + H2D/D2H excluded — the comparison the +~100-TOPS pre-swizzle baseline was measured at), H200 NVL: + +| size (M=K=N) | binary TOPS | ms/launch | +|--------------|-------------|-----------| +| 4096 | ~4,000 | 0.034 | +| 8192 | ~6,700 | 0.163 | +| 16384 | ~8,600 | 1.02 | +| 32768 | ~9,600 | 7.33 | + +i.e. roughly a **~96× kernel speedup** over the ~100-TOPS pre-swizzle state, and +the >16384 L2 cliff is **gone** — throughput now *climbs* with size. + +Getting there took closing an L2-residency-of-B cliff, then a bandwidth wall, +then a latency wall (all measured on-device, since ncu can't attach through the +CUDA-13 driver — see the standalone wgmma/L2 microbenchmark approach): +- **Persistent grid + grouped rasterization + clusters/TMA-multicast** (Phases + 8–9) keep B's reuse distance short so it stays L2-resident. `bench_shapes` now + shows equal-FLOPs B-fits and B-spills shapes running at the *same* throughput — + the cliff is closed. +- **Register-blocking** (Phase 10): each CTA computes a 192×128 output block as 3 + `m64n128` strips sharing one loaded B sub-tile, cutting operand-refill + bytes/MAC by 33% — the kernel had become L2→SMEM-refill bound (~8 TB/s + sustained vs a ~12.5k-TOPS single-warpgroup wgmma ceiling). +- **Epilogue overlap** (Phase 11): double-buffered `sC` + a deferred + `cp.async.bulk.wait_group.read 1` so tile T's output store drains during tile + T+1's compute. +- **Fence hoisting** (Phase 12): one `wgmma.fence` before the K-loop instead of + two per k-chunk (a warpgroup-wide sync). After blocking, skipping an entire + operand load barely changes throughput — the kernel is now **TMA-latency + bound** (deeper pipelines help; pipeline depth is capped by SMEM at STAGES=4). + +Run `cargo run --release -p fp-cuda --example bench_shapes` (L2-residency check) +or `--example tune` (fast throughput sweep) to reproduce. + +The end-to-end `cargo bench` figures (≤30 TOPS) are dominated by host +serialization and the TMA-layout pre-arrangement; use `cargo run --release -p +fp-cuda --example bench_kernel_only` for the kernel number. + +Reproduce in this order (each gates the next): + +1. **64×256×64 identity / small product first.** Smallest path that exercises one + swizzled tile end-to-end (the first `matmul_b1_demo` case). A failure here + points at the swizzled wgmma descriptor constants (`DESC_LBO = 16`, + `DESC_SBO = 1024`, per-k256 advance of 32 bytes), or a host-layout / TMA-box + mismatch. These derive from CUTLASS `make_gmma_desc` + (`LayoutType::B128`). +2. **Full size sweep** via `cargo run -p fp-cuda --example matmul_b1_demo` + (bit-exact CPU↔GPU for 64…8192). +3. **Kernel-only throughput + correctness** via + `cargo run --release -p fp-cuda --example bench_kernel_only`; compare binary + TOPS against the ~100-TOPS pre-swizzle baseline. + +Other points worth a trace once: the dynamic-SMEM base must be 128-byte aligned +for TMA (declared `extern __shared__ __align__(128)`), and the per-stage +`expect_tx = (TILE + TILE_B) * 8` bytes (one A tile + one B tile) must match +the `cp.async.bulk.tensor.complete_tx::bytes` notifications from the two issued +TMA loads. + +## Roadmap + +Following the optimization ladder of Pranjal Shankhdhar's "Outperforming cuBLAS +on H100" worklog, adapted to the binary (`b1`) GF(2) kernel. + +Done: + +- **128B swizzle** (Phase 3) on both operands — TMA loads with + `CU_TENSOR_MAP_SWIZZLE_128B`; wgmma descriptors set `layout_type = 1` + (LBO = 16 B, SBO = 1024 B), avoiding bank conflicts. The SMEM K-tile is 1024 + bits (a full 128B K-major swizzle atom = 4 k256 sub-chunks); per-stage wgmmas + run behind one `commit_group`/`wait_group` and accumulate in-hardware + (`scale-D = 1`). Operands moved to dynamic shared memory. Host pre-arrangement + is plain row-major tiles (no `cm()` interleave). +- **Wide binary MMA** (Phase 4) — began with one `m64n256k256` per k-step + (replacing four `m64n64k256`); the current kernel uses `m64n128k256` strips + (see Phase 10). +- **Register reallocation** (Phase 5) — `setmaxnreg.dec(40)` in the producer, + `setmaxnreg.inc(N)` in the consumer (N sized to the accumulator block). +- **Deeper pipeline** (Phase 6) — `STAGES` full/empty buffers (now 4). +- **TMA output store** (Phase 7) — the packed `sC` tile is written back with a + single `cp.async.bulk.tensor.2d.global.shared::cta`; C is padded to whole + NG-limb column groups so every stored tile is complete. +- **Persistent kernel + grouped rasterization** (Phase 8) — a 1-D grid of + ~SM-count CTAs sweeps the output tiles in a grouped-along-M order (`GROUP_M` + M-tiles per band), shortening each B-panel's reuse distance to keep it + L2-resident. Cuts B's HBM re-reads by ~`GROUP_M`. +- **Clusters + TMA multicast** (Phase 9) — `CLUSTER` CTAs along M form a + thread-block cluster and share one HBM read of each B-panel via + `cp.async.bulk.tensor…multicast::cluster` (each computes a different M-tile). + Cluster-wide empty barrier; the pipeline barriers init once and flow + continuously across tiles. Mirrors the proven `pranjalssh/fast.cu` matmul_9. +- **Register-blocking** (Phase 10) — each CTA computes a `TM×NB` output block + (`MSTRIPS` `m64n128` strips × `NB`=128 cols) whose strips all reuse one loaded + B sub-tile, cutting operand-refill bytes/MAC (the bottleneck after Phase 9). + `MSTRIPS`=3 → a 192×128 block, −33% bytes/MAC. `MSTRIPS`/`TILE_M`/`NG` in the + kernel and `src/lib.rs` must match. +- **Epilogue overlap** (Phase 11) — double-buffered `sC` + deferred + `cp.async.bulk.wait_group.read 1` so tile T's output store overlaps tile T+1's + compute. +- **Fence hoisting** (Phase 12) — one `wgmma.fence` before the K-loop instead of + two per k-chunk. + +Phases 8–9 were validated on hardware (H200 NVL) alongside 10–12. + +- **Wired into `fp`** — the `fp` crate has an optional `gpu` feature + (`cargo …--features fp/gpu`) that pulls in `fp-cuda` and dispatches large + `p = 2` products from `<&Matrix as Mul>::mul` to the GPU, falling back to the + CPU BLAS kernel when no device is present, the size is below + `FP_CUDA_THRESHOLD` (default 2048), or a launch fails. `fp-cuda`'s library API + is fp-agnostic (raw row-major limb slices) so the dependency is acyclic; the + `Matrix` glue lives on the `fp` side (`src/blas/cuda.rs`) and in the + examples/benches (a dev-dependency on `fp`). `FP_CUDA_DISABLE` forces the CPU + path. See the `gpu_dispatch_matches_cpu` integration test. + +Remaining (now TMA-latency bound — the consumer out-runs per-tile TMA latency and +pipeline depth is SMEM-capped at `STAGES`=4): + +- Keep operands resident on the device across `step_resolution`'s successive + multiplications (the current dispatch re-marshals and re-copies each product). +- Extend the parent Nix flake to provide nvcc when the user opts in. diff --git a/ext/crates/fp-cuda/benches/matmul_b1.rs b/ext/crates/fp-cuda/benches/matmul_b1.rs new file mode 100644 index 0000000000..26af1f456f --- /dev/null +++ b/ext/crates/fp-cuda/benches/matmul_b1.rs @@ -0,0 +1,110 @@ +use std::time::Instant; + +use criterion::{BatchSize, Criterion, criterion_group, criterion_main}; +use fp::{matrix::Matrix, prime::TWO}; +use fp_cuda::GpuContext; + +#[path = "../examples/common/mod.rs"] +mod common; +use common::matmul_b1; +use rand::Rng; + +const SIZES: &[usize] = &[128, 256, 512, 1024, 2048, 4096, 8192]; + +fn random_matrix(rows: usize, cols: usize) -> Matrix { + let mut rng = rand::rng(); + let data_len = rows * cols.div_ceil(64); + let data: Vec = (0..data_len).map(|_| rng.random()).collect(); + Matrix::from_data(TWO, rows, cols, data) +} + +fn assert_bit_equal(cpu: &Matrix, gpu: &Matrix) { + assert_eq!(cpu.rows(), gpu.rows(), "row count mismatch"); + assert_eq!(cpu.columns(), gpu.columns(), "column count mismatch"); + assert!( + cpu == gpu, + "CPU and GPU F_2 matmul results disagree at {}x{}", + cpu.rows(), + cpu.columns(), + ); +} + +fn binary_tops(m: usize, k: usize, n: usize, secs: f64) -> f64 { + // 2 * M * N * K binary ops (one AND + one XOR per inner-product step). + 2.0 * (m as f64) * (n as f64) * (k as f64) / secs / 1e12 +} + +fn bench_square(c: &mut Criterion, gpu: &GpuContext, size: usize) { + let mut group = c.benchmark_group(format!("matmul_b1_{size}x{size}")); + group.throughput(criterion::Throughput::Elements( + (2 * size * size * size) as u64, + )); + + // One-shot correctness check (outside the criterion timing loop) before benching. + let (a, b) = (random_matrix(size, size), random_matrix(size, size)); + let cpu_ref = &a * &b; + let gpu_ref = matmul_b1(gpu, &a, &b).expect("GPU matmul launch failed"); + assert_bit_equal(&cpu_ref, &gpu_ref); + + group.bench_function("cpu_fast_mul_concurrent", |bencher| { + bencher.iter_batched( + || (random_matrix(size, size), random_matrix(size, size)), + |(a, b)| &a * &b, + BatchSize::SmallInput, + ); + }); + + group.bench_function("gpu_matmul_b1", |bencher| { + bencher.iter_batched( + || (random_matrix(size, size), random_matrix(size, size)), + |(a, b)| matmul_b1(gpu, &a, &b).expect("GPU matmul launch failed"), + BatchSize::SmallInput, + ); + }); + + group.finish(); + + // Coarse manual TFLOPS report (criterion has its own throughput, but explicit + // logging makes binary-op throughput easy to grep from the bench output). + let runs = 5; + let start = Instant::now(); + for _ in 0..runs { + let _ = matmul_b1(gpu, &a, &b).expect("GPU matmul launch failed"); + } + let gpu_avg = start.elapsed().as_secs_f64() / runs as f64; + + let start = Instant::now(); + for _ in 0..runs { + let _ = &a * &b; + } + let cpu_avg = start.elapsed().as_secs_f64() / runs as f64; + + println!( + "[matmul_b1_{size}x{size}] CPU {:.2} TOPS ({:.2} ms) GPU {:.2} TOPS ({:.2} ms) speedup \ + {:.2}x", + binary_tops(size, size, size, cpu_avg), + cpu_avg * 1e3, + binary_tops(size, size, size, gpu_avg), + gpu_avg * 1e3, + cpu_avg / gpu_avg, + ); +} + +fn bench_all(c: &mut Criterion) { + let gpu = GpuContext::new(0).expect("failed to initialise GpuContext"); + let (major, minor) = gpu + .compute_capability() + .expect("failed to query compute capability"); + println!("fp-cuda bench running on sm_{major}{minor}"); + + for &size in SIZES { + bench_square(c, &gpu, size); + } +} + +criterion_group! { + name = matmul_b1_bench; + config = Criterion::default().measurement_time(std::time::Duration::from_secs(3)); + targets = bench_all +} +criterion_main!(matmul_b1_bench); diff --git a/ext/crates/fp-cuda/build.rs b/ext/crates/fp-cuda/build.rs new file mode 100644 index 0000000000..3f6817247e --- /dev/null +++ b/ext/crates/fp-cuda/build.rs @@ -0,0 +1,76 @@ +//! Compile the CUDA C++ kernel to PTX via nvcc. +//! +//! The emitted `matmul_b1.ptx` is picked up by `src/lib.rs` via +//! `include_bytes!(concat!(env!("OUT_DIR"), "/matmul_b1.ptx"))`. +//! +//! When `nvcc` is **not on PATH** (e.g. CI, or a contributor without the CUDA +//! Toolkit), we emit a *stub* PTX instead of failing the build. This keeps the +//! crate `cargo check`/`clippy`-able everywhere — important because the `fp` +//! crate's optional `cuda` feature pulls `fp-cuda` in, and CI runs +//! `cargo check --workspace --all-features` on a machine with no CUDA. The stub +//! carries no kernel, so at runtime `GpuContext::new` fails cleanly and callers +//! fall back to the CPU path. A `nvcc` that *is* present but fails to compile is +//! still a hard error (a real kernel bug we want to surface). + +use std::{env, fs, path::PathBuf, process::Command}; + +const KERNEL_SRC: &str = "cuda_kernels/matmul_b1.cu"; +const PTX_NAME: &str = "matmul_b1.ptx"; +const ARCH: &str = "sm_90a"; + +/// A minimal, valid PTX module with no kernel. `include_bytes!` needs a file to +/// exist; loading this at runtime fails at `load_function("matmul_b1_kernel")`, +/// which the Rust side surfaces as an error so the caller uses the CPU path. +const STUB_PTX: &str = "//\n// fp-cuda stub: nvcc was unavailable at build time; no GPU kernel \ + was\n// compiled. Install the CUDA Toolkit (12.x+) and rebuild to enable \ + the GPU\n// backend.\n//\n.version 7.8\n.target sm_90a\n.address_size 64\n"; + +fn main() { + println!("cargo:rerun-if-changed={KERNEL_SRC}"); + println!("cargo:rerun-if-changed=build.rs"); + println!("cargo:rerun-if-env-changed=NVCC"); + + let out_dir = PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR not set by cargo")); + let ptx_out = out_dir.join(PTX_NAME); + + let nvcc = env::var("NVCC").unwrap_or_else(|_| "nvcc".to_string()); + + let status = Command::new(&nvcc) + .args([ + "-ptx", + "-O3", + "-std=c++17", + "--use_fast_math", + &format!("-arch={ARCH}"), + KERNEL_SRC, + "-o", + ]) + .arg(&ptx_out) + .status(); + + match status { + // nvcc ran and compiled the kernel: the real PTX is in place. + Ok(s) if s.success() => {} + // nvcc exists but failed to compile: surface it (real kernel error). + Ok(s) => panic!( + "nvcc failed to compile {KERNEL_SRC} (exit status: {s}).\nCheck that your CUDA \ + Toolkit supports {ARCH} (Hopper sm_90a)." + ), + // nvcc genuinely absent: degrade to a stub so the crate still builds. + // The GPU backend is simply unavailable at runtime. + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + println!( + "cargo:warning=fp-cuda: nvcc not found ('{nvcc}': {e}); emitting a stub PTX. The \ + GPU backend will be unavailable at runtime. Install the CUDA Toolkit (12.x+) and \ + rebuild to enable it." + ); + fs::write(&ptx_out, STUB_PTX).expect("failed to write stub PTX to OUT_DIR"); + } + // nvcc is present but couldn't be executed (permissions, a broken exec, + // etc.): fail fast rather than silently hiding a broken CUDA setup. + Err(e) => panic!( + "failed to run nvcc ('{nvcc}': {e}). Set the NVCC env var to a working nvcc, or unset \ + it to fall back to a stub only when nvcc is absent." + ), + } +} diff --git a/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu new file mode 100644 index 0000000000..f55eee6905 --- /dev/null +++ b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu @@ -0,0 +1,463 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// +// Hopper wgmma.b1 F_2 GEMM kernel — 128B-swizzle operands, pipelined wgmmas, +// the widest binary MMA shape (m64n256k256), a persistent grid with grouped +// tile rasterization (Phase 8), and thread-block clusters + TMA B-multicast +// (Phase 9) — both target L2 residency of B at large N. +// +// Both operands are K-major. They are pre-arranged on the host as plain +// row-major tiles and loaded via TMA cp.async.bulk.tensor.2d with +// CU_TENSOR_MAP_SWIZZLE_128B: the TMA hardware applies the 128B swizzle on the +// way into SMEM, landing the data exactly where the swizzled wgmma matrix +// descriptor expects it — so the host emits the natural layout and there is no +// hand-rolled interleave. +// +// Each loaded tile spans a full 128B K-major swizzle atom (8 rows × 1024 bits), +// i.e. KSUB = 4 consecutive k256 sub-chunks. A CTA computes a register-blocked +// TM×NB output block (MSTRIPS m64 row-strips × 128 columns). Each k256 step +// loads B once and issues MSTRIPS m64n128k256 wgmmas that all reuse it — one +// L2→SMEM read of B feeds every strip, which cuts the operand-refill bytes per +// MAC (the measured bottleneck on Hopper: the tensor core out-runs L2→SMEM +// bandwidth). Each strip accumulates into its own resident 64-reg accumulator +// (scale-D = 1, popcounts summed in-hardware), all live across the whole K loop. +// +// The grid is persistent: ~SM-count CTAs (in clusters of CLUSTER along M) sweep +// the output tile grid in a grouped-along-M rasterized order so each B-panel's +// reuse distance stays short (L2-resident). Within a cluster the CLUSTER CTAs +// share one HBM read of each B-panel via TMA multicast — each computes a +// different M-tile but receives the same B into its own SMEM. The pipeline +// barriers are initialized once and flow continuously across tiles; the empty +// barrier is cluster-wide. This mirrors the proven pattern in +// pranjalssh/fast.cu matmul_9.cuh. + +#include +#include +#include + +// ── Helpers ───────────────────────────────────────────────────────────────── + +// Build a wgmma SMEM matrix descriptor. +// p : SMEM address of the operand sub-tile (already swizzled by TMA). +// lead : leading-dimension byte offset (LBO), per CUTLASS make_gmma_desc. +// stride: stride-dimension byte offset (SBO). +// swiz : layout_type — 0 = none, 1 = 128B, 2 = 64B, 3 = 32B. +// Byte offsets are stored with their low 4 bits dropped (uint128 units). +__device__ __forceinline__ uint64_t make_desc( + const void* p, uint32_t lead, uint32_t stride, uint32_t swiz) { + uint32_t a = (uint32_t)__cvta_generic_to_shared(p); + uint64_t d = 0; + d |= ((uint64_t)a >> 4) & 0x3FFFULL; + d |= ((uint64_t)(lead >> 4) & 0x3FFFULL) << 16; + d |= ((uint64_t)(stride >> 4) & 0x3FFFULL) << 32; + d |= ((uint64_t)(swiz & 0x3)) << 62; + return d; +} + +__device__ __forceinline__ void mbar_init(uint64_t* b, uint32_t cnt) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"((uint32_t)__cvta_generic_to_shared(b)), "r"(cnt)); +} +__device__ __forceinline__ void mbar_tx(uint64_t* b, uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" + :: "r"((uint32_t)__cvta_generic_to_shared(b)), "r"(bytes) : "memory"); +} +__device__ __forceinline__ void mbar_wait(uint64_t* b, uint32_t phase) { + uint32_t a = (uint32_t)__cvta_generic_to_shared(b); + asm volatile( + "{ .reg .pred p;\n" + " L: mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n" + " @!p bra L;\n" + "}\n" :: "r"(a), "r"(phase) : "memory"); +} +__device__ __forceinline__ void tma_2d( + void* dst, const CUtensorMap* tm, int x, int y, uint64_t* b) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1, {%2,%3}], [%4];\n" + :: "r"((uint32_t)__cvta_generic_to_shared(dst)), + "l"((uint64_t)tm), "r"(x), "r"(y), + "r"((uint32_t)__cvta_generic_to_shared(b)) + : "memory"); +} + +// ── Cluster helpers (Phase 9: clusters + TMA multicast) ─────────────────────── +// All mirror the proven pattern in pranjalssh/fast.cu matmul_9.cuh. + +// This CTA's rank within its cluster (0..CLUSTER-1). +__device__ __forceinline__ uint32_t cluster_ctarank() { + uint32_t r; + asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(r) :); + return r; +} + +// Cluster-wide barrier: every thread of every CTA in the cluster must arrive. +__device__ __forceinline__ void cluster_sync() { + asm volatile("barrier.cluster.arrive;\n" ::: "memory"); + asm volatile("barrier.cluster.wait;\n" ::: "memory"); +} + +// Arrive (count 1) on the mbarrier `b` located in cluster-mate CTA `cta_id`, +// using mapa to translate the local SMEM address into that CTA's window. +__device__ __forceinline__ void arrive_cluster(uint64_t* b, uint32_t cta_id) { + uint32_t local = (uint32_t)__cvta_generic_to_shared(b); + asm volatile( + "{ .reg .b32 rem;\n" + " mapa.shared::cluster.u32 rem, %0, %1;\n" + " mbarrier.arrive.shared::cluster.b64 _, [rem], 1;\n" + "}\n" :: "r"(local), "r"(cta_id) : "memory"); +} + +// TMA load with cluster multicast: one HBM read of the source tile is fanned +// out into the SMEM of every CTA whose bit is set in `mask` (same `dst` SMEM +// offset and `b` mbarrier offset in each), and counts complete_tx bytes against +// each of their barriers. Issued by a single thread of one CTA. +__device__ __forceinline__ void tma_2d_multicast( + void* dst, const CUtensorMap* tm, int x, int y, uint64_t* b, uint16_t mask) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%2,%3}], [%4], %5;\n" + :: "r"((uint32_t)__cvta_generic_to_shared(dst)), + "l"((uint64_t)tm), "r"(x), "r"(y), + "r"((uint32_t)__cvta_generic_to_shared(b)), "h"(mask) + : "memory"); +} + +// m64n128k256 binary MMA, scale-D = 1 (accumulate into the 64 s32 regs of +// `acc`). da/db are the swizzled operand descriptors. Half the N of the +// widest binary shape, so MSTRIPS of these share one B tile to cut refill BW. +__device__ __forceinline__ void wgmma_n128(int32_t acc[64], uint64_t da, uint64_t db) { + asm volatile( + "wgmma.mma_async.sync.aligned.m64n128k256.row.col.s32.b1.b1.and.popc " + "{" \ + "%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15," \ + "%16,%17,%18,%19,%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31," \ + "%32,%33,%34,%35,%36,%37,%38,%39,%40,%41,%42,%43,%44,%45,%46,%47," \ + "%48,%49,%50,%51,%52,%53,%54,%55,%56,%57,%58,%59,%60,%61,%62,%63}," \ + "%64,%65, 1;\n" + : "+r"(acc[0]),"+r"(acc[1]),"+r"(acc[2]),"+r"(acc[3]),"+r"(acc[4]),"+r"(acc[5]),"+r"(acc[6]),"+r"(acc[7]), + "+r"(acc[8]),"+r"(acc[9]),"+r"(acc[10]),"+r"(acc[11]),"+r"(acc[12]),"+r"(acc[13]),"+r"(acc[14]),"+r"(acc[15]), + "+r"(acc[16]),"+r"(acc[17]),"+r"(acc[18]),"+r"(acc[19]),"+r"(acc[20]),"+r"(acc[21]),"+r"(acc[22]),"+r"(acc[23]), + "+r"(acc[24]),"+r"(acc[25]),"+r"(acc[26]),"+r"(acc[27]),"+r"(acc[28]),"+r"(acc[29]),"+r"(acc[30]),"+r"(acc[31]), + "+r"(acc[32]),"+r"(acc[33]),"+r"(acc[34]),"+r"(acc[35]),"+r"(acc[36]),"+r"(acc[37]),"+r"(acc[38]),"+r"(acc[39]), + "+r"(acc[40]),"+r"(acc[41]),"+r"(acc[42]),"+r"(acc[43]),"+r"(acc[44]),"+r"(acc[45]),"+r"(acc[46]),"+r"(acc[47]), + "+r"(acc[48]),"+r"(acc[49]),"+r"(acc[50]),"+r"(acc[51]),"+r"(acc[52]),"+r"(acc[53]),"+r"(acc[54]),"+r"(acc[55]), + "+r"(acc[56]),"+r"(acc[57]),"+r"(acc[58]),"+r"(acc[59]),"+r"(acc[60]),"+r"(acc[61]),"+r"(acc[62]),"+r"(acc[63]) + : "l"(da), "l"(db)); +} +__device__ __forceinline__ void wgmma_fence() { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); } +__device__ __forceinline__ void wgmma_commit() { asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } +__device__ __forceinline__ void wgmma_wait() { asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); } + +// TMA bulk tensor store (SMEM → global) plus its completion group helpers and +// the async-proxy fence that makes generic-proxy SMEM writes visible to it. +__device__ __forceinline__ void tma_store_2d( + const CUtensorMap* tm, int x, int y, const void* src) { + asm volatile( + "cp.async.bulk.tensor.2d.global.shared::cta.bulk_group [%0, {%1, %2}], [%3];\n" + :: "l"((uint64_t)tm), "r"(x), "r"(y), + "r"((uint32_t)__cvta_generic_to_shared(src)) : "memory"); +} +__device__ __forceinline__ void tma_store_commit() { asm volatile("cp.async.bulk.commit_group;\n" ::: "memory"); } +__device__ __forceinline__ void tma_store_wait() { asm volatile("cp.async.bulk.wait_group 0;\n" ::: "memory"); } +// Wait until all but the most-recent store group have *read* their SMEM source +// (`.read` = don't wait for global visibility). Lets a double-buffered sC free +// the buffer from two tiles ago while the newest store drains in the background. +__device__ __forceinline__ void tma_store_wait_read1() { asm volatile("cp.async.bulk.wait_group.read 1;\n" ::: "memory"); } +__device__ __forceinline__ void fence_async_shared(){ asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } + +// Per-warpgroup register reallocation (warpgroup-aligned). The producer needs +// few registers, so it releases its surplus; the consumer (MSTRIPS*ACC_N-reg +// accumulator) claims them. Counts must be multiples of 8 in [24,256]; at +// 1 CTA/SM the SM's 64K-register budget is ample (128*(40+232) = 34816 at +// MSTRIPS=3), so the binding limit is the 255-reg-per-thread hardware cap. +#define SET_MAXNREG_DEC(N) asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" :: "n"(N)) +#define SET_MAXNREG_INC(N) asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" :: "n"(N)) + +// Output block = MSTRIPS m64 row-strips × NB columns per CTA. Each k256 step +// issues MSTRIPS m64n128 wgmmas that SHARE one B sub-tile, so a single L2→SMEM +// load of B feeds MSTRIPS strips — cutting refill bytes/MAC (the bottleneck) by +// ~1/(1+NB/BM). MSTRIPS is the block knob: 2 → 128×128 block (−20% bytes/MAC, +// 128 acc regs), 3 → 192×128 (−33%, 192 acc regs). NB is fixed at 128 (the +// wgmma_n128 shape); acc regs per thread = MSTRIPS*ACC_N ≤ 240. +constexpr int MSTRIPS = 3; // m64 row-strips per CTA (block knob) +constexpr int MW = 64; // wgmma M extent (fixed for binary wgmma) +constexpr int TK = 1024, KL = TK/64; +constexpr int TM = MW*MSTRIPS; // output rows per CTA (192 at MSTRIPS=3) +constexpr int NB = 128; // n128 output width (columns) per CTA +constexpr int NG = NB/64; // 2 output column-limbs per CTA +constexpr int ACC_N = NB/2; // 64 s32 accumulator regs per m64n128 strip +constexpr int TILE_A = TM*KL; // A tile: TM rows × 16 u64 (192 rows → 3072 u64 = 24 KB) +constexpr int TILE_B = NB*KL; // B tile: 128 cols × 16 u64 = 2048 u64 = 16 KB +constexpr int STROW = MW*KL; // u64 per m64 strip in sA (64*16 = 1024 = 8 KB) +constexpr int SC_STRIDE = NG*TM; // u64 per sC output buffer (double-buffered) +constexpr int KSUB = TK/256; // 4 k256 wgmma sub-chunks per loaded tile +constexpr int KSUB_U64 = 256/64; // 4 u64 = 32 bytes per k256 sub-chunk +// K-loop pipeline depth. Cap is 4 under CLUSTER>1: each stage is 40 KB, so +// STAGES=5 needs ~206 KB, which is under the 227 KB opt-in cap BUT a cluster +// launch reserves extra shared memory (distributed-SMEM / cluster-barrier +// bookkeeping), so 206 KB intermittently over-commits and the kernel faults +// with a flaky "unspecified launch failure" (verified 2026-07-07 H200; the +// sanitizers miss it because it is an async/resource fault). STAGES=5 is stable +// only with CLUSTER=1, and gives no large-N speedup there, so 4 it is. +constexpr int STAGES = 4; // K-loop pipeline depth (full/empty buffers) +constexpr int THREADS_PER_WG = 128; +constexpr int GROUP_M = 16; // M-tiles per rasterization group (L2 reuse knob) +constexpr int CLUSTER = 2; // CTAs per cluster along M (multicast B; reuse knob) +constexpr int PRODUCER_REGS = 40; +// Consumer holds MSTRIPS*ACC_N accumulator regs live across K, plus addressing; +// round up to a multiple of 8, ≤ 240. 1 CTA/SM so the SM reg budget is ample. +constexpr int CONSUMER_REGS = ((MSTRIPS*ACC_N + 40 + 7)/8)*8; + +// wgmma 128B K-major descriptor constants (CUTLASS make_gmma_desc, +// LayoutType::B128): LBO = 1 uint128 = 16 bytes, SBO = 8-row-brick stride = +// 1024 bytes (independent of the MN extent), swizzle = 1. A k256 sub-chunk c +// sits at byte offset c*32 within the tile (advance start_address; the +// hardware re-applies the swizzle). +constexpr uint32_t DESC_LBO = 16; +constexpr uint32_t DESC_SBO = 1024; +constexpr uint32_t DESC_SWIZ = 1; + +// ── Kernel ────────────────────────────────────────────────────────────────── + +// Producer-consumer kernel: 2 warpgroups (256 threads/CTA). +// Warpgroup 0 (t in [0, 128)) = PRODUCER: issues TMA loads in a tight +// K-loop into a STAGES-deep circular SMEM +// buffer. +// Warpgroup 1 (t in [128, 256)) = CONSUMER: waits for each stage to be full, +// runs KSUB*MSTRIPS pipelined m64n128 wgmmas +// against it, signals the stage empty so the +// producer can refill. +// +// Dynamic SMEM per CTA (carved from `smem`, 128B-aligned for TMA): +// sA[STAGES][TILE_A] = STAGES * 24576 B (TM=192 rows) +// sB[STAGES][TILE_B] = STAGES * 16384 B (NB=128 cols) +// sC[2][TM][NG] = 2 * TM * NG * 8 B (double-buffered so the output store +// of tile T overlaps tile T+1's compute) +// mbar_full[STAGES] + mbar_empty[STAGES] +// +// Per stage = sA (24 KB) + sB (16 KB) = 40 KB; STAGES=4 ≈ 163 KB total (requires +// the opt-in CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES set host-side). +// At 40 KB/stage the block runs 1 CTA/SM; STAGES is the K-pipeline-depth knob. +// +// The output block (TM rows × NG limbs) is packed row-major into sC and written +// back with a single TMA bulk store (S2G). C is padded to whole NG-limb column +// groups on the host so every stored tile is complete. +extern "C" __global__ void __cluster_dims__(CLUSTER, 1, 1) matmul_b1_kernel( + const __grid_constant__ CUtensorMap tma_a, + const __grid_constant__ CUtensorMap tma_b, + const __grid_constant__ CUtensorMap tma_c, + uint32_t m_tiles, + uint32_t n_groups, + uint32_t M, uint32_t K) +{ + extern __shared__ __align__(128) uint64_t smem[]; + uint64_t* sA = smem; // [STAGES][TILE_A] + uint64_t* sB = sA + STAGES * TILE_A; // [STAGES][TILE_B] + uint64_t* sC = sB + STAGES * TILE_B; // [2][TM][NG] row-major (double-buffered) + uint64_t* mbar_full = sC + 2 * SC_STRIDE; // [STAGES] + uint64_t* mbar_empty = mbar_full + STAGES; // [STAGES] + + const int t = threadIdx.x; + const int wg = t / THREADS_PER_WG; // 0 = producer, 1 = consumer + const int t_wg = t - wg * THREADS_PER_WG; // 0..127 within warpgroup + + const int nchunks = (K + TK - 1) / TK; + // One full A tile + one full B tile per stage (B is zero-padded on the + // host to a multiple of NB columns, so it is always a complete tile). A is + // loaded per-CTA; B arrives via multicast — both target this CTA's full + // barrier, so the expected bytes are the same as the single-CTA case. + const uint32_t expected_tx = (uint32_t)((TILE_A + TILE_B) * sizeof(uint64_t)); + + // Cluster geometry: CLUSTER CTAs along M share one B-panel via multicast, + // so the schedule walks "M-super-rows" of CLUSTER M-tiles. The host pads + // m_tiles to a multiple of CLUSTER, so m_super divides exactly. + const uint32_t rank = cluster_ctarank(); // 0..CLUSTER-1 (= M offset) + const uint32_t cluster_id = blockIdx.x / CLUSTER; + const uint32_t num_clusters = gridDim.x / CLUSTER; + const uint32_t m_super = m_tiles / CLUSTER; + const uint32_t total_cl = m_super * n_groups; + const uint16_t bmask = (uint16_t)((1u << CLUSTER) - 1u); // all ranks + + // Register reallocation is a one-time per-warpgroup action. + if (wg == 0) SET_MAXNREG_DEC(PRODUCER_REGS); + else SET_MAXNREG_INC(CONSUMER_REGS); + + // Initialize the pipeline barriers ONCE; they flow continuously across the + // persistent tile loop (no per-tile re-init, which would race with the + // cross-CTA arrivals/multicast of a cluster). The empty barrier is + // cluster-wide: it needs one arrival from every CTA's consumer. + if (t == 0) { + #pragma unroll + for (int s = 0; s < STAGES; ++s) { + mbar_init(&mbar_full[s], 1); + mbar_init(&mbar_empty[s], CLUSTER); + } + } + __syncthreads(); + cluster_sync(); // all CTAs' barriers initialized before any cross-CTA arrive + + // Pre-arrive every empty barrier cluster-wide so the producer's first + // STAGES `mbar_wait(empty, 0)` succeed immediately (stages logically free). + if (wg == 1 && t_wg < CLUSTER) { + #pragma unroll + for (int s = 0; s < STAGES; ++s) arrive_cluster(&mbar_empty[s], t_wg); + } + + // ===================== PERSISTENT CLUSTER LOOP ===================== + // A 1-D grid of clusters sweeps the M-super × N tile grid. The grouped + // rasterizer (super-row varies fastest within a GROUP_M band) keeps each + // B-panel's reuse distance short for L2 residency; the cluster additionally + // shares each B-panel HBM read across its CLUSTER CTAs via multicast. + // qidx/p are the running pipeline slot/phase, carried across tiles. + // sC is double-buffered so tile T's output store overlaps tile T+1's + // compute: cbuf ping-pongs per tile, and the store's wait is deferred one + // tile (see epilogue). titer counts this CTA's tiles for the ping-pong. + uint32_t qidx = 0, p = 0, titer = 0; + for (uint32_t ct = cluster_id; ct < total_cl; ct += num_clusters, ++titer) { + const uint32_t gid = ct / (GROUP_M * n_groups); + const uint32_t firstm = gid * GROUP_M; + const uint32_t curm = min((uint32_t)GROUP_M, m_super - firstm); + const uint32_t local = ct - gid * GROUP_M * n_groups; + const uint32_t sbi = firstm + local % curm; + const int bj = (int)(local / curm); + const int bi = (int)(sbi * CLUSTER + rank); // this CTA's M-tile + const int row0 = bi * TM, col0 = bj * NG; + uint64_t* sCb = sC + (titer & 1) * SC_STRIDE; // this tile's sC buffer + + // Before reusing this buffer, make sure the store that last used it (two + // tiles ago) has finished reading it. The deferred .read-1 wait below + // already drained it during the previous tile; this syncs all consumer + // threads to that wait before they overwrite the buffer. + if (wg == 1) { + for (int r = t_wg; r < TM; r += THREADS_PER_WG) { + #pragma unroll + for (int g = 0; g < NG; ++g) sCb[r * NG + g] = 0; + } + } + __syncthreads(); + + if (wg == 0) { + // ===================== PRODUCER ===================== + for (int kk = 0; kk < nchunks; ++kk) { + const uint32_t s = qidx; + + if (t_wg == 0) { + // Wait for all CTAs' consumers to release this stage, then + // set expected bytes (A + multicast B) and issue the loads. + mbar_wait(&mbar_empty[s], p); + mbar_tx(&mbar_full[s], expected_tx); + // A: this CTA's own TM-row block (MSTRIPS m64 strips). + tma_2d(&sA[s * TILE_A], &tma_a, 0, + (kk * m_tiles + bi) * TM, &mbar_full[s]); + // B: one HBM read, multicast into every cluster member's sB + // and counted against every member's full barrier. Issued by + // rank 0 only (its mask bit is set, so it fills itself too). + if (rank == 0) { + tma_2d_multicast(&sB[s * TILE_B], &tma_b, 0, + (kk * n_groups + bj) * NB, &mbar_full[s], + bmask); + } + } + if (++qidx == STAGES) { qidx = 0; p ^= 1; } + } + } else { + // ===================== CONSUMER ===================== + // MSTRIPS m64n128 accumulators (MSTRIPS*ACC_N s32 regs/thread), all + // resident across the whole K loop, re-zeroed per output tile. + int32_t acc[MSTRIPS][ACC_N]; + #pragma unroll + for (int si = 0; si < MSTRIPS; ++si) + #pragma unroll + for (int r = 0; r < ACC_N; ++r) acc[si][r] = 0; + + // One wgmma.fence for the whole K-loop: it orders the non-wgmma + // accumulator zeroing above before the first wgmma. Inside the loop + // the accumulators are written only by wgmma, so no per-chunk fence + // is needed (a warpgroup-wide sync we don't want 32× per tile). The + // per-chunk wgmma.wait_group 0 makes the results readable at the end. + wgmma_fence(); + for (int kk = 0; kk < nchunks; ++kk) { + const uint32_t s = qidx; + + // Wait for the producer's TMAs to finish populating this stage. + mbar_wait(&mbar_full[s], p); + + // Issue every k256 wgmma for this stage behind one commit/wait so + // they pipeline. Each k256 loads B once and reuses it across all + // MSTRIPS strips (independent accumulators → they can overlap). + // scale-D = 1 accumulates each sub-chunk in-hardware. + #pragma unroll + for (int c = 0; c < KSUB; ++c) { + uint64_t db = make_desc(&sB[s * TILE_B + c * KSUB_U64], + DESC_LBO, DESC_SBO, DESC_SWIZ); + #pragma unroll + for (int si = 0; si < MSTRIPS; ++si) { + uint64_t da = make_desc(&sA[s * TILE_A + si * STROW + c * KSUB_U64], + DESC_LBO, DESC_SBO, DESC_SWIZ); + wgmma_n128(acc[si], da, db); + } + } + wgmma_commit(); + wgmma_wait(); + + // Release this stage cluster-wide: arrive on every CTA's empty + // barrier (so rank 0 may overwrite their multicast sB). + if (t_wg < CLUSTER) arrive_cluster(&mbar_empty[s], t_wg); + if (++qidx == STAGES) { qidx = 0; p ^= 1; } + } + + // Pack each strip's NB-wide accumulator into sC's NG output limbs. + // The m64n128 fragment is the m64n64 layout tiled along N: register + // group gi (0..NB/8-1) covers output columns [gi*8, gi*8+8); within it + // this thread owns columns cb, cb+1 for rows rb and rb+8. Strip si adds + // si*MW to the row. Column c maps to limb c/64, bit c%64. + const int wid = t_wg >> 5, lane = t_wg & 31; + const int rb = wid*16 + (lane>>2), cb = (lane&3)*2; + #pragma unroll + for (int si = 0; si < MSTRIPS; ++si) { + uint64_t lo[NG] = {0}, hi[NG] = {0}; + #pragma unroll + for (int gi = 0; gi < NB/8; ++gi) { + int c0 = cb + gi*8, c1 = c0 + 1; + int l0 = c0 >> 6, b0p = c0 & 63; + int l1 = c1 >> 6, b1p = c1 & 63; + lo[l0] |= (uint64_t)(acc[si][gi*4+0]&1) << b0p; + lo[l1] |= (uint64_t)(acc[si][gi*4+1]&1) << b1p; + hi[l0] |= (uint64_t)(acc[si][gi*4+2]&1) << b0p; + hi[l1] |= (uint64_t)(acc[si][gi*4+3]&1) << b1p; + } + // Row-major sC[row*NG + limb]; padded limbs (out-of-range columns) + // get zero popcounts from the zero-padded B, so they store harmless + // zeros into C's padded region (trimmed on the host). + const int rlo = si*MW + rb, rhi = si*MW + rb + 8; + #pragma unroll + for (int g = 0; g < NG; ++g) { + uint32_t* clo = reinterpret_cast(&sCb[rlo * NG + g]); + uint32_t* chi = reinterpret_cast(&sCb[rhi * NG + g]); + atomicXor(&clo[0], (uint32_t)lo[g]); + atomicXor(&clo[1], (uint32_t)(lo[g]>>32)); + atomicXor(&chi[0], (uint32_t)hi[g]); + atomicXor(&chi[1], (uint32_t)(hi[g]>>32)); + } + } + } + + // Write the TM×NG output block back with a single TMA bulk store, then + // DEFER its wait: `.read 1` blocks only until every store but the newest + // (this one) has finished reading its sC buffer, so tile T's store drains + // during tile T+1's compute instead of stalling here. The buffer freed is + // the one from two tiles ago — safe to reuse next time cbuf lands on it. + __syncthreads(); // sC[cbuf] fully packed by the consumer + fence_async_shared(); // make the atomicXor writes visible to the async proxy + if (t == 0) { + tma_store_2d(&tma_c, col0 * 2, row0, sCb); // x in UINT32 units (2 per limb) + tma_store_commit(); + tma_store_wait_read1(); + } + __syncthreads(); // order the deferred wait before the next tile + // reuses (zeroes) an sC buffer + } + // Drain the last outstanding output store before the CTA exits. + if (t == 0) tma_store_wait(); +} diff --git a/ext/crates/fp-cuda/examples/bench_giant.rs b/ext/crates/fp-cuda/examples/bench_giant.rs new file mode 100644 index 0000000000..89e32ecdb0 --- /dev/null +++ b/ext/crates/fp-cuda/examples/bench_giant.rs @@ -0,0 +1,80 @@ +use std::time::Instant; + +use fp::{matrix::Matrix, prime::TWO}; +use fp_cuda::GpuContext; + +mod common; +use common::matmul_b1; +use rand::Rng; + +fn main() -> Result<(), Box> { + let gpu = GpuContext::new(0)?; + let (major, minor) = gpu.compute_capability()?; + println!("GPU: sm_{major}{minor} (H100)"); + println!(); + + let mut rng = rand::rng(); + + // Several GB per matrix: for an NxN binary matrix, storage = N*N/8 bytes + // N=131072: 131072^2/8 = 2 GB per matrix + // N=65536: 65536^2/8 = 512 MB per matrix + // N=32768: 32768^2/8 = 128 MB per matrix + // + // For F2 matmul: 2*N^3 bit-operations (N^3 ANDs + N^3 XORs) + // Binary TOPS = 2*N^3 / time_seconds / 1e12 + + // N*N/8 bytes per matrix: + // N=32768, K=32768: 128 MB each + // N=65536, K=65536: 512 MB each (1 GB pair) + // N=131072, K=32768: 512 MB each + // N=131072, K=65536: 1 GB each (2 GB pair) + // N=131072, K=131072: 2 GB each (4 GB pair) + // N=262144, K=65536: 2 GB each (4 GB pair) + // N=262144, K=131072: 4 GB each (8 GB pair) + for &(m, k_actual, n) in &[ + (32768usize, 32768usize, 32768usize), // 128 MB each, warmup + (65536, 65536, 65536), // 512 MB each + (131072, 65536, 131072), // 1 GB each + (131072, 131072, 131072), // 2 GB each + (262144, 65536, 262144), // 2 GB + 2 GB + (262144, 131072, 262144), // 4 GB each + ] { + let stride_a = k_actual.div_ceil(64); + let stride_c = n.div_ceil(64); + let a_elems = m * stride_a; + let b_elems = k_actual * stride_c; + let a_bytes = a_elems * 8; + let b_bytes = b_elems * 8; + + println!("=== {m} x {k_actual} x {n} ==="); + println!(" A: {m}x{k_actual} = {:.1} MB", a_bytes as f64 / 1e6); + println!(" B: {k_actual}x{n} = {:.1} MB", b_bytes as f64 / 1e6); + + let a_data: Vec = (0..a_elems).map(|_| rng.random()).collect(); + let b_data: Vec = (0..b_elems).map(|_| rng.random()).collect(); + let a = Matrix::from_data(TWO, m, k_actual, a_data); + let b = Matrix::from_data(TWO, k_actual, n, b_data); + + // Warmup + let _ = matmul_b1(&gpu, &a, &b)?; + + // Timed runs + let trials = 3; + let mut best_secs = f64::MAX; + for _ in 0..trials { + let t0 = Instant::now(); + let _ = matmul_b1(&gpu, &a, &b)?; + let elapsed = t0.elapsed().as_secs_f64(); + best_secs = best_secs.min(elapsed); + } + + // 2*M*N*K bit-ops (AND + XOR per element) + let bit_ops = 2.0 * (m as f64) * (n as f64) * (k_actual as f64); + let tops = bit_ops / best_secs / 1e12; + println!(" Time: {:.3} ms", best_secs * 1e3); + println!(" Binary TOPS: {:.2}", tops); + println!(); + } + + Ok(()) +} diff --git a/ext/crates/fp-cuda/examples/bench_kernel.rs b/ext/crates/fp-cuda/examples/bench_kernel.rs new file mode 100644 index 0000000000..173dc15222 --- /dev/null +++ b/ext/crates/fp-cuda/examples/bench_kernel.rs @@ -0,0 +1,86 @@ +use std::time::Instant; + +use fp::{matrix::Matrix, prime::TWO}; +use fp_cuda::GpuContext; + +mod common; +use common::matmul_b1; +use rand::Rng; + +fn main() -> Result<(), Box> { + let gpu = GpuContext::new(0)?; + let (major, minor) = gpu.compute_capability()?; + println!("GPU: sm_{major}{minor}"); + println!("Timing includes host serialization + H2D + kernel + D2H."); + println!(); + + let mut rng = rand::rng(); + + // Focus on compute-bound regime: large K maximizes wgmma fraction + for &(m, k, n) in &[ + (8192usize, 8192usize, 8192usize), + (16384, 16384, 16384), + (32768, 32768, 32768), + ] { + let stride_a = k.div_ceil(64); + let stride_b = n.div_ceil(64); + let a_bytes = m * stride_a * 8; + let b_bytes = k * stride_b * 8; + + println!("=== {m} x {k} x {n} ==="); + println!( + " A: {:.1} MB, B: {:.1} MB", + a_bytes as f64 / 1e6, + b_bytes as f64 / 1e6 + ); + + let a_data: Vec = (0..m * stride_a).map(|_| rng.random()).collect(); + let b_data: Vec = (0..k * stride_b).map(|_| rng.random()).collect(); + let a = Matrix::from_data(TWO, m, k, a_data); + let b = Matrix::from_data(TWO, k, n, b_data); + + // Warmup (includes compilation, allocation) + let _ = matmul_b1(&gpu, &a, &b)?; + + // Timed: end-to-end (host serial + H2D + kernel + D2H) + let trials = 3; + let mut best = f64::MAX; + for _ in 0..trials { + let t0 = Instant::now(); + let _ = matmul_b1(&gpu, &a, &b)?; + best = best.min(t0.elapsed().as_secs_f64()); + } + + let bit_ops = 2.0 * (m as f64) * (n as f64) * (k as f64); + let tops = bit_ops / best / 1e12; + + // Estimate host overhead: time just serialization + padding + let t_host = Instant::now(); + let _a_ser = { + let mut v = Vec::new(); + a.to_bytes(&mut v).unwrap(); + v + }; + let _b_ser = { + let mut v = Vec::new(); + b.to_bytes(&mut v).unwrap(); + v + }; + let host_ser_ms = t_host.elapsed().as_secs_f64() * 1e3; + + println!( + " End-to-end: {:.1} ms → {:.1} binary TOPS", + best * 1e3, + tops + ); + println!( + " Host serialization alone: {:.1} ms ({:.0}% of total)", + host_ser_ms, + host_ser_ms / (best * 1e3) * 100.0 + ); + println!(" K-chunks per CTA: {}", k.div_ceil(256)); + println!(); + } + + Ok(()) +} diff --git a/ext/crates/fp-cuda/examples/bench_kernel_only.rs b/ext/crates/fp-cuda/examples/bench_kernel_only.rs new file mode 100644 index 0000000000..133945ae87 --- /dev/null +++ b/ext/crates/fp-cuda/examples/bench_kernel_only.rs @@ -0,0 +1,69 @@ +//! Kernel-only throughput for `matmul_b1`. +//! +//! Unlike `bench_kernel` (end-to-end, host-serialization-bound) and the +//! `cargo bench` criterion harness (also end-to-end), this isolates the GPU +//! kernel: all host (de)serialization, the TMA-layout pre-arrangement, and the +//! H2D/D2H copies happen once, then only back-to-back kernel launches are +//! timed (see `matmul_b1_timed`). This is the apples-to-apples number to +//! compare against the ~100-binary-TOPS pre-swizzle kernel baseline. +//! +//! Run: `cargo run --release -p fp-cuda --example bench_kernel_only`. + +use fp::{matrix::Matrix, prime::TWO}; +use fp_cuda::GpuContext; + +mod common; +use common::{matmul_b1, matmul_b1_timed}; +use rand::Rng; + +fn binary_tops(m: usize, k: usize, n: usize, secs: f64) -> f64 { + 2.0 * (m as f64) * (n as f64) * (k as f64) / secs / 1e12 +} + +fn main() -> Result<(), Box> { + let gpu = GpuContext::new(0)?; + let (major, minor) = gpu.compute_capability()?; + println!("GPU: sm_{major}{minor}"); + println!("Kernel-only binary TOPS (host setup + H2D/D2H excluded):\n"); + + let mut rng = rand::rng(); + let mut make = |rows: usize, cols: usize| { + let data: Vec = (0..rows * cols.div_ceil(64)) + .map(|_| rng.random()) + .collect(); + Matrix::from_data(TWO, rows, cols, data) + }; + + for &(m, k, n, iters) in &[ + (4096usize, 4096, 4096, 50), + (8192, 8192, 8192, 30), + (16384, 16384, 16384, 10), + (32768, 32768, 32768, 5), + ] { + let a = make(m, k); + let b = make(k, n); + + // Bit-exact correctness check against the CPU path once per shape. + let cpu = &a * &b; + let (gpu_ref, _) = matmul_b1_timed(&gpu, &a, &b, 1)?; + let ok = cpu == gpu_ref; + // matmul_b1 (single launch) must agree with the timed multi-launch path. + let single = matmul_b1(&gpu, &a, &b)?; + let idempotent = single == gpu_ref; + + let (_, secs) = matmul_b1_timed(&gpu, &a, &b, iters)?; + println!( + " {m:>6} x {k:>6} x {n:>6}: {:>7.1} binary TOPS ({:>8.3} ms/launch, {iters} iters) \ + correct={ok} idempotent={idempotent}", + binary_tops(m, k, n, secs), + secs * 1e3, + ); + if !ok || !idempotent { + eprintln!(" CORRECTNESS FAILURE at {m}x{k}x{n}"); + std::process::exit(1); + } + } + + println!("\nH100 binary tensor-op peak is ~360,000 TOPS."); + Ok(()) +} diff --git a/ext/crates/fp-cuda/examples/bench_shapes.rs b/ext/crates/fp-cuda/examples/bench_shapes.rs new file mode 100644 index 0000000000..547546d674 --- /dev/null +++ b/ext/crates/fp-cuda/examples/bench_shapes.rs @@ -0,0 +1,92 @@ +//! Isolate *why* kernel-only throughput drops past N=16384: is it total size, or +//! the B operand spilling out of the 50 MB L2? +//! +//! Each B column-panel is reused across every M-tile, so the L2-reuse condition +//! is simply whether the whole B matrix (K*N/8 bytes) fits in L2. These shapes +//! hold FLOPs fixed while flipping "B fits in L2", which a pure size/occupancy +//! story cannot explain. +//! +//! Run: `cargo run --release -p fp-cuda --example bench_shapes`. + +use fp::{matrix::Matrix, prime::TWO}; +use fp_cuda::GpuContext; + +mod common; +use common::matmul_b1_timed; +use rand::Rng; + +fn binary_tops(m: usize, k: usize, n: usize, secs: f64) -> f64 { + 2.0 * (m as f64) * (n as f64) * (k as f64) / secs / 1e12 +} + +fn main() -> Result<(), Box> { + let gpu = GpuContext::new(0)?; + // Assumed L2 capacity (H100 / H200 NVL ≈ 50 MB); not read from the device. + // Override when profiling a card with a different L2 size. + let l2_mb = 50.0; + println!( + "Assumed GPU L2 ~= {l2_mb} MB (H100/H200 NVL). B in L2 (bytes = K*N/8) governs \ + cross-M-tile reuse.\n" + ); + + let mut rng = rand::rng(); + let mut make = |rows: usize, cols: usize| { + let data: Vec = (0..rows * cols.div_ceil(64)) + .map(|_| rng.random()) + .collect(); + Matrix::from_data(TWO, rows, cols, data) + }; + + // (M, K, N, iters, note) + let shapes = [ + (16384usize, 16384, 16384, 10, "cube, B fits"), + (32768, 32768, 32768, 5, "cube, B spills"), + // Same FLOPs (1.76e13), only B-in-L2 differs: + ( + 65536, + 16384, + 16384, + 10, + "tall: huge M, B FITS (=16384^3 x4 FLOPs)", + ), + ( + 16384, + 16384, + 65536, + 5, + "wide: huge N, B SPILLS (same FLOPs as above)", + ), + ( + 16384, + 65536, + 16384, + 5, + "deep: huge K, B SPILLS (same FLOPs as above)", + ), + // B fits even at 2x the tall-case FLOPs: + ( + 131072, + 16384, + 16384, + 5, + "taller: M=128K, B FITS (=16384^3 x8 FLOPs)", + ), + ]; + + println!( + "{:>7} {:>7} {:>7} | {:>9} {:>6} | {:>9} | note", + "M", "K", "N", "B (MB)", "fits", "TOPS" + ); + for &(m, k, n, iters, note) in &shapes { + let b_mb = (k as f64) * (n as f64) / 8.0 / 1e6; + let a = make(m, k); + let b = make(k, n); + let (_, secs) = matmul_b1_timed(&gpu, &a, &b, iters)?; + println!( + "{m:>7} {k:>7} {n:>7} | {b_mb:>9.1} {:>6} | {:>9.1} | {note}", + if b_mb <= l2_mb { "yes" } else { "NO" }, + binary_tops(m, k, n, secs), + ); + } + Ok(()) +} diff --git a/ext/crates/fp-cuda/examples/common/mod.rs b/ext/crates/fp-cuda/examples/common/mod.rs new file mode 100644 index 0000000000..67efe61580 --- /dev/null +++ b/ext/crates/fp-cuda/examples/common/mod.rs @@ -0,0 +1,54 @@ +//! Shared `fp::Matrix` glue for the examples and benches. +//! +//! The `fp-cuda` library itself is fp-agnostic (it takes raw row-major limb +//! slices) so that the `fp` crate can depend on it without a dependency cycle. +//! These thin wrappers restore the `Matrix`-typed convenience the examples want, +//! using `fp` — which `fp-cuda` pulls in only as a dev-dependency. +//! +//! `#![allow(dead_code)]` because each example/bench uses a different subset. +#![allow(dead_code)] + +use fp::{matrix::Matrix, prime::TWO}; +use fp_cuda::GpuContext; + +/// Row-major, K-major `u64` limbs, exactly as [`matmul_b1_raw`] expects: +/// `rows × columns.div_ceil(64)` limbs, one bit per entry, no inter-row padding. +/// +/// [`matmul_b1_raw`]: fp_cuda::matmul_b1_raw +pub fn to_limbs(m: &Matrix) -> Vec { + let stride = m.columns().div_ceil(64); + let mut bytes = Vec::with_capacity(m.rows() * stride * 8); + m.to_bytes(&mut bytes).expect("Vec writes never fail"); + let (chunks, _) = bytes.as_chunks::<8>(); + chunks.iter().map(|&c| u64::from_le_bytes(c)).collect() +} + +/// `Matrix`-typed wrapper over [`fp_cuda::matmul_b1_raw`]. +pub fn matmul_b1( + gpu: &GpuContext, + a: &Matrix, + b: &Matrix, +) -> Result> { + assert_eq!(a.prime(), TWO); + assert_eq!(b.prime(), TWO); + assert_eq!(a.columns(), b.rows()); + let (m, k, n) = (a.rows(), a.columns(), b.columns()); + let c = fp_cuda::matmul_b1_raw(gpu, &to_limbs(a), m, k, &to_limbs(b), n)?; + Ok(Matrix::from_data(TWO, m, n, c)) +} + +/// `Matrix`-typed wrapper over [`fp_cuda::matmul_b1_raw_timed`]. +pub fn matmul_b1_timed( + gpu: &GpuContext, + a: &Matrix, + b: &Matrix, + time_iters: usize, +) -> Result<(Matrix, f64), Box> { + assert_eq!(a.prime(), TWO); + assert_eq!(b.prime(), TWO); + assert_eq!(a.columns(), b.rows()); + let (m, k, n) = (a.rows(), a.columns(), b.columns()); + let (c, secs) = + fp_cuda::matmul_b1_raw_timed(gpu, &to_limbs(a), m, k, &to_limbs(b), n, time_iters)?; + Ok((Matrix::from_data(TWO, m, n, c), secs)) +} diff --git a/ext/crates/fp-cuda/examples/matmul_b1_demo.rs b/ext/crates/fp-cuda/examples/matmul_b1_demo.rs new file mode 100644 index 0000000000..ed59faa26c --- /dev/null +++ b/ext/crates/fp-cuda/examples/matmul_b1_demo.rs @@ -0,0 +1,61 @@ +//! Smoke test for the `fp-cuda` matmul kernel. +//! +//! Multiplies one pair of small F_2 matrices on the GPU and verifies the result against +//! `fp::blas`. Run with `cargo oxide run -p fp-cuda --example matmul_b1_demo`. + +use fp::{matrix::Matrix, prime::TWO}; +use fp_cuda::GpuContext; + +mod common; +use common::matmul_b1; +use rand::Rng; + +fn main() -> Result<(), Box> { + let gpu = GpuContext::new(0)?; + let (major, minor) = gpu.compute_capability()?; + println!("=== fp-cuda matmul_b1 demo ==="); + println!("GPU compute capability: sm_{major}{minor}"); + + let mut rng = rand::rng(); + let mut make = |rows: usize, cols: usize| { + let data: Vec = (0..rows * cols.div_ceil(64)) + .map(|_| rng.random()) + .collect(); + Matrix::from_data(TWO, rows, cols, data) + }; + + for &(m, k, n) in &[ + (64, 256, 64), + (128, 256, 128), + (256, 256, 256), + (512, 512, 512), + (1024, 1024, 1024), + (2048, 512, 2048), + (4096, 256, 4096), + (8192, 256, 8192), + ] { + let a = make(m, k); + let b = make(k, n); + let cpu = &a * &b; + let gpu_out = matmul_b1(&gpu, &a, &b)?; + let ok = cpu == gpu_out; + println!( + " {m}x{k} * {k}x{n}: {}", + if ok { "OK" } else { "MISMATCH" } + ); + if !ok { + let mut cb = Vec::new(); + cpu.to_bytes(&mut cb).unwrap(); + let mut gb = Vec::new(); + gpu_out.to_bytes(&mut gb).unwrap(); + let cv = u64::from_le_bytes(cb[..8].try_into().unwrap()); + let gv = u64::from_le_bytes(gb[..8].try_into().unwrap()); + println!(" row 0: cpu={cv:016x} gpu={gv:016x}"); + println!(" GPU all zeros: {}", gb.iter().all(|&b| b == 0)); + std::process::exit(1); + } + } + + println!("All shapes matched."); + Ok(()) +} diff --git a/ext/crates/fp-cuda/examples/prof_sizes.rs b/ext/crates/fp-cuda/examples/prof_sizes.rs new file mode 100644 index 0000000000..e2fb8846ed --- /dev/null +++ b/ext/crates/fp-cuda/examples/prof_sizes.rs @@ -0,0 +1,31 @@ +//! Minimal profiling target: exactly one `matmul_b1` kernel launch at each of +//! two sizes (16384, 32768), no CPU cross-check. Built for `ncu --launch-count` +//! so the L2-cliff hypothesis can be checked with hardware counters. +//! +//! `ncu --set basic --launch-count 2 target/release/examples/prof_sizes` + +use fp::{matrix::Matrix, prime::TWO}; +use fp_cuda::GpuContext; + +mod common; +use common::matmul_b1; +use rand::Rng; + +fn main() -> Result<(), Box> { + let gpu = GpuContext::new(0)?; + let mut rng = rand::rng(); + let mut make = |rows: usize, cols: usize| { + let data: Vec = (0..rows * cols.div_ceil(64)) + .map(|_| rng.random()) + .collect(); + Matrix::from_data(TWO, rows, cols, data) + }; + + for &n in &[16384usize, 32768] { + let a = make(n, n); + let b = make(n, n); + let _ = matmul_b1(&gpu, &a, &b)?; + eprintln!("launched {n}"); + } + Ok(()) +} diff --git a/ext/crates/fp-cuda/examples/tune.rs b/ext/crates/fp-cuda/examples/tune.rs new file mode 100644 index 0000000000..f0125e983a --- /dev/null +++ b/ext/crates/fp-cuda/examples/tune.rs @@ -0,0 +1,55 @@ +//! Fast tuning target: kernel-only TOPS at a couple of sizes, no CPU +//! cross-check, single correctness spot-check at 4096. Used by the sweep +//! driver to compare tuning-knob configurations quickly. +//! +//! Run: `cargo run --release -p fp-cuda --example tune` + +use fp::{matrix::Matrix, prime::TWO}; +use fp_cuda::GpuContext; + +mod common; +use common::{matmul_b1, matmul_b1_timed}; +use rand::Rng; + +fn binary_tops(m: usize, k: usize, n: usize, secs: f64) -> f64 { + 2.0 * (m as f64) * (n as f64) * (k as f64) / secs / 1e12 +} + +fn main() -> Result<(), Box> { + let gpu = GpuContext::new(0)?; + let mut rng = rand::rng(); + let mut make = |rows: usize, cols: usize| { + let data: Vec = (0..rows * cols.div_ceil(64)) + .map(|_| rng.random()) + .collect(); + Matrix::from_data(TWO, rows, cols, data) + }; + + // One cheap correctness spot-check so a broken config is caught fast. + { + let a = make(4096, 4096); + let b = make(4096, 4096); + let cpu = &a * &b; + let gpu_ref = matmul_b1(&gpu, &a, &b)?; + if cpu != gpu_ref { + eprintln!("CORRECTNESS FAILURE at 4096"); + std::process::exit(1); + } + } + + for &(m, k, n, iters) in &[ + (8192usize, 8192, 8192, 30), + (16384, 16384, 16384, 15), + (32768, 32768, 32768, 8), + ] { + let a = make(m, k); + let b = make(k, n); + let (_, secs) = matmul_b1_timed(&gpu, &a, &b, iters)?; + println!( + " {m:>6} x {k:>6} x {n:>6}: {:>7.1} TOPS ({:>8.3} ms)", + binary_tops(m, k, n, secs), + secs * 1e3, + ); + } + Ok(()) +} diff --git a/ext/crates/fp-cuda/src/lib.rs b/ext/crates/fp-cuda/src/lib.rs new file mode 100644 index 0000000000..dc34a0d172 --- /dev/null +++ b/ext/crates/fp-cuda/src/lib.rs @@ -0,0 +1,403 @@ +//! CUDA backend for `fp::blas` F_2 matrix multiplication on Hopper. +//! +//! Both operands are pre-arranged on the host as plain row-major K-major tiles +//! and loaded via TMA with 128B swizzle, which lands them in the SMEM layout the +//! swizzled wgmma matrix descriptors expect. The kernel register-blocks a +//! TILE_M×(NG*64) output tile per CTA out of MSTRIPS m64n128 wgmma.b1 strips +//! that share each loaded B tile (cuts operand-refill bandwidth, the bottleneck). + +use std::{ffi::c_void, mem::MaybeUninit, sync::Arc, time::Instant}; + +use cudarc::{ + driver::{ + CudaContext, CudaFunction, CudaModule, CudaStream, DevicePtr, DeviceRepr, LaunchConfig, + PushKernelArg, sys, + }, + nvrtc::Ptx, +}; + +static PTX_IMAGE: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/matmul_b1.ptx")); + +const TILE_M: usize = 192; // MW*MSTRIPS in the kernel; must match +const TILE_K: usize = 1024; +const KL: usize = TILE_K / 64; // 16 +const THREADS: u32 = 256; // 2 warpgroups: producer (0..128) + consumer (128..256) +const NG: u32 = 2; // output column-limbs per CTA (NB/64 = 128/64); must match the kernel +const STAGES: usize = 4; // K-loop pipeline depth; must match the kernel +const CLUSTER: usize = 2; // CTAs per cluster along M (multicast B); must match the kernel + +/// Lets us pass a `CUtensorMap` by value as a (grid-constant) kernel argument +/// through cudarc's typed launch builder. `repr(transparent)` so the pointer +/// cudarc pushes is the address of the 128-byte descriptor itself. +#[repr(transparent)] +struct TmaArg(sys::CUtensorMap); +unsafe impl DeviceRepr for TmaArg {} + +pub struct GpuContext { + ctx: Arc, + #[allow(dead_code)] + module: Arc, + kernel: CudaFunction, +} + +impl GpuContext { + pub fn new(device_id: usize) -> Result> { + let ctx = CudaContext::new(device_id)?; + let ptx = Ptx::from_src(String::from_utf8(PTX_IMAGE.to_vec())?); + let module = ctx.load_module(ptx)?; + let kernel = module.load_function("matmul_b1_kernel")?; + Ok(Self { + ctx, + module, + kernel, + }) + } + + pub fn compute_capability(&self) -> Result<(i32, i32), Box> { + let major = self.ctx.attribute( + sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, + )?; + let minor = self.ctx.attribute( + sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, + )?; + Ok((major, minor)) + } + + pub fn default_stream(&self) -> Arc { + self.ctx.default_stream() + } + + pub fn kernel(&self) -> &CudaFunction { + &self.kernel + } +} + +/// Multiply two bit-packed F₂ matrices on the GPU. +/// +/// Operands are plain **row-major, K-major** limb arrays — the exact layout +/// `fp::Matrix::to_bytes` produces (little-endian `u64` limbs, one bit per +/// entry, `columns.div_ceil(64)` limbs per row, no inter-row padding): +/// +/// - `a`: the `m`×`k` left operand, `m * k.div_ceil(64)` limbs. +/// - `b`: the `k`×`n` right operand, `k * n.div_ceil(64)` limbs. +/// +/// Returns C = A·B as `m * n.div_ceil(64)` limbs in the same layout, ready to +/// hand to `fp::Matrix::from_data`. This keeps `fp-cuda` free of any dependency +/// on `fp` (the higher-level crate depends on this one, not the reverse). +pub fn matmul_b1_raw( + gpu: &GpuContext, + a: &[u64], + m: usize, + k: usize, + b: &[u64], + n: usize, +) -> Result, Box> { + Ok(matmul_b1_inner(gpu, a, m, k, b, n, 1)?.0) +} + +/// Like [`matmul_b1_raw`], but also returns the average **kernel-only** wall +/// time (seconds) over `time_iters` back-to-back launches, excluding host +/// (de)serialization, the TMA-layout pre-arrangement, and the H2D/D2H copies. +/// +/// The kernel zeroes its SMEM accumulator and writes C with a bulk-tensor +/// *store* (overwrite, not accumulate), so repeated launches against the same +/// device buffers are idempotent and the returned limbs are the correct +/// product. Use this to compare against the ~100-binary-TOPS pre-swizzle +/// kernel baseline; the end-to-end `cargo bench` figures are dominated by host +/// serialization and understate kernel throughput. +pub fn matmul_b1_raw_timed( + gpu: &GpuContext, + a: &[u64], + m: usize, + k: usize, + b: &[u64], + n: usize, + time_iters: usize, +) -> Result<(Vec, f64), Box> { + matmul_b1_inner(gpu, a, m, k, b, n, time_iters.max(1)) +} + +#[allow(clippy::too_many_arguments)] +fn matmul_b1_inner( + gpu: &GpuContext, + a: &[u64], + m: usize, + k: usize, + b: &[u64], + n: usize, + time_iters: usize, +) -> Result<(Vec, f64), Box> { + let n_lim = n.div_ceil(64); + assert_eq!(a.len(), m * k.div_ceil(64), "A limb count mismatch"); + assert_eq!(b.len(), k * n_lim, "B limb count mismatch"); + + let k_padded = k.next_multiple_of(TILE_K); + // Pad M to a whole number of clusters (CLUSTER M-tiles) so every cluster + // rank has a valid M-tile; the extra padded rows produce zeros that the + // `take(m)` readback trims. + let m_padded = m.next_multiple_of(TILE_M * CLUSTER); + let m_tiles = m_padded / TILE_M; + let k_chunks = k_padded / TILE_K; + // Each CTA computes a TILE_M×(NG*64) output block via MSTRIPS m64n128 wgmmas, + // so B (and the C output) are grouped/padded to whole NG-limb column tiles. + let n_groups = n_lim.div_ceil(NG as usize); + let n_padded_lim = n_groups * NG as usize; + + let stream = gpu.ctx.default_stream(); + + let a_padded = pad_2d(a, m, k.div_ceil(64), m_padded, k_padded / 64); + let b_padded = pad_2d(b, k, n_lim, k_padded, n_lim); + + // Gather A into row-major K-major tiles; the TMA applies the 128B swizzle. + let a_interleaved = interleave_a(&a_padded, m_padded, k_padded); + // Pre-transpose B into row-major K-major tiles (swizzled by the TMA). + let bt = transpose_b(&b_padded, k_padded, n_lim); + + let a_dev = stream.clone_htod(&a_interleaved)?; + let bt_dev = stream.clone_htod(&bt)?; + let c_dev = stream.alloc_zeros::(m_padded * n_padded_lim)?; + + // Raw device addresses for the TMA descriptors. The returned guards keep the + // reads ordered on the stream; hold them until after the launch. + let (a_ptr, _ga) = a_dev.device_ptr(&stream); + let (b_ptr, _gb) = bt_dev.device_ptr(&stream); + let (c_ptr, _gc) = c_dev.device_ptr(&stream); + + // TMA tensor maps. A: TILE_M-row block per (k_chunk, M-block), split into + // MSTRIPS m64 strips by the consumer. B: (NG*64)-column tile per (k_chunk, + // column group), reused across the strips. Both have a 128-byte inner dim + // (= the 128B swizzle width). C: TILE_M-row × NG-limb output blocks, no + // swizzle, for the bulk store. + let tma_a = encode_tma( + a_ptr, + [32, (k_chunks * m_tiles * TILE_M) as u64], + [32, TILE_M as u32], + 128, + sys::CUtensorMapSwizzle_enum::CU_TENSOR_MAP_SWIZZLE_128B, + )?; + let tma_b = encode_tma( + b_ptr, + [32, (k_chunks * n_groups * NG as usize * 64) as u64], + [32, (NG as usize * 64) as u32], + 128, + sys::CUtensorMapSwizzle_enum::CU_TENSOR_MAP_SWIZZLE_128B, + )?; + let tma_c = encode_tma( + c_ptr, + [(n_padded_lim * 2) as u64, m_padded as u64], + [(NG as usize * 2) as u32, TILE_M as u32], + (n_padded_lim * 8) as u64, + sys::CUtensorMapSwizzle_enum::CU_TENSOR_MAP_SWIZZLE_NONE, + )?; + + // Dynamic SMEM per CTA: sA + sB + 2*sC (double-buffered) + 2*STAGES mbarriers. + let tile_a = TILE_M * KL; // TILE_M-row A block + let tile_b = NG as usize * 64 * KL; // (NG*64)-col B tile + let smem_u64 = STAGES * tile_a + STAGES * tile_b + 2 * NG as usize * TILE_M + 2 * STAGES; + let smem_bytes = (smem_u64 * std::mem::size_of::()) as u32; + + // Opt in to >48 KB shared memory (Hopper static default cap). + gpu.kernel.set_attribute( + sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + smem_bytes as i32, + )?; + + // Persistent grid: co-resident CTAs = (occupancy per SM) × SM count, so the + // grid exactly fills the machine and the kernel's persistent loop sweeps all + // output tiles in grouped-rasterized order. Rounded to a whole number of + // clusters (`__cluster_dims__` requires gridDim.x % CLUSTER == 0); surplus + // clusters run an empty loop. + // + // This is 1 CTA/SM at the register-blocked config, and that is optimal: + // 2 CTAs/SM was measured (2026-07-07 H200) and loses badly. Two CTAs need the + // compiled register count ≤128/thread (2·256·128 = the 64K reg file), but the + // resident accumulator that gives the kernel its arithmetic intensity is + // 192 regs/thread at MSTRIPS=3. Shrinking it to fit two CTAs (MSTRIPS=1, + // 64-reg acc → occ=2) collapses AI and drops 16384 from ~8,600 to ~5,500 + // TOPS. High AI (big accumulator) and high occupancy compete for the same + // register file and AI wins, so we run 1 CTA/SM with the largest accumulator + // that fits under the 255-reg cap. + let sms = gpu + .ctx + .attribute(sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)? + as u32; + let occ = gpu + .kernel + .occupancy_max_active_blocks_per_multiprocessor(THREADS, smem_bytes as usize, None)? + .max(1); + let num_ctas = (occ * sms / CLUSTER as u32).max(1) * CLUSTER as u32; + if std::env::var("FP_CUDA_DEBUG").is_ok() { + eprintln!("[fp-cuda] occ={occ}/SM sms={sms} num_ctas={num_ctas} smem={smem_bytes}B"); + } + + let ta = TmaArg(tma_a); + let tb = TmaArg(tma_b); + let tc = TmaArg(tma_c); + let mt = m_tiles as u32; + let ng = n_groups as u32; + let m_val = m_padded as u32; + let k_val = k_padded as u32; + + let launch = || -> Result<(), cudarc::driver::DriverError> { + let cfg = LaunchConfig { + grid_dim: (num_ctas, 1, 1), + block_dim: (THREADS, 1, 1), + shared_mem_bytes: smem_bytes, + }; + let mut lb = stream.launch_builder(&gpu.kernel); + lb.arg(&ta) + .arg(&tb) + .arg(&tc) + .arg(&mt) + .arg(&ng) + .arg(&m_val) + .arg(&k_val); + unsafe { lb.launch(cfg) }?; + Ok(()) + }; + + // Warm up once (untimed) when measuring, so the timed loop excludes any + // first-launch JIT/allocation costs. + if time_iters > 1 { + launch()?; + stream.synchronize()?; + } + + let start = Instant::now(); + for _ in 0..time_iters { + launch()?; + } + stream.synchronize()?; + let kernel_secs = start.elapsed().as_secs_f64() / time_iters as f64; + + let c_all = stream.clone_dtoh(&c_dev)?; + let c_limbs: Vec = c_all + .chunks_exact(n_padded_lim) + .take(m) + .flat_map(|row| row[..n_lim].iter().copied()) + .collect(); + Ok((c_limbs, kernel_secs)) +} + +/// Encode a 2D row-major TMA tensor map of UINT32 elements. +fn encode_tma( + dev_ptr: sys::CUdeviceptr, + gdim: [u64; 2], + boxdim: [u32; 2], + row_stride_bytes: u64, + swizzle: sys::CUtensorMapSwizzle_enum, +) -> Result> { + let gstride = [row_stride_bytes]; + let elemstride = [1u32, 1u32]; + let mut tmap = MaybeUninit::::uninit(); + unsafe { + sys::cuTensorMapEncodeTiled( + tmap.as_mut_ptr(), + sys::CUtensorMapDataType_enum::CU_TENSOR_MAP_DATA_TYPE_UINT32, + 2, + dev_ptr as *mut c_void, + gdim.as_ptr(), + gstride.as_ptr(), + boxdim.as_ptr(), + elemstride.as_ptr(), + sys::CUtensorMapInterleave_enum::CU_TENSOR_MAP_INTERLEAVE_NONE, + swizzle, + sys::CUtensorMapL2promotion_enum::CU_TENSOR_MAP_L2_PROMOTION_NONE, + sys::CUtensorMapFloatOOBfill_enum::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE, + ) + .result()?; + Ok(tmap.assume_init()) + } +} + +/// Gather A into plain row-major K-major tiles for TMA 128B swizzle. +/// +/// Output: contiguous tiles, each TILE_M rows × KL u64s (64 × 128 bytes). The +/// TMA applies the 128B swizzle on load, so the host layout is the natural +/// row-major sub-block: tile row `row` holds K bits `kk*TILE_K .. +TILE_K` of +/// global row `bi*TILE_M + row`, zero-padded out of bounds. +/// +/// Tiles are ordered: for K-chunk kk=0..k_chunks-1, then M-tile bi=0..m_tiles-1. +fn interleave_a(a: &[u64], m: usize, k: usize) -> Vec { + let sa = k / 64; + let k_chunks = k / TILE_K; + let m_tiles = m / TILE_M; + let tile_u64s = TILE_M * KL; + let mut out = vec![0u64; k_chunks * m_tiles * tile_u64s]; + + for kk in 0..k_chunks { + for bi in 0..m_tiles { + let base = (kk * m_tiles + bi) * tile_u64s; + for row in 0..TILE_M { + for kl in 0..KL { + let global_row = bi * TILE_M + row; + let global_kl = kk * KL + kl; + let val = if global_row < m && global_kl < sa { + a[global_row * sa + global_kl] + } else { + 0 + }; + out[base + row * KL + kl] = val; + } + } + } + } + out +} + +/// Pre-transpose B into plain row-major K-major tiles for TMA 128B swizzle. +/// +/// Each (k_chunk, column group) tile is NB = NG*64 rows (= the NG*64 output +/// columns of the group) × KL u64s (= TILE_K K bits); the consumer feeds it to +/// MSTRIPS m64n128 wgmmas that share it. Operand row `lg*64 + jj` is output +/// column `cg*NG*64 + lg*64 + jj`; element `[..][kl] bit` is bit `jj` of +/// `B[k_chunk*TILE_K + kl*64 + bit][cg*NG + lg]`. +/// Groups whose limb runs past `n_lim` are left zero-padded. Output is +/// row-major; the TMA applies the swizzle on load. +fn transpose_b(b: &[u64], k: usize, n_lim: usize) -> Vec { + let k_chunks = k / TILE_K; + let ng = NG as usize; + let n_groups = n_lim.div_ceil(ng); + let tile = ng * 64 * KL; // 256 rows × KL u64 + let mut out = vec![0u64; k_chunks * n_groups * tile]; + let mut buf = [0u64; TILE_K]; + + for kk in 0..k_chunks { + for cg in 0..n_groups { + let base = (kk * n_groups + cg) * tile; + for lg in 0..ng { + let limb = cg * ng + lg; + if limb >= n_lim { + continue; // padded column group → leave zeros + } + for (i, slot) in buf.iter_mut().enumerate() { + let br = kk * TILE_K + i; + *slot = if br < k { b[br * n_lim + limb] } else { 0 }; + } + for jj in 0..64usize { + let j = lg * 64 + jj; // operand row within the 256-col tile + for kl in 0..KL { + let mut val: u64 = 0; + for bit in 0..64usize { + val |= ((buf[kl * 64 + bit] >> jj) & 1) << bit; + } + out[base + j * KL + kl] = val; + } + } + } + } + } + out +} + +fn pad_2d(src: &[u64], rows: usize, stride: usize, nr: usize, ns: usize) -> Vec { + if rows == nr && stride == ns { + return src.to_vec(); + } + let mut out = vec![0u64; nr * ns]; + for r in 0..rows { + let n = stride.min(ns); + out[r * ns..r * ns + n].copy_from_slice(&src[r * stride..r * stride + n]); + } + out +} diff --git a/ext/crates/fp/Cargo.toml b/ext/crates/fp/Cargo.toml index b0123aba05..fec9c6765f 100644 --- a/ext/crates/fp/Cargo.toml +++ b/ext/crates/fp/Cargo.toml @@ -21,6 +21,12 @@ serde_json = "1.0.141" maybe-rayon = { path = "../maybe-rayon" } query = { path = "../query" } +# GPU backend (Hopper wgmma.b1). Optional and off by default: gated behind the +# `gpu` feature. `fp-cuda`'s library is fp-agnostic (raw limbs), so this +# dependency does not form a cycle. Building it requires nvcc at build time (a +# stub PTX is emitted otherwise) and a Hopper GPU + driver at runtime. +fp-cuda = { path = "../fp-cuda", optional = true } + [dev-dependencies] # We use the proptest harness for our own tests fp = { path = ".", default-features = false, features = ["proptest"] } @@ -38,6 +44,8 @@ build_const = "0.2.2" default = ["odd-primes"] concurrent = ["maybe-rayon/concurrent"] odd-primes = [] +# Dispatch large p=2 matrix products to the Hopper GPU backend (`fp-cuda`). +gpu = ["dep:fp-cuda"] [[bench]] name = "mul" diff --git a/ext/crates/fp/src/blas/cuda.rs b/ext/crates/fp/src/blas/cuda.rs new file mode 100644 index 0000000000..070e1a77c0 --- /dev/null +++ b/ext/crates/fp/src/blas/cuda.rs @@ -0,0 +1,89 @@ +//! GPU dispatch for F₂ matrix multiplication (Hopper `wgmma.b1`). +//! +//! Compiled only under the `gpu` feature. [`try_mul`] is consulted by +//! `<&Matrix as Mul>::mul` before the CPU BLAS path: for large enough `p = 2` +//! products it converts the operands to the raw row-major limb layout +//! `fp-cuda` expects, runs the kernel, and rebuilds a [`Matrix`]. Anything that +//! makes the GPU path unavailable or unsuitable — no device, a launch error, or +//! a below-threshold size — returns `None`, and the caller falls back to the +//! (bit-identical) CPU kernel. +//! +//! Tuning knobs (environment variables, read once): +//! - `FP_CUDA_DISABLE` — set to any value to force the CPU path. +//! - `FP_CUDA_THRESHOLD` — minimum of `m`, `k`, `n` (in bits) below which the +//! CPU path is used. Defaults to 2048; the GPU only wins once the kernel work +//! dwarfs the H2D/D2H + TMA-layout marshalling, which dominates small sizes. + +use std::sync::{Mutex, OnceLock}; + +use fp_cuda::GpuContext; + +use crate::{matrix::Matrix, prime::TWO}; + +/// Smallest `min(m, k, n)` for which we attempt the GPU. Below this the host +/// marshalling (bit-repack into TMA tiles + copies) costs more than it saves. +const DEFAULT_THRESHOLD: usize = 2048; + +fn threshold() -> usize { + std::env::var("FP_CUDA_THRESHOLD") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(DEFAULT_THRESHOLD) +} + +/// The process-wide GPU context, created lazily on first use. `None` if no +/// usable device is present (no driver, no Hopper GPU, or the kernel PTX is the +/// nvcc-absent build stub). Wrapped in a `Mutex` because a single CUDA context +/// serialises submission anyway and `GpuContext` is not shared concurrently. +fn context() -> Option<&'static Mutex> { + static GPU: OnceLock>> = OnceLock::new(); + GPU.get_or_init(|| { + if std::env::var_os("FP_CUDA_DISABLE").is_some() { + return None; + } + match GpuContext::new(0) { + Ok(ctx) => Some(Mutex::new(ctx)), + Err(_) => None, + } + }) + .as_ref() +} + +/// Row-major, K-major `u64` limbs — the exact layout `fp_cuda::matmul_b1_raw` +/// expects (`rows × columns.div_ceil(64)` limbs, no inter-row padding). Uses +/// `Matrix::to_bytes`, which already strips the physical row stride. +fn to_limbs(m: &Matrix) -> Vec { + let stride = m.columns().div_ceil(64); + let mut bytes = Vec::with_capacity(m.rows() * stride * 8); + m.to_bytes(&mut bytes).expect("Vec writes never fail"); + let (chunks, _) = bytes.as_chunks::<8>(); + chunks.iter().map(|&c| u64::from_le_bytes(c)).collect() +} + +/// Try to compute `a · b` on the GPU. Returns `None` (and the caller uses the +/// CPU path) if the GPU is unavailable, the product is below the size +/// threshold, or the launch fails. The result is bit-identical to the CPU path. +/// +/// Assumes `a.prime() == b.prime() == 2` and `a.columns() == b.rows()` — the +/// same preconditions the caller has already checked. +pub(super) fn try_mul(a: &Matrix, b: &Matrix) -> Option { + debug_assert_eq!(a.prime(), TWO); + debug_assert_eq!(b.prime(), TWO); + debug_assert_eq!(a.columns(), b.rows()); + + let (m, k, n) = (a.rows(), a.columns(), b.columns()); + let t = threshold(); + if m < t || k < t || n < t { + return None; + } + + let ctx = context()?; + let a_limbs = to_limbs(a); + let b_limbs = to_limbs(b); + + let c = { + let guard = ctx.lock().ok()?; + fp_cuda::matmul_b1_raw(&guard, &a_limbs, m, k, &b_limbs, n).ok()? + }; + Some(Matrix::from_data(TWO, m, n, c)) +} diff --git a/ext/crates/fp/src/blas/mod.rs b/ext/crates/fp/src/blas/mod.rs index b5487628a8..73e5976bb6 100644 --- a/ext/crates/fp/src/blas/mod.rs +++ b/ext/crates/fp/src/blas/mod.rs @@ -31,6 +31,9 @@ use crate::matrix::Matrix; pub mod block; pub mod tile; +#[cfg(feature = "gpu")] +mod cuda; + impl std::ops::Mul for &Matrix { type Output = Matrix; @@ -44,6 +47,10 @@ impl std::ops::Mul for &Matrix { { // Can use optimized BLAS operations (matrix rows are padded to multiple of 64) // TODO: Use different block sizes and loop orders based on the size of the matrices + #[cfg(feature = "gpu")] + if let Some(result) = cuda::try_mul(self, rhs) { + return result; + } self.fast_mul_concurrent(rhs) } else { // Use naive multiplication for: diff --git a/ext/crates/fp/tests/cuda_dispatch.rs b/ext/crates/fp/tests/cuda_dispatch.rs new file mode 100644 index 0000000000..09e1db72ea --- /dev/null +++ b/ext/crates/fp/tests/cuda_dispatch.rs @@ -0,0 +1,45 @@ +//! GPU-dispatch correctness for `<&Matrix as Mul>::mul` under the `gpu` +//! feature. Only compiled when the feature is on; if no usable GPU is present +//! the dispatch transparently falls back to the CPU kernel, so this test still +//! passes (it just doesn't exercise the device). Run with `FP_CUDA_DEBUG=1` to +//! see the `[fp-cuda]` launch line and confirm the GPU path was taken. +#![cfg(feature = "gpu")] + +use fp::{matrix::Matrix, prime::TWO}; +use rand::Rng; + +fn random_matrix(rows: usize, cols: usize) -> Matrix { + let mut rng = rand::rng(); + let limbs = cols.div_ceil(64); + let data: Vec = (0..rows * limbs).map(|_| rng.random()).collect(); + Matrix::from_data(TWO, rows, cols, data) +} + +/// The dispatched product must be bit-identical to the CPU BLAS kernel. Sizes +/// are chosen above the default 2048 threshold so the GPU path is attempted. +#[test] +fn gpu_dispatch_matches_cpu() { + for &(m, k, n) in &[ + (2048, 2048, 2048), + (4096, 2048, 3072), + (3072, 4096, 2048), + // Non-tile-aligned dims (not multiples of the kernel's 192/128/1024 + // tiles; rows still pad to a multiple of 64 so dispatch fires): + // exercise edge masks, partial limbs, and raster tails. + (2049, 2051, 2053), + (3000, 2112, 4097), + ] { + let a = random_matrix(m, k); + let b = random_matrix(k, n); + + // `&a * &b` consults the GPU (feature on); `fast_mul_concurrent` is the + // reference CPU kernel. They must agree bit-for-bit. + let dispatched = &a * &b; + let reference = a.fast_mul_concurrent(&b); + + assert_eq!( + dispatched, reference, + "GPU/CPU mismatch at {m}x{k} * {k}x{n}" + ); + } +} diff --git a/ext/flake.nix b/ext/flake.nix index 35e0d02fc8..0858136926 100644 --- a/ext/flake.nix +++ b/ext/flake.nix @@ -7,7 +7,23 @@ outputs = {super, ...}: super.flake-utils.lib.eachDefaultSystem (system: let - pkgs = import super.nixpkgs {inherit system;}; + # Allow CUDA (unfree in nixpkgs). Scoped to the CUDA / NVIDIA prefix so + # we don't accidentally unfree-allow anything else. nixpkgs splits the + # toolkit into many sub-derivations (cuda_nvcc, cuda_cudart, cuda-merged, + # cuda_cuobjdump, libcublas, ...) — listing them individually is whack- + # a-mole, so we match by prefix. + pkgs = import super.nixpkgs { + inherit system; + config.allowUnfreePredicate = pkg: + let + lib = super.nixpkgs.lib; + name = lib.getName pkg; + in + lib.hasPrefix "cuda" name + || lib.hasPrefix "libcu" name + || lib.hasPrefix "libnv" name + || lib.hasPrefix "libnpp" name; + }; pythonEnv = pkgs.python3.withPackages (ps: [ ps.black @@ -27,6 +43,13 @@ pkgs.perf ] ++ super.defaultPackages.devTools.${system}; + + # CUDA toolkit for building fp-cuda's Hopper wgmma.b1 kernel: nvcc + headers + # at build time. Kept out of `commonPackages` (and the default shell) so + # contributors and the `apps.test`/CI closure don't fetch the multi-GB unfree + # CUDA tree for the opt-in backend. cudarc dlopens libcuda at runtime, so + # only running — not building the Rust — needs the host driver. + cudatoolkit = pkgs.cudaPackages.cudatoolkit; in { devShells.default = pkgs.mkShell { packages = commonPackages; @@ -35,6 +58,17 @@ ''; }; + # GPU dev shell: `nix develop .#gpu`. Adds the CUDA toolkit (nvcc + headers) + # and points the loader at both it and the host driver's libcuda. + devShells.gpu = pkgs.mkShell { + packages = commonPackages ++ [cudatoolkit]; + shellHook = '' + export RUST_LOG=info + export CUDA_PATH="${cudatoolkit}" + export LD_LIBRARY_PATH="${cudatoolkit}/lib:/run/opengl-driver/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + ''; + }; + apps.test = { type = "app"; packages = commonPackages;