wgblas 2.1.0 → 2.2.0

package/src/shaders/drotm.wgsl

@@ -0,0 +1,99 @@
+// drotm: applies a modified Givens rotation H to vectors x and y — double-
+// double (Dekker) f64 emulation of srotm. paramHi/paramLo[0] = flag: -1
+// (full H), 0 (unit diagonal), 1 (unit off-diagonal). param = [ flag, h11,
+// h21, h12, h22 ], each entry an f32 (hi, lo) pair.
+// flag == -2 (identity/no-op) is handled in JS before dispatch reaches here.
+//
+// h11/h12/h21/h22 are resolved once, outside the loop, from the (uniform
+// across every thread) flag — same shape as srot's c/s, so no barrier is
+// needed for that selection itself. Each element then costs four
+// ddMulProtected + two ddAddProtected, same as drot.
+
+@group(0) @binding(0) var<storage, read_write> xHi: array<f32>;
+@group(0) @binding(1) var<storage, read_write> xLo: array<f32>;
+@group(0) @binding(2) var<storage, read_write> yHi: array<f32>;
+@group(0) @binding(3) var<storage, read_write> yLo: array<f32>;
+@group(0) @binding(4) var<storage, read> paramHi: array<f32>;
+@group(0) @binding(5) var<storage, read> paramLo: array<f32>;
+@group(0) @binding(6) var<uniform> params: Params;
+
+struct Params {
+  n:     u32,
+  x_inc: u32,
+  y_inc: u32,
+}
+
+const WGS: u32 = 64;
+const ONE: DD = DD(1.0, 0.0);
+const NEG_ONE: DD = DD(-1.0, 0.0);
+
+@compute @workgroup_size(64)
+fn drotm_main(
+  @builtin(global_invocation_id) gid:    vec3u,
+  @builtin(local_invocation_id)  lid:    vec3u,
+  @builtin(workgroup_id)         wgid:   vec3u,
+  @builtin(num_workgroups)       num_wg: vec3u,
+) {
+  let flag = paramHi[0]; // exact small integer (-1, 0, or 1) — lo is always 0
+
+  var h11: DD; var h12: DD;
+  var h21: DD; var h22: DD;
+
+  if (flag == -1.0) {
+    // full 2x2 matrix
+    h11 = DD(paramHi[1], paramLo[1]); h21 = DD(paramHi[2], paramLo[2]);
+    h12 = DD(paramHi[3], paramLo[3]); h22 = DD(paramHi[4], paramLo[4]);
+  } else if (flag == 0.0) {
+    // diagonal fixed at 1
+    h11 = ONE;                        h21 = DD(paramHi[2], paramLo[2]);
+    h12 = DD(paramHi[3], paramLo[3]); h22 = ONE;
+  } else {
+    // flag == 1.0: off-diagonal fixed at +1 / -1
+    h11 = DD(paramHi[1], paramLo[1]); h21 = NEG_ONE;
+    h12 = ONE;                        h22 = DD(paramHi[4], paramLo[4]);
+  }
+
+  let stride = num_wg.x * WGS;
+
+  let n_floor = (params.n / stride) * stride;
+  let mainIters = n_floor / stride;
+  for (var iter = 0u; iter < mainIters; iter++) {
+    let id = gid.x + iter * stride;
+    let ix = id * params.x_inc;
+    let iy = id * params.y_inc;
+    let xi = DD(xHi[ix], xLo[ix]);
+    let yi = DD(yHi[iy], yLo[iy]);
+    let xNew = ddAddProtected(ddMulProtected(h11, xi, lid.x), ddMulProtected(h12, yi, lid.x), lid.x);
+    let yNew = ddAddProtected(ddMulProtected(h21, xi, lid.x), ddMulProtected(h22, yi, lid.x), lid.x);
+    xHi[ix] = xNew.hi;
+    xLo[ix] = xNew.lo;
+    yHi[iy] = yNew.hi;
+    yLo[iy] = yNew.lo;
+  }
+
+  // Tail is ragged (0 or 1 extra per thread) — pad to this workgroup's worst
+  // case so every thread in the workgroup still calls ddMulProtected/
+  // ddAddProtected the same number of times (their barriers need that),
+  // masking only the write.
+  let wgBaseGid = wgid.x * WGS;
+  var tailIters = 0u;
+  if (n_floor + wgBaseGid < params.n) {
+    tailIters = (params.n - 1u - n_floor - wgBaseGid) / stride + 1u;
+  }
+  for (var iter = 0u; iter < tailIters; iter++) {
+    let id = n_floor + gid.x + iter * stride;
+    let valid = id < params.n;
+    let ix = select(0u, id * params.x_inc, valid); // index 0 always in-bounds
+    let iy = select(0u, id * params.y_inc, valid);
+    let xi = DD(xHi[ix], xLo[ix]);
+    let yi = DD(yHi[iy], yLo[iy]);
+    let xNew = ddAddProtected(ddMulProtected(h11, xi, lid.x), ddMulProtected(h12, yi, lid.x), lid.x);
+    let yNew = ddAddProtected(ddMulProtected(h21, xi, lid.x), ddMulProtected(h22, yi, lid.x), lid.x);
+    if (valid) {
+      xHi[ix] = xNew.hi;
+      xLo[ix] = xNew.lo;
+      yHi[iy] = yNew.hi;
+      yLo[iy] = yNew.lo;
+    }
+  }
+}

