wgblas 2.0.0 → 2.2.0

package/src/shaders/index.mjs

@@ -3,25 +3,244 @@
  *
  * `shaders/*.wgsl` — one WGSL compute shader per BLAS routine (sscal, saxpy, sdot, …).
  *
- * `shaders/browser-shaders.mjs` — the browser's runtime shader source. In Node.js, shaders are
- * read directly from disk via `readFileSync`. In the browser there is no filesystem, so this file
- * provides all shader strings inline. Vite bundles it by importing each `.wgsl` file as a string.
+ * `routineShaders` below is the single source of truth: routine name → the WGSL source(s)
+ * its `getPipeline()` calls actually reference, verified against every `src/<routine>/<routine>.mjs`
+ * rather than inferred from naming convention (see its doc comment for the exceptions). Each
+ * shader is imported right above the line that adds it — the import *is* the mapping entry, no
+ * separate block to cross-reference. `shaderSources`, the flat name → source registry the
+ * browser bundle's runtime lookup needs, is *derived* from `routineShaders` rather than
+ * hand-duplicated, so the two can never drift apart. In Node.js neither is read — shaders are
+ * `readFileSync` from disk directly; `scripts/build-browser.mjs` inlines this module into the
+ * browser's IIFE bundle via esbuild instead.
  *
  * ## Cross-shader patterns
  *
- * **Fixed workgroup size of 64.** Every shader declares `const WGS: u32 = 64` and
- * `@workgroup_size(64)`. 64 is the minimum `maxComputeInvocationsPerWorkgroup` guaranteed across
- * all WebGPU devices, so this works everywhere without querying device limits.
+ * **Single bind group.** Every shader with bindings uses `@group(0)` only — the JS side always
+ * calls `pipeline.getBindGroupLayout(0)`, no secondary groups to track. Binding order is
+ * consistent too: any read-only storage buffers come before read_write ones, with the
+ * `uniform Params` struct always last. `@binding` indices match the position of each resource in
+ * the array passed to `createBindGroup`, which appends `resultBuffer` last.
  *
- * **Single bind group.** All bindings use `@group(0)`. This means the JS side always calls
- * `pipeline.getBindGroupLayout(0)` — no secondary groups to track.
+ * **All counts and strides are `u32`.** `n`, `x_inc`, `y_inc`, and every other index/count field
+ * in a `Params` struct is unsigned, avoiding implicit sign-extension in index expressions like
+ * `id * params.x_inc`.
  *
- * The `@binding` indices must match the position of each resource in the array passed to
- * `createBindGroup` — it assigns `binding: 0, 1, 2 …` sequentially, with `resultBuffer` appended last.
- *
- * **All counts and strides are `u32`.** `n`, `x_inc`, `y_inc`, and any other index fields in the
- * `Params` uniform struct are unsigned. This avoids implicit sign-extension when they appear in
- * index expressions like `id * params.x_inc`.
+ * **Entry points don't have to be named `main`.** `loadShader` (`util/pipeline.mjs`)
+ * auto-detects the sole `@compute` function in a module instead of requiring a fixed name, so
+ * `dasum_main`, `strsv_invert_block_main`, etc. work without renaming.
  *
  * @module devdocs/shaders
  */
