wgblas 2.0.0 → 2.1.0

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/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);
         }
       }
     }

package/src/shaders/sgemm_small.wgsl

@@ -5,6 +5,13 @@
 // workgroups to fill the GPU. Same structure as sgemm_large.wgsl (2D
 // register-blocked, shared-memory-tiled), just smaller.
 //
+// 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).
+// NUM_THREADS (256) exceeds some small-tile load shapes, so the vectorized
+// paths whose lane count doesn't tile exactly guard their As/Bs stores.
+//
 // col mapped to gid.x for coalesced B/C access (row-major: col contiguous).
 
 const BM: u32 = 32u;
@@ -18,9 +25,11 @@ const NUM_THREADS: u32 = THREADS_X * THREADS_Y; // 256
 const STRIDE_A: u32 = NUM_THREADS / BK;
 const STRIDE_B: u32 = NUM_THREADS / BN;
 
-@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,
@@ -33,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>;
@@ -65,17 +76,103 @@ 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, one vec4 per thread. NUM_THREADS
+      // spans BM/4×(BK/4) several times over — guard the store.
+      let r4 = tid / (BK / 4u);
+      let c4 = tid % (BK / 4u);
+      if (r4 < BM) {
+        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. NUM_THREADS over-spans
+      // the BK-column tile — guard the store.
+      let r4 = tid % (BM / 4u);
+      let c  = tid / (BM / 4u);
+      if (c < BK) {
+        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, one vec4 per thread. NUM_THREADS
+      // over-spans the BK-row tile — guard the store.
+      let r  = tid / (BN / 4u);
+      let c4 = tid % (BN / 4u);
+      if (r < BK) {
+        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, one vec4 per thread —
+      // NUM_THREADS over-spans the 32-column tile, so guard the store.
+      let r4 = tid % (BK / 4u);
+      let c  = tid / (BK / 4u);
+      if (c < BN) {
+        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();
@@ -104,7 +201,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);
         }
       }
     }

package/src/shaders/sgemmtr_large.wgsl

@@ -109,7 +109,10 @@ fn main(
         let inTriangle = select(col >= row, col <= row, params.uplo == 0u);
         if (col < params.n && inTriangle) {
           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);
         }
       }
     }

package/src/shaders/sgemmtr_small.wgsl

@@ -102,7 +102,10 @@ fn main(
         let inTriangle = select(col >= row, col <= row, params.uplo == 0u);
         if (col < params.n && inTriangle) {
           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);
         }
       }
     }

package/src/shaders/sgemv_n.wgsl

@@ -67,7 +67,9 @@ fn main(
 
     if lid.x == 0u {
       let yi = row * params.incy;
-      y[yi] = params.alpha * scratch[0] + params.beta * y[yi];
+      // BLAS beta==0 semantics: y is written, not accumulated — must not read y.
+      let acc = params.alpha * scratch[0];
+      y[yi] = select(acc, acc + params.beta * y[yi], params.beta != 0.0);
     }
     // All 64 threads must agree before the next row reuses scratch[].
     workgroupBarrier();

package/src/shaders/sgemv_t.wgsl

@@ -60,6 +60,8 @@ fn main(
       acc0 += A[k * params.lda + col] * x[k * params.incx];
     }
     let yi = col * params.incy;
-    y[yi] = params.alpha * (acc0 + acc1 + acc2 + acc3) + params.beta * y[yi];
+    // BLAS beta==0 semantics: y is written, not accumulated — must not read y.
+    let acc = params.alpha * (acc0 + acc1 + acc2 + acc3);
+    y[yi] = select(acc, acc + params.beta * y[yi], params.beta != 0.0);
   }
 }

package/src/shaders/snrm2.wgsl

@@ -1,9 +1,22 @@
-// snrm2: result = sqrt(sum(x[i] * x[i]))
-// pass 1 dispatches exactly 2 * WGS workgroups; pass 2 uses reduction/sqsum.wgsl.
+// snrm2: result = sqrt(sum(x[i] * x[i])), computed via scaled accumulation
+// (Blue's algorithm / reference BLAS's SLASSQ) rather than naive squaring —
+// naive `sum += x_i * x_i` overflows to inf for |x_i| ≳ 1.8e19 (f32's
+// squaring range is only sqrt(f32_max)) and loses precision on tiny
+// magnitudes squaring into the denormal range. Running state is (scale,
+// ssq) with true-sum-of-squares == scale² · ssq: scale tracks the largest
+// |x_i| seen so far, and every other contribution is expressed *relative
+// to* scale (never squared in absolute terms), so ssq stays near 1
+// regardless of x's magnitude range. Merging two independent partials
+// (ssqMerge) is associative, so this composes with the same 4-way-ILP +
+// tree-reduction shape every other Level 1 reduction here uses — see
+// reduction/scaledSum.wgsl for the pass-2 counterpart, which finishes with
+// scale·sqrt(ssq).
+// pass 1 dispatches exactly 2 * WGS workgroups; pass 2 uses reduction/scaledSum.wgsl.
 
