wgblas 0.1.2 → 1.1.0

package/src/strmv/strmv.mjs

@@ -0,0 +1,138 @@
+import {
+  uploadBuffer,
+  createParamsBuffer,
+  stageReadback,
+  destroyBuffers,
+} from "../util/buffer.mjs";
+import { createBindGroup } from "../util/bindgroup.mjs";
+import { runComputePass, submit } from "../util/compute.mjs";
+import { extractResult } from "../util/result.mjs";
+import { extractTimestamp } from "../util/benchmark.mjs";
+import { getPipeline } from "../util/pipeline.mjs";
+import { GpuVector } from "../classes/GpuVector.mjs";
+import { GpuMatrix } from "../classes/GpuMatrix.mjs";
+
+export async function strmv(device, uplo, trans, diag, n, A, lda, x, incx, y, incy, layout = "row-major") {
+  const xIsGpu = x instanceof GpuVector;
+  const yIsGpu = y instanceof GpuVector;
+  const AIsGpu = A instanceof GpuMatrix;
+  const isUnit = diag === "unit";
+
+  if (!(device instanceof GPUDevice))
+    throw new Error("device must be a GPUDevice.");
+  if (uplo !== "lower" && uplo !== "upper")
+    throw new Error("uplo must be 'lower' or 'upper'.");
+  if (trans !== "no-transpose" && trans !== "transpose")
+    throw new Error("trans must be 'no-transpose' or 'transpose'.");
+  if (!isUnit && diag !== "non-unit")
+    throw new Error("diag must be 'unit' or 'non-unit'.");
+  if (layout !== "row-major" && layout !== "column-major")
+    throw new Error("layout must be 'row-major' or 'column-major'.");
+  if (
+    !Number.isInteger(n) ||
+    !Number.isInteger(incx) ||
+    !Number.isInteger(incy) ||
+    !Number.isInteger(lda)
+  )
+    throw new Error("n, incx, incy, and lda must be integers.");
+  if (incx <= 0 || incy <= 0)
+    throw new Error("incx and incy must be positive.");
+  if (lda < n) throw new Error("lda must be >= n.");
+  if (!AIsGpu && !(A instanceof Float32Array))
+    throw new Error("A must be a Float32Array or GpuMatrix.");
+  if (!xIsGpu && !(x instanceof Float32Array))
+    throw new Error("x must be a Float32Array or GpuVector.");
+  if (!yIsGpu && !(y instanceof Float32Array))
+    throw new Error("y must be a Float32Array or GpuVector.");
+  if (xIsGpu !== yIsGpu)
+    throw new Error(
+      "x and y must be the same type (both Float32Array or both GpuVector).",
+    );
+  if (xIsGpu && x._buf === y._buf)
+    throw new Error(
+      "x and y must not reference the same GPU buffer when both are GpuVectors.",
+    );
+  if (xIsGpu && !AIsGpu)
+    throw new Error("A must be a GpuMatrix when x and y are GpuVectors.");
+  if (AIsGpu && yIsGpu && A._buf === y._buf)
+    throw new Error("A and y must not reference the same GPU buffer.");
+  if (AIsGpu && lda !== A.lda)
+    throw new Error("lda must match A.lda when A is a GpuMatrix.");
+  if (AIsGpu && (A.rows < n || A.cols < n))
+    throw new Error("A is too small for the given n.");
+  if (n < 0) throw new Error("n must be non-negative.");
+  if (n === 0) return yIsGpu ? {} : { y };
+
+  if (!AIsGpu && A.length < (n - 1) * lda + n)
+    throw new Error(
+      "A does not have enough elements for the given n and lda.",
+    );
+  if (x.length < (n - 1) * incx + 1)
+    throw new Error(
+      "x does not have enough elements for the given n and incx.",
+    );
+  if (y.length < (n - 1) * incy + 1)
+    throw new Error(
+      "y does not have enough elements for the given n and incy.",
+    );
+
+  // GpuMatrix's own layout wins over the argument; column-major A reinterpreted row-major is A^T, so flip both uplo and trans to reproduce the requested op(A).
+  const effLayout = AIsGpu ? A.layout : layout;
+  const isColMajor = effLayout === "column-major";
+  const isLower = isColMajor ? uplo === "upper" : uplo === "lower";
+  const isNoTrans = isColMajor ? trans === "transpose" : trans === "no-transpose";
+
+  const pipeline = await getPipeline(device, "strmv");
+
+  let ABuffer = null;
+  let xBuffer = null;
+  let yBuffer = null;
+  let paramsBuffer = null;
+
+  try {
+    ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "strmv-A", false);
+    xBuffer = xIsGpu ? x._buf : uploadBuffer(x, "strmv-x", false);
+    yBuffer = yIsGpu ? y._buf : uploadBuffer(y, "strmv-y", true);
+    paramsBuffer = createParamsBuffer(
+      [
+        { value: n,             type: "u32" },
+        { value: incx,          type: "u32" },
+        { value: incy,          type: "u32" },
+        { value: lda,           type: "u32" },
+        { value: isNoTrans ? 0 : 1, type: "u32" },
+        { value: isLower ? 0 : 1,   type: "u32" },
+        { value: isUnit ? 1 : 0,    type: "u32" },
+      ],
+      "strmv-params",
+    );
+
+    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+      ABuffer,
+      xBuffer,
+      yBuffer,
+      paramsBuffer,
+    ]);
+
+    const wgCount = Math.min(n, device.limits.maxComputeWorkgroupsPerDimension);
+    const { commandEncoder, ts } = runComputePass(pipeline, bindGroup, wgCount);
+    const readBuffer = yIsGpu ? null : stageReadback(commandEncoder, yBuffer);
+
+    submit(commandEncoder);
+
+    const gpuTimeMs = await extractTimestamp(ts);
+
+    if (yIsGpu) {
+      if (gpuTimeMs !== undefined) return { gpuTimeMs };
+      return {};
+    }
+
+    const result = await extractResult(readBuffer, Float32Array);
+    if (gpuTimeMs !== undefined) return { y: result, gpuTimeMs };
+    return { y: result };
+  } finally {
+    if (!AIsGpu && ABuffer) destroyBuffers(ABuffer);
+    if (!xIsGpu && xBuffer) destroyBuffers(xBuffer);
+    if (!yIsGpu && yBuffer) destroyBuffers(yBuffer);
+    if (paramsBuffer) destroyBuffers(paramsBuffer);
+  }
+}