+
+/**
+ * Routine name → the WGSL source(s) its `getPipeline()` calls reference. Keys are the exact
+ * shader names `getPipeline(device, name)` is called with — most routines have one, some pick
+ * one of several conditionally (e.g. sgemv's `sgemv_n`/`sgemv_t`, by `trans`), and some have no
+ * dedicated shader at all:
+ *
+ * - `sgemmtr`/`ssyrk`/`ssyr2k` all dispatch through `sgemmtr_small`/`sgemmtr_large`.
+ * - `strsm` reuses `strsv_invert_block` and `sscal`, plus its own `block_transfer` and the
+ *   shared `sgemm_small`/`sgemm_large`.
+ * - `dasum`/`idamax` concatenate several f64 utility shaders with their own — see
+ *   `getPipeline`'s `shaderName: string[]` behaviour.
+ * - `random` has no entry — CPU-only, no `getPipeline()` call.
+ *
+ * Built up entry by entry so each import sits next to the mapping entry that uses it.
+ * @public
+ */
+export const routineShaders = {};
+
+import sscal from "./sscal.wgsl";
+routineShaders.sscal = { sscal };
+
+import cscal from "./cscal.wgsl";
+routineShaders.cscal = { cscal };
+
+import sswap from "./sswap.wgsl";
+routineShaders.sswap = { sswap };
+
+import dswap from "./dswap.wgsl"; // f64 sibling of sswap — pure data movement, no arithmetic, no reduction/barrier shader needed
+routineShaders.dswap = { dswap };
+
+import saxpy from "./saxpy.wgsl";
+routineShaders.saxpy = { saxpy };
+
+import scopy from "./scopy.wgsl";
+routineShaders.scopy = { scopy };
+
+import dcopy from "./dcopy.wgsl"; // f64 sibling of scopy — pure data movement, no arithmetic, no reduction/barrier shader needed
+routineShaders.dcopy = { dcopy };
+
+import sdot from "./sdot.wgsl";
+import sum from "./reduction/sum.wgsl";
+routineShaders.sdot = { sdot, "reduction/sum": sum };
+
+import sasum from "./sasum.wgsl";
+routineShaders.sasum = { sasum, "reduction/sum": sum };
+
+import snrm2 from "./snrm2.wgsl";
+import scaledSum from "./reduction/scaledSum.wgsl";
+routineShaders.snrm2 = { snrm2, "reduction/scaledSum": scaledSum };
+
+import isamax from "./isamax.wgsl";
+import argmax from "./reduction/argmax.wgsl";
+routineShaders.isamax = { isamax, "reduction/argmax": argmax };
+
+import dekker from "./f64/dekker.wgsl";
+import ddAbs from "./f64/utils/abs.wgsl";
+import ddAddUtil from "./f64/utils/add.wgsl";
+import dasum from "./dasum.wgsl";
+import sumF64 from "./reduction/sumF64.wgsl";
+routineShaders.dasum = {
+  "f64/dekker": dekker,
+  "f64/utils/abs": ddAbs,
+  "f64/utils/add": ddAddUtil,
+  dasum,
+  "reduction/sumF64": sumF64,
+};
+
+import ddMulUtil from "./f64/utils/multiply.wgsl";
+import ddot from "./ddot.wgsl";
+// multiply.wgsl needs dekker's DD struct and add.wgsl's fsub/negf and
+// fastTwoSumProtected, so those two precede it here.
+routineShaders.ddot = {
+  "f64/dekker": dekker,
+  "f64/utils/add": ddAddUtil,
+  "f64/utils/multiply": ddMulUtil,
+  ddot,
+  "reduction/sumF64": sumF64,
+};
+
+import dscal from "./dscal.wgsl"; // f64 sibling of sscal — no reduction shader needed, unlike dasum/ddot
+routineShaders.dscal = {
+  "f64/dekker": dekker,
+  "f64/utils/add": ddAddUtil,
+  "f64/utils/multiply": ddMulUtil,
+  dscal,
+};
+
+import daxpy from "./daxpy.wgsl"; // f64 sibling of saxpy — one ddMulProtected + one ddAddProtected per element, no reduction shader needed
+routineShaders.daxpy = {
+  "f64/dekker": dekker,
+  "f64/utils/add": ddAddUtil,
+  "f64/utils/multiply": ddMulUtil,
+  daxpy,
+};
+
+import ddGreater from "./f64/utils/greater.wgsl";
+import ddEqual from "./f64/utils/equal.wgsl";
+import idamax from "./idamax.wgsl";
+import argmaxF64 from "./reduction/argmaxF64.wgsl";
+routineShaders.idamax = {
+  "f64/dekker": dekker,
+  "f64/utils/abs": ddAbs,
+  "f64/utils/greater": ddGreater,
+  "f64/utils/equal": ddEqual,
+  idamax,
+  "reduction/argmaxF64": argmaxF64,
+};
+
+import srot from "./srot.wgsl";
+routineShaders.srot = { srot };
+
+import drot from "./drot.wgsl"; // f64 sibling of srot — four ddMulProtected + two ddAddProtected per element, no reduction shader needed
+routineShaders.drot = {
+  "f64/dekker": dekker,
+  "f64/utils/add": ddAddUtil,
+  "f64/utils/multiply": ddMulUtil,
+  drot,
+};
+
+import srotm from "./srotm.wgsl";
+routineShaders.srotm = { srotm };
+
+import drotm from "./drotm.wgsl"; // f64 sibling of srotm — four ddMulProtected + two ddAddProtected per element, no reduction shader needed
+routineShaders.drotm = {
+  "f64/dekker": dekker,
+  "f64/utils/add": ddAddUtil,
+  "f64/utils/multiply": ddMulUtil,
+  drotm,
+};
+
+import ddDivUtil from "./f64/utils/divide.wgsl";
+import ddSqrtUtil from "./f64/utils/sqrt.wgsl";
+import dnrm2 from "./dnrm2.wgsl"; // f64 sibling of snrm2 — scaled accumulation (Blue's algorithm) ported to double-double, branch-free (select()) since ddDivProtected/ddMulProtected/ddAddProtected's barriers need every thread to take the same path
+import scaledSumF64 from "./reduction/scaledSumF64.wgsl";
+routineShaders.dnrm2 = {
+  "f64/dekker": dekker,
+  "f64/utils/abs": ddAbs,
+  "f64/utils/greater": ddGreater,
+  "f64/utils/add": ddAddUtil,
+  "f64/utils/multiply": ddMulUtil,
+  "f64/utils/divide": ddDivUtil,
+  "f64/utils/sqrt": ddSqrtUtil,
+  dnrm2,
+  "reduction/scaledSumF64": scaledSumF64,
+};
+
+import sgemv_n from "./sgemv_n.wgsl";
+import sgemv_t from "./sgemv_t.wgsl";
+routineShaders.sgemv = { sgemv_n, sgemv_t }; // one or the other, picked by trans
+
+import ssymv from "./ssymv.wgsl";
+routineShaders.ssymv = { ssymv };
+
+import strmv from "./strmv.wgsl";
+routineShaders.strmv = { strmv };
+
+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";
+routineShaders.strsv = {
+  strsv_invert_block,
+  strsv_apply_inverse,
+  strsv_update,
+};
+
+import sger from "./sger.wgsl";
+routineShaders.sger = { sger };
+
+import ssyr from "./ssyr.wgsl";
+routineShaders.ssyr = { ssyr };
+
+import ssyr2 from "./ssyr2.wgsl";
+routineShaders.ssyr2 = { ssyr2 };
+
+import sgemm_small from "./sgemm_small.wgsl";
+import sgemm_large from "./sgemm_large.wgsl";
+routineShaders.sgemm = { sgemm_small, sgemm_large }; // one or the other, picked by a tile-size threshold
+
+import sgemmtr_small from "./sgemmtr_small.wgsl";
+import sgemmtr_large from "./sgemmtr_large.wgsl";
+routineShaders.sgemmtr = { sgemmtr_small, sgemmtr_large };
+
+routineShaders.ssyrk = { sgemmtr_small, sgemmtr_large }; // no shader of its own — rides on sgemmtr's
+routineShaders.ssyr2k = { sgemmtr_small, sgemmtr_large }; // no shader of its own — rides on sgemmtr's
+
+import symmetrize from "./symmetrize.wgsl";
+routineShaders.ssymm = { sgemm_small, sgemm_large, symmetrize };
+
+import triangularize from "./triangularize.wgsl";
+routineShaders.strmm = { sgemm_small, sgemm_large, triangularize };
+
+import blockTransfer from "./block_transfer.wgsl";
+routineShaders.strsm = {
+  strsv_invert_block,
+  block_transfer: blockTransfer,
+  sscal,
+  sgemm_small,
+  sgemm_large,
+};
+
+/**
+ * Flat shader-name → WGSL source-string registry — what `getPipeline()`/`loadShader()` (see
+ * `util/pipeline.mjs`) actually look shaders up in, in the browser. Derived from
+ * `routineShaders` by merging every routine's shaders together; shared shaders (e.g.
+ * `"reduction/sum"`, used by two different routines above) collapse harmlessly here since
+ * every routine's copy is the same imported string, never independently authored text.
+ * @public
+ */
+export const shaderSources = Object.assign(
+  {},
+  ...Object.values(routineShaders),
+);