package/src/shaders/dscal.wgsl

@@ -0,0 +1,60 @@
+// dscal: x := alpha * x, double-double (Dekker) f64 emulation of sscal.
+// alpha and x are each an f32 (hi, lo) pair. See f64/utils/multiply.wgsl for
+// ddMulProtected and why plain ddMulRaw isn't safe without a renormalizing
+// barrier — that barrier needs a provably uniform loop trip count across
+// every thread in the workgroup, so (like dasum.wgsl's reduction loop) this
+// splits into a uniform main pass plus a ragged, select-masked tail rather
+// than a plain `id < params.n` grid-stride loop.
+
+@group(0) @binding(0) var<storage, read_write> xHi: array<f32>;
+@group(0) @binding(1) var<storage, read_write> xLo: array<f32>;
+@group(0) @binding(2) var<uniform> params: Params;
+
+struct Params {
+  n:       u32,
+  alphaHi: f32,
+  alphaLo: f32,
+  x_inc:   u32,
+}
+
+const WGS: u32 = 64;
+
+@compute @workgroup_size(64)
+fn dscal_main(
+  @builtin(global_invocation_id) gid:    vec3u,
+  @builtin(local_invocation_id)  lid:    vec3u,
+  @builtin(workgroup_id)         wgid:   vec3u,
+  @builtin(num_workgroups)       num_wg: vec3u,
+) {
+  let alpha = DD(params.alphaHi, params.alphaLo);
+  let stride = num_wg.x * WGS;
+
+  let n_floor = (params.n / stride) * stride;
+  let mainIters = n_floor / stride;
+  for (var iter = 0u; iter < mainIters; iter++) {
+    let id = gid.x + iter * stride;
+    let i = id * params.x_inc;
+    let result = ddMulProtected(alpha, DD(xHi[i], xLo[i]), lid.x);
+    xHi[i] = result.hi;
+    xLo[i] = result.lo;
+  }
+
+  // Tail is ragged (0 or 1 extra per thread) — pad to this workgroup's worst
+  // case so every thread in the workgroup still calls ddMulProtected the
+  // same number of times (its barrier needs that), masking only the write.
+  let wgBaseGid = wgid.x * WGS;
+  var tailIters = 0u;
+  if (n_floor + wgBaseGid < params.n) {
+    tailIters = (params.n - 1u - n_floor - wgBaseGid) / stride + 1u;
+  }
+  for (var iter = 0u; iter < tailIters; iter++) {
+    let id = n_floor + gid.x + iter * stride;
+    let valid = id < params.n;
+    let i = select(0u, id * params.x_inc, valid); // index 0 always in-bounds
+    let result = ddMulProtected(alpha, DD(xHi[i], xLo[i]), lid.x);
+    if (valid) {
+      xHi[i] = result.hi;
+      xLo[i] = result.lo;
+    }
+  }
+}

package/src/shaders/dswap.wgsl

