wgblas 2.0.0 → 2.2.0

package/src/sdot/sdot.mjs

@@ -12,15 +12,15 @@ import { extractTimestamp } from "../util/benchmark.mjs";
 import { extractResult } from "../util/result.mjs";
 import { getPipeline } from "../util/pipeline.mjs";
 import { GpuVector } from "../classes/GpuVector.mjs";
-
-const WGS = 64; //workgroup size
+import { WGS } from "../util/constants.mjs";
+import { requireGpuDevice, requireSameDevice } from "../util/device.mjs";
 
 export async function sdot(device, n, x, incx, y, incy) {
   const xIsGpu = x instanceof GpuVector;
   const yIsGpu = y instanceof GpuVector;
 
-  if (!(device instanceof GPUDevice))
-    throw new Error("device must be a GPUDevice.");
+  requireGpuDevice(device);
+  requireSameDevice(device, "sdot", { x, y });
   if (
     !Number.isInteger(n) ||
     !Number.isInteger(incx) ||
@@ -58,11 +58,12 @@ export async function sdot(device, n, x, incx, y, incy) {
   let readBuffer = null;
 
   try {
-    xBuffer = xIsGpu ? x._buf : uploadBuffer(x, "sdot-x", false);
-    yBuffer = yIsGpu ? y._buf : uploadBuffer(y, "sdot-y", false);
-    partialsBuffer = createStorageBuffer(2 * WGS * 4, "sdot-partials"); //to hold 2*WGS partial sums of f32
-    resultBuffer = createResultBuffer(4, "sdot-result"); //to hold the final float32 dot product
+    xBuffer = xIsGpu ? x._buf : uploadBuffer(device, x, "sdot-x", false);
+    yBuffer = yIsGpu ? y._buf : uploadBuffer(device, y, "sdot-y", false);
+    partialsBuffer = createStorageBuffer(device, 2 * WGS * 4, "sdot-partials"); //to hold 2*WGS partial sums of f32
+    resultBuffer = createResultBuffer(device, 4, "sdot-result"); //to hold the final float32 dot product
     paramsBuffer = createParamsBuffer(
+      device,
       [
         { value: n, type: "u32" },
         { value: incx, type: "u32" },
@@ -71,32 +72,35 @@ export async function sdot(device, n, x, incx, y, incy) {
       "sdot-params",
     );
 
-    const bgMain = createBindGroup(pipelineMain.getBindGroupLayout(0), [
+    const bgMain = createBindGroup(device, pipelineMain.getBindGroupLayout(0), [
       xBuffer,
       yBuffer,
       partialsBuffer,
       paramsBuffer,
     ]);
     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,
-      resultBuffer,
-    ]);
+    const bgReduce = createBindGroup(
+      device,
+      pipelineReduce.getBindGroupLayout(0),
+      [partialsBuffer, resultBuffer],
+    );
     const { commandEncoder: enc2, ts: ts2 } = runComputePass(
+      device,
       pipelineReduce,
       bgReduce,
       1,
     ); // dispatch 1 workgroup to reduce the partial sums to a 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
@@ -117,7 +121,7 @@ export async function sdot(device, n, x, incx, y, incy) {
     if (partialsBuffer) destroyBuffers(partialsBuffer);
     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/sgemm/sgemm.d.mts

@@ -1,7 +1,7 @@
 import { GpuMatrix } from "../classes/GpuMatrix.mjs";
 
 /**
- * Performs the matrix-matrix operation C = alpha * op(A) * op(B) + beta * C
+ * Performs the matrix-matrix operation $$C \leftarrow \alpha \mathrm{op}(A) \mathrm{op}(B) + \beta C$$
  *
  * - `transA/transB='no-transpose'`: op(A) = A (m×k), op(B) = B (k×n)
  * - `transA/transB='transpose'`:    op(A) = A^T,     op(B) = B^T
@@ -57,7 +57,7 @@ export declare function sgemm(
 ): Promise<{ C: Float32Array; gpuTimeMs?: number }>;
 
 /**
- * Performs the matrix-matrix operation C = alpha * op(A) * op(B) + beta * C
+ * Performs the matrix-matrix operation $$C \leftarrow \alpha \mathrm{op}(A) \mathrm{op}(B) + \beta C$$
  *
  * A, B, and C are all kept GPU-resident. Each matrix's own `layout` (set at
  * `GpuMatrix.from` time) determines the operation — there is no separate

package/src/sgemm/sgemm.mjs

@@ -3,6 +3,8 @@ import {
   createParamsBuffer,
   stageReadback,
   destroyBuffers,
+  vec4ViewBinding,
+  vec4Usable,
 } from "../util/buffer.mjs";
 import { createBindGroup } from "../util/bindgroup.mjs";
 import { runComputePass, submit } from "../util/compute.mjs";
@@ -10,32 +12,49 @@ import { extractResult } from "../util/result.mjs";
 import { extractTimestamp } from "../util/benchmark.mjs";
 import { getPipeline } from "../util/pipeline.mjs";
 import { GpuMatrix } from "../classes/GpuMatrix.mjs";
-
-const BM_SMALL = 32, BN_SMALL = 32; // sgemm_small.wgsl's block tile
-const BM_LARGE = 64, BN_LARGE = 64; // sgemm_large.wgsl's block tile
-const LARGE_TILE_WORKGROUP_THRESHOLD = 36; // large tile needs >= a 6x6 grid of its own tiles to beat the small tile
+import { requireWorkgroupCount } from "../util/workgroup.mjs";
+import {
+  BM_SMALL,
+  BN_SMALL,
+  BM_LARGE,
+  BN_LARGE,
+  LARGE_TILE_WORKGROUP_THRESHOLD,
+} from "../util/constants.mjs";
+import { requireGpuDevice, requireSameDevice } from "../util/device.mjs";
 
 export async function sgemm(
-  device, transA, transB, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc, layout = "row-major",
+  device,
+  transA,
+  transB,
+  m,
+  n,
+  k,
+  alpha,
+  A,
+  lda,
+  B,
+  ldb,
+  beta,
+  C,
+  ldc,
+  layout = "row-major",
 ) {
   let AIsGpu = A instanceof GpuMatrix;
   let BIsGpu = B instanceof GpuMatrix;
   const CIsGpu = C instanceof GpuMatrix;
 
-  if (!(device instanceof GPUDevice))
-    throw new Error("device must be a GPUDevice.");
+  requireGpuDevice(device);
+  requireSameDevice(device, "sgemm", { A, B, C });
   if (transA !== "no-transpose" && transA !== "transpose")
     throw new Error("transA must be 'no-transpose' or 'transpose'.");
   if (transB !== "no-transpose" && transB !== "transpose")
     throw new Error("transB must be 'no-transpose' or 'transpose'.");
   if (layout !== "row-major" && layout !== "column-major")
     throw new Error("layout must be 'row-major' or 'column-major'.");
-  if (typeof alpha !== "number")
-    throw new Error("alpha must be a number.");
+  if (typeof alpha !== "number") throw new Error("alpha must be a number.");
   if (Number.isNaN(alpha)) throw new Error("alpha must not be NaN.");
   if (!Number.isFinite(alpha)) throw new Error("alpha must be finite.");
-  if (typeof beta !== "number")
-    throw new Error("beta must be a number.");
+  if (typeof beta !== "number") throw new Error("beta must be a number.");
   if (Number.isNaN(beta)) throw new Error("beta must not be NaN.");
   if (!Number.isFinite(beta)) throw new Error("beta must be finite.");
   if (
@@ -57,7 +76,10 @@ export async function sgemm(
     throw new Error("C must be a GpuMatrix when A or B is a GpuMatrix.");
   if (CIsGpu && (!AIsGpu || !BIsGpu))
     throw new Error("A and B must be GpuMatrix when C is a GpuMatrix.");
-  if (m < 0 || n < 0 || k < 0) throw new Error("m, n, and k must be non-negative.");
+  if (m < 0 || n < 0 || k < 0)
+    throw new Error("m, n, and k must be non-negative.");
+  if (lda <= 0 || ldb <= 0 || ldc <= 0)
+    throw new Error("lda, ldb, and ldc must be positive.");
   if (m === 0 || n === 0) return CIsGpu ? {} : { C };
 
   // GpuMatrix's own .layout wins over the shared `layout` argument.
@@ -72,14 +94,19 @@ export async function sgemm(
   const aOuter = transA === "no-transpose" ? aRows : aCols; // # of stored chunks
   const aInner = transA === "no-transpose" ? aCols : aRows; // required chunk length
   if (lda < aInner)
-    throw new Error(`lda must be >= ${effLayoutA === "column-major" ? "rows" : "cols"} of A as stored.`);
+    throw new Error(
+      `lda must be >= ${effLayoutA === "column-major" ? "rows" : "cols"} of A as stored.`,
+    );
   if (AIsGpu) {
-    if (lda !== A.lda) throw new Error("lda must match A.lda when A is a GpuMatrix.");
+    if (lda !== A.lda)
+      throw new Error("lda must match A.lda when A is a GpuMatrix.");
     const [aLogRows, aLogCols] = transA === "no-transpose" ? [m, k] : [k, m];
     if (A.rows < aLogRows || A.cols < aLogCols)
       throw new Error("A is too small for the given m, k, and transA.");
   } else if (A.length < (aOuter - 1) * lda + aInner) {
-    throw new Error("A does not have enough elements for the given dimensions and lda.");
+    throw new Error(
+      "A does not have enough elements for the given dimensions and lda.",
+    );
   }
 
   // B: same reasoning as A, with op(B) = k x n.
@@ -88,26 +115,37 @@ export async function sgemm(
   const bOuter = transB === "no-transpose" ? bRows : bCols;
   const bInner = transB === "no-transpose" ? bCols : bRows;
   if (ldb < bInner)
-    throw new Error(`ldb must be >= ${effLayoutB === "column-major" ? "rows" : "cols"} of B as stored.`);
+    throw new Error(
+      `ldb must be >= ${effLayoutB === "column-major" ? "rows" : "cols"} of B as stored.`,
+    );
   if (BIsGpu) {
-    if (ldb !== B.lda) throw new Error("ldb must match B.lda when B is a GpuMatrix.");
+    if (ldb !== B.lda)
+      throw new Error("ldb must match B.lda when B is a GpuMatrix.");
     const [bLogRows, bLogCols] = transB === "no-transpose" ? [k, n] : [n, k];
     if (B.rows < bLogRows || B.cols < bLogCols)
       throw new Error("B is too small for the given n, k, and transB.");
   } else if (B.length < (bOuter - 1) * ldb + bInner) {
-    throw new Error("B does not have enough elements for the given dimensions and ldb.");
+    throw new Error(
+      "B does not have enough elements for the given dimensions and ldb.",
+    );
   }
 
   // C: always m x n (no trans flag) — layout only affects lda/storage order.
   const cOuter = effLayoutC === "column-major" ? n : m;
   const cInner = effLayoutC === "column-major" ? m : n;
   if (ldc < cInner)
-    throw new Error(`ldc must be >= ${effLayoutC === "column-major" ? "rows" : "cols"} of C as stored.`);
+    throw new Error(
+      `ldc must be >= ${effLayoutC === "column-major" ? "rows" : "cols"} of C as stored.`,
+    );
   if (CIsGpu) {
-    if (ldc !== C.lda) throw new Error("ldc must match C.lda when C is a GpuMatrix.");
-    if (C.rows < m || C.cols < n) throw new Error("C is too small for the given m and n.");
+    if (ldc !== C.lda)
+      throw new Error("ldc must match C.lda when C is a GpuMatrix.");
+    if (C.rows < m || C.cols < n)
+      throw new Error("C is too small for the given m and n.");
   } else if (C.length < (cOuter - 1) * ldc + cInner) {
-    throw new Error("C does not have enough elements for the given dimensions and ldc.");
+    throw new Error(
+      "C does not have enough elements for the given dimensions and ldc.",
+    );
   }
 
   // Column-major A/B reinterpreted row-major is A^T/B^T — flip the trans flag.
@@ -133,48 +171,79 @@ export async function sgemm(
   const largeWgY = Math.ceil(m / BM_LARGE);
   const useLargeTile = largeWgX * largeWgY >= LARGE_TILE_WORKGROUP_THRESHOLD;
 
-  const pipeline = await getPipeline(device, useLargeTile ? "sgemm_large" : "sgemm_small");
+  const pipeline = await getPipeline(
+    device,
+    useLargeTile ? "sgemm_large" : "sgemm_small",
+  );
 
-  const ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "sgemm-A", false);
-  const BBuffer = BIsGpu ? B._buf : uploadBuffer(B, "sgemm-B", false);
-  const CBuffer = CIsGpu ? C._buf : uploadBuffer(C, "sgemm-C", true);
+  const ABuffer = AIsGpu ? A._buf : uploadBuffer(device, A, "sgemm-A", false);
+  const BBuffer = BIsGpu ? B._buf : uploadBuffer(device, B, "sgemm-B", false);
+  const CBuffer = CIsGpu ? C._buf : uploadBuffer(device, C, "sgemm-C", true);
+  // Vectorized-load enablement — kernel-side view after the column-major swap.
+  // op(A) is m×k (no-trans) or k×m (trans); op(B) is k×n or n×k.
+  const aNot = transA === "no-transpose";
+  const bNot = transB === "no-transpose";
+  const useVecA = aNot && vec4Usable(ABuffer, lda, m, k); // transposed-A vec stores bank-conflict smem — measured slower than scalar
+  const useVecB = vec4Usable(BBuffer, ldb, bNot ? k : n, bNot ? n : k);
   const paramsBuffer = createParamsBuffer(
+    device,
     [
-      { value: m,   type: "u32" },
-      { value: n,   type: "u32" },
-      { value: k,   type: "u32" },
+      { value: m, type: "u32" },
+      { value: n, type: "u32" },
+      { value: k, type: "u32" },
       { value: alpha, type: "f32" },
-      { value: beta,  type: "f32" },
+      { value: beta, type: "f32" },
       { value: lda, type: "u32" },
       { value: ldb, type: "u32" },
       { value: ldc, type: "u32" },
       { value: transA === "transpose" ? 1 : 0, type: "u32" },
       { value: transB === "transpose" ? 1 : 0, type: "u32" },
+      { value: useVecA ? 1 : 0, type: "u32" },
+      { value: useVecB ? 1 : 0, type: "u32" },
     ],
     "sgemm-params",
   );
 
   try {
-    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+    const bindGroup = createBindGroup(device, pipeline.getBindGroupLayout(0), [
       ABuffer,
+      vec4ViewBinding(device, ABuffer),
       BBuffer,
+      vec4ViewBinding(device, BBuffer),
       CBuffer,
       paramsBuffer,
     ]);
 
     const wgCount = useLargeTile
       ? {
-        x: Math.min(largeWgX, device.limits.maxComputeWorkgroupsPerDimension),
-        y: Math.min(largeWgY, device.limits.maxComputeWorkgroupsPerDimension),
-      }
+          x: requireWorkgroupCount(device, largeWgX, "sgemm", "x"),
+          y: requireWorkgroupCount(device, largeWgY, "sgemm", "y"),
+        }
       : {
-        x: Math.min(Math.ceil(n / BN_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
-        y: Math.min(Math.ceil(m / BM_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
-      };
-    const { commandEncoder, ts } = runComputePass(pipeline, bindGroup, wgCount);
-    const readBuffer = CIsGpu ? null : stageReadback(commandEncoder, CBuffer);
+          x: requireWorkgroupCount(
+            device,
+            Math.ceil(n / BN_SMALL),
+            "sgemm",
+            "x",
+          ),
+          y: requireWorkgroupCount(
+            device,
+            Math.ceil(m / BM_SMALL),
+            "sgemm",
+            "y",
+          ),
+        };
+    const { commandEncoder, ts } = runComputePass(
+      device,
+      pipeline,
+      bindGroup,
+      wgCount,
+    );
+    const readBuffer = CIsGpu
+      ? null
+      : stageReadback(device, commandEncoder, CBuffer);
 
-    submit(commandEncoder);
+    submit(device, commandEncoder);
 
     const gpuTimeMs = await extractTimestamp(ts);
 

package/src/sgemmtr/sgemmtr.d.mts

@@ -1,7 +1,8 @@
 import { GpuMatrix } from "../classes/GpuMatrix.mjs";
 
 /**
- * Performs the matrix-matrix operation C := uplo(alpha * op(A) * op(B) + beta * C) —
+ * Performs the matrix-matrix operation $$C \leftarrow \mathrm{uplo}(\alpha \mathrm{op}(A) \mathrm{op}(B) + \beta C)$$
+ *
  * `sgemm`'s operation, but only the triangle of C named by `uplo` is read or
  * written (`'lower'`: `col <= row`, `'upper'`: `col >= row`; C need not be
  * square — the test applies over the full m×n grid).
@@ -57,7 +58,7 @@ export declare function sgemmtr(
 ): Promise<{ C: Float32Array; gpuTimeMs?: number }>;
 
 /**
- * Performs the matrix-matrix operation C := uplo(alpha * op(A) * op(B) + beta * C)
+ * Performs the matrix-matrix operation $$C \leftarrow \mathrm{uplo}(\alpha \mathrm{op}(A) \mathrm{op}(B) + \beta C)$$
  *
  * A, B, and C are all kept GPU-resident. Each matrix's own `layout` (set at
  * `GpuMatrix.from` time) determines the operation — there is no separate

package/src/sgemmtr/sgemmtr.mjs

@@ -10,20 +10,40 @@ import { extractResult } from "../util/result.mjs";
 import { extractTimestamp } from "../util/benchmark.mjs";
 import { getPipeline } from "../util/pipeline.mjs";
 import { GpuMatrix } from "../classes/GpuMatrix.mjs";
-
-const BM_SMALL = 32, BN_SMALL = 32; // sgemmtr_small.wgsl's block tile
-const BM_LARGE = 64, BN_LARGE = 64; // sgemmtr_large.wgsl's block tile
-const LARGE_TILE_WORKGROUP_THRESHOLD = 36; // same threshold sgemm uses — see sgemm.mjs
+import { requireWorkgroupCount } from "../util/workgroup.mjs";
+import {
+  BM_SMALL,
+  BN_SMALL,
+  BM_LARGE,
+  BN_LARGE,
+  LARGE_TILE_WORKGROUP_THRESHOLD,
+} from "../util/constants.mjs";
+import { requireGpuDevice, requireSameDevice } from "../util/device.mjs";
 
 export async function sgemmtr(
-  device, uplo, transA, transB, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc, layout = "row-major",
+  device,
+  uplo,
+  transA,
+  transB,
+  m,
+  n,
+  k,
+  alpha,
+  A,
+  lda,
+  B,
+  ldb,
+  beta,
+  C,
+  ldc,
+  layout = "row-major",
 ) {
   let AIsGpu = A instanceof GpuMatrix;
   let BIsGpu = B instanceof GpuMatrix;
   const CIsGpu = C instanceof GpuMatrix;
 
-  if (!(device instanceof GPUDevice))
-    throw new Error("device must be a GPUDevice.");
+  requireGpuDevice(device);
+  requireSameDevice(device, "sgemmtr", { A, B, C });
   if (uplo !== "lower" && uplo !== "upper")
     throw new Error("uplo must be 'lower' or 'upper'.");
   if (transA !== "no-transpose" && transA !== "transpose")
@@ -32,12 +52,10 @@ export async function sgemmtr(
     throw new Error("transB must be 'no-transpose' or 'transpose'.");
   if (layout !== "row-major" && layout !== "column-major")
     throw new Error("layout must be 'row-major' or 'column-major'.");
-  if (typeof alpha !== "number")
-    throw new Error("alpha must be a number.");
+  if (typeof alpha !== "number") throw new Error("alpha must be a number.");
   if (Number.isNaN(alpha)) throw new Error("alpha must not be NaN.");
   if (!Number.isFinite(alpha)) throw new Error("alpha must be finite.");
-  if (typeof beta !== "number")
-    throw new Error("beta must be a number.");
+  if (typeof beta !== "number") throw new Error("beta must be a number.");
   if (Number.isNaN(beta)) throw new Error("beta must not be NaN.");
   if (!Number.isFinite(beta)) throw new Error("beta must be finite.");
   if (
@@ -59,7 +77,10 @@ export async function sgemmtr(
     throw new Error("C must be a GpuMatrix when A or B is a GpuMatrix.");
   if (CIsGpu && (!AIsGpu || !BIsGpu))
     throw new Error("A and B must be GpuMatrix when C is a GpuMatrix.");
-  if (m < 0 || n < 0 || k < 0) throw new Error("m, n, and k must be non-negative.");
+  if (m < 0 || n < 0 || k < 0)
+    throw new Error("m, n, and k must be non-negative.");
+  if (lda <= 0 || ldb <= 0 || ldc <= 0)
+    throw new Error("lda, ldb, and ldc must be positive.");
   if (m === 0 || n === 0) return CIsGpu ? {} : { C };
 
   // GpuMatrix's own .layout wins over the shared `layout` argument.
@@ -74,14 +95,19 @@ export async function sgemmtr(
   const aOuter = transA === "no-transpose" ? aRows : aCols; // # of stored chunks
   const aInner = transA === "no-transpose" ? aCols : aRows; // required chunk length
   if (lda < aInner)
-    throw new Error(`lda must be >= ${effLayoutA === "column-major" ? "rows" : "cols"} of A as stored.`);
+    throw new Error(
+      `lda must be >= ${effLayoutA === "column-major" ? "rows" : "cols"} of A as stored.`,
+    );
   if (AIsGpu) {
-    if (lda !== A.lda) throw new Error("lda must match A.lda when A is a GpuMatrix.");
+    if (lda !== A.lda)
+      throw new Error("lda must match A.lda when A is a GpuMatrix.");
     const [aLogRows, aLogCols] = transA === "no-transpose" ? [m, k] : [k, m];
     if (A.rows < aLogRows || A.cols < aLogCols)
       throw new Error("A is too small for the given m, k, and transA.");
   } else if (A.length < (aOuter - 1) * lda + aInner) {
-    throw new Error("A does not have enough elements for the given dimensions and lda.");
+    throw new Error(
+      "A does not have enough elements for the given dimensions and lda.",
+    );
   }
 
   // B: same reasoning as A, with op(B) = k x n.
@@ -90,26 +116,37 @@ export async function sgemmtr(
   const bOuter = transB === "no-transpose" ? bRows : bCols;
   const bInner = transB === "no-transpose" ? bCols : bRows;
   if (ldb < bInner)
-    throw new Error(`ldb must be >= ${effLayoutB === "column-major" ? "rows" : "cols"} of B as stored.`);
+    throw new Error(
+      `ldb must be >= ${effLayoutB === "column-major" ? "rows" : "cols"} of B as stored.`,
+    );
   if (BIsGpu) {
-    if (ldb !== B.lda) throw new Error("ldb must match B.lda when B is a GpuMatrix.");
+    if (ldb !== B.lda)
+      throw new Error("ldb must match B.lda when B is a GpuMatrix.");
     const [bLogRows, bLogCols] = transB === "no-transpose" ? [k, n] : [n, k];
     if (B.rows < bLogRows || B.cols < bLogCols)
       throw new Error("B is too small for the given n, k, and transB.");
   } else if (B.length < (bOuter - 1) * ldb + bInner) {
-    throw new Error("B does not have enough elements for the given dimensions and ldb.");
+    throw new Error(
+      "B does not have enough elements for the given dimensions and ldb.",
+    );
   }
 
   // C: always m x n (no trans flag) — layout only affects lda/storage order.
   const cOuter = effLayoutC === "column-major" ? n : m;
   const cInner = effLayoutC === "column-major" ? m : n;
   if (ldc < cInner)
-    throw new Error(`ldc must be >= ${effLayoutC === "column-major" ? "rows" : "cols"} of C as stored.`);
+    throw new Error(
+      `ldc must be >= ${effLayoutC === "column-major" ? "rows" : "cols"} of C as stored.`,
+    );
   if (CIsGpu) {
-    if (ldc !== C.lda) throw new Error("ldc must match C.lda when C is a GpuMatrix.");
-    if (C.rows < m || C.cols < n) throw new Error("C is too small for the given m and n.");
+    if (ldc !== C.lda)
+      throw new Error("ldc must match C.lda when C is a GpuMatrix.");
+    if (C.rows < m || C.cols < n)
+      throw new Error("C is too small for the given m and n.");
   } else if (C.length < (cOuter - 1) * ldc + cInner) {
-    throw new Error("C does not have enough elements for the given dimensions and ldc.");
+    throw new Error(
+      "C does not have enough elements for the given dimensions and ldc.",
+    );
   }
 
   // Column-major A/B reinterpreted row-major is A^T/B^T — flip the trans flag.
@@ -140,18 +177,22 @@ export async function sgemmtr(
   const largeWgY = Math.ceil(m / BM_LARGE);
   const useLargeTile = largeWgX * largeWgY >= LARGE_TILE_WORKGROUP_THRESHOLD;
 
-  const pipeline = await getPipeline(device, useLargeTile ? "sgemmtr_large" : "sgemmtr_small");
+  const pipeline = await getPipeline(
+    device,
+    useLargeTile ? "sgemmtr_large" : "sgemmtr_small",
+  );
 
-  const ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "sgemmtr-A", false);
-  const BBuffer = BIsGpu ? B._buf : uploadBuffer(B, "sgemmtr-B", false);
-  const CBuffer = CIsGpu ? C._buf : uploadBuffer(C, "sgemmtr-C", true);
+  const ABuffer = AIsGpu ? A._buf : uploadBuffer(device, A, "sgemmtr-A", false);
+  const BBuffer = BIsGpu ? B._buf : uploadBuffer(device, B, "sgemmtr-B", false);
+  const CBuffer = CIsGpu ? C._buf : uploadBuffer(device, C, "sgemmtr-C", true);
   const paramsBuffer = createParamsBuffer(
+    device,
     [
-      { value: m,   type: "u32" },
-      { value: n,   type: "u32" },
-      { value: k,   type: "u32" },
+      { value: m, type: "u32" },
+      { value: n, type: "u32" },
+      { value: k, type: "u32" },
       { value: alpha, type: "f32" },
-      { value: beta,  type: "f32" },
+      { value: beta, type: "f32" },
       { value: lda, type: "u32" },
       { value: ldb, type: "u32" },
       { value: ldc, type: "u32" },
@@ -163,7 +204,7 @@ export async function sgemmtr(
   );
 
   try {
-    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+    const bindGroup = createBindGroup(device, pipeline.getBindGroupLayout(0), [
       ABuffer,
       BBuffer,
       CBuffer,
@@ -172,17 +213,34 @@ export async function sgemmtr(
 
     const wgCount = useLargeTile
       ? {
-        x: Math.min(largeWgX, device.limits.maxComputeWorkgroupsPerDimension),
-        y: Math.min(largeWgY, device.limits.maxComputeWorkgroupsPerDimension),
-      }
+          x: requireWorkgroupCount(device, largeWgX, "sgemmtr", "x"),
+          y: requireWorkgroupCount(device, largeWgY, "sgemmtr", "y"),
+        }
       : {
-        x: Math.min(Math.ceil(n / BN_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
-        y: Math.min(Math.ceil(m / BM_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
-      };
-    const { commandEncoder, ts } = runComputePass(pipeline, bindGroup, wgCount);
-    const readBuffer = CIsGpu ? null : stageReadback(commandEncoder, CBuffer);
+          x: requireWorkgroupCount(
+            device,
+            Math.ceil(n / BN_SMALL),
+            "sgemmtr",
+            "x",
+          ),
+          y: requireWorkgroupCount(
+            device,
+            Math.ceil(m / BM_SMALL),
+            "sgemmtr",
+            "y",
+          ),
+        };
+    const { commandEncoder, ts } = runComputePass(
+      device,
+      pipeline,
+      bindGroup,
+      wgCount,
+    );
+    const readBuffer = CIsGpu
+      ? null
+      : stageReadback(device, commandEncoder, CBuffer);
 
-    submit(commandEncoder);
+    submit(device, commandEncoder);
 
     const gpuTimeMs = await extractTimestamp(ts);
 

package/src/sgemv/sgemv.d.mts

@@ -2,7 +2,7 @@ import { GpuVector } from "../classes/GpuVector.mjs";
 import { GpuMatrix } from "../classes/GpuMatrix.mjs";
 
 /**
- * Performs the matrix-vector operation y = alpha * op(A) * x + beta * y
+ * Performs the matrix-vector operation $$y \leftarrow \alpha \mathrm{op}(A) x + \beta y$$
  *
  * - `trans='no-transpose'`: op(A) = A,   x is length n, y is length m
  * - `trans='transpose'`:    op(A) = A^T, x is length m, y is length n
@@ -51,7 +51,7 @@ export declare function sgemv(
 ): Promise<{ y: Float32Array; gpuTimeMs?: number }>;
 
 /**
- * Performs the matrix-vector operation y = alpha * op(A) * x + beta * y
+ * Performs the matrix-vector operation $$y \leftarrow \alpha \mathrm{op}(A) x + \beta y$$
  *
  * x and y are kept resident on the GPU. A must be a GpuMatrix; its own
  * `layout` (set at `GpuMatrix.from` time) determines the operation — there is