wgblas 2.0.0 → 2.2.0

package/src/util/compute.mjs

@@ -1,17 +1,12 @@
 /** @module devdocs/utility-functions/compute */
-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
  * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUQueue/submit GPUQueue.submit()}
  */
-export function submit(commandEncoder) {
-  const device = getDevice();
+export function submit(device, commandEncoder) {
   device.queue.submit([commandEncoder.finish()]);
 }
 
@@ -23,9 +18,8 @@ export function submit(commandEncoder) {
  * @returns {{ commandEncoder: GPUCommandEncoder, querySet: GPUQuerySet|null, passDescriptor: object|undefined }}
  * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUDevice/createCommandEncoder GPUDevice.createCommandEncoder()}
  */
-export function beginTimedEncoder() {
-  const device = getDevice();
-  const { querySet, passDescriptor } = beginTimestamp();
+export function beginTimedEncoder(device) {
+  const { querySet, passDescriptor } = beginTimestamp(device);
   const commandEncoder = device.createCommandEncoder();
   return { commandEncoder, querySet, passDescriptor };
 }
@@ -43,7 +37,13 @@ export function beginTimedEncoder() {
  * @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) {
+export function encodePass(
+  commandEncoder,
+  pipeline,
+  bindGroup,
+  workgroups,
+  passDescriptor,
+) {
   const passEncoder = commandEncoder.beginComputePass(passDescriptor);
 
   passEncoder.setPipeline(pipeline);
@@ -53,12 +53,14 @@ export function encodePass(commandEncoder, pipeline, bindGroup, workgroups, pass
     passEncoder.dispatchWorkgroups(workgroups);
   } else {
     // `?? 1` is load-bearing — confirmed passing undefined (from an {x,y}-only caller) crashes the process, not just no-ops.
-    passEncoder.dispatchWorkgroups(workgroups.x, workgroups.y, workgroups.z ?? 1);
+    passEncoder.dispatchWorkgroups(
+      workgroups.x,
+      workgroups.y,
+      workgroups.z ?? 1,
+    );
   }
 
   passEncoder.end();
-
-  _passEncoders.set(commandEncoder, passEncoder);
 }
 
 /**
@@ -71,11 +73,12 @@ export function encodePass(commandEncoder, pipeline, bindGroup, workgroups, pass
  * @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();
+export function runComputePass(device, pipeline, bindGroup, workgroups) {
+  const { commandEncoder, querySet, passDescriptor } =
+    beginTimedEncoder(device);
   encodePass(commandEncoder, pipeline, bindGroup, workgroups, passDescriptor);
 
-  const ts = resolveTimestamp(commandEncoder, querySet);
+  const ts = resolveTimestamp(device, commandEncoder, querySet);
 
   return { commandEncoder, ts };
 }

package/src/util/constants.mjs

@@ -0,0 +1,57 @@
+/** @module devdocs/utility-functions/constants */
+
+/**
+ * Single source of truth for every constant the JS dispatch side shares with a
+ * shader.
+ *
+ * WGSL has no import mechanism, so each shader necessarily declares its own
+ * copy of these values. That makes them a silent-corruption hazard: change
+ * `BM` in `sgemm_large.wgsl` without changing `BM_LARGE` here and the host
+ * dispatches the wrong grid — too few workgroups computes part of the matrix
+ * and reports success. Hoisting them here removes the JS-to-JS duplication;
+ * `tests/utils/test.constants.js` closes the remaining JS-to-WGSL gap
+ * by parsing the shader sources and asserting they still agree.
+ *
+ * Every export below names the shader declaration it mirrors. Changing one
+ * means changing both, and the test will tell you if you forget.
+ */
+
+// --- gemm block tiles ------------------------------------------------------
+
+/** `BM` in sgemm_small.wgsl / sgemmtr_small.wgsl. */
+export const BM_SMALL = 32;
+/** `BN` in sgemm_small.wgsl / sgemmtr_small.wgsl. */
+export const BN_SMALL = 32;
+/** `BM` in sgemm_large.wgsl / sgemmtr_large.wgsl. */
+export const BM_LARGE = 64;
+/** `BN` in sgemm_large.wgsl / sgemmtr_large.wgsl. */
+export const BN_LARGE = 64;
+
+/**
+ * The large tile only pays for its bigger workgroups once the problem needs at
+ * least a 6x6 grid of them; below that the small tile wins. JS-only (no shader
+ * counterpart) — it selects *which* shader runs.
+ */
+export const LARGE_TILE_WORKGROUP_THRESHOLD = 36;
+
+// --- 1D / reduction kernels ------------------------------------------------
+
+/**
+ * `const WGS: u32 = 64` — declared by every 1D and reduction shader, and the
+ * workgroup size `calcWorkgroups` divides by for a 1D dispatch.
+ */
+export const WGS = 64;
+
+// --- 2D helper kernels -----------------------------------------------------
+
+/**
+ * `@workgroup_size(8, 8)` in symmetrize.wgsl, triangularize.wgsl and
+ * block_transfer.wgsl, and the size `calcWorkgroups` divides by per dimension
+ * for a 2D dispatch.
+ */
+export const TILE_WG_2D = 8;
+
+// --- triangular solve ------------------------------------------------------
+
+/** `BLOCK_SIZE` in strsv_invert_block.wgsl — the diagonal block order. */
+export const BLOCK_SIZE = 64;

package/src/util/device.mjs

@@ -0,0 +1,49 @@
+/** @module devdocs/utility-functions/device */
+import { GpuVector } from "../classes/GpuVector.mjs";
+import { GpuMatrix } from "../classes/GpuMatrix.mjs";
+
+/**
+ * Throws if `device` is not a `GPUDevice`.
+ *
+ * Every routine's first guard, extracted since the check and message are
+ * identical across all of them.
+ *
+ * @param {GPUDevice} device - the value to check
+ * @throws {Error} if `device` is not a `GPUDevice`
+ */
+export function requireGpuDevice(device) {
+  if (!(device instanceof GPUDevice))
+    throw new Error("device must be a GPUDevice.");
+}
+
+/**
+ * Throws if any GPU-resident operand belongs to a device other than the one
+ * the routine was called with.
+ *
+ * A `GPUBuffer` is bound to the device that created it, and WebGPU has no way
+ * to share one across devices. Handing a routine a `GpuMatrix` from device A
+ * while passing device B fails deep inside bind-group creation as a
+ * `GPUValidationError` with no indication that two devices are involved —
+ * this turns it into a named, actionable error at the call boundary.
+ *
+ * Scalars, plain typed arrays and `undefined` entries are ignored, so callers
+ * can pass their whole operand set without filtering.
+ *
+ * @param {GPUDevice} device - the device the routine will dispatch on
+ * @param {string} routine - routine name, for the error message
+ * @param {Record<string, unknown>} operands - operand name -> value
+ * @throws {Error} if an operand is GPU-resident on a different device
+ */
+export function requireSameDevice(device, routine, operands) {
+  for (const [name, value] of Object.entries(operands)) {
+    if (!(value instanceof GpuVector) && !(value instanceof GpuMatrix))
+      continue;
+    if (value.device !== device) {
+      throw new Error(
+        `${routine}: ${name} belongs to a different GPUDevice than the one passed in. ` +
+          "GPU buffers cannot be shared across devices — recreate the operand on this " +
+          "device, or call the routine with the device that owns it.",
+      );
+    }
+  }
+}

package/src/util/pipeline.mjs

@@ -1,5 +1,4 @@
 /** @module devdocs/utility-functions/pipeline */
-import { getDevice } from "../init.mjs";
 
 // WeakMap keyed by GPUDevice so pipelines are released automatically when the device is destroyed.
 const _pipelines = new WeakMap();
@@ -24,23 +23,31 @@ export async function getPipeline(device, shaderName, entryPoint = "main") {
   const names = Array.isArray(shaderName) ? shaderName : [shaderName];
   const key = `${names.join("+")}::${entryPoint}`;
   if (!byName.has(key)) {
-    byName.set(key, await loadShader(names, entryPoint));
+    // Cache the in-flight promise (not its resolved value) so a concurrent
+    // call awaits the same compile instead of starting a duplicate one;
+    // drop the entry on failure so a later call can retry.
+    const pending = loadShader(device, names, entryPoint).catch((err) => {
+      byName.delete(key);
+      throw err;
+    });
+    byName.set(key, pending);
   }
   return byName.get(key);
 }
 
 /**
  * Loads WGSL source for `shaderName`. In the browser, reads from the inline bundle
- * (`browser-shaders.mjs`); in Node.js, reads the `.wgsl` file directly from disk.
+ * (`shaders/index.mjs`'s `shaderSources`); in Node.js, reads the `.wgsl` file directly from disk.
  * @param {string} shaderName
  * @returns {Promise<string>}
  */
 async function loadCode(shaderName) {
   // Check for Node.js explicitly — `window` is undefined in Web Workers too, so it's not a reliable signal.
   if (typeof process === "undefined" || !process.versions?.node) {
-    const { shaderSources } = await import("../shaders/browser-shaders.mjs");
+    const { shaderSources } = await import("../shaders/index.mjs");
     const src = shaderSources[shaderName];
-    if (!src) throw new Error(`Shader "${shaderName}" not found in browser bundle.`);
+    if (!src)
+      throw new Error(`Shader "${shaderName}" not found in browser bundle.`);
     return src;
   } else {
     const { readFileSync } = await import("fs");
@@ -56,6 +63,7 @@ async function loadCode(shaderName) {
  * `GPUComputePipeline`. Throws with line-level detail if compilation fails, rather than
  * surfacing a raw GPU error. Uses `layout: "auto"` so the pipeline derives its bind group
  * layout from the shader — no manual layout definition needed.
+ * @param {GPUDevice} device
  * @param {string[]} shaderNames - filenames without `.wgsl`, concatenated in array order
  * @param {string} [entryPoint="main"] - which `@compute` function in the combined module to run
  * @returns {Promise<GPUComputePipeline>}
@@ -64,11 +72,34 @@ async function loadCode(shaderName) {
  * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUShaderModule/getCompilationInfo GPUShaderModule.getCompilationInfo()}
  * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUDevice/createComputePipeline GPUDevice.createComputePipeline()}
  */
-export async function loadShader(shaderNames, entryPoint = "main") {
-  const device = getDevice();
+export async function loadShader(device, shaderNames, entryPoint = "main") {
   const label = shaderNames.join("+");
-  const code = (await Promise.all(shaderNames.map(loadCode))).join("\n");
+  const codes = await Promise.all(shaderNames.map(loadCode));
+
+  // Per-file line ranges within the concatenated module, so a compile error
+  // reports as "<file>.wgsl:<local line>" instead of a whole-module line
+  // number — confusing for multi-file pipelines like dasum's.
+  let offset = 0;
+  const ranges = codes.map((c, i) => {
+    const lineCount = c.split("\n").length;
+    const range = {
+      name: shaderNames[i],
+      startLine: offset + 1,
+      endLine: offset + lineCount,
+    };
+    offset += lineCount;
+    return range;
+  });
+  const locate = (lineNum) => {
+    const range =
+      lineNum &&
+      ranges.find((r) => lineNum >= r.startLine && lineNum <= r.endLine);
+    return range
+      ? `${range.name}.wgsl:${lineNum - range.startLine + 1}`
+      : `line ${lineNum}`;
+  };
 
+  const code = codes.join("\n");
   const shaderModule = device.createShaderModule({ label, code });
 
   const info = await shaderModule.getCompilationInfo();
@@ -76,7 +107,7 @@ export async function loadShader(shaderNames, entryPoint = "main") {
   const errors = info.messages.filter((m) => m.type === "error");
   if (errors.length > 0) {
     throw new Error(
-      `Shader "${label}" compilation failed:\n${errors.map((m) => `  line ${m.lineNum}: ${m.message}`).join("\n")}`,
+      `Shader "${label}" compilation failed:\n${errors.map((m) => `  ${locate(m.lineNum)}: ${m.message}`).join("\n")}`,
     );
   }
 
@@ -84,7 +115,10 @@ export async function loadShader(shaderNames, entryPoint = "main") {
   // this project's WebGPU backend is unstable (intermittent multi-minute hangs and wrong
   // results, confirmed by bisection) when entryPoint is set explicitly, even to the shader's
   // only/correct entry point. Auto-detecting the single entry point is the stable path.
-  const compute = entryPoint === "main" ? { module: shaderModule } : { module: shaderModule, entryPoint };
+  const compute =
+    entryPoint === "main"
+      ? { module: shaderModule }
+      : { module: shaderModule, entryPoint };
   const pipeline = device.createComputePipeline({
     label,
     layout: "auto",

package/src/util/workgroup.mjs

@@ -1,14 +1,23 @@
 /** @module devdocs/utility-functions/workgroup */
-import { getDevice } from "../init.mjs";
-
-// Fixed sizes match the shader declarations (WGS = 64 for 1D, 8×8 = 64 threads for 2D).
-const WORKGROUP_SIZE_1D = 64;
-const WORKGROUP_SIZE_2D = 8;
+// Fixed sizes match the shader declarations (WGS = 64 for 1D, 8×8 = 64 threads
+// for 2D) — see constants.mjs, which is where both values are defined and
+// where the WGSL cross-check hangs off.
+import {
+  WGS as WORKGROUP_SIZE_1D,
+  TILE_WG_2D as WORKGROUP_SIZE_2D,
+} from "./constants.mjs";
 
 /**
  * Calculates the number of workgroups to dispatch, clamped to the device's
  * `maxComputeWorkgroupsPerDimension` limit (default 65535 across most devices).
  *
+ * ONLY for shaders whose kernel is a grid-stride loop driven by
+ * `num_workgroups` — those re-walk the whole domain regardless of how many
+ * workgroups actually launch, so clamping costs a little parallelism and
+ * nothing else. A shader that indexes straight off `workgroup_id` or
+ * `global_invocation_id` silently drops every row past the clamp; those must
+ * use {@link requireWorkgroups} / {@link requireWorkgroupCount} instead.
+ *
  * - 1D (pass only `rows`): returns a single count for `dispatchWorkgroups(n)`.
  * - 2D (pass both `rows` and `cols`): returns `{ x, y }` for `dispatchWorkgroups(x, y)`.
  *   `rows` maps to the y dimension and `cols` maps to the x dimension.
@@ -18,8 +27,8 @@ const WORKGROUP_SIZE_2D = 8;
  * @returns {number | { x: number, y: number }}
  * @see {@link https://developer.mozilla.org/en-US/docs/Web/API/GPUSupportedLimits GPUSupportedLimits} (`maxComputeWorkgroupsPerDimension`)
  */
-export function calcWorkgroups(rows, cols) {
-  const max = getDevice().limits.maxComputeWorkgroupsPerDimension;
+export function calcWorkgroups(device, rows, cols) {
+  const max = device.limits.maxComputeWorkgroupsPerDimension;
   if (cols === undefined) {
     return Math.min(Math.ceil(rows / WORKGROUP_SIZE_1D), max);
   } else {
@@ -29,3 +38,59 @@ export function calcWorkgroups(rows, cols) {
     };
   }
 }
+
+/**
+ * Returns `count` unchanged if it fits the device's dispatch limit, and throws
+ * otherwise. The counterpart to {@link calcWorkgroups} for shaders with no
+ * grid-stride fallback: silently clamping those computes only part of the
+ * result and reports success, so a refusal is the safer failure.
+ *
+ * @param {number} count - workgroups this dispatch requires in one dimension
+ * @param {string} routine - routine name, for the error message
+ * @param {string} [dim] - dimension label ("x"/"y"), for the error message
+ * @returns {number} `count`
+ * @throws {Error} when `count` exceeds `maxComputeWorkgroupsPerDimension`
+ */
+export function requireWorkgroupCount(device, count, routine, dim = "x") {
+  const max = device.limits.maxComputeWorkgroupsPerDimension;
+  if (count > max)
+    throw new Error(
+      `${routine}: this problem needs ${count} workgroups in ${dim}, but the device allows ` +
+        `${max} (maxComputeWorkgroupsPerDimension). The operands are too large for this device — ` +
+        `split the operation into smaller blocks.`,
+    );
+  return count;
+}
+
+/**
+ * {@link calcWorkgroups} with the clamp replaced by a throw — same arguments
+ * and same return shape, for shaders without a grid-stride fallback.
+ *
+ * @param {string} routine - routine name, for the error message
+ * @param {number} rows - row count (1D: element count)
+ * @param {number} [cols] - column count; omit for a 1D dispatch
+ * @returns {number | { x: number, y: number }}
+ * @throws {Error} when either dimension exceeds `maxComputeWorkgroupsPerDimension`
+ */
+export function requireWorkgroups(device, routine, rows, cols) {
+  if (cols === undefined)
+    return requireWorkgroupCount(
+      device,
+      Math.ceil(rows / WORKGROUP_SIZE_1D),
+      routine,
+    );
+  return {
+    x: requireWorkgroupCount(
+      device,
+      Math.ceil(cols / WORKGROUP_SIZE_2D),
+      routine,
+      "x",
+    ),
+    y: requireWorkgroupCount(
+      device,
+      Math.ceil(rows / WORKGROUP_SIZE_2D),
+      routine,
+      "y",
+    ),
+  };
+}

package/src/shaders/browser-shaders.mjs

@@ -1,81 +0,0 @@
-import argmax from "./reduction/argmax.wgsl";
-import argmaxF64 from "./reduction/argmaxF64.wgsl";
-import sum from "./reduction/sum.wgsl";
-import sumF64 from "./reduction/sumF64.wgsl";
-import sscal from "./sscal.wgsl";
-import sswap from "./sswap.wgsl";
-import saxpy from "./saxpy.wgsl";
-import scopy from "./scopy.wgsl";
-import sdot from "./sdot.wgsl";
-import sasum from "./sasum.wgsl";
-import snrm2 from "./snrm2.wgsl";
-import srot from "./srot.wgsl";
-import srotm from "./srotm.wgsl";
-import isamax from "./isamax.wgsl";
-import sgemv_n from "./sgemv_n.wgsl";
-import sgemv_t from "./sgemv_t.wgsl";
-import ssymv from "./ssymv.wgsl";
-import strmv from "./strmv.wgsl";
-import sger from "./sger.wgsl";
-import ssyr from "./ssyr.wgsl";
-import ssyr2 from "./ssyr2.wgsl";
-import f64add from "./f64add.wgsl";
-import dekker from "./f64/dekker.wgsl";
-import ddAbs from "./f64/utils/abs.wgsl";
-import ddAddUtil from "./f64/utils/add.wgsl";
-import ddGreater from "./f64/utils/greater.wgsl";
-import ddEqual from "./f64/utils/equal.wgsl";
-import dasum from "./dasum.wgsl";
-import idamax from "./idamax.wgsl";
-import strsv_invert_block from "./strsv_invert_block.wgsl";
-import strsv_apply_inverse from "./strsv_apply_inverse.wgsl";
-import strsv_update from "./strsv_update.wgsl";
-import sgemm_small from "./sgemm_small.wgsl";
-import sgemm_large from "./sgemm_large.wgsl";
-import sgemmtr_small from "./sgemmtr_small.wgsl";
-import sgemmtr_large from "./sgemmtr_large.wgsl";
-import symmetrize from "./symmetrize.wgsl";
-import triangularize from "./triangularize.wgsl";
-import blockTransfer from "./block_transfer.wgsl";
-
-export const shaderSources = {
-  "reduction/argmax": argmax,
-  "reduction/argmaxF64": argmaxF64,
-  "reduction/sum": sum,
-  "reduction/sumF64": sumF64,
-  sscal,
-  sswap,
-  saxpy,
-  scopy,
-  sdot,
-  sasum,
-  snrm2,
-  srot,
-  srotm,
-  isamax,
-  sgemv_n,
-  sgemv_t,
-  ssymv,
-  strmv,
-  sger,
-  ssyr,
-  ssyr2,
-  f64add,
-  "f64/dekker": dekker,
-  "f64/utils/abs": ddAbs,
-  "f64/utils/add": ddAddUtil,
-  "f64/utils/greater": ddGreater,
-  "f64/utils/equal": ddEqual,
-  dasum,
-  idamax,
-  strsv_invert_block,
-  strsv_apply_inverse,
-  strsv_update,
-  sgemm_small,
-  sgemm_large,
-  sgemmtr_small,
-  sgemmtr_large,
-  symmetrize,
-  triangularize,
-  block_transfer: blockTransfer,
-};