qualia_core_db/wgsl_forge/emit/
coopmat.rs

1//! Cooperative-matrix (tensor-core) WGSL emission (plan §18).
2//!
3//! Emits a single 8x8x8 GEMM tile `C = A * B` using naga's WGSL cooperative
4//! matrix extension (`enable wgpu_cooperative_matrix`, `coop_mat8x8<T, role>`,
5//! `coopLoadT`/`coopMultiplyAdd`/`coopStoreT`). One subgroup cooperatively
6//! computes the tile. Gated on [`crate::wgsl_forge::AdapterConstraints::supports_coopmat`].
7//!
8//! ## Why all-f32, and why the multiply is *execution*-blocked on wgpu 29.0.3
9//!
10//! This emitter is all-f32 8x8x8 because that is the only configuration wgpu/naga
11//! 29.0.3 even claims to support: `wgpu-types-29.0.3/features.rs:1375` states
12//! "The implementation currently only supports 8x8 **f32** matrices", and on Vulkan
13//! it gates the feature on `vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR`
14//! matching 8x8x8 f32. (Mixed-precision f16-in/f32-acc MulAdd — the canonical
15//! reduced-precision tensor-core config — was added to wgpu *after* the 29 line,
16//! gfx-rs/wgpu#9629, MSL-first; an earlier f16 revision of this file requested that
17//! unimplemented config.) The participation set is correct at `@workgroup_size(32)`:
18//! naga declares the coop-matrix SPIR-V type with `Scope::Subgroup`
19//! (`naga-29.0.3/back/spv/writer.rs`: `get_index_constant(spirv::Scope::Subgroup)`),
20//! so on NVIDIA one 32-lane warp is exactly one subgroup — verified empirically,
21//! `@workgroup_size(8,8,1)` gives the identical result.
22//!
23//! Even so, the all-f32 `coopMultiplyAdd` returns **all-zeros** when executed on the
24//! 29.0.3 Vulkan path (the `coopLoadT`/`coopStoreT` round-trip works — only the
25//! multiply fails). This matches gfx-rs/wgpu#9729/#9741: coopmat emits Device-scope
26//! SPIR-V memory ops that are invalid/no-op'd unless `vulkanMemoryModelDeviceScope`
27//! is auto-enabled at device creation — a fix that landed on git `main` **after**
28//! 29.0.3. 29.0.3 is the newest published wgpu (crates.io), so there is no released
29//! fix; the WGSL multiply will start working when a wgpu release carries #9741 (or
30//! by pinning wgpu to a git commit — a core-dependency decision left to the human).
31//! naga's own coopmat test is a WGSL→SPIR-V *translation* test, not a GPU-execution
32//! test, so its passing never implied the multiply runs.
33//!
34//! The kernel below is therefore correct and naga-validated, and `evaluate_matmul_tc`
35//! is kept ready to assert it the moment the upstream fix ships. Until then the
36//! genuine tensor-core multiply is delivered + hardware-verified via the CUDA WMMA
37//! path (`emit::cuda_c::WMMA_GEMM_16X16_SRC`, `oracle::evaluate_matmul_tc_cuda`),
38//! which uses NVIDIA's mature `nvcuda::wmma` API and is unaffected by this wgpu bug.
39
40/// The fixed tile dimension (rows == columns == K) of the emitted GEMM.
41pub const TILE: u32 = 8;
42
43/// Emits the cooperative-matrix 8x8 GEMM tile `C = A * B`, all-f32 (the only
44/// configuration wgpu/naga 29 implements). Row-major loads/stores (`coopLoadT`/
45/// `coopStoreT`, stride = TILE) reproduce a standard row-major reference
46/// `c[i][j] = sum_k a[i][k] * b[k][j]`, so it verifies against
47/// [`crate::wgsl_forge::oracle::matmul_cpu`]. The accumulator is seeded by loading
48/// `c`, which the caller zero-fills.
49pub fn matmul_tc_wgsl() -> String {
50    format!(
51        r#"enable wgpu_cooperative_matrix;
52
53@group(0) @binding(0) var<storage, read> a: array<f32>;
54@group(0) @binding(1) var<storage, read> b: array<f32>;
55@group(0) @binding(2) var<storage, read_write> c: array<f32>;
56
57@compute @workgroup_size(32)
58fn matmul_tc() {{
59    // `T` (row-major) loads + a row-major store match a host row-major matmul.
60    let a_frag = coopLoadT<coop_mat8x8<f32, A>>(&a[0], {TILE}u);
61    let b_frag = coopLoadT<coop_mat8x8<f32, B>>(&b[0], {TILE}u);
62    let acc_in = coopLoadT<coop_mat8x8<f32, C>>(&c[0], {TILE}u);
63    let acc = coopMultiplyAdd(a_frag, b_frag, acc_in);
64    coopStoreT(acc, &c[0], {TILE}u);
65}}"#,
66        TILE = TILE,
67    )
68}
69
70/// Entry point of the tiled cooperative-matrix GEMM ([`matmul_tc_wgsl_tiled`]).
71pub const MATMUL_TC_TILED_ENTRY: &str = "matmul_tc_tiled";
72
73/// Emits a **tiled** cooperative-matrix GEMM `C[m×n] = A[m×k]·B[k×n]` (row-major,
74/// all-f32) that loops the proven single-8×8×8-tile primitive over arbitrary `m`, `n`,
75/// `k` (each a multiple of [`TILE`]). One workgroup (== one subgroup == one warp on
76/// NVIDIA, `@workgroup_size(32)`) computes one 8×8 output tile, accumulating across the
77/// K dimension in a cooperative-matrix register fragment — exactly the structure of the
78/// CUDA WMMA tiled kernel ([`crate::wgsl_forge::emit::cuda_c::WMMA_GEMM_TILED_SRC`]), so
79/// the two backends mirror each other for the same DAG node.
80///
81/// Bindings: `a`(0, read) `b`(1, read) `c`(2, read_write, zero-seeded) and
82/// `dims`(3, read) = `[m, n, k]`. Dispatch one workgroup per output tile:
83/// `num_tiles = (m/8)·(n/8)`, each picked by `@builtin(workgroup_id).x`. Row-major
84/// `coopLoadT`/`coopStoreT` with the runtime leading dimensions (`k` for A, `n` for B/C)
85/// reproduce the standard reference, so it verifies against
86/// [`crate::wgsl_forge::oracle::matmul_cpu`].
87///
88/// **Dormant on wgpu 29.0.3**: the `coopMultiplyAdd` returns zeros there (gfx-rs/wgpu
89/// #9741, merged upstream but unreleased — see the module header). The kernel is correct
90/// and naga-validated; [`crate::wgsl_forge::dispatch::coopmat_usable`] probes the multiply
91/// at runtime so this path stays gated off until a wgpu release (or soft-fork) carries the
92/// fix, then self-activates. Until then the genuine tensor-core GEMM ships via CUDA WMMA.
93pub fn matmul_tc_wgsl_tiled() -> String {
94    format!(
95        r#"enable wgpu_cooperative_matrix;
96
97@group(0) @binding(0) var<storage, read> a: array<f32>;
98@group(0) @binding(1) var<storage, read> b: array<f32>;
99@group(0) @binding(2) var<storage, read_write> c: array<f32>;
100@group(0) @binding(3) var<storage, read> dims: array<u32>;
101
102@compute @workgroup_size(32)
103fn {ENTRY}(@builtin(workgroup_id) wid: vec3<u32>) {{
104    let m = dims[0];
105    let n = dims[1];
106    let k = dims[2];
107    let tiles_n = n / {TILE}u;
108    let num_tiles = (m / {TILE}u) * tiles_n;
109    let tile = wid.x;
110    if (tile >= num_tiles) {{ return; }}
111    let tile_row = tile / tiles_n;
112    let tile_col = tile % tiles_n;
113    // Output-tile base offset in the row-major C[m×n].
114    let c_off = (tile_row * {TILE}u) * n + tile_col * {TILE}u;
115    var acc = coopLoadT<coop_mat8x8<f32, C>>(&c[c_off], n);
116    // Accumulate the 8-wide K tiles: A row-major leading dim k, B leading dim n.
117    for (var kt: u32 = 0u; kt < k; kt = kt + {TILE}u) {{
118        let a_off = (tile_row * {TILE}u) * k + kt;
119        let b_off = kt * n + tile_col * {TILE}u;
120        let a_frag = coopLoadT<coop_mat8x8<f32, A>>(&a[a_off], k);
121        let b_frag = coopLoadT<coop_mat8x8<f32, B>>(&b[b_off], n);
122        acc = coopMultiplyAdd(a_frag, b_frag, acc);
123    }}
124    coopStoreT(acc, &c[c_off], n);
125}}"#,
126        ENTRY = MATMUL_TC_TILED_ENTRY,
127        TILE = TILE,
128    )
129}
qualia_core_db/wgsl_forge/emit/coopmat.rs

qualia_core_db/wgsl_forge/emit/
coopmat.rs