package/src/shaders/reduction/scaledSum.wgsl

@@ -0,0 +1,65 @@
+// scaledSum reduction: collapses 2*WGS (scale, ssq) partials from
+// snrm2.wgsl into the final norm — sqrt(scale² · ssq) == scale · sqrt(ssq).
+// Mirrors reduction/sum.wgsl's shape exactly, merging via ssqMerge (see
+// snrm2.wgsl for the derivation) instead of plain `+`, and taking the final
+// sqrt here rather than on the CPU — unlike sasum/sdot's plain sum, "sum of
+// squares" isn't a meaningful standalone value to hand back, only
+// scale·sqrt(ssq) is.
+// dispatch: 1 workgroup of WGS threads.
+// partialsScale/partialsSsq must have exactly 2*WGS entries each.
+
+@group(0) @binding(0) var<storage, read>       partialsScale: array<f32>;
+@group(0) @binding(1) var<storage, read>       partialsSsq:   array<f32>;
+@group(0) @binding(2) var<storage, read_write> result:        array<f32>;
+
+const WGS: u32 = 64;
+
+// True sum-of-squares represented so far == scale² · ssq — see snrm2.wgsl.
+struct ScaleSsq {
+  scale: f32,
+  ssq:   f32,
+}
+
+// Associative merge of two independent (scale, ssq) partials.
+fn ssqMerge(a: ScaleSsq, b: ScaleSsq) -> ScaleSsq {
+  if (a.scale == 0.0 && b.scale == 0.0) { return ScaleSsq(0.0, 1.0); }
+  if (a.scale >= b.scale) {
+    let r = b.scale / a.scale;
+    return ScaleSsq(a.scale, a.ssq + b.ssq * r * r);
+  }
+  let r = a.scale / b.scale;
+  return ScaleSsq(b.scale, b.ssq + a.ssq * r * r);
+}
+
+var<workgroup> tileScale: array<f32, 64>;
+var<workgroup> tileSsq:   array<f32, 64>;
+
+@compute @workgroup_size(64)
+fn reduce_scaled(
+  @builtin(local_invocation_id) lid: vec3u,
+) {
+  let i = lid.x;
+  let merged0 = ssqMerge(
+    ScaleSsq(partialsScale[i], partialsSsq[i]),
+    ScaleSsq(partialsScale[i + WGS], partialsSsq[i + WGS]),
+  );
+  tileScale[i] = merged0.scale;
+  tileSsq[i]   = merged0.ssq;
+  workgroupBarrier();
+
+  for (var s = WGS / 2u; s > 0u; s >>= 1u) {
+    if (i < s) {
+      let merged = ssqMerge(
+        ScaleSsq(tileScale[i], tileSsq[i]),
+        ScaleSsq(tileScale[i + s], tileSsq[i + s]),
+      );
+      tileScale[i] = merged.scale;
+      tileSsq[i]   = merged.ssq;
+    }
+    workgroupBarrier();
+  }
+
+  if (i == 0u) {
+    result[0] = tileScale[0] * sqrt(tileSsq[0]);
+  }
+}

