wgblas 1.2.1 → 2.1.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/shaders/symmetrize.wgsl

@@ -0,0 +1,31 @@
+// symmetrize: Adense := full dense expansion of a symmetric matrix stored
+// with only its `uplo` triangle meaningful (the other triangle is implied
+// by symmetry: A[i,j] = A[j,i]). A plain element-wise pass, no tiling or
+// shared memory needed — used to materialize a dense operand for routines
+// that read a symmetric matrix as a normal dense gemm input (e.g. ssymm),
+// rather than teaching the tiled gemm kernel itself to mirror-read.
+
+@group(0) @binding(0) var<storage, read>       A: array<f32>;
+@group(0) @binding(1) var<storage, read_write> Adense: array<f32>;
+
+struct Params {
+  n:    u32,
+  lda:  u32,
+  ldd:  u32, // leading dimension of Adense
+  uplo: u32, // 0 = lower (stored where col <= row), 1 = upper (col >= row)
+}
+
+@group(0) @binding(2) var<uniform> params: Params;
+
+@compute @workgroup_size(8, 8)
+fn main(@builtin(global_invocation_id) gid: vec3u) {
+  let row = gid.y;
+  let col = gid.x;
+  if (row >= params.n || col >= params.n) {
+    return;
+  }
+
+  let isStored = select(col >= row, col <= row, params.uplo == 0u);
+  let srcIdx = select(col * params.lda + row, row * params.lda + col, isStored);
+  Adense[row * params.ldd + col] = A[srcIdx];
+}

package/src/shaders/triangularize.wgsl

@@ -0,0 +1,44 @@
+// triangularize: Adense := dense expansion of op(A) (A or A^T per `trans`),
+// zero-filling the unstored triangle (exact for a matmul) so strmm can reuse
+// sgemm's kernel unchanged. `diag=1` substitutes 1.0 on the diagonal.
+
+@group(0) @binding(0) var<storage, read>       A: array<f32>;
+@group(0) @binding(1) var<storage, read_write> Adense: array<f32>;
+
+struct Params {
+  n:     u32,
+  lda:   u32,
+  ldd:   u32, // leading dimension of Adense
+  uplo:  u32, // 0 = lower, 1 = upper
+  trans: u32, // 0 = no-transpose (op(A) = A), 1 = transpose (op(A) = A^T)
+  diag:  u32, // 0 = non-unit, 1 = unit
+}
+
+@group(0) @binding(2) var<uniform> params: Params;
+
+@compute @workgroup_size(8, 8)
+fn main(@builtin(global_invocation_id) gid: vec3u) {
+  let row = gid.y;
+  let col = gid.x;
+  if (row >= params.n || col >= params.n) {
+    return;
+  }
+
+  if (row == col) {
+    Adense[row * params.ldd + col] = select(A[row * params.lda + row], 1.0, params.diag == 1u);
+    return;
+  }
+
+  var isMeaningful: bool;
+  var srcRow: u32;
+  var srcCol: u32;
+  if (params.trans == 0u) {
+    isMeaningful = select(col >= row, col <= row, params.uplo == 0u);
+    srcRow = row; srcCol = col;
+  } else {
+    isMeaningful = select(col <= row, col >= row, params.uplo == 0u);
+    srcRow = col; srcCol = row;
+  }
+
+  Adense[row * params.ldd + col] = select(0.0, A[srcRow * params.lda + srcCol], isMeaningful);
+}

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,9 +24,9 @@ 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/gpuvec.snrm2.js}
+ * {@includeCode ../../examples/snrm2/gpu.snrm2.js}
  *
  * @param device - GPUDevice from `init()`
  * @param n - number of elements (must be a positive integer)

package/src/snrm2/snrm2.mjs

@@ -12,14 +12,16 @@ import { extractTimestamp } from "../util/benchmark.mjs";
 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";
 
-const WGS = 64; // workgroup size
 
 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.");
