Skip to main content

Module dispatch

Module dispatch 

Source
Expand description

Capability-aware “best path on this machine” compute dispatcher, keystoned on GEMM.

The forge proves and tunes individual kernels; this module is the layer that, at runtime, picks the best compute path actually available on this machine for a given call, while keeping a CPU floor so the call is never broken. All backends are present in the code; which one activates is decided by the machine’s probed ComputeCaps, not by the call site.

§Why f32 and f64 take different best-paths

WGSL has no f64 — only f32/f16/i32/u32. So the GPU best-path for single precision is the certified WGSL GEMM (via ForgeRuntime); but for double precision there is no WGSL analogue at all. Native f64 on the GPU has to come from CUDA/PTX, which has a real double type and fma.rn.f64. That is the whole reason gemm_f32 and gemm_f64 resolve to different backends:

dtypebest path (if available)floor (always present)
f32WGSL GEMM (ForgeRuntime::gemm)gemm_cpu (f32)
f64native CUDA-f64 → df64-WGSL (double-single)gemm_cpu_f64

§f64 on every GPU: the 3-tier chain (native CUDA → df64-WGSL → CPU)

f64 now has a GPU path on every machine, not just NVIDIA. The chain in gemm_f64 is three tiers:

  1. native CUDA-f64 ([gemm_f64_cuda]) — exact double via PTX fma.rn.f64, NVIDIA only (cuda feature + a CUDA device).
  2. df64 / double-single WGSL (gemm_f64_df64) — emulated double on any other wgpu adapter (AMD, Intel, Apple, mobile). Each f64 is a hi/lo pair of f32 and the accumulation uses error-free transforms (Dekker/TwoSum/TwoProd), giving ~44–48 effective mantissa bits — well beyond a single f32’s 24. The kernel is the raw WGSL GEMM_DF64_WGSL.
  3. CPU floor (gemm_cpu_f64) — exact double, always present, never broken.

So a non-NVIDIA GPU can get real f64 acceleration (tier 2) instead of dropping straight to the CPU — but only where the adapter’s WGSL float arithmetic preserves the df64 error-free transforms. Many drivers (incl. the naga→SPIR-V→NVIDIA-Vulkan path) reassociate floats (c - (c - a)a, fma(x,y,-(x*y))0), which silently collapses df64 to f32 precision. WGSL exposes no portable way to forbid that, so tier 2 is gated on a runtime precision probe ([df64_usable]): df64 runs only where it actually delivers ~double precision; elsewhere the chain uses native CUDA (if present) or the exact CPU floor — never a degraded df64 masquerading as f64. (GEMV’s f64 chain is CUDA-f64 → CPU — the df64 path is GEMM-only today.)

Structs§

ComputeCaps
Probed compute capabilities of this machine. Every flag reflects what was actually constructible at probe time, not what the build was compiled with — a cuda-feature build on a machine with no NVIDIA device still reports cuda == false.

Constants§

GEMM_GPU_THRESHOLD
Problem-size threshold (in m * n * k multiply-adds) below which GEMM stays on the CPU regardless of available accelerators. Small GEMMs are dominated by dispatch/transfer overhead, so the GPU path only earns its keep above this size. 1 << 15 (32768 FMAs, e.g. a 32×32×32 GEMM) is a conservative crossover that keeps the hand-checked unit tests (well below it) firmly on the CPU floor.

Functions§