package/src/strsv/strsv.d.mts

@@ -0,0 +1,106 @@
+import { GpuVector } from "../classes/GpuVector.mjs";
+import { GpuMatrix } from "../classes/GpuMatrix.mjs";
+
+/**
+ * Solves the triangular system op(A) * x = b for x, in place (x holds b on
+ * input, the solution on output).
+ *
+ * A is an n×n triangular matrix stored in row-major order. Only the triangle
+ * specified by `uplo` is referenced; the other triangle is not accessed.
+ *
+ * {@includeCode ../../examples/strsv/strsv.js}
+ *
+ * **Browser (standalone HTML):**
+ * {@includeCode ../../examples/strsv/web/strsv.html}
+ *
+ * @param device - GPUDevice from `init()`
+ * @param uplo   - `'lower'` to use the lower triangle, `'upper'` to use the upper triangle
+ * @param trans  - `'no-transpose'` to solve A*x=b, `'transpose'` to solve A^T*x=b
+ * @param diag   - `'unit'` to treat the diagonal as all-ones (A's diagonal is not read), `'non-unit'` to read it
+ * @param n      - order of the matrix A (number of rows and columns)
+ * @param A      - Float32Array, row-major or column-major (see `layout`), at least (n-1)*lda+n elements
+ * @param lda    - leading dimension of A (>= n either way — A is square)
+ * @param x      - Float32Array holding b on input, the solution on output; length at least (n-1)*incx+1
+ * @param incx   - stride for x (must be a positive integer)
+ * @param layout - storage layout of `A` (default: `'row-major'`); column-major
+ *   flips both the stored triangle and the effective `trans` (the system
+ *   being solved stays what you asked for either way)
+ * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/strsv/strsv.mjs#L15">Source code: strsv.mjs (L15)</a>
+ * @category BLAS Level 2
+ */
+export declare function strsv(
+  device: GPUDevice,
+  uplo: 'lower' | 'upper',
+  trans: 'no-transpose' | 'transpose',
+  diag: 'unit' | 'non-unit',
+  n: number,
+  A: Float32Array,
+  lda: number,
+  x: Float32Array,
+  incx: number,
+  layout?: 'row-major' | 'column-major',
+): Promise<{ x: Float32Array; gpuTimeMs?: number }>;
+
+/**
+ * Solves the triangular system op(A) * x = b for x, in place.
+ *
+ * A is kept GPU-resident; x is a CPU Float32Array. `A`'s own `layout` (set at
+ * `GpuMatrix.from` time) determines the operation — there is no separate
+ * `layout` argument here.
+ *
+ * @param device - GPUDevice from `init()`
+ * @param uplo   - `'lower'` to use the lower triangle, `'upper'` to use the upper triangle
+ * @param trans  - `'no-transpose'` to solve A*x=b, `'transpose'` to solve A^T*x=b
+ * @param diag   - `'unit'` to treat the diagonal as all-ones (A's diagonal is not read), `'non-unit'` to read it
+ * @param n      - order of the matrix A
+ * @param A      - GpuMatrix, GPU-resident
+ * @param lda    - leading dimension of A (must equal A.lda)
+ * @param x      - Float32Array holding b on input, the solution on output
+ * @param incx   - stride for x (must be a positive integer)
+ * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/strsv/strsv.mjs#L15">Source code: strsv.mjs (L15)</a>
+ * @category BLAS Level 2
+ */
+export declare function strsv(
+  device: GPUDevice,
+  uplo: 'lower' | 'upper',
+  trans: 'no-transpose' | 'transpose',
+  diag: 'unit' | 'non-unit',
+  n: number,
+  A: GpuMatrix,
+  lda: number,
+  x: Float32Array,
+  incx: number,
+): Promise<{ x: Float32Array; gpuTimeMs?: number }>;
+
+/**
+ * Solves the triangular system op(A) * x = b for x, in place.
+ *
+ * x is kept resident on the GPU (mutated in place). A must be a GpuMatrix;
+ * its own `layout` (set at `GpuMatrix.from` time) determines the operation —
+ * there is no separate `layout` argument here.
+ *
+ * {@includeCode ../../examples/strsv/gpuvec.strsv.js}
+ *
+ * @param device - GPUDevice from `init()`
+ * @param uplo   - `'lower'` to use the lower triangle, `'upper'` to use the upper triangle
+ * @param trans  - `'no-transpose'` to solve A*x=b, `'transpose'` to solve A^T*x=b
+ * @param diag   - `'unit'` to treat the diagonal as all-ones (A's diagonal is not read), `'non-unit'` to read it
+ * @param n      - order of the matrix A
+ * @param A      - GpuMatrix, GPU-resident
+ * @param lda    - leading dimension of A (must equal A.lda)
+ * @param x      - GpuVector holding b on input, the solution on output (mutated in place)
+ * @param incx   - stride for x (must be a positive integer)
+ * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/strsv/strsv.mjs#L15">Source code: strsv.mjs (L15)</a>
+ * @category BLAS Level 2
+ */
+export declare function strsv(
+  device: GPUDevice,
+  uplo: 'lower' | 'upper',
+  trans: 'no-transpose' | 'transpose',
+  diag: 'unit' | 'non-unit',
+  n: number,
+  A: GpuMatrix,
+  lda: number,
+  x: GpuVector,
+  incx: number,
+): Promise<{ gpuTimeMs?: number }>;

