wgblas 2.0.0 → 2.1.0

package/src/ssyr/ssyr.mjs

@@ -11,6 +11,7 @@ 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 { requireSameDevice } from "../util/device.mjs";
 
 export async function ssyr(device, uplo, n, alpha, x, incx, A, lda, layout = "row-major") {
   const xIsGpu = x instanceof GpuVector;
@@ -18,6 +19,7 @@ export async function ssyr(device, uplo, n, alpha, x, incx, A, lda, layout = "ro
 
   if (!(device instanceof GPUDevice))
     throw new Error("device must be a GPUDevice.");
+  requireSameDevice(device, "ssyr", { A, x });
   if (uplo !== "lower" && uplo !== "upper")
     throw new Error("uplo must be 'lower' or 'upper'.");
   if (layout !== "row-major" && layout !== "column-major")
@@ -63,9 +65,9 @@ export async function ssyr(device, uplo, n, alpha, x, incx, A, lda, layout = "ro
   let paramsBuffer = null;
 
   try {
-    xBuffer = xIsGpu ? x._buf : uploadBuffer(x, "ssyr-x", false);
-    ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "ssyr-A", true);
-    paramsBuffer = createParamsBuffer(
+    xBuffer = xIsGpu ? x._buf : uploadBuffer(device, x, "ssyr-x", false);
+    ABuffer = AIsGpu ? A._buf : uploadBuffer(device, A, "ssyr-A", true);
+    paramsBuffer = createParamsBuffer(device,
       [
         { value: n,               type: "u32" },
         { value: alpha,           type: "f32" },
@@ -76,7 +78,7 @@ export async function ssyr(device, uplo, n, alpha, x, incx, A, lda, layout = "ro
       "ssyr-params",
     );
 
-    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+    const bindGroup = createBindGroup(device, pipeline.getBindGroupLayout(0), [
       xBuffer,
       ABuffer,
       paramsBuffer,
@@ -85,10 +87,10 @@ export async function ssyr(device, uplo, n, alpha, x, incx, A, lda, layout = "ro
     // One workgroup per row of A; clamped to device limit — the shader's
     // grid-stride loop handles remaining rows when n > dispatch count.
     const wgCount = Math.min(n, device.limits.maxComputeWorkgroupsPerDimension);
-    const { commandEncoder, ts } = runComputePass(pipeline, bindGroup, wgCount);
-    const readBuffer = AIsGpu ? null : stageReadback(commandEncoder, ABuffer);
+    const { commandEncoder, ts } = runComputePass(device, pipeline, bindGroup, wgCount);
+    const readBuffer = AIsGpu ? null : stageReadback(device, commandEncoder, ABuffer);
 
-    submit(commandEncoder);
+    submit(device, commandEncoder);
 
     const gpuTimeMs = await extractTimestamp(ts);
 

package/src/ssyr2/ssyr2.mjs

@@ -11,6 +11,7 @@ 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 { requireSameDevice } from "../util/device.mjs";
 
 export async function ssyr2(device, uplo, n, alpha, x, incx, y, incy, A, lda, layout = "row-major") {
   const xIsGpu = x instanceof GpuVector;
@@ -19,6 +20,7 @@ export async function ssyr2(device, uplo, n, alpha, x, incx, y, incy, A, lda, la
 
   if (!(device instanceof GPUDevice))
     throw new Error("device must be a GPUDevice.");
+  requireSameDevice(device, "ssyr2", { A, x, y });
   if (uplo !== "lower" && uplo !== "upper")
     throw new Error("uplo must be 'lower' or 'upper'.");
   if (layout !== "row-major" && layout !== "column-major")
@@ -83,10 +85,10 @@ export async function ssyr2(device, uplo, n, alpha, x, incx, y, incy, A, lda, la
   let paramsBuffer = null;
 
   try {
-    xBuffer = xIsGpu ? x._buf : uploadBuffer(x, "ssyr2-x", false);
-    yBuffer = yIsGpu ? y._buf : uploadBuffer(y, "ssyr2-y", false);
-    ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "ssyr2-A", true);
-    paramsBuffer = createParamsBuffer(
+    xBuffer = xIsGpu ? x._buf : uploadBuffer(device, x, "ssyr2-x", false);
+    yBuffer = yIsGpu ? y._buf : uploadBuffer(device, y, "ssyr2-y", false);
+    ABuffer = AIsGpu ? A._buf : uploadBuffer(device, A, "ssyr2-A", true);
+    paramsBuffer = createParamsBuffer(device,
       [
         { value: n,               type: "u32" },
         { value: alpha,           type: "f32" },
@@ -98,7 +100,7 @@ export async function ssyr2(device, uplo, n, alpha, x, incx, y, incy, A, lda, la
       "ssyr2-params",
     );
 
-    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+    const bindGroup = createBindGroup(device, pipeline.getBindGroupLayout(0), [
       xBuffer,
       yBuffer,
       ABuffer,
@@ -108,10 +110,10 @@ export async function ssyr2(device, uplo, n, alpha, x, incx, y, incy, A, lda, la
     // One workgroup per row of A; clamped to device limit — the shader's
     // grid-stride loop handles remaining rows when n > dispatch count.
     const wgCount = Math.min(n, device.limits.maxComputeWorkgroupsPerDimension);
-    const { commandEncoder, ts } = runComputePass(pipeline, bindGroup, wgCount);
-    const readBuffer = AIsGpu ? null : stageReadback(commandEncoder, ABuffer);
+    const { commandEncoder, ts } = runComputePass(device, pipeline, bindGroup, wgCount);
+    const readBuffer = AIsGpu ? null : stageReadback(device, commandEncoder, ABuffer);
 
-    submit(commandEncoder);
+    submit(device, commandEncoder);
 
     const gpuTimeMs = await extractTimestamp(ts);
 

package/src/ssyr2k/ssyr2k.mjs

@@ -10,10 +10,10 @@ 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";
+import { requireWorkgroupCount } from "../util/workgroup.mjs";
+import { BM_SMALL, BN_SMALL, BM_LARGE, BN_LARGE, LARGE_TILE_WORKGROUP_THRESHOLD } from "../util/constants.mjs";
+import { requireSameDevice } from "../util/device.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
 
 // 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.
@@ -26,6 +26,7 @@ export async function ssyr2k(
 
   if (!(device instanceof GPUDevice))
     throw new Error("device must be a GPUDevice.");
+  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")
@@ -131,24 +132,24 @@ export async function ssyr2k(
   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);
+  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(
+    const makeParams = (p, betaVal) => createParamsBuffer(device,
       [
         { value: n,   type: "u32" },
         { value: n,   type: "u32" },
@@ -167,19 +168,19 @@ export async function ssyr2k(
     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 { 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.mjs

@@ -11,10 +11,10 @@ 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";
+import { requireWorkgroupCount } from "../util/workgroup.mjs";
+import { BM_SMALL, BN_SMALL, BM_LARGE, BN_LARGE, LARGE_TILE_WORKGROUP_THRESHOLD } from "../util/constants.mjs";
+import { requireSameDevice } from "../util/device.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
 
 // 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).
@@ -26,6 +26,7 @@ export async function ssyrk(
 
   if (!(device instanceof GPUDevice))
     throw new Error("device must be a GPUDevice.");
+  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")
@@ -106,14 +107,14 @@ export async function ssyrk(
 
   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);
-  const paramsBuffer = createParamsBuffer(
+    ? 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
@@ -131,7 +132,7 @@ export async function ssyrk(
   );
 
   try {
-    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+    const bindGroup = createBindGroup(device, pipeline.getBindGroupLayout(0), [
       ABuffer,
       BBuffer,
       CBuffer,
@@ -140,22 +141,22 @@ 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();
+    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.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,11 +12,11 @@ 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";
+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 { requireSameDevice } from "../util/device.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)
 
 // 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
@@ -30,6 +31,7 @@ export async function strmm(
 
   if (!(device instanceof GPUDevice))
     throw new Error("device must be a GPUDevice.");
+  requireSameDevice(device, "strmm", { A, B });
   if (side !== "left" && side !== "right")
     throw new Error("side must be 'left' or 'right'.");
   if (uplo !== "lower" && uplo !== "upper")
@@ -110,29 +112,35 @@ export async function strmm(
   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),
+      x: requireWorkgroupCount(device, Math.ceil(ng / BN_SMALL), "strmm", "x"),
+      y: requireWorkgroupCount(device, Math.ceil(mg / BM_SMALL), "strmm", "y"),
     };
 
-  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,
-  );
+  // 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 {
-    triParams = createParamsBuffer(
+    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 +151,7 @@ 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;
@@ -151,7 +159,7 @@ export async function strmm(
     const YBuffer = swapXY ? AdenseBuffer : BBuffer;
     const ldY = swapXY ? ldDense : ldb;
 
-    gemmParams = createParamsBuffer(
+    gemmParams = createParamsBuffer(device,
       [
         { value: mg,  type: "u32" },
         { value: ng,  type: "u32" },
@@ -166,9 +174,16 @@ export async function strmm(
       ],
       "strmm-gemm-params",
     );
-    const gemmBindGroup = createBindGroup(gemmPipeline.getBindGroupLayout(0), [XBuffer, YBuffer, outBuffer, gemmParams]);
-
-    const { commandEncoder, querySet } = beginTimedEncoder();
+    const gemmBindGroup = createBindGroup(device, gemmPipeline.getBindGroupLayout(0), [
+      XBuffer,
+      vec4ViewBinding(device, XBuffer),
+      YBuffer,
+      vec4ViewBinding(device, YBuffer),
+      outBuffer,
+      gemmParams,
+    ]);
+
+    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
@@ -176,13 +191,13 @@ export async function strmm(
     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, 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 +216,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.mjs

@@ -11,6 +11,7 @@ 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 { requireSameDevice } from "../util/device.mjs";
 
 export async function strmv(device, uplo, trans, diag, n, A, lda, x, incx, y, incy, layout = "row-major") {
   const xIsGpu = x instanceof GpuVector;
@@ -20,6 +21,7 @@ export async function strmv(device, uplo, trans, diag, n, A, lda, x, incx, y, in
 
   if (!(device instanceof GPUDevice))
     throw new Error("device must be a GPUDevice.");
+  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")
@@ -92,10 +94,10 @@ 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);
-    paramsBuffer = createParamsBuffer(
+    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" },
@@ -108,7 +110,7 @@ export async function strmv(device, uplo, trans, diag, n, A, lda, x, incx, y, in
       "strmv-params",
     );
 
-    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+    const bindGroup = createBindGroup(device, pipeline.getBindGroupLayout(0), [
       ABuffer,
       xBuffer,
       yBuffer,
@@ -116,10 +118,10 @@ 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);