wgblas 1.2.0 → 2.0.0

package/src/shaders/browser-shaders.mjs

@@ -1,4 +1,5 @@
 import argmax from "./reduction/argmax.wgsl";
+import argmaxF64 from "./reduction/argmaxF64.wgsl";
 import sum from "./reduction/sum.wgsl";
 import sumF64 from "./reduction/sumF64.wgsl";
 import sscal from "./sscal.wgsl";
@@ -20,13 +21,26 @@ import ssyr from "./ssyr.wgsl";
 import ssyr2 from "./ssyr2.wgsl";
 import f64add from "./f64add.wgsl";
 import dekker from "./f64/dekker.wgsl";
+import ddAbs from "./f64/utils/abs.wgsl";
+import ddAddUtil from "./f64/utils/add.wgsl";
+import ddGreater from "./f64/utils/greater.wgsl";
+import ddEqual from "./f64/utils/equal.wgsl";
 import dasum from "./dasum.wgsl";
+import idamax from "./idamax.wgsl";
 import strsv_invert_block from "./strsv_invert_block.wgsl";
 import strsv_apply_inverse from "./strsv_apply_inverse.wgsl";
 import strsv_update from "./strsv_update.wgsl";
+import sgemm_small from "./sgemm_small.wgsl";
+import sgemm_large from "./sgemm_large.wgsl";
+import sgemmtr_small from "./sgemmtr_small.wgsl";
+import sgemmtr_large from "./sgemmtr_large.wgsl";
+import symmetrize from "./symmetrize.wgsl";
+import triangularize from "./triangularize.wgsl";
+import blockTransfer from "./block_transfer.wgsl";
 
 export const shaderSources = {
   "reduction/argmax": argmax,
+  "reduction/argmaxF64": argmaxF64,
   "reduction/sum": sum,
   "reduction/sumF64": sumF64,
   sscal,
@@ -48,8 +62,20 @@ export const shaderSources = {
   ssyr2,
   f64add,
   "f64/dekker": dekker,
+  "f64/utils/abs": ddAbs,
+  "f64/utils/add": ddAddUtil,
+  "f64/utils/greater": ddGreater,
+  "f64/utils/equal": ddEqual,
   dasum,
+  idamax,
   strsv_invert_block,
   strsv_apply_inverse,
   strsv_update,
+  sgemm_small,
+  sgemm_large,
+  sgemmtr_small,
+  sgemmtr_large,
+  symmetrize,
+  triangularize,
+  block_transfer: blockTransfer,
 };

package/src/shaders/dasum.wgsl

@@ -1,6 +1,7 @@
 // dasum: sum(|x[i]|), double-double (Dekker). Same ILP=4 shape as sasum.wgsl;
-// see f64/dekker.wgsl for ddAddProtected and why plain ddAdd isn't safe.
-// GpuVector input isn't pre-abs'd, so ddAbs() applies unconditionally below.
+// see f64/utils/add.wgsl for ddAddProtected and why plain ddAdd isn't safe.
+// GpuVector input isn't pre-abs'd, so ddAbs() (f64/utils/abs.wgsl) applies
+// unconditionally below.
 
 @group(0) @binding(0) var<storage, read>       xHi:        array<f32>;
 @group(0) @binding(1) var<storage, read>       xLo:        array<f32>;

package/src/shaders/f64/dekker.wgsl

@@ -7,93 +7,12 @@
 //
 // No bindings, no entry point — a helper library, concatenated with a
 // consumer's own bindings/entry point by getPipeline (WGSL has no #include).
+// The DD struct lives here — abs.wgsl/add.wgsl/greater.wgsl/equal.wgsl all
+// use it but don't redefine it (WGSL errors on duplicate struct definitions
+// once concatenated), so any consumer using those must concatenate this
+// file too, first.
 
 struct DD {
   hi: f32,
   lo: f32,
 }
-
-// |a| for a double-double pair. Negation is exact (no rounding), so this is
-// just a sign flip on both components — hi alone determines the pair's sign.
-fn ddAbs(a: DD) -> DD {
-  if (a.hi < 0.0) {
-    return DD(-a.hi, -a.lo);
-  }
-  return a;
-}
-
-// ── A real compiler bug — read before touching anything below ──────────────
-//
-// twoSum/fastTwoSum's error term `e` should be nonzero (that's the point —
-// `s` is rounded). A front-end-optimizer bug zeros it anyway, on both NVIDIA
-// and Mesa (ANV + llvmpipe), via two different mechanisms needing two fixes:
-// bitcast-based subtraction (`fsub`/`negf`, fixes NVIDIA) and materializing
-// the sum through workgroup memory + workgroupBarrier() (fixes Mesa). Only
-// both together (ddAddProtected) is verified correct everywhere — the plain
-// twoSum/fastTwoSum/ddAdd below are reference-only, not safe to use.
-fn negf(x: f32) -> f32 {
-  return bitcast<f32>(bitcast<u32>(x) ^ 0x80000000u);
-}
-fn fsub(a: f32, b: f32) -> f32 {
-  return a + negf(b);
-}
-
-// Knuth/Møller's TwoSum: s = fl(a+b), e = exact rounding error, a+b == s+e.
-// Works for any a, b. UNPROTECTED — see header above.
-fn twoSum(a: f32, b: f32) -> DD {
-  let s = a + b;
-  let v = s - a;
-  let e = (a - (s - v)) + (b - v);
-  return DD(s, e);
-}
-
-// Dekker's Fast-Two-Sum: same contract, but only correct when |a| >= |b|.
-// UNPROTECTED — see header above.
-fn fastTwoSum(a: f32, b: f32) -> DD {
-  let s = a + b;
-  let e = b - (s - a);
-  return DD(s, e);
-}
-
-// Double-double addition (Dekker's Add2). UNPROTECTED — see header above.
-fn ddAdd(a: DD, b: DD) -> DD {
-  let s = twoSum(a.hi, b.hi);
-  let loSum = a.lo + b.lo;
-  return fastTwoSum(s.hi, s.lo + loSum);
-}
-
-// ── Protected variants — use these ──────────────────────────────────────────
-//
-// Bitcast subtraction + workgroup-barrier materialization, verified correct
-// on all three backends tested. Costs a real barrier: fine for O(1)-per-
-// thread or O(log n) reduction use, not a long per-element loop. A
-// workgroupBarrier() requires uniform control flow, so:
-//   - `threadSlot` must be unique per concurrent caller (e.g. local_invocation_index).
-//   - Every thread in the workgroup must call this the same number of times
-//     — including ones whose result gets discarded. Compute unconditionally;
-//     only the write-back should be conditional.
-var<workgroup> dekkerScratch: array<f32, 64>;
-
-fn twoSumProtected(a: f32, b: f32, threadSlot: u32) -> DD {
-  dekkerScratch[threadSlot] = a + b;
-  workgroupBarrier();
-  let s = dekkerScratch[threadSlot];
-  let v = fsub(s, a);
-  let e = fsub(a, fsub(s, v)) + fsub(b, v);
-  return DD(s, e);
-}
-
-fn fastTwoSumProtected(a: f32, b: f32, threadSlot: u32) -> DD {
-  dekkerScratch[threadSlot] = a + b;
-  workgroupBarrier();
-  let s = dekkerScratch[threadSlot];
-  let e = fsub(b, fsub(s, a));
-  return DD(s, e);
-}
-
-// Protected double-double addition — same contract as ddAdd, but exact.
-fn ddAddProtected(a: DD, b: DD, threadSlot: u32) -> DD {
-  let s = twoSumProtected(a.hi, b.hi, threadSlot);
-  let loSum = a.lo + b.lo;
-  return fastTwoSumProtected(s.hi, s.lo + loSum, threadSlot);
-}

package/src/shaders/f64/utils/abs.wgsl

@@ -0,0 +1,10 @@
+// Requires f64/dekker.wgsl concatenated first for the DD struct.
+
+// |a| for a double-double pair. Negation is exact (no rounding), so this is
+// just a sign flip on both components — hi alone determines the pair's sign.
+fn ddAbs(a: DD) -> DD {
+  if (a.hi < 0.0) {
+    return DD(-a.hi, -a.lo);
+  }
+  return a;
+}

package/src/shaders/f64/utils/add.wgsl

@@ -0,0 +1,77 @@
+// Requires f64/dekker.wgsl concatenated first for the DD struct.
+
+// ── A real compiler bug — read before touching anything below ──────────────
+//
+// twoSum/fastTwoSum's error term `e` should be nonzero (that's the point —
+// `s` is rounded). A front-end-optimizer bug zeros it anyway, on both NVIDIA
+// and Mesa (ANV + llvmpipe), via two different mechanisms needing two fixes:
+// bitcast-based subtraction (`fsub`/`negf`, fixes NVIDIA) and materializing
+// the sum through workgroup memory + workgroupBarrier() (fixes Mesa). Only
+// both together (ddAddProtected) is verified correct everywhere — the plain
+// twoSum/fastTwoSum/ddAdd below are reference-only, not safe to use.
+fn negf(x: f32) -> f32 {
+  return bitcast<f32>(bitcast<u32>(x) ^ 0x80000000u);
+}
+fn fsub(a: f32, b: f32) -> f32 {
+  return a + negf(b);
+}
+
+// Knuth/Møller's TwoSum: s = fl(a+b), e = exact rounding error, a+b == s+e.
+// Works for any a, b. UNPROTECTED — see header above.
+fn twoSum(a: f32, b: f32) -> DD {
+  let s = a + b;
+  let v = s - a;
+  let e = (a - (s - v)) + (b - v);
+  return DD(s, e);
+}
+
+// Dekker's Fast-Two-Sum: same contract, but only correct when |a| >= |b|.
+// UNPROTECTED — see header above.
+fn fastTwoSum(a: f32, b: f32) -> DD {
+  let s = a + b;
+  let e = b - (s - a);
+  return DD(s, e);
+}
+
+// Double-double addition (Dekker's Add2). UNPROTECTED — see header above.
+fn ddAdd(a: DD, b: DD) -> DD {
+  let s = twoSum(a.hi, b.hi);
+  let loSum = a.lo + b.lo;
+  return fastTwoSum(s.hi, s.lo + loSum);
+}
+
+// ── Protected variants — use these ──────────────────────────────────────────
+//
+// Bitcast subtraction + workgroup-barrier materialization, verified correct
+// on all three backends tested. Costs a real barrier: fine for O(1)-per-
+// thread or O(log n) reduction use, not a long per-element loop. A
+// workgroupBarrier() requires uniform control flow, so:
+//   - `threadSlot` must be unique per concurrent caller (e.g. local_invocation_index).
+//   - Every thread in the workgroup must call this the same number of times
+//     — including ones whose result gets discarded. Compute unconditionally;
+//     only the write-back should be conditional.
+var<workgroup> dekkerScratch: array<f32, 64>;
+
+fn twoSumProtected(a: f32, b: f32, threadSlot: u32) -> DD {
+  dekkerScratch[threadSlot] = a + b;
+  workgroupBarrier();
+  let s = dekkerScratch[threadSlot];
+  let v = fsub(s, a);
+  let e = fsub(a, fsub(s, v)) + fsub(b, v);
+  return DD(s, e);
+}
+
+fn fastTwoSumProtected(a: f32, b: f32, threadSlot: u32) -> DD {
+  dekkerScratch[threadSlot] = a + b;
+  workgroupBarrier();
+  let s = dekkerScratch[threadSlot];
+  let e = fsub(b, fsub(s, a));
+  return DD(s, e);
+}
+
+// Protected double-double addition — same contract as ddAdd, but exact.
+fn ddAddProtected(a: DD, b: DD, threadSlot: u32) -> DD {
+  let s = twoSumProtected(a.hi, b.hi, threadSlot);
+  let loSum = a.lo + b.lo;
+  return fastTwoSumProtected(s.hi, s.lo + loSum, threadSlot);
+}

package/src/shaders/f64/utils/equal.wgsl

@@ -0,0 +1,7 @@
+// Requires f64/dekker.wgsl concatenated first for the DD struct.
+
+// a == b for double-double pairs — exact field equality, no rounding
+// involved, so (like ddGreater) this needs no protection.
+fn ddEqual(a: DD, b: DD) -> bool {
+  return a.hi == b.hi && a.lo == b.lo;
+}

package/src/shaders/f64/utils/greater.wgsl

@@ -0,0 +1,12 @@
+// Requires f64/dekker.wgsl concatenated first for the DD struct.
+
+// a > b for double-double pairs. hi dominates (|lo| <= ulp(hi)/2 always), so
+// comparing hi alone is correct except on an exact hi tie, when lo breaks it.
+// A plain comparison, not a rounding-identity subtraction — no reassociation
+// risk, so unlike twoSum/fastTwoSum this needs no protection.
+fn ddGreater(a: DD, b: DD) -> bool {
+  if (a.hi != b.hi) {
+    return a.hi > b.hi;
+  }
+  return a.lo > b.lo;
+}

package/src/shaders/f64/utils/multiply.wgsl

@@ -0,0 +1,81 @@
+// Requires f64/dekker.wgsl concatenated first for the DD struct, and
+// f64/utils/add.wgsl for fsub/negf (bitcast-based subtraction/negation) and,
+// for ddMulProtected at the bottom, fastTwoSumProtected.
+//
+// Use twoProdBit — verified universal (0 corrupting failures across 3000+
+// random trials on NVIDIA/Intel-Mesa-ANV/llvmpipe), no barrier protection
+// needed. The classic approaches below (twoProd, twoProdFma) each fail on
+// one backend in a way barrier materialization doesn't fix; twoProdBit
+// sidesteps the bug instead by deriving the split via bitcast+bitmask
+// rather than an arithmetic identity, leaving nothing for a reassociating
+// compiler to fold. Intel Mesa ANV shows frequent last-bit-only diffs from
+// strict ground truth (never data-corrupting) — consistent with the driver
+// legitimately auto-fusing `x - y*z` into hardware FMA.
+const SPLIT_CONST: f32 = 4097.0;
+
+fn bitSplit(a: f32) -> DD {
+  let bits = bitcast<u32>(a);
+  let hiBits = bits & 0xFFFFF800u; // keep sign+exponent+top 12 mantissa bits
+  let hi = bitcast<f32>(hiBits);
+  let lo = fsub(a, hi); // exact by Sterbenz's lemma (hi, a share an exponent, are close)
+  return DD(hi, lo);
+}
+
+fn twoProdBit(a: f32, b: f32) -> DD {
+  let s = a * b;
+  let aSplit = bitSplit(a);
+  let bSplit = bitSplit(b);
+  let e = fsub(fsub(fsub(fsub(s, aSplit.hi * bSplit.hi), aSplit.lo * bSplit.hi), aSplit.hi * bSplit.lo), aSplit.lo * bSplit.lo);
+  return DD(s, negf(e));
+}
+
+// ── Unsafe historical reference — do not use ────────────────────────────
+// Both broken on one backend, confirmed via isolated cross-driver testing,
+// NOT fixed by barrier materialization (unlike addition's bug):
+//   - veltkampSplit/twoProd (Dekker's original): fails on NVIDIA — compiler
+//     folds `hi = c - (c - a)` to `= a` straight through fsub/negf, even
+//     with every intermediate barrier-materialized (11/11 fail, worse than
+//     unprotected's 6/11).
+//   - twoProdFma (Ogita/Rump/Oishi): fails on llvmpipe — its software fma()
+//     likely isn't genuinely fused, making `fma(a,b,-(a*b))` correctly (not
+//     buggily) zero. Materializing `s` doesn't change this.
+fn veltkampSplit(a: f32) -> DD {
+  let c = SPLIT_CONST * a;
+  let big = fsub(c, a);
+  let hi = fsub(c, big);
+  let lo = fsub(a, hi);
+  return DD(hi, lo);
+}
+
+fn twoProd(a: f32, b: f32) -> DD {
+  let s = a * b;
+  let aSplit = veltkampSplit(a);
+  let bSplit = veltkampSplit(b);
+  let e = fsub(fsub(fsub(fsub(s, aSplit.hi * bSplit.hi), aSplit.lo * bSplit.hi), aSplit.hi * bSplit.lo), aSplit.lo * bSplit.lo);
+  return DD(s, negf(e));
+}
+
+fn twoProdFma(a: f32, b: f32) -> DD {
+  let s = a * b;
+  let e = fma(a, b, negf(s));
+  return DD(s, e);
+}
+
+// DD × DD product (Dekker/Bailey): twoProdBit(a.hi, b.hi) already captures
+// the dominant term to full DD precision, and the cross terms are below the
+// ~48-bit floor anyway, so folding them in with plain f32 loses nothing —
+// only the final renormalization needs barrier protection. Split into
+// ddMulRaw (unprotected) and ddMulProtected (renormalizes via
+// fastTwoSumProtected) so callers with several products can batch them
+// through one shared barrier. ddMulRaw's result isn't a valid DD pair on
+// its own — it must be renormalized before use.
+fn ddMulRaw(a: DD, b: DD) -> DD {
+  let p = twoProdBit(a.hi, b.hi);
+  let crossAndLo = p.lo + (a.hi * b.lo + a.lo * b.hi);
+  return DD(p.hi, crossAndLo);
+}
+
+fn ddMulProtected(a: DD, b: DD, threadSlot: u32) -> DD {
+  let raw = ddMulRaw(a, b);
+  return fastTwoSumProtected(raw.hi, raw.lo, threadSlot);
+}

package/src/shaders/idamax.wgsl

@@ -0,0 +1,96 @@
+// idamax: returns index of element with largest absolute value (f64, double-double)
+// pass 1 dispatches exactly 2 * WGS workgroups; pass 2 uses reduction/argmaxF64.wgsl.
+// Concatenated after f64/dekker.wgsl (DD struct), f64/utils/abs.wgsl (ddAbs),
+// f64/utils/greater.wgsl (ddGreater), and f64/utils/equal.wgsl (ddEqual).
+
+@group(0) @binding(0) var<storage, read>       xHi:            array<f32>;
+@group(0) @binding(1) var<storage, read>       xLo:            array<f32>;
+@group(0) @binding(2) var<storage, read_write> partialsValHi:  array<f32>;
+@group(0) @binding(3) var<storage, read_write> partialsValLo:  array<f32>;
+@group(0) @binding(4) var<storage, read_write> partialsIdx:    array<u32>;
+@group(0) @binding(5) var<uniform>             params:         Params;
+
+struct Params {
+  n:     u32,
+  x_inc: u32,
+}
+
+const WGS: u32 = 64;
+
+var<workgroup> tile_val: array<DD, 64>;
+var<workgroup> tile_idx: array<u32, 64>;
+
+@compute @workgroup_size(64)
+fn idamax_main(
+  @builtin(global_invocation_id) gid:    vec3u,
+  @builtin(local_invocation_id)  lid:    vec3u,
+  @builtin(workgroup_id)         wgid:   vec3u,
+  @builtin(num_workgroups)       num_wg: vec3u,
+) {
+  // DD(-1.0, 0.0) is a safe sentinel: any |x[i]| >= 0 beats it,
+  // so workgroups with no elements lose gracefully in the epilogue.
+  var best_val0 = DD(-1.0, 0.0); var best_idx0: u32 = 0u;
+  var best_val1 = DD(-1.0, 0.0); var best_idx1: u32 = 0u;
+  var best_val2 = DD(-1.0, 0.0); var best_idx2: u32 = 0u;
+  var best_val3 = DD(-1.0, 0.0); var best_idx3: u32 = 0u;
+
+  let stride   = num_wg.x * WGS;
+  let n4_floor = (params.n / (4u * stride)) * (4u * stride);
+
+  for (var id = gid.x; id < n4_floor; id += 4u * stride) {
+    let i0 =  id                * params.x_inc;
+    let i1 = (id +      stride) * params.x_inc;
+    let i2 = (id + 2u * stride) * params.x_inc;
+    let i3 = (id + 3u * stride) * params.x_inc;
+    let v0 = ddAbs(DD(xHi[i0], xLo[i0]));
+    let v1 = ddAbs(DD(xHi[i1], xLo[i1]));
+    let v2 = ddAbs(DD(xHi[i2], xLo[i2]));
+    let v3 = ddAbs(DD(xHi[i3], xLo[i3]));
+    if (ddGreater(v0, best_val0)) { best_val0 = v0; best_idx0 = id; }
+    if (ddGreater(v1, best_val1)) { best_val1 = v1; best_idx1 = id +      stride; }
+    if (ddGreater(v2, best_val2)) { best_val2 = v2; best_idx2 = id + 2u * stride; }
+    if (ddGreater(v3, best_val3)) { best_val3 = v3; best_idx3 = id + 3u * stride; }
+  }
+  for (var id = n4_floor + gid.x; id < params.n; id += stride) {
+    let i = id * params.x_inc;
+    let v = ddAbs(DD(xHi[i], xLo[i]));
+    if (ddGreater(v, best_val0)) { best_val0 = v; best_idx0 = id; }
+  }
+
+  // merge 4 independent lanes; prefer lower index on tie (first occurrence wins)
+  if (ddGreater(best_val1, best_val0) ||
+      (ddEqual(best_val1, best_val0) && best_idx1 < best_idx0)) {
+    best_val0 = best_val1; best_idx0 = best_idx1;
+  }
+  if (ddGreater(best_val2, best_val0) ||
+      (ddEqual(best_val2, best_val0) && best_idx2 < best_idx0)) {
+    best_val0 = best_val2; best_idx0 = best_idx2;
+  }
+  if (ddGreater(best_val3, best_val0) ||
+      (ddEqual(best_val3, best_val0) && best_idx3 < best_idx0)) {
+    best_val0 = best_val3; best_idx0 = best_idx3;
+  }
+
+  tile_val[lid.x] = best_val0;
+  tile_idx[lid.x] = best_idx0;
+  workgroupBarrier();
+
+  for (var s = WGS / 2u; s > 0u; s >>= 1u) {
+    if (lid.x < s) {
+      let a_val = tile_val[lid.x];
+      let b_val = tile_val[lid.x + s];
+      if (ddGreater(b_val, a_val) ||
+          (ddEqual(b_val, a_val) && tile_idx[lid.x + s] < tile_idx[lid.x])) {
+        tile_val[lid.x] = b_val;
+        tile_idx[lid.x] = tile_idx[lid.x + s];
+      }
+    }
+    workgroupBarrier();
+  }
+
+  if (lid.x == 0u) {
+    partialsValHi[wgid.x] = tile_val[0].hi;
+    partialsValLo[wgid.x] = tile_val[0].lo;
+    partialsIdx[wgid.x]   = tile_idx[0];
+  }
+}

package/src/shaders/reduction/argmaxF64.wgsl

@@ -0,0 +1,50 @@
+// amax reduction (f64, double-double): collapses 2*WGS (value, index) pairs
+// into one index, using ddGreater/ddEqual instead of plain f32 `>`/`==` (see
+// reduction/argmax.wgsl for the f32 original this mirrors).
+// dispatch: 1 workgroup of WGS threads. partialsValHi/partialsValLo/
+// partialsIdx must have exactly 2*WGS entries each. Concatenated after
+// f64/dekker.wgsl (DD struct), f64/utils/greater.wgsl (ddGreater), and
+// f64/utils/equal.wgsl (ddEqual).
+
+@group(0) @binding(0) var<storage, read>       partialsValHi: array<f32>;
+@group(0) @binding(1) var<storage, read>       partialsValLo: array<f32>;
+@group(0) @binding(2) var<storage, read>       partialsIdx:   array<u32>;
+@group(0) @binding(3) var<storage, read_write> result:        array<u32>;
+
+const WGS: u32 = 64;
+
+var<workgroup> tile_val: array<DD, 64>;
+var<workgroup> tile_idx: array<u32, 64>;
+
+@compute @workgroup_size(64)
+fn reduce_f64(
+  @builtin(local_invocation_id) lid: vec3u,
+) {
+  let i = lid.x;
+  let a_val = DD(partialsValHi[i], partialsValLo[i]);
+  let b_val = DD(partialsValHi[i + WGS], partialsValLo[i + WGS]);
+  if (ddGreater(b_val, a_val) ||
+      (ddEqual(b_val, a_val) && partialsIdx[i + WGS] < partialsIdx[i])) {
+    tile_val[i] = b_val;
+    tile_idx[i] = partialsIdx[i + WGS];
+  } else {
+    tile_val[i] = a_val;
+    tile_idx[i] = partialsIdx[i];
+  }
+  workgroupBarrier();
+
+  for (var s = WGS / 2u; s > 0u; s >>= 1u) {
+    if (i < s) {
+      let c_val = tile_val[i];
+      let d_val = tile_val[i + s];
+      if (ddGreater(d_val, c_val) ||
+          (ddEqual(d_val, c_val) && tile_idx[i + s] < tile_idx[i])) {
+        tile_val[i] = d_val;
+        tile_idx[i] = tile_idx[i + s];
+      }
+    }
+    workgroupBarrier();
+  }
+
+  if (i == 0u) { result[0] = tile_idx[0]; }
+}

package/src/shaders/reduction/sumF64.wgsl

@@ -2,8 +2,8 @@
 // into one, using ddAddProtected instead of plain f32 `+` (see
 // reduction/sum.wgsl for the f32 original this mirrors).
 // dispatch: 1 workgroup of WGS threads. partialsHi/partialsLo must have
-// exactly 2*WGS entries each. Concatenated after f64/dekker.wgsl for
-// DD/ddAddProtected (see it for why plain ddAdd isn't safe).
+// exactly 2*WGS entries each. Concatenated after f64/dekker.wgsl (DD struct)
+// and f64/utils/add.wgsl (ddAddProtected — see it for why plain ddAdd isn't safe).
 
 @group(0) @binding(0) var<storage, read>       partialsHi: array<f32>;
 @group(0) @binding(1) var<storage, read>       partialsLo: array<f32>;

package/src/shaders/sgemm_large.wgsl

@@ -0,0 +1,117 @@
+// sgemm_large: C = alpha * op(A) * op(B) + beta * C — large-tile half of
+// the two-tier autotuned dispatch (see sgemm.mjs and sgemm_small.wgsl).
+// BM=BN=64, BK=8, TM=8, TN=4 (128 threads/workgroup) — the kernel 9
+// autotuning winner (temp/autotune_sweep.mjs, temp/gen_sweep_kernel.mjs,
+// swept BM/BN/BK/TM/TN and warp-tiled variants), +69% over the old BM=32
+// single-tier baseline at n=512, +84% at n=1024. But BM=64 loses to BM=32
+// below a 6x6=36 workgroup grid (not enough workgroups to fill the GPU at
+// that tile size), hence the two-tier split rather than one global config.
+// Neither vectorized loads (kernel 6) nor warp-tiling (kernel 10) beat this
+// at the sizes tried, including warp-tiled variants in the same sweep at
+// BM=64/128.
+
+const BM: u32 = 64u;
+const BN: u32 = 64u;
+const BK: u32 = 8u;
+const TM: u32 = 8u;
+const TN: u32 = 4u;
+const THREADS_X: u32 = BN / TN;
+const THREADS_Y: u32 = BM / TM;
+const NUM_THREADS: u32 = THREADS_X * THREADS_Y; // 128
+const STRIDE_A: u32 = NUM_THREADS / BK; // rows of As covered per load-loop step
+const STRIDE_B: u32 = NUM_THREADS / BN; // rows of Bs covered per load-loop step
+
+@group(0) @binding(0) var<storage, read>       A: array<f32>;
+@group(0) @binding(1) var<storage, read>       B: array<f32>;
+@group(0) @binding(2) var<storage, read_write> C: array<f32>;
+
+struct Params {
+  m:      u32,
+  n:      u32,
+  k:      u32,
+  alpha:  f32,
+  beta:   f32,
+  lda:    u32,
+  ldb:    u32,
+  ldc:    u32,
+  transA: u32, // 0 = no-transpose, 1 = transpose
+  transB: u32,
+}
+
+@group(0) @binding(3) var<uniform> params: Params;
+
+var<workgroup> As: array<f32, BM * BK>;
+var<workgroup> Bs: array<f32, BK * BN>;
+
+@compute @workgroup_size(THREADS_X, THREADS_Y)
+fn main(
+  @builtin(workgroup_id) wid: vec3u,
+  @builtin(local_invocation_id) lid: vec3u,
+  @builtin(local_invocation_index) tid: u32,
+) {
+  let blockRow = wid.y * BM;
+  let blockCol = wid.x * BN;
+  let threadCol = lid.x;
+  let threadRow = lid.y;
+
+  // Load indices, independent of the compute thread shape — a loop since
+  // NUM_THREADS doesn't match the tile size 1:1 at this config.
+  let innerRowA = tid / BK;
+  let innerColA = tid % BK;
+  let innerRowB = tid / BN;
+  let innerColB = tid % BN;
+
+  var threadResults: array<f32, TM * TN>;
+  for (var i = 0u; i < TM * TN; i++) {
+    threadResults[i] = 0.0;
+  }
+  var regM: array<f32, TM>;
+  var regN: array<f32, TN>;
+
+  let numTiles = (params.k + BK - 1u) / BK;
+  for (var t = 0u; t < numTiles; t++) {
+    for (var loadOffset = 0u; loadOffset < BM; loadOffset += STRIDE_A) {
+      let gRowA = blockRow + innerRowA + loadOffset;
+      let gColA = t * BK + innerColA;
+      let aIdx = select(gRowA * params.lda + gColA, gColA * params.lda + gRowA, params.transA != 0u);
+      As[(innerRowA + loadOffset) * BK + innerColA] = select(0.0, A[aIdx], gRowA < params.m && gColA < params.k);
+    }
+    for (var loadOffset = 0u; loadOffset < BK; loadOffset += STRIDE_B) {
+      let gRowB = t * BK + innerRowB + loadOffset;
+      let gColB = blockCol + innerColB;
+      let bIdx = select(gRowB * params.ldb + gColB, gColB * params.ldb + gRowB, params.transB != 0u);
+      Bs[(innerRowB + loadOffset) * BN + innerColB] = select(0.0, B[bIdx], gRowB < params.k && gColB < params.n);
+    }
+
+    workgroupBarrier();
+
+    for (var dotIdx = 0u; dotIdx < BK; dotIdx++) {
+      for (var i = 0u; i < TM; i++) {
+        regM[i] = As[(threadRow * TM + i) * BK + dotIdx];
+      }
+      for (var i = 0u; i < TN; i++) {
+        regN[i] = Bs[dotIdx * BN + threadCol * TN + i];
+      }
+      for (var resIdxM = 0u; resIdxM < TM; resIdxM++) {
+        for (var resIdxN = 0u; resIdxN < TN; resIdxN++) {
+          threadResults[resIdxM * TN + resIdxN] += regM[resIdxM] * regN[resIdxN];
+        }
+      }
+    }
+
+    workgroupBarrier();
+  }
+
+  for (var resIdxM = 0u; resIdxM < TM; resIdxM++) {
+    let row = blockRow + threadRow * TM + resIdxM;
+    if (row < params.m) {
+      for (var resIdxN = 0u; resIdxN < TN; resIdxN++) {
+        let col = blockCol + threadCol * TN + resIdxN;
+        if (col < params.n) {
+          let cIdx = row * params.ldc + col;
+          C[cIdx] = params.alpha * threadResults[resIdxM * TN + resIdxN] + params.beta * C[cIdx];
+        }
+      }
+    }
+  }
+}