package/src/strsv/strsv.mjs

@@ -0,0 +1,207 @@
+import {
+  uploadBuffer,
+  createParamsBuffer,
+  createStorageBuffer,
+  stageReadback,
+  destroyBuffers,
+} from "../util/buffer.mjs";
+import { createBindGroup } from "../util/bindgroup.mjs";
+import { beginTimedEncoder, encodePass, submit } from "../util/compute.mjs";
+import { extractResult } from "../util/result.mjs";
+import { resolveTimestamp, extractTimestamp } from "../util/benchmark.mjs";
+import { getPipeline } from "../util/pipeline.mjs";
+import { GpuVector } from "../classes/GpuVector.mjs";
+import { GpuMatrix } from "../classes/GpuMatrix.mjs";
+
+// Blocked triangular solve via explicit block inversion (invert/apply/update passes) instead of barrier-per-row substitution.
+const BLOCK_SIZE = 64;
+
+// One shared buffer holds all blocks' params (offset blockIndex*stride) instead of one buffer per block — avoids the O(numBlocks) createBuffer/writeBuffer calls that dominated CPU time.
+function packBlockParams(numBlocks, stride, fieldsPerBlock) {
+  const data = new ArrayBuffer(numBlocks * stride);
+  const view = new DataView(data);
+  for (let blockIndex = 0; blockIndex < numBlocks; blockIndex++) {
+    const fields = fieldsPerBlock(blockIndex);
+    const base = blockIndex * stride;
+    fields.forEach((value, i) => view.setUint32(base + i * 4, value, true));
+  }
+  return data;
+}
+
+function createSharedParamsBuffer(device, data, label) {
+  const buffer = device.createBuffer({
+    label,
+    size: data.byteLength,
+    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
+  });
+  device.queue.writeBuffer(buffer, 0, data);
+  return buffer;
+}
+
+export async function strsv(device, uplo, trans, diag, n, A, lda, x, incx, layout = "row-major") {
+  const xIsGpu = x instanceof GpuVector;
+  const AIsGpu = A instanceof GpuMatrix;
+  const isUnit = diag === "unit";
+
+  if (!(device instanceof GPUDevice))
+    throw new Error("device must be a GPUDevice.");
+  if (uplo !== "lower" && uplo !== "upper")
+    throw new Error("uplo must be 'lower' or 'upper'.");
+  if (trans !== "no-transpose" && trans !== "transpose")
+    throw new Error("trans must be 'no-transpose' or 'transpose'.");
+  if (!isUnit && diag !== "non-unit")
+    throw new Error("diag must be 'unit' or 'non-unit'.");
+  if (layout !== "row-major" && layout !== "column-major")
+    throw new Error("layout must be 'row-major' or 'column-major'.");
+  if (!Number.isInteger(n) || !Number.isInteger(incx) || !Number.isInteger(lda))
+    throw new Error("n, incx, and lda must be integers.");
+  if (incx <= 0) throw new Error("incx must be positive.");
+  if (lda < n) throw new Error("lda must be >= n.");
+  if (!AIsGpu && !(A instanceof Float32Array))
+    throw new Error("A must be a Float32Array or GpuMatrix.");
+  if (!xIsGpu && !(x instanceof Float32Array))
+    throw new Error("x must be a Float32Array or GpuVector.");
+  if (xIsGpu && !AIsGpu)
+    throw new Error("A must be a GpuMatrix when x is a GpuVector.");
+  if (AIsGpu && lda !== A.lda)
+    throw new Error("lda must match A.lda when A is a GpuMatrix.");
+  if (AIsGpu && (A.rows < n || A.cols < n))
+    throw new Error("A is too small for the given n.");
+  if (n < 0) throw new Error("n must be non-negative.");
+  if (n === 0) return xIsGpu ? {} : { x };
+
+  if (!AIsGpu && A.length < (n - 1) * lda + n)
+    throw new Error(
+      "A does not have enough elements for the given n and lda.",
+    );
+  if (x.length < (n - 1) * incx + 1)
+    throw new Error(
+      "x does not have enough elements for the given n and incx.",
+    );
+
+  // GpuMatrix's own layout wins over the argument; column-major A reinterpreted row-major is A^T, so flip both uplo and trans to reproduce the requested system.
+  const effLayout = AIsGpu ? A.layout : layout;
+  const isColMajor = effLayout === "column-major";
+  const isLower = isColMajor ? uplo === "upper" : uplo === "lower";
+  const isNoTrans = isColMajor ? trans === "transpose" : trans === "no-transpose";
+
+  const invertPipeline = await getPipeline(device, "strsv_invert_block");
+  const applyPipeline = await getPipeline(device, "strsv_apply_inverse");
+  const updatePipeline = await getPipeline(device, "strsv_update");
+
+  // Same forward/backward pairing the shaders use.
+  const forward = isNoTrans === isLower;
+  const blockStarts = [];
+  for (let s = 0; s < n; s += BLOCK_SIZE) blockStarts.push(s);
+  if (!forward) blockStarts.reverse();
+  const numBlocks = blockStarts.length;
+
+  const maxWg = device.limits.maxComputeWorkgroupsPerDimension;
+  const stride = device.limits.minUniformBufferOffsetAlignment;
+
+  let ABuffer = null;
+  let xBuffer = null;
+  let AinvBuffer = null;
+  let applyParamsBuffer = null;
+  let updateParamsBuffer = null;
+  let invertParams = null;
+
+  try {
+    ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "strsv-A", false);
+    xBuffer = xIsGpu ? x._buf : uploadBuffer(x, "strsv-x", true);
+    // One BLOCK_SIZE x BLOCK_SIZE dense region per block (row-major), even
+    // though only a triangular half is ever nonzero — see strsv_invert_block.wgsl.
+    AinvBuffer = createStorageBuffer(
+      numBlocks * BLOCK_SIZE * BLOCK_SIZE * 4,
+      "strsv-Ainv",
+    );
+
+    // Every block's {blockStart, blockEnd} is fixed by its natural index
+    // regardless of traversal direction, so both packed buffers are indexed
+    // by blockIndex (0..numBlocks-1), not by loop position.
+    const applyData = packBlockParams(numBlocks, stride, (blockIndex) => {
+      const blockStart = blockIndex * BLOCK_SIZE;
+      const blockEnd = Math.min(blockStart + BLOCK_SIZE, n);
+      return [incx, blockIndex, blockStart, blockEnd];
+    });
+    applyParamsBuffer = createSharedParamsBuffer(device, applyData, "strsv-apply-params");
+
+    const updateData = packBlockParams(numBlocks, stride, (blockIndex) => {
+      const blockStart = blockIndex * BLOCK_SIZE;
+      const blockEnd = Math.min(blockStart + BLOCK_SIZE, n);
+      return [n, incx, lda, isNoTrans ? 0 : 1, isLower ? 0 : 1, blockStart, blockEnd];
+    });
+    updateParamsBuffer = createSharedParamsBuffer(device, updateData, "strsv-update-params");
+
+    const { commandEncoder, querySet } = beginTimedEncoder();
+
+    // Pre-pass: every block's inverse, fully parallel, one dispatch.
+    invertParams = createParamsBuffer(
+      [
+        { value: n,                 type: "u32" },
+        { value: lda,               type: "u32" },
+        { value: isNoTrans ? 0 : 1, type: "u32" },
+        { value: isLower ? 0 : 1,   type: "u32" },
+        { value: isUnit ? 1 : 0,    type: "u32" },
+      ],
+      "strsv-invert-params",
+    );
+    const invertBindGroup = createBindGroup(invertPipeline.getBindGroupLayout(0), [
+      ABuffer, AinvBuffer, invertParams,
+    ]);
+    const invertDesc = querySet
+      ? { timestampWrites: { querySet, beginningOfPassWriteIndex: 0 } }
+      : undefined;
+    encodePass(commandEncoder, invertPipeline, invertBindGroup, { x: BLOCK_SIZE, y: numBlocks }, invertDesc);
+
+    for (let bi = 0; bi < blockStarts.length; bi++) {
+      const blockStart = blockStarts[bi];
+      const blockEnd = Math.min(blockStart + BLOCK_SIZE, n);
+      const blockIndex = blockStart / BLOCK_SIZE;
+      const isLastPass = bi === blockStarts.length - 1;
+      const paramsOffset = blockIndex * stride;
+
+      const applyBindGroup = createBindGroup(applyPipeline.getBindGroupLayout(0), [
+        AinvBuffer, xBuffer, { buffer: applyParamsBuffer, offset: paramsOffset, size: 16 },
+      ]);
+
+      const applyDesc = isLastPass && querySet
+        ? { timestampWrites: { querySet, endOfPassWriteIndex: 1 } }
+        : undefined;
+      encodePass(commandEncoder, applyPipeline, applyBindGroup, 1, applyDesc);
+
+      const remaining = forward ? n - blockEnd : blockStart;
+      if (remaining === 0) continue;
+
+      const updateBindGroup = createBindGroup(updatePipeline.getBindGroupLayout(0), [
+        ABuffer, xBuffer, { buffer: updateParamsBuffer, offset: paramsOffset, size: 32 },
+      ]);
+
+      const wgCount = Math.min(remaining, maxWg);
+      encodePass(commandEncoder, updatePipeline, updateBindGroup, wgCount);
+    }
+
+    const ts = resolveTimestamp(commandEncoder, querySet);
+    const readBuffer = xIsGpu ? null : stageReadback(commandEncoder, xBuffer);
+
+    submit(commandEncoder);
+
+    const gpuTimeMs = await extractTimestamp(ts);
+
+    if (xIsGpu) {
+      if (gpuTimeMs !== undefined) return { gpuTimeMs };
+      return {};
+    }
+
+    const result = await extractResult(readBuffer, Float32Array);
+    if (gpuTimeMs !== undefined) return { x: result, gpuTimeMs };
+    return { x: result };
+  } finally {
+    if (!AIsGpu && ABuffer) destroyBuffers(ABuffer);
+    if (!xIsGpu && xBuffer) destroyBuffers(xBuffer);
+    if (AinvBuffer) destroyBuffers(AinvBuffer);
+    if (applyParamsBuffer) destroyBuffers(applyParamsBuffer);
+    if (updateParamsBuffer) destroyBuffers(updateParamsBuffer);
+    if (invertParams) destroyBuffers(invertParams);
+  }
+}