package/src/shaders/reduction/scaledSumF64.wgsl

@@ -0,0 +1,93 @@
+// scaledSum reduction (f64, double-double): collapses 2*WGS (scale, ssq) DD
+// partials from dnrm2.wgsl into the final norm — sqrt(scale² · ssq) ==
+// scale · sqrt(ssq), via ddMulProtected/ddSqrtProtected. Mirrors
+// reduction/scaledSum.wgsl's shape exactly; ssqMergeProtected is duplicated
+// from dnrm2.wgsl rather than shared via f64/utils/ — see that file's own
+// header for why (same convention the f32 pair already uses).
+// dispatch: 1 workgroup of WGS threads.
+// partialsScale*/partialsSsq* must have exactly 2*WGS entries each.
+
+@group(0) @binding(0) var<storage, read>       partialsScaleHi: array<f32>;
+@group(0) @binding(1) var<storage, read>       partialsScaleLo: array<f32>;
+@group(0) @binding(2) var<storage, read>       partialsSsqHi:   array<f32>;
+@group(0) @binding(3) var<storage, read>       partialsSsqLo:   array<f32>;
+@group(0) @binding(4) var<storage, read_write> resultHi:        array<f32, 1>;
+@group(0) @binding(5) var<storage, read_write> resultLo:        array<f32, 1>;
+
+const WGS: u32 = 64;
+
+struct ScaleSsq {
+  scale: DD,
+  ssq:   DD,
+}
+
+fn ddSelect(a: DD, b: DD, cond: bool) -> DD {
+  return DD(select(a.hi, b.hi, cond), select(a.lo, b.lo, cond));
+}
+
+// Associative merge of two independent (scale, ssq) partials — see
+// dnrm2.wgsl for the derivation and why this is branch-free.
+fn ssqMergeProtected(a: ScaleSsq, b: ScaleSsq, threadSlot: u32) -> ScaleSsq {
+  let isBigger = !ddGreater(b.scale, a.scale); // a.scale >= b.scale
+  let bigger = ddSelect(b.scale, a.scale, isBigger);
+  let smaller = ddSelect(a.scale, b.scale, isBigger);
+  let biggerSsq = ddSelect(b.ssq, a.ssq, isBigger);
+  let smallerSsq = ddSelect(a.ssq, b.ssq, isBigger);
+  let biggerIsZero = bigger.hi == 0.0;
+  let safeBigger = ddSelect(bigger, DD(1.0, 0.0), biggerIsZero);
+  let r = ddDivProtected(smaller, safeBigger, threadSlot);
+  let rsq = ddMulProtected(r, r, threadSlot);
+  let smallerSsqTimesRsq = ddMulProtected(smallerSsq, rsq, threadSlot);
+  let newSsq = ddAddProtected(biggerSsq, smallerSsqTimesRsq, threadSlot);
+  return ScaleSsq(bigger, newSsq);
+}
+
+var<workgroup> tileScaleHi: array<f32, 64>;
+var<workgroup> tileScaleLo: array<f32, 64>;
+var<workgroup> tileSsqHi:   array<f32, 64>;
+var<workgroup> tileSsqLo:   array<f32, 64>;
+
+@compute @workgroup_size(64)
+fn reduce_scaled_f64(
+  @builtin(local_invocation_id) lid: vec3u,
+) {
+  let i = lid.x;
+  let a = ScaleSsq(DD(partialsScaleHi[i], partialsScaleLo[i]), DD(partialsSsqHi[i], partialsSsqLo[i]));
+  let b = ScaleSsq(DD(partialsScaleHi[i + WGS], partialsScaleLo[i + WGS]), DD(partialsSsqHi[i + WGS], partialsSsqLo[i + WGS]));
+  let merged0 = ssqMergeProtected(a, b, i);
+  tileScaleHi[i] = merged0.scale.hi;
+  tileScaleLo[i] = merged0.scale.lo;
+  tileSsqHi[i]   = merged0.ssq.hi;
+  tileSsqLo[i]   = merged0.ssq.lo;
+  workgroupBarrier();
+
+  for (var s = WGS / 2u; s > 0u; s >>= 1u) {
+    let partner = select(i, i + s, i < s);
+    let ai = ScaleSsq(DD(tileScaleHi[i], tileScaleLo[i]), DD(tileSsqHi[i], tileSsqLo[i]));
+    let bi = ScaleSsq(DD(tileScaleHi[partner], tileScaleLo[partner]), DD(tileSsqHi[partner], tileSsqLo[partner]));
+    let merged = ssqMergeProtected(ai, bi, i);
+    workgroupBarrier();
+    if (i < s) {
+      tileScaleHi[i] = merged.scale.hi;
+      tileScaleLo[i] = merged.scale.lo;
+      tileSsqHi[i]   = merged.ssq.hi;
+      tileSsqLo[i]   = merged.ssq.lo;
+    }
+    workgroupBarrier();
+  }
+
+  // ddSqrtProtected/ddMulProtected's own workgroupBarrier()s need every
+  // thread to call them — every thread redundantly computes the same final
+  // scale·sqrt(ssq) from tile[0] (still visible to all after the reduction
+  // above), and only the write-back is conditional. Guarding the calls
+  // themselves behind `if (i == 0u)` (as the plain-f32 original safely
+  // does with its unprotected `sqrt()`) would leave 63 threads never
+  // reaching a barrier the one remaining thread still needs.
+  let scale = DD(tileScaleHi[0], tileScaleLo[0]);
+  let ssq = DD(tileSsqHi[0], tileSsqLo[0]);
+  let result = ddMulProtected(scale, ddSqrtProtected(ssq, i), i);
+  if (i == 0u) {
+    resultHi[0] = result.hi;
+    resultLo[0] = result.lo;
+  }
+}

