wgblas 2.0.0 → 2.2.0

package/src/ssyr2k/ssyr2k.mjs

@@ -10,39 +10,58 @@ import { extractResult } from "../util/result.mjs";
 import { resolveTimestamp, 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/sgemmtr/ssyrk use
+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";
 
 // ssyr2k: C := uplo(alpha*op(A)*op(B)^T + alpha*op(B)*op(A)^T + beta*C). No
 // dedicated shader — two sgemmtr passes on one encoder, second with beta=1.
 export async function ssyr2k(
-  device, uplo, trans, n, k, alpha, A, lda, B, ldb, beta, C, ldc, layout = "row-major",
+  device,
+  uplo,
+  trans,
+  n,
+  k,
+  alpha,
+  A,
+  lda,
+  B,
+  ldb,
+  beta,
+  C,
+  ldc,
+  layout = "row-major",
 ) {
   const AIsGpu = A instanceof GpuMatrix;
   const 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, "ssyr2k", { A, B, C });
   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 (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 (
-    !Number.isInteger(n) || !Number.isInteger(k) ||
-    !Number.isInteger(lda) || !Number.isInteger(ldb) || !Number.isInteger(ldc)
+    !Number.isInteger(n) ||
+    !Number.isInteger(k) ||
+    !Number.isInteger(lda) ||
+    !Number.isInteger(ldb) ||
+    !Number.isInteger(ldc)
   )
     throw new Error("n, k, lda, ldb, and ldc must be integers.");
   if (!AIsGpu && !(A instanceof Float32Array))
@@ -56,6 +75,8 @@ export async function ssyr2k(
   if (CIsGpu && (!AIsGpu || !BIsGpu))
     throw new Error("A and B must be GpuMatrix when C is a GpuMatrix.");
   if (n < 0 || k < 0) throw new Error("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 (n === 0) return CIsGpu ? {} : { C };
 
   const effLayoutA = AIsGpu ? A.layout : layout;
@@ -68,14 +89,19 @@ export async function ssyr2k(
   const aOuter = trans === "no-transpose" ? aRows : aCols;
   const aInner = trans === "no-transpose" ? aCols : aRows;
   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] = trans === "no-transpose" ? [n, k] : [k, n];
     if (A.rows < aLogRows || A.cols < aLogCols)
       throw new Error("A is too small for the given n, k, and trans.");
   } 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 shape rule as A — netlib's syr2k shares one trans across both operands.
@@ -84,23 +110,32 @@ export async function ssyr2k(
   const bOuter = trans === "no-transpose" ? bRows : bCols;
   const bInner = trans === "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] = trans === "no-transpose" ? [n, k] : [k, n];
     if (B.rows < bLogRows || B.cols < bLogCols)
       throw new Error("B is too small for the given n, k, and trans.");
   } 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 n x n symmetric — layout only affects storage order, not size.
   if (ldc < n) throw new Error("ldc must be >= n.");
   if (CIsGpu) {
-    if (ldc !== C.lda) throw new Error("ldc must match C.lda when C is a GpuMatrix.");
-    if (C.rows < n || C.cols < n) throw new Error("C is too small for the given n.");
+    if (ldc !== C.lda)
+      throw new Error("ldc must match C.lda when C is a GpuMatrix.");
+    if (C.rows < n || C.cols < n)
+      throw new Error("C is too small for the given n.");
   } else if (C.length < (n - 1) * ldc + n) {
-    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 trans per operand.
@@ -111,75 +146,118 @@ export async function ssyr2k(
   if (effLayoutB === "column-major")
     effTransB = effTransB === "no-transpose" ? "transpose" : "no-transpose";
 
-  const uploEff = effLayoutC === "column-major"
-    ? (uplo === "lower" ? "upper" : "lower")
-    : uplo;
+  const uploEff =
+    effLayoutC === "column-major"
+      ? uplo === "lower"
+        ? "upper"
+        : "lower"
+      : uplo;
   const flip = (t) => (t === "no-transpose" ? "transpose" : "no-transpose");
 
   // One pass: op(X) as-is, op(Y) transposed. transOwnY is explicit, not inferred.
   function passShape(transOwnX, X, ldX, transOwnY, Y, ldY) {
     const transX = transOwnX;
     const transY = flip(transOwnY);
-    if (effLayoutC !== "column-major") return { transX, X, ldX, transY, Y, ldY };
-    return { transX: flip(transY), X: Y, ldX: ldY, transY: flip(transX), Y: X, ldY: ldX };
+    if (effLayoutC !== "column-major")
+      return { transX, X, ldX, transY, Y, ldY };
+    return {
+      transX: flip(transY),
+      X: Y,
+      ldX: ldY,
+      transY: flip(transX),
+      Y: X,
+      ldY: ldX,
+    };
   }
 
   // Shape-based auto-select — see sgemmtr_small.wgsl/sgemmtr_large.wgsl. m=n=n here (square C).
   const largeWgX = Math.ceil(n / BN_LARGE);
   const largeWgY = Math.ceil(n / 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 wgCount = useLargeTile
     ? {
-      x: Math.min(largeWgX, device.limits.maxComputeWorkgroupsPerDimension),
-      y: Math.min(largeWgY, device.limits.maxComputeWorkgroupsPerDimension),
-    }
+        x: requireWorkgroupCount(device, largeWgX, "ssyr2k", "x"),
+        y: requireWorkgroupCount(device, largeWgY, "ssyr2k", "y"),
+      }
     : {
-      x: Math.min(Math.ceil(n / BN_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
-      y: Math.min(Math.ceil(n / BM_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
-    };
+        x: requireWorkgroupCount(
+          device,
+          Math.ceil(n / BN_SMALL),
+          "ssyr2k",
+          "x",
+        ),
+        y: requireWorkgroupCount(
+          device,
+          Math.ceil(n / BM_SMALL),
+          "ssyr2k",
+          "y",
+        ),
+      };
 
-  const ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "ssyr2k-A", false);
-  const BBuffer = BIsGpu ? B._buf : uploadBuffer(B, "ssyr2k-B", false);
-  const CBuffer = CIsGpu ? C._buf : uploadBuffer(C, "ssyr2k-C", true);
-  let paramsBuffer1 = null, paramsBuffer2 = null;
+  const ABuffer = AIsGpu ? A._buf : uploadBuffer(device, A, "ssyr2k-A", false);
+  const BBuffer = BIsGpu ? B._buf : uploadBuffer(device, B, "ssyr2k-B", false);
+  const CBuffer = CIsGpu ? C._buf : uploadBuffer(device, C, "ssyr2k-C", true);
+  let paramsBuffer1 = null,
+    paramsBuffer2 = null;
 
   try {
     const pass1 = passShape(effTransA, ABuffer, lda, effTransB, BBuffer, ldb);
     const pass2 = passShape(effTransB, BBuffer, ldb, effTransA, ABuffer, lda);
 
-    const makeParams = (p, betaVal) => createParamsBuffer(
-      [
-        { value: n,   type: "u32" },
-        { value: n,   type: "u32" },
-        { value: k,   type: "u32" },
-        { value: alpha, type: "f32" },
-        { value: betaVal, type: "f32" },
-        { value: p.ldX, type: "u32" },
-        { value: p.ldY, type: "u32" },
-        { value: ldc, type: "u32" },
-        { value: p.transX === "transpose" ? 1 : 0, type: "u32" },
-        { value: p.transY === "transpose" ? 1 : 0, type: "u32" },
-        { value: uploEff === "upper" ? 1 : 0, type: "u32" },
-      ],
-      "ssyr2k-params",
-    );
+    const makeParams = (p, betaVal) =>
+      createParamsBuffer(
+        device,
+        [
+          { value: n, type: "u32" },
+          { value: n, type: "u32" },
+          { value: k, type: "u32" },
+          { value: alpha, type: "f32" },
+          { value: betaVal, type: "f32" },
+          { value: p.ldX, type: "u32" },
+          { value: p.ldY, type: "u32" },
+          { value: ldc, type: "u32" },
+          { value: p.transX === "transpose" ? 1 : 0, type: "u32" },
+          { value: p.transY === "transpose" ? 1 : 0, type: "u32" },
+          { value: uploEff === "upper" ? 1 : 0, type: "u32" },
+        ],
+        "ssyr2k-params",
+      );
     paramsBuffer1 = makeParams(pass1, beta);
     paramsBuffer2 = makeParams(pass2, 1.0);
 
-    const bindGroup1 = createBindGroup(pipeline.getBindGroupLayout(0), [pass1.X, pass1.Y, CBuffer, paramsBuffer1]);
-    const bindGroup2 = createBindGroup(pipeline.getBindGroupLayout(0), [pass2.X, pass2.Y, CBuffer, paramsBuffer2]);
+    const bindGroup1 = createBindGroup(device, pipeline.getBindGroupLayout(0), [
+      pass1.X,
+      pass1.Y,
+      CBuffer,
+      paramsBuffer1,
+    ]);
+    const bindGroup2 = createBindGroup(device, pipeline.getBindGroupLayout(0), [
+      pass2.X,
+      pass2.Y,
+      CBuffer,
+      paramsBuffer2,
+    ]);
 
-    const { commandEncoder, querySet } = beginTimedEncoder();
-    const desc1 = querySet ? { timestampWrites: { querySet, beginningOfPassWriteIndex: 0 } } : undefined;
-    const desc2 = querySet ? { timestampWrites: { querySet, endOfPassWriteIndex: 1 } } : undefined;
+    const { commandEncoder, querySet } = beginTimedEncoder(device);
+    const desc1 = querySet
+      ? { timestampWrites: { querySet, beginningOfPassWriteIndex: 0 } }
+      : undefined;
+    const desc2 = querySet
+      ? { timestampWrites: { querySet, endOfPassWriteIndex: 1 } }
+      : undefined;
     encodePass(commandEncoder, pipeline, bindGroup1, wgCount, desc1);
     encodePass(commandEncoder, pipeline, bindGroup2, wgCount, desc2);
 
-    const ts = resolveTimestamp(commandEncoder, querySet);
-    const readBuffer = CIsGpu ? null : stageReadback(commandEncoder, CBuffer);
+    const ts = resolveTimestamp(device, commandEncoder, querySet);
+    const readBuffer = CIsGpu
+      ? null
+      : stageReadback(device, commandEncoder, CBuffer);
 
-    submit(commandEncoder);
+    submit(device, commandEncoder);
 
     const gpuTimeMs = await extractTimestamp(ts);
 

package/src/ssyrk/ssyrk.d.mts

@@ -1,7 +1,8 @@
 import { GpuMatrix } from "../classes/GpuMatrix.mjs";
 
 /**
- * Performs the symmetric rank-k update C := uplo(alpha * op(A) * op(A)^T + beta * C) —
+ * Performs the symmetric rank-k update $$C \leftarrow \mathrm{uplo}(\alpha \mathrm{op}(A) \mathrm{op}(A)^{T} + \beta C)$$
+ *
  * only the triangle of C named by `uplo` is read or written (`'lower'`: `col <= row`,
  * `'upper'`: `col >= row`). C is always n×n.
  *
@@ -51,7 +52,7 @@ export declare function ssyrk(
 ): Promise<{ C: Float32Array; gpuTimeMs?: number }>;
 
 /**
- * Performs the symmetric rank-k update C := uplo(alpha * op(A) * op(A)^T + beta * C)
+ * Performs the symmetric rank-k update $$C \leftarrow \mathrm{uplo}(\alpha \mathrm{op}(A) \mathrm{op}(A)^{T} + \beta C)$$
  *
  * A and C are both kept GPU-resident. Each matrix's own `layout` (set at
  * `GpuMatrix.from` time) determines the operation — there is no separate

package/src/ssyrk/ssyrk.mjs

@@ -11,36 +11,55 @@ import { extractResult } from "../util/result.mjs";
 import { resolveTimestamp, 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/sgemmtr use
+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";
 
 // ssyrk: C := uplo(alpha*op(A)*op(A)^T + beta*C). No dedicated shader —
 // sgemmtr's kernel with A duplicated into a separate B buffer (B := A).
 export async function ssyrk(
-  device, uplo, trans, n, k, alpha, A, lda, beta, C, ldc, layout = "row-major",
+  device,
+  uplo,
+  trans,
+  n,
+  k,
+  alpha,
+  A,
+  lda,
+  beta,
+  C,
+  ldc,
+  layout = "row-major",
 ) {
   const AIsGpu = A instanceof GpuMatrix;
   const CIsGpu = C instanceof GpuMatrix;
 
-  if (!(device instanceof GPUDevice))
-    throw new Error("device must be a GPUDevice.");
+  requireGpuDevice(device);
+  requireSameDevice(device, "ssyrk", { A, C });
   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 (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 (!Number.isInteger(n) || !Number.isInteger(k) || !Number.isInteger(lda) || !Number.isInteger(ldc))
+  if (
+    !Number.isInteger(n) ||
+    !Number.isInteger(k) ||
+    !Number.isInteger(lda) ||
+    !Number.isInteger(ldc)
+  )
     throw new Error("n, k, lda, and ldc must be integers.");
   if (!AIsGpu && !(A instanceof Float32Array))
     throw new Error("A must be a Float32Array or GpuMatrix.");
@@ -51,6 +70,7 @@ export async function ssyrk(
   if (CIsGpu && !AIsGpu)
     throw new Error("A must be a GpuMatrix when C is a GpuMatrix.");
   if (n < 0 || k < 0) throw new Error("n and k must be non-negative.");
+  if (lda <= 0 || ldc <= 0) throw new Error("lda and ldc must be positive.");
   if (n === 0) return CIsGpu ? {} : { C };
 
   // GpuMatrix's own .layout wins over the shared `layout` argument.
@@ -63,23 +83,32 @@ export async function ssyrk(
   const aOuter = trans === "no-transpose" ? aRows : aCols;
   const aInner = trans === "no-transpose" ? aCols : aRows;
   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] = trans === "no-transpose" ? [n, k] : [k, n];
     if (A.rows < aLogRows || A.cols < aLogCols)
       throw new Error("A is too small for the given n, k, and trans.");
   } 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.",
+    );
   }
 
   // C: always n x n symmetric — layout only affects storage order, not size.
   if (ldc < n) throw new Error("ldc must be >= n.");
   if (CIsGpu) {
-    if (ldc !== C.lda) throw new Error("ldc must match C.lda when C is a GpuMatrix.");
-    if (C.rows < n || C.cols < n) throw new Error("C is too small for the given n.");
+    if (ldc !== C.lda)
+      throw new Error("ldc must match C.lda when C is a GpuMatrix.");
+    if (C.rows < n || C.cols < n)
+      throw new Error("C is too small for the given n.");
   } else if (C.length < (n - 1) * ldc + n) {
-    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 reinterpreted row-major is A^T — flip trans, same trick sgemm/sgemmtr use.
@@ -104,22 +133,31 @@ export async function ssyrk(
   const largeWgY = Math.ceil(n / 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, "ssyrk-A", false);
-  const CBuffer = CIsGpu ? C._buf : uploadBuffer(C, "ssyrk-C", true);
+  const ABuffer = AIsGpu ? A._buf : uploadBuffer(device, A, "ssyrk-A", false);
+  const CBuffer = CIsGpu ? C._buf : uploadBuffer(device, C, "ssyrk-C", true);
   // B := A, but as a genuinely separate buffer — re-uploaded for Float32Array,
   // GPU-copied (see below) for GpuMatrix, since the caller owns ABuffer.
   const BBuffer = AIsGpu
-    ? createStorageBuffer(ABuffer.size, "ssyrk-B", GPUBufferUsage.COPY_DST)
-    : uploadBuffer(A, "ssyrk-B", false);
+    ? createStorageBuffer(
+        device,
+        ABuffer.size,
+        "ssyrk-B",
+        GPUBufferUsage.COPY_DST,
+      )
+    : uploadBuffer(device, A, "ssyrk-B", false);
   const paramsBuffer = createParamsBuffer(
+    device,
     [
-      { value: n,   type: "u32" }, // gemmtr's m
-      { value: n,   type: "u32" }, // gemmtr's n
-      { value: k,   type: "u32" },
+      { value: n, type: "u32" }, // gemmtr's m
+      { value: n, type: "u32" }, // gemmtr's n
+      { value: k, type: "u32" },
       { value: alpha, type: "f32" },
-      { value: beta,  type: "f32" },
+      { value: beta, type: "f32" },
       { value: lda, type: "u32" }, // gemmtr's lda
       { value: lda, type: "u32" }, // gemmtr's ldb — B := A, same lda
       { value: ldc, type: "u32" },
@@ -131,7 +169,7 @@ export async function ssyrk(
   );
 
   try {
-    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+    const bindGroup = createBindGroup(device, pipeline.getBindGroupLayout(0), [
       ABuffer,
       BBuffer,
       CBuffer,
@@ -140,22 +178,36 @@ export async function ssyrk(
 
     const wgCount = useLargeTile
       ? {
-        x: Math.min(largeWgX, device.limits.maxComputeWorkgroupsPerDimension),
-        y: Math.min(largeWgY, device.limits.maxComputeWorkgroupsPerDimension),
-      }
+          x: requireWorkgroupCount(device, largeWgX, "ssyrk", "x"),
+          y: requireWorkgroupCount(device, largeWgY, "ssyrk", "y"),
+        }
       : {
-        x: Math.min(Math.ceil(n / BN_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
-        y: Math.min(Math.ceil(n / BM_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
-      };
+          x: requireWorkgroupCount(
+            device,
+            Math.ceil(n / BN_SMALL),
+            "ssyrk",
+            "x",
+          ),
+          y: requireWorkgroupCount(
+            device,
+            Math.ceil(n / BM_SMALL),
+            "ssyrk",
+            "y",
+          ),
+        };
     // Manual encoder (not runComputePass) so the A->B duplicate copy lands
     // on the same command encoder, strictly before the compute pass reads B.
-    const { commandEncoder, querySet, passDescriptor } = beginTimedEncoder();
-    if (AIsGpu) commandEncoder.copyBufferToBuffer(ABuffer, 0, BBuffer, 0, ABuffer.size);
+    const { commandEncoder, querySet, passDescriptor } =
+      beginTimedEncoder(device);
+    if (AIsGpu)
+      commandEncoder.copyBufferToBuffer(ABuffer, 0, BBuffer, 0, ABuffer.size);
     encodePass(commandEncoder, pipeline, bindGroup, wgCount, passDescriptor);
-    const ts = resolveTimestamp(commandEncoder, querySet);
-    const readBuffer = CIsGpu ? null : stageReadback(commandEncoder, CBuffer);
+    const ts = resolveTimestamp(device, commandEncoder, querySet);
+    const readBuffer = CIsGpu
+      ? null
+      : stageReadback(device, commandEncoder, CBuffer);
 
-    submit(commandEncoder);
+    submit(device, commandEncoder);
 
     const gpuTimeMs = await extractTimestamp(ts);
 

package/src/strmm/strmm.d.mts

@@ -2,8 +2,9 @@ import { GpuMatrix } from "../classes/GpuMatrix.mjs";
 
 /**
  * Performs the triangular matrix-matrix operation
- * B := alpha * op(A) * B (`side='left'`) or
- * B := alpha * B * op(A) (`side='right'`) — `A` is triangular, only its
+ * $$B \leftarrow \alpha \mathrm{op}(A) B \quad (\texttt{side='left'})$$
+ * $$B \leftarrow \alpha B \mathrm{op}(A) \quad (\texttt{side='right'})$$
+ * `A` is triangular, only its
  * `uplo` triangle stored; `B` is a general m×n matrix, overwritten in place.
  *
  * - `side='left'`:  `A` is m×m — `A` premultiplies `B`
@@ -58,8 +59,8 @@ export declare function strmm(
 
 /**
  * Performs the triangular matrix-matrix operation
- * B := alpha * op(A) * B (`side='left'`) or
- * B := alpha * B * op(A) (`side='right'`)
+ * $$B \leftarrow \alpha \mathrm{op}(A) B \quad (\texttt{side='left'})$$
+ * $$B \leftarrow \alpha B \mathrm{op}(A) \quad (\texttt{side='right'})$$
  *
  * A and B are both kept GPU-resident; B is mutated in place. Each matrix's
  * own `layout` (set at `GpuMatrix.from` time) determines the operation —