package/src/util/benchmark.mjs

@@ -99,5 +99,5 @@ export async function extractTimestamp(ts) {
   tsReadBuffer.destroy();
   resolveBuffer.destroy(); // never mapped — no unmap() needed
   querySet.destroy();
-  return Number(timestamps[1] - timestamps[0]) / 1e6;
+  return Math.max(0, Number(timestamps[1] - timestamps[0])) / 1e6;
 }

package/src/util/bindgroup.mjs

@@ -2,21 +2,25 @@
 import { getDevice } from "../init.mjs";
 
 /**
- * Creates a `GPUBindGroup` by mapping each buffer to sequential binding indices (0, 1, 2 …).
- * `resultBuffer` is appended last so its binding index follows all input buffers.
- * The order of `buffers` must match the `@binding` indices declared in the shader.
+ * Creates a `GPUBindGroup` by mapping each buffer to sequential binding indices
+ * (startBinding, startBinding + 1, …). The order of `buffers` must match the `@binding`s
+ * declared in the shader.
  * @param {GPUBindGroupLayout} layout - from `pipeline.getBindGroupLayout(0)`; the layout is derived automatically from the shader because `loadShader` uses `layout: "auto"`
- * @param {GPUBuffer[]} buffers - input and intermediate buffers in binding order
- * @param {GPUBuffer|null} [resultBuffer=null] - appended as the final binding if provided
+ * @param {(GPUBuffer|{buffer: GPUBuffer, offset: number, size: number})[]} buffers - input and
+ *   intermediate buffers in binding order; a plain `GPUBuffer` binds the whole buffer, or pass
+ *   `{buffer, offset, size}` to bind a sub-range — e.g. many small per-call param structs packed
+ *   into one shared buffer instead of allocating a separate tiny buffer per call
+ * @param {number} [startBinding=0] - binding index of `buffers[0]`; nonzero only when a shader's
+ *   own bindings don't start at 0 — e.g. a module built by concatenating two `.wgsl` files where
+ *   the second file's bindings continue after the first's own bindings (if any)
  * @returns {GPUBindGroup}
  * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUDevice/createBindGroup GPUDevice.createBindGroup()}
  */