@@ -0,0 +1,38 @@
+// dswap: x <-> y, double-double (Dekker) f64 emulation of sswap. A swap is
+// pure data movement — hi and lo are exchanged verbatim, with no arithmetic
+// at all — so (unlike dscal/daxpy/ddot) this needs no
+// ddMulProtected/ddAddProtected renormalizing barrier, and so no
+// ragged-tail-with-barrier split; a plain grid-stride loop is safe, same
+// shape as sswap.wgsl itself.
+
+@group(0) @binding(0) var<storage, read_write> xHi: array<f32>;
+@group(0) @binding(1) var<storage, read_write> xLo: array<f32>;
+@group(0) @binding(2) var<storage, read_write> yHi: array<f32>;
+@group(0) @binding(3) var<storage, read_write> yLo: array<f32>;
+
+struct Params {
+  n:     u32,
+  x_inc: u32,
+  y_inc: u32,
+}
+
+@group(0) @binding(4) var<uniform> params: Params;
+
+const WGS: u32 = 64;
+
+@compute @workgroup_size(64)
+fn main(
+  @builtin(global_invocation_id) gid: vec3u,
+  @builtin(num_workgroups) num_wg: vec3u,
+) {
+  for (var id = gid.x; id < params.n; id += num_wg.x * WGS) {
+    let ix = id * params.x_inc;
+    let iy = id * params.y_inc;
+    let tempHi = xHi[ix];
+    let tempLo = xLo[ix];
+    xHi[ix] = yHi[iy];
+    xLo[ix] = yLo[iy];
+    yHi[iy] = tempHi;
+    yLo[iy] = tempLo;
+  }
+}

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

@@ -75,3 +75,9 @@ fn ddAddProtected(a: DD, b: DD, threadSlot: u32) -> DD {
   let loSum = a.lo + b.lo;
   return fastTwoSumProtected(s.hi, s.lo + loSum, threadSlot);
 }
+
+// Double-double subtraction — a - b, via exact negation (a sign-bit flip,
+// no rounding) then ddAddProtected. Same protection contract.
+fn ddSubProtected(a: DD, b: DD, threadSlot: u32) -> DD {
+  return ddAddProtected(a, DD(negf(b.hi), negf(b.lo)), threadSlot);
+}

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

@@ -0,0 +1,45 @@
+// Requires f64/dekker.wgsl concatenated first for the DD struct, and
+// f64/utils/add.wgsl (ddSubProtected/ddAddProtected/negf) and
+// f64/utils/multiply.wgsl (ddMulProtected).
+//
+// Verified empirically on real hardware (NVIDIA GTX 1650 + Intel Mesa
+// integrated GPU, ~8000 random trials each plus explicit edge cases): no
+// compiler-reassociation-style corruption of the kind that broke twoSum/
+// twoProd (see add.wgsl's/multiply.wgsl's own headers) — every failure
+// found was a genuine algorithm gap, not a driver miscompile, and both are
+// fixed below (the b.hi==0.0 guard). One real, expected-shape difference
+// from every other protected op here: the low-power backend's observed
+// forward-error factor for this op specifically runs noticeably higher
+// (~7-14000x eps, vs ~3-9x on high-performance) than ddSqrtProtected's
+// (~3-4x on both) — division inherently amplifies input imprecision more
+// than a sum/product does, so a real routine built on this needs its own
+// backend-calibrated threshold, same as every other f64 arithmetic routine
+// in this codebase (see e.g. tests/drot/src/test.drot.js's THRESHOLDS).
+//
+// One Newton-style long-division refinement (Bailey/QD-style): q1 = a.hi /
+// b.hi is a plain f32 quotient, accurate to ~24 bits. Computing the residual
+// a - q1*b in DD arithmetic (not f32) recovers the bits q1 lost, and a
+// second plain division of that residual resolves them into a correction
+// term — combining q1 + q2 gives roughly double a lone f32 divide's
+// precision, matching this scheme's ~48-bit double-double target (already
+// short of real f64's 52 bits, so a second refinement step would chase
+// precision this representation has no room for).
+// b.hi == 0.0 makes q1 = a.hi/0.0 already the IEEE-754-correct answer
+// (±Infinity, or NaN for 0/0) via plain float division, but the refinement
+// below would corrupt it: p1 = q1*b multiplies that Infinity by a zero
+// divisor, and Infinity*0 is NaN by definition, poisoning everything after.
+// Substituting a safe non-zero denominator via select() — rather than
+// branching/returning early — keeps every thread calling ddMulProtected/
+// ddSubProtected/ddAddProtected unconditionally, which their internal
+// workgroupBarrier() requires; only the final result is selected between
+// the refined value and q1's own already-correct answer.
+fn ddDivProtected(a: DD, b: DD, threadSlot: u32) -> DD {
+  let bIsZero = b.hi == 0.0;
+  let q1 = a.hi / b.hi;
+  let safeB = DD(select(b.hi, 1.0, bIsZero), select(b.lo, 0.0, bIsZero));
+  let p1 = ddMulProtected(DD(q1, 0.0), safeB, threadSlot);
+  let r1 = ddSubProtected(a, p1, threadSlot);
+  let q2 = r1.hi / safeB.hi;
+  let refined = ddAddProtected(DD(q1, 0.0), DD(q2, 0.0), threadSlot);
+  return DD(select(refined.hi, q1, bIsZero), select(refined.lo, 0.0, bIsZero));
+}

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