package/src/shaders/sgemm_large.wgsl

@@ -6,9 +6,13 @@
 // 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.
+//
+// A and B are bound twice — scalar array<f32> and array<vec4<f32>> views of
+// the same GPUBuffer (see vec4ViewBinding) — so each tile load can issue
+// 16-byte vector reads along op(A)/op(B)'s contiguous dimension when the
+// stride allows it (stride % 4 == 0 keeps every row base 16-byte aligned).
+// Transposed or odd-stride operands take the scalar path; both paths
+// zero-fill out-of-bounds components identically.
 
 const BM: u32 = 64u;
 const BN: u32 = 64u;
@@ -21,9 +25,11 @@ 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>;
+@group(0) @binding(0) var<storage, read>       A:  array<f32>;
+@group(0) @binding(1) var<storage, read>       A4: array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read>       B:  array<f32>;
+@group(0) @binding(3) var<storage, read>       B4: array<vec4<f32>>;
+@group(0) @binding(4) var<storage, read_write> C:  array<f32>;
 
 struct Params {
   m:      u32,
@@ -36,9 +42,11 @@ struct Params {
   ldc:    u32,
   transA: u32, // 0 = no-transpose, 1 = transpose
   transB: u32,
+  useVecA: u32, // 1 = A's vec4 view covers every in-bounds element (see vec4Usable)
+  useVecB: u32, // 1 = B's vec4 view covers every in-bounds element
 }
 
-@group(0) @binding(3) var<uniform> params: Params;
+@group(0) @binding(5) var<uniform> params: Params;
 
 var<workgroup> As: array<f32, BM * BK>;
 var<workgroup> Bs: array<f32, BK * BN>;
@@ -70,17 +78,95 @@ fn main(
 
   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);
+    // ── Load the BM×BK A tile into As (vectorized along op(A)'s fast dim
+    // when lda allows; every branch here is dispatch-uniform) ──
+    if (params.useVecA == 1u && params.transA == 0u) {
+      // No-transpose: columns contiguous. Each thread loads one vec4 of 4
+      // columns; 64 rows × 2 column-lanes = NUM_THREADS exactly, single pass.
+      let r4 = tid / (BK / 4u);
+      let c4 = tid % (BK / 4u);
+      let gRow = blockRow + r4;
+      let gCol = t * BK + c4 * 4u;
+      var v = A4[(gRow * params.lda + gCol) / 4u];
+      let rowOK = gRow < params.m;
+      v.x = select(0.0, v.x, rowOK &&  gCol            < params.k);
+      v.y = select(0.0, v.y, rowOK && (gCol + 1u) < params.k);
+      v.z = select(0.0, v.z, rowOK && (gCol + 2u) < params.k);
+      v.w = select(0.0, v.w, rowOK && (gCol + 3u) < params.k);
+      As[r4 * BK + c4 * 4u]      = v.x;
+      As[r4 * BK + c4 * 4u + 1u] = v.y;
+      As[r4 * BK + c4 * 4u + 2u] = v.z;
+      As[r4 * BK + c4 * 4u + 3u] = v.w;
+    } else if (params.useVecA == 1u && params.transA != 0u) {
+      // Transpose: rows contiguous within a column. Each thread loads one
+      // vec4 of 4 rows; 16 row-lanes × 8 columns = NUM_THREADS, single pass.
+      let r4 = tid % (BM / 4u);
+      let c  = tid / (BM / 4u);
+      let gRow = blockRow + r4 * 4u;
+      let gCol = t * BK + c;
+      var v = A4[(gCol * params.lda + gRow) / 4u];
+      let colOK = gCol < params.k;
+      v.x = select(0.0, v.x, colOK &&  gRow            < params.m);
+      v.y = select(0.0, v.y, colOK && (gRow + 1u) < params.m);
+      v.z = select(0.0, v.z, colOK && (gRow + 2u) < params.m);
+      v.w = select(0.0, v.w, colOK && (gRow + 3u) < params.m);
+      As[(r4 * 4u) * BK + c]      = v.x;
+      As[(r4 * 4u + 1u) * BK + c] = v.y;
+      As[(r4 * 4u + 2u) * BK + c] = v.z;
+      As[(r4 * 4u + 3u) * BK + c] = v.w;
+    } else {
+      // Scalar fallback: odd stride or unhandled orientation.
+      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);
+
+    // ── Load the BK×BN B tile into Bs ──
+    if (params.useVecB == 1u && params.transB == 0u) {
+      // No-transpose: columns contiguous. 8 rows × 16 column-lanes cover the
+      // tile in one pass (BK = NUM_THREADS / (BN/4)).
+      let r  = tid / (BN / 4u);
+      let c4 = tid % (BN / 4u);
+      let gRow = t * BK + r;
+      let gCol = blockCol + c4 * 4u;
+      var v = B4[(gRow * params.ldb + gCol) / 4u];
+      let rowOK = gRow < params.k;
+      v.x = select(0.0, v.x, rowOK &&  gCol            < params.n);
+      v.y = select(0.0, v.y, rowOK && (gCol + 1u) < params.n);
+      v.z = select(0.0, v.z, rowOK && (gCol + 2u) < params.n);
+      v.w = select(0.0, v.w, rowOK && (gCol + 3u) < params.n);
+      Bs[r * BN + c4 * 4u]      = v.x;
+      Bs[r * BN + c4 * 4u + 1u] = v.y;
+      Bs[r * BN + c4 * 4u + 2u] = v.z;
+      Bs[r * BN + c4 * 4u + 3u] = v.w;
+    } else if (params.useVecB == 1u && params.transB != 0u) {
+      // Transpose: rows contiguous within a column. 2 row-lanes × 64 columns
+      // cover the tile in one pass (BN = NUM_THREADS / (BK/4)).
+      let r4 = tid % (BK / 4u);
+      let c  = tid / (BK / 4u);
+      let gRow = t * BK + r4 * 4u;
+      let gCol = blockCol + c;
+      var v = B4[(gCol * params.ldb + gRow) / 4u];
+      let colOK = gCol < params.n;
+      v.x = select(0.0, v.x, colOK &&  gRow            < params.k);
+      v.y = select(0.0, v.y, colOK && (gRow + 1u) < params.k);
+      v.z = select(0.0, v.z, colOK && (gRow + 2u) < params.k);
+      v.w = select(0.0, v.w, colOK && (gRow + 3u) < params.k);
+      Bs[(r4 * 4u) * BN + c]     = v.x;
+      Bs[(r4 * 4u + 1u) * BN + c] = v.y;
+      Bs[(r4 * 4u + 2u) * BN + c] = v.z;
+      Bs[(r4 * 4u + 3u) * BN + c] = v.w;
+    } else {
+      // Scalar fallback.
+      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();
@@ -109,7 +195,10 @@ fn main(
         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];
+          // BLAS beta==0 semantics: C is written, not accumulated — must not
+          // read C (stale NaN/Inf bits would survive 0 * C as NaN).
+          let acc = params.alpha * threadResults[resIdxM * TN + resIdxN];
+          C[cIdx] = select(acc, acc + params.beta * C[cIdx], params.beta != 0.0);
         }
       }
     }