-export function createBindGroup(layout, buffers, resultBuffer = null) {
+export function createBindGroup(layout, buffers, startBinding = 0) {
   const device = getDevice();
-  const allBuffers = resultBuffer ? [...buffers, resultBuffer] : [...buffers];
-  const entries = allBuffers.map((buffer, i) => ({
-    binding: i,
-    resource: { buffer },
+  const entries = buffers.map((buffer, i) => ({
+    binding: startBinding + i,
+    resource: buffer instanceof GPUBuffer ? { buffer } : buffer,
   }));
   return device.createBindGroup({ layout, entries });
 }

package/src/util/buffer.mjs

@@ -12,7 +12,11 @@ export function destroyBuffers(...buffers) {
 
 /**
  * Creates a GPU storage buffer and uploads `data` into it via mapped-at-creation.
- * @param {Float32Array} data
+ * The mapped view is constructed from `data`'s own typed-array constructor, so
+ * bits are copied as-is regardless of element type (e.g. a Uint32Array's raw
+ * bit patterns are preserved — critical for dasum's aux half, which must
+ * never pass through a Float32Array view and risk NaN-bit-pattern canonicalization).
+ * @param {Float32Array|Uint32Array|Int32Array} data
  * @param {string} [label] - debug label visible in browser DevTools GPU inspection
  * @param {boolean} [readback=false] - add `COPY_SRC` so the buffer can be copied to a readback buffer
  * @throws {Error} if `data.byteLength` exceeds the device's `maxStorageBufferBindingSize`
@@ -45,7 +49,8 @@ export function uploadBuffer(data, label = "blas-input", readback = false) {
     mappedAtCreation: true,
   });
 
-  const mappedArray = new Float32Array(buffer.getMappedRange());
+  const ViewCtor = data.constructor;
+  const mappedArray = new ViewCtor(buffer.getMappedRange());
   mappedArray.set(data);
   buffer.unmap();
 

package/src/util/compute.mjs

@@ -2,6 +2,9 @@
 import { getDevice } from "../init.mjs";
 import { beginTimestamp, resolveTimestamp } from "./benchmark.mjs";
 
+// Anchors the pass encoder to its command encoder to prevent premature GC.
+const _passEncoders = new WeakMap();
+
 /**
  * Finalises `commandEncoder` into a command buffer and submits it to the GPU queue.
  * @param {GPUCommandEncoder} commandEncoder
@@ -13,24 +16,33 @@ export function submit(commandEncoder) {
 }
 
 /**
- * Encodes and submits a single compute pass: sets the pipeline and bind group,
- * dispatches workgroups, and optionally wraps the pass in GPU timestamp queries.
- * @param {GPUComputePipeline} pipeline
- * @param {GPUBindGroup} bindGroup
- * @param {number | { x: number, y: number }} workgroups - workgroup count; number for 1D dispatch, `{x, y}` for 2D
- * @returns {{ commandEncoder: GPUCommandEncoder, ts: any }} encoded commands and timestamp handle
+ * Starts a new command encoder together with its GPU timestamp query (if
+ * benchmarking is enabled) — the pairing every routine needs before encoding
+ * its first pass, whether that's `runComputePass`'s single pass or a
+ * multi-pass routine (e.g. strsv's blocked solve) encoding several by hand.
+ * @returns {{ commandEncoder: GPUCommandEncoder, querySet: GPUQuerySet|null, passDescriptor: object|undefined }}
  * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUDevice/createCommandEncoder GPUDevice.createCommandEncoder()}
- * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUCommandEncoder/beginComputePass GPUCommandEncoder.beginComputePass()}
- * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUComputePassEncoder/setPipeline GPUComputePassEncoder.setPipeline()}
- * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUComputePassEncoder/setBindGroup GPUComputePassEncoder.setBindGroup()}
- * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUComputePassEncoder/dispatchWorkgroups GPUComputePassEncoder.dispatchWorkgroups()}
  */
-export function runComputePass(pipeline, bindGroup, workgroups) {
+export function beginTimedEncoder() {
   const device = getDevice();
-
   const { querySet, passDescriptor } = beginTimestamp();
-
   const commandEncoder = device.createCommandEncoder();
+  return { commandEncoder, querySet, passDescriptor };
+}
+
+/**
+ * Encodes one compute pass (set pipeline, set bind group, dispatch, end) onto an
+ * existing command encoder — the shared building block behind `runComputePass`,
+ * also used directly by routines that need several passes on one encoder (e.g.
+ * strsv's blocked solve, where each pass depends on the previous one completing).
+ * @param {GPUCommandEncoder} commandEncoder
+ * @param {GPUComputePipeline} pipeline
+ * @param {GPUBindGroup} bindGroup
+ * @param {number | { x: number, y: number }} workgroups - workgroup count; number for 1D dispatch, `{x, y}` for 2D
+ * @param {object} [passDescriptor] - passed to `beginComputePass`, e.g. for timestamp writes
+ * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUCommandEncoder/beginComputePass GPUCommandEncoder.beginComputePass()}
+ */
+export function encodePass(commandEncoder, pipeline, bindGroup, workgroups, passDescriptor) {
   const passEncoder = commandEncoder.beginComputePass(passDescriptor);
 
   passEncoder.setPipeline(pipeline);
@@ -44,9 +56,23 @@ export function runComputePass(pipeline, bindGroup, workgroups) {
 
   passEncoder.end();
 
-  const ts = resolveTimestamp(commandEncoder, querySet);
+  _passEncoders.set(commandEncoder, passEncoder);
+}
 
-  commandEncoder._passEncoder = passEncoder; // anchor — GC'd passEncoder may crash native encoder
+/**
+ * Encodes and submits a single compute pass: sets the pipeline and bind group,
+ * dispatches workgroups, and optionally wraps the pass in GPU timestamp queries.
+ * @param {GPUComputePipeline} pipeline
+ * @param {GPUBindGroup} bindGroup
+ * @param {number | { x: number, y: number }} workgroups - workgroup count; number for 1D dispatch, `{x, y}` for 2D
+ * @returns {{ commandEncoder: GPUCommandEncoder, ts: any }} encoded commands and timestamp handle
+ * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUDevice/createCommandEncoder GPUDevice.createCommandEncoder()}
+ */
+export function runComputePass(pipeline, bindGroup, workgroups) {
+  const { commandEncoder, querySet, passDescriptor } = beginTimedEncoder();
+  encodePass(commandEncoder, pipeline, bindGroup, workgroups, passDescriptor);
+
+  const ts = resolveTimestamp(commandEncoder, querySet);
 
   return { commandEncoder, ts };
 }