pub const TERNARY_GEMM_2BIT_WGSL: &str = "// \u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\n// QualiaDB \u{2014} BitNet b1.58 ternary GEMM, **2-bit branchless** variant (STELLAR \u{a7}A).\n//\n// out[m][i] = scale \u{b7} \u{3a3}_j trit(W[i][j]) \u{b7} act[m][j] W \u{2208} {-1, 0, +1}\n//\n// Optimised for GPU execution (external-review-driven):\n// \u{2022} 2-bit packing (4 trits/byte: 00=0, 01=+1, 10=-1) \u{2192} unpack with SHIFT + MASK only,\n// no integer `/` or `%` (base-3 needs both \u{2014} dozens of cycles on Ampere).\n// \u{2022} BRANCHLESS accumulation: the trit becomes a float multiplier `f32(c==1) - f32(c==2)`\n// and the body is a single FMA. Every thread in the warp runs identical instructions \u{2014}\n// no divergence (vs the `if trit>0 \u{2026} else if trit<0` form in `ternary_gemm.wgsl`).\n// On a GPU the multiply is free (FMA); the ternary win here is bandwidth + occupancy.\n//\n// Bindings + params are identical to `ternary_gemm.wgsl`; CPU oracle: `ternary::ternary_gemm_cpu_2bit`.\n// \u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\n\nstruct TernaryParams {\n n_in: u32,\n n_out: u32,\n n_batch: u32,\n in_row_stride: u32,\n out_row_stride: u32,\n scale: f32,\n _pad0: u32,\n _pad1: u32,\n}\n\n@group(0) @binding(0) var<storage, read> activations: array<f32>;\n@group(0) @binding(1) var<storage, read> trit_words: array<u32>;\n@group(0) @binding(2) var<uniform> params: TernaryParams;\n@group(0) @binding(3) var<storage, read_write> ternary_output: array<f32>;\n\n// One packed byte from the u32 word array (little-endian within the word).\nfn read_byte(idx: u32) -> u32 {\n let w = trit_words[idx >> 2u];\n return (w >> ((idx & 3u) * 8u)) & 0xFFu;\n}\n\n// 2-bit code {0,1,2} at linear weight index k (4 codes/byte) \u{2192} trit {0,+1,-1} as f32, branchless.\nfn trit_f32(k: u32) -> f32 {\n let byte = read_byte(k >> 2u); // k / 4\n let code = (byte >> ((k & 3u) * 2u)) & 3u; // (k % 4) * 2 bits\n return f32(code == 1u) - f32(code == 2u); // +1, -1, or 0 \u{2014} no branch\n}\n\n@compute @workgroup_size(64)\nfn ternary_gemm(@builtin(global_invocation_id) gid: vec3<u32>) {\n let i = gid.x; // output feature\n let m = gid.y; // batch row\n let batch = max(params.n_batch, 1u);\n if (m >= batch || i >= params.n_out) {\n return;\n }\n\n let in_stride = select(params.n_in, params.in_row_stride, params.in_row_stride > 0u);\n let out_stride = select(params.n_out, params.out_row_stride, params.out_row_stride > 0u);\n let in_base = m * in_stride;\n let row0 = i * params.n_in;\n\n var acc = 0.0;\n for (var j = 0u; j < params.n_in; j = j + 1u) {\n // single FMA, identical across the warp\n acc = acc + trit_f32(row0 + j) * activations[in_base + j];\n }\n\n ternary_output[m * out_stride + i] = params.scale * acc;\n}\n";Expand description
The branchless 2-bit WGSL ternary-GEMM kernel; CPU oracle is ternary_gemm_cpu_2bit.