wgblas 2.0.0 → 2.2.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}
  *

package/src/snrm2/snrm2.mjs

@@ -12,14 +12,14 @@ import { extractTimestamp } from "../util/benchmark.mjs";
 import { extractResult } from "../util/result.mjs";
 import { getPipeline } from "../util/pipeline.mjs";
 import { GpuVector } from "../classes/GpuVector.mjs";
-
-const WGS = 64; // workgroup size
+import { WGS } from "../util/constants.mjs";
+import { requireGpuDevice, requireSameDevice } from "../util/device.mjs";
 
 export async function snrm2(device, n, x, incx) {
   const xIsGpu = x instanceof GpuVector;
 
-  if (!(device instanceof GPUDevice))
-    throw new Error("device must be a GPUDevice.");
+  requireGpuDevice(device);
+  requireSameDevice(device, "snrm2", { x });
   if (!Number.isInteger(n) || !Number.isInteger(incx))
     throw new Error("n and incx must be integers.");
   if (incx <= 0) throw new Error("incx must be positive.");
@@ -32,19 +32,31 @@ export async function snrm2(device, n, x, incx) {
     );
 
   const pipelineMain = await getPipeline(device, "snrm2");
-  const pipelineReduce = await getPipeline(device, "reduction/sum");
+  const pipelineReduce = await getPipeline(device, "reduction/scaledSum");
 
   let xBuffer = null;
-  let partialsBuffer = null;
+  let partialsScaleBuffer = null;
+  let partialsSsqBuffer = null;
   let resultBuffer = null;
   let paramsBuffer = null;
   let readBuffer = null;
 
   try {
-    xBuffer = xIsGpu ? x._buf : uploadBuffer(x, "snrm2-x", false);
-    partialsBuffer = createStorageBuffer(2 * WGS * 4, "snrm2-partials"); // 2*WGS partial sums of f32
-    resultBuffer = createResultBuffer(4, "snrm2-result"); // final f32 scalar
+    xBuffer = xIsGpu ? x._buf : uploadBuffer(device, x, "snrm2-x", false);
+    // 2*WGS partial (scale, ssq) pairs — see snrm2.wgsl for what they represent.
+    partialsScaleBuffer = createStorageBuffer(
+      device,
+      2 * WGS * 4,
+      "snrm2-partials-scale",
+    );
+    partialsSsqBuffer = createStorageBuffer(
+      device,
+      2 * WGS * 4,
+      "snrm2-partials-ssq",
+    );
+    resultBuffer = createResultBuffer(device, 4, "snrm2-result"); // final f32 scalar
     paramsBuffer = createParamsBuffer(
+      device,
       [
         { value: n, type: "u32" },
         { value: incx, type: "u32" },
@@ -52,53 +64,59 @@ export async function snrm2(device, n, x, incx) {
       "snrm2-params",
     );
 
-    const bgMain = createBindGroup(pipelineMain.getBindGroupLayout(0), [
+    const bgMain = createBindGroup(device, pipelineMain.getBindGroupLayout(0), [
       xBuffer,
-      partialsBuffer,
+      partialsScaleBuffer,
+      partialsSsqBuffer,
       paramsBuffer,
     ]);
     const { commandEncoder: enc1, ts: ts1 } = runComputePass(
+      device,
       pipelineMain,
       bgMain,
       2 * WGS,
     ); // dispatch 2*WGS workgroups
 
-    submit(enc1);
+    submit(device, enc1);
 
-    const bgReduce = createBindGroup(pipelineReduce.getBindGroupLayout(0), [
-      partialsBuffer,
-      resultBuffer,
-    ]);
+    const bgReduce = createBindGroup(
+      device,
+      pipelineReduce.getBindGroupLayout(0),
+      [partialsScaleBuffer, partialsSsqBuffer, resultBuffer],
+    );
     const { commandEncoder: enc2, ts: ts2 } = runComputePass(
+      device,
       pipelineReduce,
       bgReduce,
       1,
     ); // reduce partials to single result
-    readBuffer = stageReadback(enc2, resultBuffer);
+    readBuffer = stageReadback(device, enc2, resultBuffer);
 
-    submit(enc2);
+    submit(device, enc2);
 
     const resultPromise = extractResult(readBuffer, Float32Array);
     readBuffer = null; // ownership transferred — extractResult's own finally destroys it
 
-    const [gpuTime1, gpuTime2, sqsumArr] = await Promise.all([
+    const [gpuTime1, gpuTime2, resultArr] = await Promise.all([
       extractTimestamp(ts1),
       extractTimestamp(ts2),
       resultPromise,
     ]);
 
-    // sqrt is taken on CPU after the GPU sum-of-squares reduction
-    const nrm2 = Math.sqrt(sqsumArr[0]);
+    // reduction/scaledSum.wgsl already computes scale·sqrt(ssq) on the GPU —
+    // unlike the old naive-sum version, there's no separate sqrt step here.
+    const nrm2 = resultArr[0];
 
     if (gpuTime1 !== undefined && gpuTime2 !== undefined)
       return { nrm2, gpuTimeMs: gpuTime1 + gpuTime2 };
     return { nrm2 };
   } finally {
     if (!xIsGpu && xBuffer) destroyBuffers(xBuffer);
-    if (partialsBuffer) destroyBuffers(partialsBuffer);
+    if (partialsScaleBuffer) destroyBuffers(partialsScaleBuffer);
+    if (partialsSsqBuffer) destroyBuffers(partialsSsqBuffer);
     if (resultBuffer) destroyBuffers(resultBuffer);
     if (paramsBuffer) destroyBuffers(paramsBuffer);
-    // Only reached if submit(enc2) threw before ownership was transferred above.
+    // Only reached if submit(device, enc2) threw before ownership was transferred above.
     if (readBuffer) destroyBuffers(readBuffer);
   }
 }

package/src/srot/srot.d.mts

@@ -2,8 +2,7 @@ import { GpuVector } from "../classes/GpuVector.mjs";
 
 /**
  * Applies a Givens plane rotation to vectors x and y:
- *   x = c*x + s*y
- *   y = -s*x + c*y
+ * $$\begin{aligned} x &\leftarrow cx + sy \\\\ y &\leftarrow -sx + cy \end{aligned}$$
  *
  * {@includeCode ../../examples/srot/srot.js}
  *
@@ -34,8 +33,7 @@ export declare function srot(
 
 /**
  * Applies a Givens plane rotation to vectors x and y:
- *   x = c*x + s*y
- *   y = -s*x + c*y
+ * $$\begin{aligned} x &\leftarrow cx + sy \\\\ y &\leftarrow -sx + cy \end{aligned}$$
  *
  * {@includeCode ../../examples/srot/gpu.srot.js}
  *

package/src/srot/srot.mjs

@@ -11,13 +11,14 @@ import { extractResult } from "../util/result.mjs";
 import { getPipeline } from "../util/pipeline.mjs";
 import { calcWorkgroups } from "../util/workgroup.mjs";
 import { GpuVector } from "../classes/GpuVector.mjs";
+import { requireGpuDevice, requireSameDevice } from "../util/device.mjs";
 
 export async function srot(device, n, x, incx, y, incy, c, s) {
   const xIsGpu = x instanceof GpuVector;
   const yIsGpu = y instanceof GpuVector;
 
-  if (!(device instanceof GPUDevice))
-    throw new Error("device must be a GPUDevice.");
+  requireGpuDevice(device);
+  requireSameDevice(device, "srot", { x, y });
   if (
     !Number.isInteger(n) ||
     !Number.isInteger(incx) ||
@@ -26,7 +27,8 @@ export async function srot(device, n, x, incx, y, incy, c, s) {
     throw new Error("n, incx, and incy must be integers.");
   if (typeof c !== "number") throw new Error("c must be a number.");
   if (typeof s !== "number") throw new Error("s must be a number.");
-  if (Number.isNaN(c) || Number.isNaN(s)) throw new Error("c and s must not be NaN.");
+  if (Number.isNaN(c) || Number.isNaN(s))
+    throw new Error("c and s must not be NaN.");
   if (!Number.isFinite(c)) throw new Error("c must be finite.");
   if (!Number.isFinite(s)) throw new Error("s must be finite.");
   if (incx <= 0 || incy <= 0)
@@ -58,9 +60,10 @@ export async function srot(device, n, x, incx, y, incy, c, s) {
   let readY = null;
 
   try {
-    xBuffer = xIsGpu ? x._buf : uploadBuffer(x, "srot-x", true);
-    yBuffer = yIsGpu ? y._buf : uploadBuffer(y, "srot-y", true);
+    xBuffer = xIsGpu ? x._buf : uploadBuffer(device, x, "srot-x", true);
+    yBuffer = yIsGpu ? y._buf : uploadBuffer(device, y, "srot-y", true);
     paramsBuffer = createParamsBuffer(
+      device,
       [
         { value: n, type: "u32" },
         { value: c, type: "f32" },
@@ -71,23 +74,25 @@ export async function srot(device, n, x, incx, y, incy, c, s) {
       "srot-params",
     );
 
-    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+    const bindGroup = createBindGroup(device, pipeline.getBindGroupLayout(0), [
       xBuffer,
       yBuffer,
       paramsBuffer,
     ]);
     const { commandEncoder, ts } = runComputePass(
+      device,
       pipeline,
       bindGroup,
-      calcWorkgroups(n),
+      calcWorkgroups(device, n),
     );
-    readX = xIsGpu ? null : stageReadback(commandEncoder, xBuffer);
-    readY = yIsGpu ? null : stageReadback(commandEncoder, yBuffer);
-    submit(commandEncoder);
+    readX = xIsGpu ? null : stageReadback(device, commandEncoder, xBuffer);
+    readY = yIsGpu ? null : stageReadback(device, commandEncoder, yBuffer);
+    submit(device, commandEncoder);
 
     const gpuTimeMs = await extractTimestamp(ts);
 
-    if (xIsGpu && yIsGpu) {
+    if (xIsGpu) {
+      // xIsGpu === yIsGpu, enforced above
       if (gpuTimeMs !== undefined) return { gpuTimeMs };
       return {};
     }

package/src/srotm/srotm.d.mts

@@ -2,8 +2,7 @@ import { GpuVector } from "../classes/GpuVector.mjs";
 
 /**
  * Applies a modified Givens plane rotation H to vectors x and y:
- *   x = H[0][0]*x + H[0][1]*y
- *   y = H[1][0]*x + H[1][1]*y
+ * $$\begin{pmatrix} x \\\\ y \end{pmatrix} \leftarrow \begin{pmatrix} h_{11} & h_{12} \\\\ h_{21} & h_{22} \end{pmatrix} \begin{pmatrix} x \\\\ y \end{pmatrix}$$
  *
  * {@includeCode ../../examples/srotm/srotm.js}
  *
@@ -33,8 +32,7 @@ export declare function srotm(
 
 /**
  * Applies a modified Givens plane rotation H to vectors x and y:
- *   x = H[0][0]*x + H[0][1]*y
- *   y = H[1][0]*x + H[1][1]*y
+ * $$\begin{pmatrix} x \\\\ y \end{pmatrix} \leftarrow \begin{pmatrix} h_{11} & h_{12} \\\\ h_{21} & h_{22} \end{pmatrix} \begin{pmatrix} x \\\\ y \end{pmatrix}$$
  *
  * {@includeCode ../../examples/srotm/gpu.srotm.js}
  *