@@ -15,7 +15,11 @@ 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
+  // Top 11 mantissa bits, so hi carries 12 significant bits with the implicit
+  // leading 1 — the halves are multiplied pairwise and f32 holds 24, so a
+  // wider split rounds those products and the "exact" error term goes wrong.
+  // Matches SPLIT_CONST = 2^12+1 used by the Veltkamp path below.
+  let hiBits = bits & 0xFFFFF000u;
   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);
@@ -63,19 +67,24 @@ fn twoProdFma(a: f32, b: f32) -> DD {
 
 // 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 {
+// ~48-bit floor anyway, so folding them in with plain f32 loses nothing.
+//
+// Another real compiler bug, distinct from add.wgsl's twoSum one — confirmed
+// on Intel Mesa ANV: when p.lo feeds straight into `crossAndLo` unobserved,
+// the compiler folds it away entirely. Materializing p.lo itself through
+// workgroup memory + workgroupBarrier() (like twoSumProtected does for its
+// sum) is what fixes it, so ddMulRaw now takes threadSlot and always pays
+// that barrier — no longer a plain unprotected batchable helper.
+fn ddMulRaw(a: DD, b: DD, threadSlot: u32) -> DD {
   let p = twoProdBit(a.hi, b.hi);
-  let crossAndLo = p.lo + (a.hi * b.lo + a.lo * b.hi);
+  dekkerScratch[threadSlot] = p.lo;
+  workgroupBarrier();
+  let pLo = dekkerScratch[threadSlot];
+  let crossAndLo = pLo + (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);
+  let raw = ddMulRaw(a, b, threadSlot);
   return fastTwoSumProtected(raw.hi, raw.lo, threadSlot);
 }

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

@@ -0,0 +1,45 @@
+// Requires f64/dekker.wgsl concatenated first for the DD struct, and
+// f64/utils/add.wgsl (ddSubProtected/ddAddProtected) and
+// f64/utils/multiply.wgsl (twoProdBit — squaring a plain f32 needs no
+// barrier, per multiply.wgsl's own note that twoProdBit is universally safe
+// unprotected).
+//
+// Verified empirically on real hardware (NVIDIA GTX 1650 + Intel Mesa
+// integrated GPU, ~8000 random trials each plus explicit edge cases): no
+// compiler-reassociation-style corruption of the kind that broke twoSum/
+// twoProd (see add.wgsl's/multiply.wgsl's own headers) — every failure
+// found was a genuine algorithm gap (the a.hi==0.0 case below), not a
+// driver miscompile, and is fixed. Observed forward-error factor against a
+// true f64 reference stayed ~3-4x eps on both backends across every random
+// trial — noticeably tighter than ddDivProtected's own low-power spread
+// (see divide.wgsl's header), since sqrt has no denominator to be unlucky
+// about.
+//
+// One Newton refinement step (the classic extended-precision sqrt trick):
+// x0 = sqrt(a.hi) is a plain f32 approximation; the residual a - x0^2,
+// computed in DD arithmetic, recovers what x0 lost, and linearizing sqrt
+// around x0 (dividing that residual by 2*x0) gives a correction term
+// roughly doubling the precision — same ~48-bit target as ddDivProtected,
+// so one step is enough.
+//
+// Undefined for a.hi < 0.0, same as plain sqrt() — callers must guard
+// themselves; this never checks.
+//
+// a.hi == 0.0 (a genuinely zero input, not an underflowed one — zero is
+// exactly representable in f32, unlike this scheme's real range limits;
+// see splitDoubleDouble's own doc comment) makes x0 = sqrt(0) = 0, and the
+// correction step would divide by 2*x0 = 0. Substituting a safe non-zero
+// denominator via select() — rather than branching/returning early — keeps
+// every thread calling ddSubProtected/ddAddProtected unconditionally, which
+// their internal workgroupBarrier() requires; only the final result is
+// selected between the computed value and the exact DD(0,0) answer.
+fn ddSqrtProtected(a: DD, threadSlot: u32) -> DD {
+  let isZero = a.hi == 0.0;
+  let x0 = sqrt(a.hi);
+  let x0sq = twoProdBit(x0, x0);
+  let r = ddSubProtected(a, x0sq, threadSlot);
+  let safeDenom = select(2.0 * x0, 1.0, isZero);
+  let correction = r.hi / safeDenom;
+  let result = ddAddProtected(DD(x0, 0.0), DD(correction, 0.0), threadSlot);
+  return DD(select(result.hi, 0.0, isZero), select(result.lo, 0.0, isZero));
+}

package/src/shaders/index.mjs

@@ -53,15 +53,24 @@ 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 };
@@ -90,6 +99,34 @@ routineShaders.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";
@@ -106,9 +143,41 @@ routineShaders.idamax = {
 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

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/snrm2/snrm2.mjs

@@ -13,14 +13,12 @@ import { extractResult } from "../util/result.mjs";
 import { getPipeline } from "../util/pipeline.mjs";
 import { GpuVector } from "../classes/GpuVector.mjs";
 import { WGS } from "../util/constants.mjs";
-import { requireSameDevice } from "../util/device.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.");
@@ -46,13 +44,19 @@ export async function snrm2(device, n, x, incx) {
   try {
     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,
+    partialsScaleBuffer = createStorageBuffer(
+      device,
       2 * WGS * 4,
       "snrm2-partials-scale",
     );
-    partialsSsqBuffer = createStorageBuffer(device, 2 * WGS * 4, "snrm2-partials-ssq");
+    partialsSsqBuffer = createStorageBuffer(
+      device,
+      2 * WGS * 4,
+      "snrm2-partials-ssq",
+    );
     resultBuffer = createResultBuffer(device, 4, "snrm2-result"); // final f32 scalar
-    paramsBuffer = createParamsBuffer(device,
+    paramsBuffer = createParamsBuffer(
+      device,
       [
         { value: n, type: "u32" },
         { value: incx, type: "u32" },
@@ -66,7 +70,8 @@ export async function snrm2(device, n, x, incx) {
       partialsSsqBuffer,
       paramsBuffer,
     ]);
-    const { commandEncoder: enc1, ts: ts1 } = runComputePass(device,
+    const { commandEncoder: enc1, ts: ts1 } = runComputePass(
+      device,
       pipelineMain,
       bgMain,
       2 * WGS,
@@ -74,12 +79,13 @@ export async function snrm2(device, n, x, incx) {
 
     submit(device, enc1);
 
-    const bgReduce = createBindGroup(device, pipelineReduce.getBindGroupLayout(0), [
-      partialsScaleBuffer,
-      partialsSsqBuffer,
-      resultBuffer,
-    ]);
-    const { commandEncoder: enc2, ts: ts2 } = runComputePass(device,
+    const bgReduce = createBindGroup(
+      device,
+      pipelineReduce.getBindGroupLayout(0),
+      [partialsScaleBuffer, partialsSsqBuffer, resultBuffer],
+    );
+    const { commandEncoder: enc2, ts: ts2 } = runComputePass(
+      device,
       pipelineReduce,
       bgReduce,
       1,