+  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 +34,25 @@ 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
-    paramsBuffer = createParamsBuffer(
+    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 +60,57 @@ 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(
+    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,
+    const bgReduce = createBindGroup(device, pipelineReduce.getBindGroupLayout(0), [
+      partialsScaleBuffer,
+      partialsSsqBuffer,
       resultBuffer,
     ]);
-    const { commandEncoder: enc2, ts: ts2 } = runComputePass(
+    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,10 +33,9 @@ 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/gpuvec.srot.js}
+ * {@includeCode ../../examples/srot/gpu.srot.js}
  *
  * @param device - GPUDevice from `init()`
  * @param n - number of elements (must be a positive integer)

package/src/srot/srot.mjs

@@ -11,6 +11,7 @@ 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 { requireSameDevice } from "../util/device.mjs";
 
 export async function srot(device, n, x, incx, y, incy, c, s) {
   const xIsGpu = x instanceof GpuVector;
@@ -18,6 +19,7 @@ export async function srot(device, n, x, incx, y, incy, c, s) {
 
   if (!(device instanceof GPUDevice))
     throw new Error("device must be a GPUDevice.");
+  requireSameDevice(device, "srot", { x, y });
   if (
     !Number.isInteger(n) ||
     !Number.isInteger(incx) ||
@@ -58,9 +60,9 @@ 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);
-    paramsBuffer = createParamsBuffer(
+    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,19 +73,19 @@ 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(
+    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);
 

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,10 +32,9 @@ 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/gpuvec.srotm.js}
+ * {@includeCode ../../examples/srotm/gpu.srotm.js}
  *
  * @param device - GPUDevice from `init()`
  * @param n - number of elements (must be a positive integer)

package/src/srotm/srotm.mjs

@@ -11,6 +11,7 @@ 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 { requireSameDevice } from "../util/device.mjs";
 
 export async function srotm(device, n, x, incx, y, incy, param) {
   const xIsGpu = x instanceof GpuVector;
@@ -18,6 +19,7 @@ export async function srotm(device, n, x, incx, y, incy, param) {
 
   if (!(device instanceof GPUDevice))
     throw new Error("device must be a GPUDevice.");
+  requireSameDevice(device, "srotm", { x, y });
   if (
     !Number.isInteger(n) ||
     !Number.isInteger(incx) ||
@@ -26,6 +28,13 @@ export async function srotm(device, n, x, incx, y, incy, param) {
     throw new Error("n, incx, and incy must be integers.");
   if (!(param instanceof Float32Array) || param.length !== 5)
     throw new Error("param must be a Float32Array of length 5.");
+  if (
+    param[0] !== -2 &&
+    param[0] !== -1 &&
+    param[0] !== 0 &&
+    param[0] !== 1
+  )
+    throw new Error("param[0] (flag) must be one of -2, -1, 0, or 1.");
   if (incx <= 0 || incy <= 0)
     throw new Error("incx and incy must be positive.");
   if (!xIsGpu && !(x instanceof Float32Array))
@@ -56,10 +65,10 @@ export async function srotm(device, n, x, incx, y, incy, param) {
   let readY = null;
 
   try {
-    xBuffer = xIsGpu ? x._buf : uploadBuffer(x, "srotm-x", true);
-    yBuffer = yIsGpu ? y._buf : uploadBuffer(y, "srotm-y", true);
-    paramBuffer = uploadBuffer(param, "srotm-param", false);
-    paramsBuffer = createParamsBuffer(
+    xBuffer = xIsGpu ? x._buf : uploadBuffer(device, x, "srotm-x", true);
+    yBuffer = yIsGpu ? y._buf : uploadBuffer(device, y, "srotm-y", true);
+    paramBuffer = uploadBuffer(device, param, "srotm-param", false);
+    paramsBuffer = createParamsBuffer(device,
       [
         { value: n, type: "u32" },
         { value: incx, type: "u32" },
@@ -68,20 +77,20 @@ export async function srotm(device, n, x, incx, y, incy, param) {
       "srotm-params",
     );
 
-    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+    const bindGroup = createBindGroup(device, pipeline.getBindGroupLayout(0), [
       xBuffer,
       yBuffer,
       paramBuffer,
       paramsBuffer,
     ]);
-    const { commandEncoder, ts } = runComputePass(
+    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);
 

package/src/sscal/sscal.d.mts

@@ -1,7 +1,7 @@
 import { GpuVector } from "../classes/GpuVector.mjs";
 
 /**
- * Scales a single-precision vector by a constant: x = alpha * x
+ * Scales a single-precision vector by a constant: $$x \leftarrow \alpha x$$
  *
  * {@includeCode ../../examples/sscal/sscal.js}
  *
@@ -22,12 +22,12 @@ export declare function sscal(
   alpha: number,
   x: Float32Array,
   incx: number,
-): Promise<Float32Array | { result: Float32Array; gpuTimeMs: number }>;
+): Promise<{ x: Float32Array } | { x: Float32Array; gpuTimeMs: number }>;
 
 /**
- * Scales a single-precision vector by a constant: x = alpha * x
+ * Scales a single-precision vector by a constant: $$x \leftarrow \alpha x$$
  *
- * {@includeCode ../../examples/sscal/gpuvec.sscal.js}
+ * {@includeCode ../../examples/sscal/gpu.sscal.js}
  *
  * @param device - GPUDevice from `init()`
  * @param n - number of elements to scale (must be a positive integer)

package/src/sscal/sscal.mjs

@@ -11,12 +11,14 @@ import { extractTimestamp } from "../util/benchmark.mjs";
 import { getPipeline } from "../util/pipeline.mjs";
 import { calcWorkgroups } from "../util/workgroup.mjs";
 import { GpuVector } from "../classes/GpuVector.mjs";
+import { requireSameDevice } from "../util/device.mjs";
 
 export async function sscal(device, n, alpha, x, incx) {
   const xIsGpu = x instanceof GpuVector;
-  
+
   if (!(device instanceof GPUDevice))
     throw new Error("device must be a GPUDevice.");
+  requireSameDevice(device, "sscal", { x });
   if (!Number.isInteger(n) || !Number.isInteger(incx))
     throw new Error("n and incx must be integers.");
   if (typeof alpha !== "number")
@@ -26,7 +28,7 @@ export async function sscal(device, n, alpha, x, incx) {
   if (incx <= 0) throw new Error("incx must be positive.");
   if (!(x instanceof Float32Array) && !(x instanceof GpuVector))
     throw new Error("x must be a Float32Array or GpuVector.");
-  if (n <= 0) return xIsGpu ? {} : x;
+  if (n <= 0) return xIsGpu ? {} : { x };
   if (x.length < (n - 1) * incx + 1)
     throw new Error(
       "x does not have enough elements for the given n and incx.",
@@ -39,8 +41,8 @@ export async function sscal(device, n, alpha, x, incx) {
   let readBuffer = null;
 
   try {
-    xBuffer = xIsGpu ? x._buf : uploadBuffer(x, "sscal-x", true);
-    paramsBuffer = createParamsBuffer(
+    xBuffer = xIsGpu ? x._buf : uploadBuffer(device, x, "sscal-x", true);
+    paramsBuffer = createParamsBuffer(device,
       [
         { value: n, type: "u32" },
         { value: alpha, type: "f32" },
@@ -49,18 +51,18 @@ export async function sscal(device, n, alpha, x, incx) {
       "sscal-params",
     );
 
-    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+    const bindGroup = createBindGroup(device, pipeline.getBindGroupLayout(0), [
       xBuffer,
       paramsBuffer,
     ]);
-    const { commandEncoder, ts } = runComputePass(
+    const { commandEncoder, ts } = runComputePass(device,
       pipeline,
       bindGroup,
-      calcWorkgroups(n),
+      calcWorkgroups(device, n),
     );
-    readBuffer = xIsGpu ? null : stageReadback(commandEncoder, xBuffer);
+    readBuffer = xIsGpu ? null : stageReadback(device, commandEncoder, xBuffer);
 
-    submit(commandEncoder);
+    submit(device, commandEncoder);
 
     const gpuTimeMs = await extractTimestamp(ts);
 
@@ -72,7 +74,7 @@ export async function sscal(device, n, alpha, x, incx) {
     const result = await extractResult(readBuffer, Float32Array);
     readBuffer = null; // extractResult already destroyed it
     if (gpuTimeMs !== undefined) return { x: result, gpuTimeMs };
-    return result;
+    return { x: result };
   } finally {
     if (!xIsGpu && xBuffer) destroyBuffers(xBuffer);
     if (paramsBuffer) destroyBuffers(paramsBuffer);

package/src/sswap/sswap.d.mts

@@ -1,7 +1,7 @@
 import { GpuVector } from "../classes/GpuVector.mjs";
 
 /**
- * Swaps the elements of two single-precision vectors: x <-> y
+ * Swaps the elements of two single-precision vectors: $$x \leftrightarrow y$$
  *
  * {@includeCode ../../examples/sswap/sswap.js}
  *
@@ -27,9 +27,9 @@ export declare function sswap(
 ): Promise<{ x: Float32Array; y: Float32Array } | { x: Float32Array; y: Float32Array; gpuTimeMs: number }>;
 
 /**
- * Swaps the elements of two single-precision vectors: x <-> y
+ * Swaps the elements of two single-precision vectors: $$x \leftrightarrow y$$
  *
- * {@includeCode ../../examples/sswap/gpuvec.sswap.js}
+ * {@includeCode ../../examples/sswap/gpu.sswap.js}
  *
  * @param device - GPUDevice from `init()`
  * @param n - number of elements to swap (must be a positive integer)