-@group(0) @binding(0) var<storage, read>       x:        array<f32>;
-@group(0) @binding(1) var<storage, read_write> partials: array<f32>;
-@group(0) @binding(2) var<uniform>             params:   Params;
+@group(0) @binding(0) var<storage, read>       x:             array<f32>;
+@group(0) @binding(1) var<storage, read_write> partialsScale: array<f32>;
+@group(0) @binding(2) var<storage, read_write> partialsSsq:   array<f32>;
+@group(0) @binding(3) var<uniform>             params:        Params;
 
 struct Params {
   n:     u32,
@@ -12,7 +25,36 @@ struct Params {
 
 const WGS: u32 = 64;
 
-var<workgroup> tile: array<f32, 64>;
+struct ScaleSsq {
+  scale: f32,
+  ssq:   f32,
+}
+
+// Folds one more |value| into a running (scale, ssq) pair.
+fn ssqAccum(acc: ScaleSsq, absxi: f32) -> ScaleSsq {
+  if (absxi == 0.0) { return acc; }
+  if (absxi > acc.scale) {
+    let r = acc.scale / absxi; // 0/absxi == 0 on the first nonzero value — safe
+    return ScaleSsq(absxi, 1.0 + acc.ssq * r * r);
+  }
+  let r = absxi / acc.scale; // reached only once acc.scale > 0 (absxi <= acc.scale and absxi > 0)
+  return ScaleSsq(acc.scale, acc.ssq + r * r);
+}
+
+// Associative merge of two independent (scale, ssq) partials — lets this
+// compose with a tree reduction exactly like a plain sum would.
+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 main(
@@ -21,36 +63,43 @@ fn main(
   @builtin(workgroup_id)         wgid:   vec3u,
   @builtin(num_workgroups)       num_wg: vec3u,
 ) {
-  var acc0: f32 = 0.0;
-  var acc1: f32 = 0.0;
-  var acc2: f32 = 0.0;
-  var acc3: f32 = 0.0;
+  var acc0 = ScaleSsq(0.0, 1.0);
+  var acc1 = ScaleSsq(0.0, 1.0);
+  var acc2 = ScaleSsq(0.0, 1.0);
+  var acc3 = ScaleSsq(0.0, 1.0);
 
   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 v0 = x[ id                * params.x_inc];
-    let v1 = x[(id +      stride) * params.x_inc];
-    let v2 = x[(id + 2u * stride) * params.x_inc];
-    let v3 = x[(id + 3u * stride) * params.x_inc];
-    acc0 += v0 * v0;
-    acc1 += v1 * v1;
-    acc2 += v2 * v2;
-    acc3 += v3 * v3;
+    acc0 = ssqAccum(acc0, abs(x[ id                * params.x_inc]));
+    acc1 = ssqAccum(acc1, abs(x[(id +      stride) * params.x_inc]));
+    acc2 = ssqAccum(acc2, abs(x[(id + 2u * stride) * params.x_inc]));
+    acc3 = ssqAccum(acc3, abs(x[(id + 3u * stride) * params.x_inc]));
   }
   for (var id = n4_floor + gid.x; id < params.n; id += stride) {
-    let v = x[id * params.x_inc];
-    acc0 += v * v;
+    acc0 = ssqAccum(acc0, abs(x[id * params.x_inc]));
   }
 
-  tile[lid.x] = acc0 + acc1 + acc2 + acc3;
+  let combined = ssqMerge(ssqMerge(acc0, acc1), ssqMerge(acc2, acc3));
+  tileScale[lid.x] = combined.scale;
+  tileSsq[lid.x]   = combined.ssq;
   workgroupBarrier();
 
   for (var s = WGS / 2u; s > 0u; s >>= 1u) {
-    if (lid.x < s) { tile[lid.x] += tile[lid.x + s]; }
+    if (lid.x < s) {
+      let merged = ssqMerge(
+        ScaleSsq(tileScale[lid.x], tileSsq[lid.x]),
+        ScaleSsq(tileScale[lid.x + s], tileSsq[lid.x + s]),
+      );
+      tileScale[lid.x] = merged.scale;
+      tileSsq[lid.x]   = merged.ssq;
+    }
     workgroupBarrier();
   }
 
-  if (lid.x == 0u) { partials[wgid.x] = tile[0]; }
+  if (lid.x == 0u) {
+    partialsScale[wgid.x] = tileScale[0];
+    partialsSsq[wgid.x]   = tileSsq[0];
+  }
 }

package/src/shaders/ssymv.wgsl

@@ -63,7 +63,9 @@ fn main(
     }
 
     if lid.x == 0u {
-      y[i * params.incy] = params.alpha * scratch[0] + params.beta * y[i * params.incy];
+      // BLAS beta==0 semantics: y is written, not accumulated — must not read y.
+      let acc = params.alpha * scratch[0];
+      y[i * params.incy] = select(acc, acc + params.beta * y[i * params.incy], params.beta != 0.0);
     }
   }
 }

package/src/snrm2/snrm2.d.mts

@@ -1,7 +1,7 @@
 import { GpuVector } from "../classes/GpuVector.mjs";
 
 /**
- * Computes the Euclidean norm of a vector: result = sqrt(sum(x[i] * x[i]))
+ * Computes the Euclidean norm of a vector: $$\text{result} = \sqrt{\sum_{i} x_i^2}$$
  *
  * {@includeCode ../../examples/snrm2/snrm2.js}
  *
@@ -24,7 +24,7 @@ export declare function snrm2(
 ): Promise<{ nrm2: number } | { nrm2: number; gpuTimeMs: number }>;
 
 /**
- * Computes the Euclidean norm of a vector: result = sqrt(sum(x[i] * x[i]))
+ * Computes the Euclidean norm of a vector: $$\text{result} = \sqrt{\sum_{i} x_i^2}$$
  *
  * {@includeCode ../../examples/snrm2/gpu.snrm2.js}
  *