wgblas 2.0.0 → 2.2.0

package/src/strmm/strmm.mjs

@@ -4,6 +4,7 @@ import {
   createStorageBuffer,
   stageReadback,
   destroyBuffers,
+  vec4ViewBinding,
 } from "../util/buffer.mjs";
 import { createBindGroup } from "../util/bindgroup.mjs";
 import { beginTimedEncoder, encodePass, submit } from "../util/compute.mjs";
@@ -11,25 +12,42 @@ 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; // 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; // same threshold sgemm/sgemmtr/ssyrk/ssymm use
-const TRI_WG = 8; // triangularize.wgsl's @workgroup_size(8, 8)
+import { requireWorkgroupCount } from "../util/workgroup.mjs";
+import {
+  BM_SMALL,
+  BN_SMALL,
+  BM_LARGE,
+  BN_LARGE,
+  LARGE_TILE_WORKGROUP_THRESHOLD,
+} from "../util/constants.mjs";
+import { TILE_WG_2D } from "../util/constants.mjs";
+import { requireGpuDevice, requireSameDevice } from "../util/device.mjs";
 
 // strmm: B := alpha*op(A)*B (side='left') or alpha*B*op(A) (side='right'), A
 // triangular. Triangularize then sgemm, one command encoder. B is both
 // input and output, so gemm writes to a fresh buffer (no aliasing race),
 // copied back into B (GpuMatrix) or read back directly (Float32Array).
 export async function strmm(
-  device, side, uplo, transA, diag, m, n, alpha, A, lda, B, ldb, layout = "row-major",
+  device,
+  side,
+  uplo,
+  transA,
+  diag,
+  m,
+  n,
+  alpha,
+  A,
+  lda,
+  B,
+  ldb,
+  layout = "row-major",
 ) {
   const AIsGpu = A instanceof GpuMatrix;
   const BIsGpu = B instanceof GpuMatrix;
   const isUnit = diag === "unit";
 
-  if (!(device instanceof GPUDevice))
-    throw new Error("device must be a GPUDevice.");
+  requireGpuDevice(device);
+  requireSameDevice(device, "strmm", { A, B });
   if (side !== "left" && side !== "right")
     throw new Error("side must be 'left' or 'right'.");
   if (uplo !== "lower" && uplo !== "upper")
@@ -40,11 +58,15 @@ export async function strmm(
     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 (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 (!Number.isInteger(m) || !Number.isInteger(n) || !Number.isInteger(lda) || !Number.isInteger(ldb))
+  if (
+    !Number.isInteger(m) ||
+    !Number.isInteger(n) ||
+    !Number.isInteger(lda) ||
+    !Number.isInteger(ldb)
+  )
     throw new Error("m, n, lda, and ldb must be integers.");
   if (!AIsGpu && !(A instanceof Float32Array))
     throw new Error("A must be a Float32Array or GpuMatrix.");
@@ -60,37 +82,59 @@ export async function strmm(
 
   // A: triangular, order = m (side='left') or n (side='right').
   const aOrder = side === "left" ? m : n;
-  if (lda < aOrder) throw new Error("lda must be >= " + (side === "left" ? "m" : "n") + ".");
+  if (lda < aOrder)
+    throw new Error("lda must be >= " + (side === "left" ? "m" : "n") + ".");
   if (AIsGpu) {
-    if (lda !== A.lda) throw new Error("lda must match A.lda when A is a GpuMatrix.");
-    if (A.rows < aOrder || A.cols < aOrder) throw new Error("A is too small for the given m/n and side.");
+    if (lda !== A.lda)
+      throw new Error("lda must match A.lda when A is a GpuMatrix.");
+    if (A.rows < aOrder || A.cols < aOrder)
+      throw new Error("A is too small for the given m/n and side.");
   } else if (A.length < (aOrder - 1) * lda + aOrder) {
-    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: always m x n, overwritten in place with the same ldb.
   const bOuter = effLayoutB === "column-major" ? n : m;
   const bInner = effLayoutB === "column-major" ? m : n;
   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 (B.rows < m || B.cols < n) throw new Error("B is too small for the given m and n.");
+    if (ldb !== B.lda)
+      throw new Error("ldb must match B.lda when B is a GpuMatrix.");
+    if (B.rows < m || B.cols < n)
+      throw new Error("B is too small for the given m and n.");
   } 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.",
+    );
   }
 
   // A isn't symmetric: column-major = genuine transpose, so flip transA;
   // transposing also swaps which triangle looks stored, so flip uplo too.
-  const uploEffA = effLayoutA === "column-major" ? (uplo === "lower" ? "upper" : "lower") : uplo;
-  const transEffA = effLayoutA === "column-major" ? (transA === "no-transpose" ? "transpose" : "no-transpose") : transA;
+  const uploEffA =
+    effLayoutA === "column-major"
+      ? uplo === "lower"
+        ? "upper"
+        : "lower"
+      : uplo;
+  const transEffA =
+    effLayoutA === "column-major"
+      ? transA === "no-transpose"
+        ? "transpose"
+        : "no-transpose"
+      : transA;
 
   const transB = effLayoutB === "column-major" ? "transpose" : "no-transpose";
   const transDense = "no-transpose"; // Adense already embodies op(A)
 
   // X*Y (X=Adense,Y=B for side='left', swapped for 'right'). Column-major
   // output: compute (B_out)^T instead — sgemm's own trick, same as ssymm's.
-  let mg = m, ng = n;
+  let mg = m,
+    ng = n;
   const kg = aOrder;
   let transX = side === "left" ? transDense : transB;
   let transY = side === "left" ? transB : transDense;
@@ -106,33 +150,63 @@ export async function strmm(
   const largeWgX = Math.ceil(ng / BN_LARGE);
   const largeWgY = Math.ceil(mg / BM_LARGE);
   const useLargeTile = largeWgX * largeWgY >= LARGE_TILE_WORKGROUP_THRESHOLD;
-  const gemmPipeline = await getPipeline(device, useLargeTile ? "sgemm_large" : "sgemm_small");
+  const gemmPipeline = await getPipeline(
+    device,
+    useLargeTile ? "sgemm_large" : "sgemm_small",
+  );
   const triPipeline = await getPipeline(device, "triangularize");
   const gemmWgCount = useLargeTile
     ? {
-      x: Math.min(largeWgX, device.limits.maxComputeWorkgroupsPerDimension),
-      y: Math.min(largeWgY, device.limits.maxComputeWorkgroupsPerDimension),
-    }
+        x: requireWorkgroupCount(device, largeWgX, "strmm", "x"),
+        y: requireWorkgroupCount(device, largeWgY, "strmm", "y"),
+      }
     : {
-      x: Math.min(Math.ceil(ng / BN_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
-      y: Math.min(Math.ceil(mg / BM_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
-    };
-
-  const ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "strmm-A", false);
-  // readback=true (COPY_SRC): BBuffer is also the source that seeds outBuffer.
-  const BBuffer = BIsGpu ? B._buf : uploadBuffer(B, "strmm-B", true);
-  const AdenseBuffer = createStorageBuffer(aOrder * ldDense * 4, "strmm-Adense");
-  // COPY_DST: seeded from B's own content before gemm runs, so stride-padding
-  // gaps (never written by gemm's tight m x n loop) keep B's original bytes
-  // instead of reading back as zero. COPY_SRC: read back / adopted by B after.
-  const outBuffer = createStorageBuffer(
-    bOuter * ldb * 4, "strmm-out", GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
-  );
-  let triParams = null, gemmParams = null;
+        x: requireWorkgroupCount(
+          device,
+          Math.ceil(ng / BN_SMALL),
+          "strmm",
+          "x",
+        ),
+        y: requireWorkgroupCount(
+          device,
+          Math.ceil(mg / BM_SMALL),
+          "strmm",
+          "y",
+        ),
+      };
+
+  // Null-init here and allocate inside the try below, so a throw partway
+  // through the sequence still reaches finally with every handle visible
+  // (strsv.mjs is the reference for this pattern).
+  let ABuffer = null,
+    BBuffer = null;
+  let AdenseBuffer = null,
+    outBuffer = null;
+  let triParams = null,
+    gemmParams = null;
   let outBufferAdopted = false; // true once B._buf is repointed at outBuffer
 
   try {
+    ABuffer = AIsGpu ? A._buf : uploadBuffer(device, A, "strmm-A", false);
+    // readback=true (COPY_SRC): BBuffer is also the source that seeds outBuffer.
+    BBuffer = BIsGpu ? B._buf : uploadBuffer(device, B, "strmm-B", true);
+    AdenseBuffer = createStorageBuffer(
+      device,
+      aOrder * ldDense * 4,
+      "strmm-Adense",
+    );
+    // COPY_DST: seeded from B's own content before gemm runs, so stride-padding
+    // gaps (never written by gemm's tight m x n loop) keep B's original bytes
+    // instead of reading back as zero. COPY_SRC: read back / adopted by B after.
+    outBuffer = createStorageBuffer(
+      device,
+      bOuter * ldb * 4,
+      "strmm-out",
+      GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+    );
+
     triParams = createParamsBuffer(
+      device,
       [
         { value: aOrder, type: "u32" },
         { value: lda, type: "u32" },
@@ -143,7 +217,11 @@ export async function strmm(
       ],
       "strmm-tri-params",
     );
-    const triBindGroup = createBindGroup(triPipeline.getBindGroupLayout(0), [ABuffer, AdenseBuffer, triParams]);
+    const triBindGroup = createBindGroup(
+      device,
+      triPipeline.getBindGroupLayout(0),
+      [ABuffer, AdenseBuffer, triParams],
+    );
 
     // X/Y buffers and their own ld, matching swapXY above.
     const XBuffer = swapXY ? BBuffer : AdenseBuffer;
@@ -152,12 +230,13 @@ export async function strmm(
     const ldY = swapXY ? ldDense : ldb;
 
     gemmParams = createParamsBuffer(
+      device,
       [
-        { value: mg,  type: "u32" },
-        { value: ng,  type: "u32" },
-        { value: kg,  type: "u32" },
+        { value: mg, type: "u32" },
+        { value: ng, type: "u32" },
+        { value: kg, type: "u32" },
         { value: alpha, type: "f32" },
-        { value: 0.0,   type: "f32" }, // beta — strmm has no C accumulation term
+        { value: 0.0, type: "f32" }, // beta — strmm has no C accumulation term
         { value: ldX, type: "u32" },
         { value: ldY, type: "u32" },
         { value: ldb, type: "u32" },
@@ -166,23 +245,58 @@ export async function strmm(
       ],
       "strmm-gemm-params",
     );
-    const gemmBindGroup = createBindGroup(gemmPipeline.getBindGroupLayout(0), [XBuffer, YBuffer, outBuffer, gemmParams]);
+    const gemmBindGroup = createBindGroup(
+      device,
+      gemmPipeline.getBindGroupLayout(0),
+      [
+        XBuffer,
+        vec4ViewBinding(device, XBuffer),
+        YBuffer,
+        vec4ViewBinding(device, YBuffer),
+        outBuffer,
+        gemmParams,
+      ],
+    );
 
-    const { commandEncoder, querySet } = beginTimedEncoder();
+    const { commandEncoder, querySet } = beginTimedEncoder(device);
     // Seed outBuffer with B's own bytes first, so gemm's tight m x n write
     // leaves stride-padding gaps holding B's original content, not zero.
     // BBuffer may be larger than outBuffer (e.g. a validation-test baseline
     // over-provisioned for a bigger ldb it might later be substituted with).
-    commandEncoder.copyBufferToBuffer(BBuffer, 0, outBuffer, 0, Math.min(BBuffer.size, outBuffer.size));
-    const triDesc = querySet ? { timestampWrites: { querySet, beginningOfPassWriteIndex: 0 } } : undefined;
-    const gemmDesc = querySet ? { timestampWrites: { querySet, endOfPassWriteIndex: 1 } } : undefined;
-    encodePass(commandEncoder, triPipeline, triBindGroup, { x: Math.ceil(aOrder / TRI_WG), y: Math.ceil(aOrder / TRI_WG) }, triDesc);
-    encodePass(commandEncoder, gemmPipeline, gemmBindGroup, gemmWgCount, gemmDesc);
+    commandEncoder.copyBufferToBuffer(
+      BBuffer,
+      0,
+      outBuffer,
+      0,
+      Math.min(BBuffer.size, outBuffer.size),
+    );
+    const triDesc = querySet
+      ? { timestampWrites: { querySet, beginningOfPassWriteIndex: 0 } }
+      : undefined;
+    const gemmDesc = querySet
+      ? { timestampWrites: { querySet, endOfPassWriteIndex: 1 } }
+      : undefined;
+    encodePass(
+      commandEncoder,
+      triPipeline,
+      triBindGroup,
+      { x: Math.ceil(aOrder / TILE_WG_2D), y: Math.ceil(aOrder / TILE_WG_2D) },
+      triDesc,
+    );
+    encodePass(
+      commandEncoder,
+      gemmPipeline,
+      gemmBindGroup,
+      gemmWgCount,
+      gemmDesc,
+    );
 
-    const ts = resolveTimestamp(commandEncoder, querySet);
-    const readBuffer = BIsGpu ? null : stageReadback(commandEncoder, outBuffer);
+    const ts = resolveTimestamp(device, commandEncoder, querySet);
+    const readBuffer = BIsGpu
+      ? null
+      : stageReadback(device, commandEncoder, outBuffer);
 
-    submit(commandEncoder);
+    submit(device, commandEncoder);
 
     const gpuTimeMs = await extractTimestamp(ts);
 
@@ -201,10 +315,10 @@ export async function strmm(
     if (gpuTimeMs !== undefined) return { B: result, gpuTimeMs };
     return { B: result };
   } finally {
-    if (!AIsGpu) destroyBuffers(ABuffer);
-    if (!BIsGpu) destroyBuffers(BBuffer);
-    destroyBuffers(AdenseBuffer);
-    if (!outBufferAdopted) destroyBuffers(outBuffer);
+    if (!AIsGpu && ABuffer) destroyBuffers(ABuffer);
+    if (!BIsGpu && BBuffer) destroyBuffers(BBuffer);
+    if (AdenseBuffer) destroyBuffers(AdenseBuffer);
+    if (outBuffer && !outBufferAdopted) destroyBuffers(outBuffer);
     if (triParams) destroyBuffers(triParams);
     if (gemmParams) destroyBuffers(gemmParams);
   }

package/src/strmv/strmv.d.mts

@@ -2,7 +2,7 @@ import { GpuVector } from "../classes/GpuVector.mjs";
 import { GpuMatrix } from "../classes/GpuMatrix.mjs";
 
 /**
- * Performs the triangular matrix-vector operation y = op(A) * x
+ * Performs the triangular matrix-vector operation $$y \leftarrow \mathrm{op}(A) x$$
  *
  * 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.
@@ -45,7 +45,7 @@ export declare function strmv(
 ): Promise<{ y: Float32Array; gpuTimeMs?: number }>;
 
 /**
- * Performs the triangular matrix-vector operation y = op(A) * x
+ * Performs the triangular matrix-vector operation $$y \leftarrow \mathrm{op}(A) x$$
  *
  * 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

package/src/strmv/strmv.mjs

@@ -11,15 +11,29 @@ import { extractTimestamp } from "../util/benchmark.mjs";
 import { getPipeline } from "../util/pipeline.mjs";
 import { GpuVector } from "../classes/GpuVector.mjs";
 import { GpuMatrix } from "../classes/GpuMatrix.mjs";
+import { requireGpuDevice, requireSameDevice } from "../util/device.mjs";
 
-export async function strmv(device, uplo, trans, diag, n, A, lda, x, incx, y, incy, layout = "row-major") {
+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.");
+  requireGpuDevice(device);
+  requireSameDevice(device, "strmv", { A, x, y });
   if (uplo !== "lower" && uplo !== "upper")
     throw new Error("uplo must be 'lower' or 'upper'.");
   if (trans !== "no-transpose" && trans !== "transpose")
@@ -66,9 +80,7 @@ export async function strmv(device, uplo, trans, diag, n, A, lda, x, incx, y, in
   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.",
-    );
+    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.",
@@ -82,7 +94,9 @@ export async function strmv(device, uplo, trans, diag, n, A, lda, x, incx, y, in
   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 isNoTrans = isColMajor
+    ? trans === "transpose"
+    : trans === "no-transpose";
 
   const pipeline = await getPipeline(device, "strmv");
 
@@ -92,23 +106,24 @@ export async function strmv(device, uplo, trans, diag, n, A, lda, x, incx, y, in
   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);
+    ABuffer = AIsGpu ? A._buf : uploadBuffer(device, A, "strmv-A", false);
+    xBuffer = xIsGpu ? x._buf : uploadBuffer(device, x, "strmv-x", false);
+    yBuffer = yIsGpu ? y._buf : uploadBuffer(device, y, "strmv-y", true);
     paramsBuffer = createParamsBuffer(
+      device,
       [
-        { value: n,             type: "u32" },
-        { value: incx,          type: "u32" },
-        { value: incy,          type: "u32" },
-        { value: lda,           type: "u32" },
+        { 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" },
+        { value: isLower ? 0 : 1, type: "u32" },
+        { value: isUnit ? 1 : 0, type: "u32" },
       ],
       "strmv-params",
     );
 
-    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+    const bindGroup = createBindGroup(device, pipeline.getBindGroupLayout(0), [
       ABuffer,
       xBuffer,
       yBuffer,
@@ -116,10 +131,17 @@ export async function strmv(device, uplo, trans, diag, n, A, lda, x, incx, y, in
     ]);
 
     const wgCount = Math.min(n, device.limits.maxComputeWorkgroupsPerDimension);
-    const { commandEncoder, ts } = runComputePass(pipeline, bindGroup, wgCount);
-    const readBuffer = yIsGpu ? null : stageReadback(commandEncoder, yBuffer);
+    const { commandEncoder, ts } = runComputePass(
+      device,
+      pipeline,
+      bindGroup,
+      wgCount,
+    );
+    const readBuffer = yIsGpu
+      ? null
+      : stageReadback(device, commandEncoder, yBuffer);
 
-    submit(commandEncoder);
+    submit(device, commandEncoder);
 
     const gpuTimeMs = await extractTimestamp(ts);
 

package/src/strsm/strsm.d.mts

@@ -2,8 +2,9 @@ import { GpuMatrix } from "../classes/GpuMatrix.mjs";
 
 /**
  * Solves the triangular matrix equation
- * op(A) * X = alpha * B (`side='left'`) or
- * X * op(A) = alpha * B (`side='right'`), overwriting `B` with the
+ * $$\mathrm{op}(A) X = \alpha B \quad (\texttt{side='left'})$$
+ * $$X \mathrm{op}(A) = \alpha B \quad (\texttt{side='right'})$$
+ * overwriting `B` with the
  * solution `X` — `A` is triangular, only its `uplo` triangle stored; `B` is
  * a general m×n matrix.
  *
@@ -57,8 +58,9 @@ export declare function strsm(
 
 /**
  * Solves the triangular matrix equation
- * op(A) * X = alpha * B (`side='left'`) or
- * X * op(A) = alpha * B (`side='right'`), overwriting `B` in place with `X`.
+ * $$\mathrm{op}(A) X = \alpha B \quad (\texttt{side='left'})$$
+ * $$X \mathrm{op}(A) = \alpha B \quad (\texttt{side='right'})$$
+ * overwriting `B` in place with `X`.
  *
  * A and B are both kept GPU-resident. Each matrix's own `layout` (set at
  * `GpuMatrix.from` time) determines the operation — there is no separate