Skip to main content

qualia_core_db/wgsl_forge/emit/
df64.rs

1//! Double-single (`df64`) emulated double-precision GEMM in **raw WGSL**.
2//!
3//! WGSL has no `f64` — only `f32`/`f16`/`i32`/`u32`. The native exact-double GPU
4//! path is therefore CUDA-only ([`crate::wgsl_forge::emit::cuda_c::GEMM_F64_SRC`]),
5//! which covers NVIDIA. This module fills the *other* half of the "best f64 path on
6//! every machine" story: on **any** wgpu-capable GPU (AMD, Intel, Apple, Adreno,
7//! Mali, llvmpipe, …) it emulates each `f64` as a hi/lo pair of `f32` — a
8//! *double-single* number carrying ~44–48 effective mantissa bits — and does the
9//! GEMM accumulation with **error-free transforms** (Dekker `two_prod`, Knuth
10//! `two_sum`, `quick_two_sum`). This is the df64/double-single technique.
11//!
12//! # Why this is a RAW WGSL string, not the portable IR
13//!
14//! The df64 arithmetic (the `two_sum`/`two_prod`/`df_add`/`df_mul` helpers operating
15//! on `vec2<f32>` hi/lo pairs, with a Veltkamp-split Dekker `two_prod`) does not map onto the
16//! forge's portable scalar-op IR. So — exactly like the cooperative-matrix kernel
17//! ([`crate::wgsl_forge::oracle::evaluate_coopmat_loadstore`]) — it is shipped as a
18//! hand-written WGSL source string, compiled via
19//! [`crate::wgsl_forge::execute::WgpuPipeline::compile`] and dispatched directly,
20//! rather than emitted from a [`crate::wgsl_forge::KernelSpec`].
21//!
22//! # Buffer ABI (matches the CUDA-f64 / forge GEMM binding layout)
23//!
24//! | binding | name   | usage              | layout                                   |
25//! |---------|--------|--------------------|------------------------------------------|
26//! | 0       | `a`    | storage, read      | `M*K` df64 = `2*M*K` f32, `[hi,lo,hi,lo…]`|
27//! | 1       | `b`    | storage, read      | `K*N` df64 = `2*K*N` f32, `[hi,lo,…]`     |
28//! | 2       | `c`    | storage, read_write| `M*N` df64 = `2*M*N` f32, `[hi,lo,…]`     |
29//! | 3       | `dims` | storage, read      | `[m, n, k]` as `u32`                      |
30//!
31//! Row-major `C[M×N] = A[M×K] · B[K×N]`. One invocation computes one output element
32//! (`@workgroup_size(64)`, `gid.x` over `m*n`). The inner `kk` accumulation order
33//! matches the f64 CPU reference [`crate::wgsl_forge::dispatch::gemm_cpu_f64`], so the
34//! two agree to df64 precision (~1e-12 for K≈64 O(1) data) **on adapters whose WGSL
35//! float arithmetic is not reassociated by the driver**.
36//!
37//! # Correctness depends on per-op IEEE rounding (probed at runtime)
38//!
39//! Every df64 algorithm (Veltkamp split, Dekker `two_prod`, Knuth `two_sum`) relies on
40//! each f32 `+`/`-`/`*` rounding exactly once, with no algebraic reassociation. Some GPU
41//! shader toolchains break that: on the naga→SPIR-V→NVIDIA-Vulkan path on this hardware,
42//! the compiler simplifies `c - (c - a)` to `a` (and an fma-based residual `fma(x,y,-(x*y))`
43//! to `0`), which collapses the residual (lo) terms — df64 then silently degrades to f32
44//! (~2e-7 error) instead of ~double precision. Switching `two_prod` from `fma` to the
45//! Veltkamp split gave a **byte-identical** wrong result, confirming reassociation (not a
46//! missing fma) as the cause; WGSL has no portable pragma to disable it. The
47//! [`crate::wgsl_forge::dispatch::gemm_f64`] chain therefore **probes** this kernel at
48//! runtime ([`df64_usable`](crate::wgsl_forge::dispatch)) and uses it only on adapters
49//! where it actually delivers f64 precision; elsewhere it falls to native CUDA-f64 or the
50//! exact CPU floor. The kernel below is correct on a faithful-IEEE adapter and is kept as
51//! the portable f64-GPU path for those.
52
53/// Entry-point name of [`GEMM_DF64_WGSL`].
54pub const GEMM_DF64_ENTRY: &str = "gemm_df64";
55
56/// Raw WGSL source for the df64 (double-single) emulated-f64 GEMM. See the module
57/// docs for the binding ABI and precision contract. The error-free-transform helpers
58/// are transcribed verbatim — `quick_two_sum`/`two_sum`/`two_prod` and the
59/// `df_add`/`df_mul` pair-arithmetic are subtle and must not be "simplified".
60pub const GEMM_DF64_WGSL: &str = r#"@group(0) @binding(0) var<storage, read> a: array<f32>;
61@group(0) @binding(1) var<storage, read> b: array<f32>;
62@group(0) @binding(2) var<storage, read_write> c: array<f32>;
63@group(0) @binding(3) var<storage, read> dims: array<u32>;
64fn quick_two_sum(x: f32, y: f32) -> vec2<f32> { let s = x + y; let e = y - (s - x); return vec2<f32>(s, e); }
65fn two_sum(x: f32, y: f32) -> vec2<f32> { let s = x + y; let v = s - x; let e = (x - (s - v)) + (y - v); return vec2<f32>(s, e); }
66// Veltkamp split of an f32 into two 12-bit halves (factor 2^12+1 = 4097). Uses only
67// IEEE +/-/* that naga preserves, so it does NOT depend on a fused fma (WGSL `fma`
68// does not reliably lower to a true single-rounding FMA on the naga->SPIR-V->NVIDIA
69// path, which silently collapses an fma-based two_prod residual to ~0 = f32 precision).
70fn split(a: f32) -> vec2<f32> { let c = 4097.0 * a; let hi = c - (c - a); let lo = a - hi; return vec2<f32>(hi, lo); }
71// Dekker TwoProduct: exact product as (p, e) with no fma. p = round(x*y); e = the
72// rounding error reconstructed from the split partial products.
73fn two_prod(x: f32, y: f32) -> vec2<f32> { let p = x * y; let xs = split(x); let ys = split(y); let e = ((xs.x * ys.x - p) + xs.x * ys.y + xs.y * ys.x) + xs.y * ys.y; return vec2<f32>(p, e); }
74fn df_add(x: vec2<f32>, y: vec2<f32>) -> vec2<f32> { var s = two_sum(x.x, y.x); s.y = s.y + x.y + y.y; return quick_two_sum(s.x, s.y); }
75fn df_mul(x: vec2<f32>, y: vec2<f32>) -> vec2<f32> { var p = two_prod(x.x, y.x); p.y = p.y + (x.x * y.y + x.y * y.x); return quick_two_sum(p.x, p.y); }
76@compute @workgroup_size(64)
77fn gemm_df64(@builtin(global_invocation_id) gid: vec3<u32>) {
78    let m = dims[0]; let n = dims[1]; let k = dims[2];
79    let o = gid.x;
80    if (o >= m * n) { return; }
81    let row = o / n; let col = o % n;
82    var acc = vec2<f32>(0.0, 0.0);
83    for (var kk: u32 = 0u; kk < k; kk = kk + 1u) {
84        let ai = (row * k + kk) * 2u;
85        let bi = (kk * n + col) * 2u;
86        let av = vec2<f32>(a[ai], a[ai + 1u]);
87        let bv = vec2<f32>(b[bi], b[bi + 1u]);
88        acc = df_add(acc, df_mul(av, bv));
89    }
90    c[o * 2u] = acc.x;
91    c[o * 2u + 1u] = acc.y;
92}
93"#;
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use crate::wgsl_forge::validate::validate_wgsl;
99
100    /// The df64 GEMM source must naga-validate with `gemm_df64` as its entry point.
101    /// This pins that the `vec2<f32>` hi/lo helpers (`two_sum`/`two_prod` and the
102    /// `df_add`/`df_mul` pair-arithmetic, including the `fma` error term) parse and
103    /// type-check under naga — independently of any GPU adapter being present.
104    #[test]
105    fn gemm_df64_wgsl_validates() {
106        let report = validate_wgsl(GEMM_DF64_WGSL).expect("df64 GEMM WGSL must naga-validate");
107        assert!(
108            report.entry_points.iter().any(|e| e == GEMM_DF64_ENTRY),
109            "validated module must expose the {GEMM_DF64_ENTRY} entry point; got {:?}",
110            report.entry_points
111        );
112    }
113}