caps
The probed ComputeCaps for this machine, computed once and cached for the process lifetime. The probe never panics: each backend is tried with ..::new(_).is_ok(), and any failure (no adapter, no driver, no device) is simply recorded as the corresponding flag being false.
coopmat_usable
Runtime probe: does this adapter’s WGSL cooperative-matrix (tensor-core) multiply actually compute, or does it return zeros? Measured once, then cached — the f32 mirror of [df64_usable].
ensure_cuda_runtime_path
Prepend CUDA_PATH/bin/x64 (and bin) to PATH so cudarc can dlopen NVRTC. CUDA 12+/13 ships nvrtc64_*.dll under bin\x64, not bin — without this, gemm_f32_tc always soft-falls to plain f32 even when the toolkit is installed. Idempotent; safe to call from any thread (best-effort env mutation).
fft_f32
Best-path forward FFT: out = DFT(in) over n = complex_interleaved.len()/2 complex points, input and output interleaved f32 ([re0, im0, re1, im1, …], length 2*n). The transform is un-normalized and uses the forward sign convention X[k] = Σ_j x[j] · e^{−2πi kj/N}, identical on both paths.
gemm_cpu_f64
CPU reference for the double-precision dense GEMM — the f64 mirror of gemm_cpu: row-major C[M×N] = A[M×K] · B[K×N], C[i][j] = sum_{k<K} A[i*K + k] * B[k*N + j]. The inner kk sum order matches the CUDA-f64 kernel so the two agree to f64 summation precision. This is the always-present f64 floor.
gemm_f32
Best-path single-precision dense GEMM: row-major C[M×N] = A[M×K] · B[K×N].
gemm_f64
Best-path double-precision dense GEMM: row-major C[M×N] = A[M×K] · B[K×N], all f64.
gemm_f32_tc
f32-faithful tensor-core GEMM — row-major C[m×n] = A[m×k]·B[k×n], f32 in/out, computed to full f32 accuracy. This is the entry point for callers who want tensor-core throughput without trading precision.
gemm_f32_tc_coopmat
Cooperative-matrix (tensor-core) f32 GEMM on a wgpu adapter: row-major C[m×n] = A[m×k]·B[k×n], all f32, computed by the tiled coopmat kernel (matmul_tc_wgsl_tiled). m, n, k must be non-zero multiples of 8 (the 8×8×8 cooperative-matrix tile).
gemm_f32_tc_reduced
Reduced-precision tensor-core GEMM — the explicit opt-in for callers that are precision-tolerant (LLM matmuls are already f16-tolerant) and want maximum tensor-core throughput. Row-major C[m×n] = A[m×k]·B[k×n], f32 in/out, but the result may be f16-precision when the CUDA WMMA tier fires.
gemm_f64_df64
Emulated double-precision (df64 / double-single) dense GEMM on any wgpu adapter: row-major C[M×N] = A[M×K] · B[K×N], all f64.
gemm_tc_cuda
Tensor-core GEMM via the tiled CUDA WMMA kernel: row-major C[m×n] = A[m×k]·B[k×n], with A/B rounded to f16 and accumulated in f32 on NVIDIA tensor cores. This is the genuine reduced-precision tensor-core path — the throughput win that the plain f32 GEMM cannot get — exposed as an opt-in (MatMul.tc) because it trades f32 precision for speed. m, n, k must be non-zero multiples of 16 (the WMMA tile); callers with other shapes pad or fall back to the plain path.
gemv_cpu_f64
CPU reference for the double-precision dense GEMV — the f64 mirror of gemv_cpu: row-major y[M] = A[M×N] · x[N], y[i] = sum_{j<N} A[i*N + j] * x[j]. The inner j sum order matches the CUDA-f64 kernel so the two agree to f64 summation precision. This is the always-present f64 floor.
gemv_f32
Best-path single-precision dense GEMV: row-major y[M] = A[M×N] · x[N].
gemv_f64
Best-path double-precision dense GEMV: row-major y[M] = A[M×N] · x[N], all f64.
pairwise_sq_dist_cpu_f64
Exact direct reference for pairwise_sq_dist_f64: D[i][j] = Σ_d (a[i][d] − b[j][d])² computed without the ‖·‖² identity, so there is no cancellation. Always on the CPU, always correct; this is the always-present floor and the differential oracle for the accelerated form. Returns a zero-filled n×m for any shape mismatch.
pairwise_sq_dist_f64
All-pairs squared Euclidean distance D[i][j] = ‖a_i − b_j‖² between the rows of a (n×p, row-major) and b (m×p, row-major), returned row-major n×m.