wgblas 2.0.0 → 2.1.0

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,34 @@
+/** @module devdocs/utility-functions/device */
+import { GpuVector } from "../classes/GpuVector.mjs";
+import { GpuMatrix } from "../classes/GpuMatrix.mjs";
+
+/**
+ * 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,21 +23,21 @@ 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));
+    byName.set(key, await loadShader(device, names, entryPoint));
   }
   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.`);
     return src;
@@ -56,6 +55,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,8 +64,7 @@ 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");
 

package/src/util/workgroup.mjs

@@ -1,14 +1,20 @@
 /** @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 +24,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 +35,45 @@ 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,
-};