wgblas 1.2.1 → 2.0.0

package/src/ssyr2k/ssyr2k.mjs

@@ -0,0 +1,201 @@
+import {
+  uploadBuffer,
+  createParamsBuffer,
+  stageReadback,
+  destroyBuffers,
+} from "../util/buffer.mjs";
+import { createBindGroup } from "../util/bindgroup.mjs";
+import { beginTimedEncoder, encodePass, submit } from "../util/compute.mjs";
+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
+
+// 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",
+) {
+  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.");
+  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 (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 (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)
+  )
+    throw new Error("n, k, lda, ldb, and ldc must be integers.");
+  if (!AIsGpu && !(A instanceof Float32Array))
+    throw new Error("A must be a Float32Array or GpuMatrix.");
+  if (!BIsGpu && !(B instanceof Float32Array))
+    throw new Error("B must be a Float32Array or GpuMatrix.");
+  if (!CIsGpu && !(C instanceof Float32Array))
+    throw new Error("C must be a Float32Array or GpuMatrix.");
+  if ((AIsGpu || BIsGpu) && !CIsGpu)
+    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 (n < 0 || k < 0) throw new Error("n and k must be non-negative.");
+  if (n === 0) return CIsGpu ? {} : { C };
+
+  const effLayoutA = AIsGpu ? A.layout : layout;
+  const effLayoutB = BIsGpu ? B.layout : layout;
+  const effLayoutC = CIsGpu ? C.layout : layout;
+
+  // A: op(A) is n x k; A itself is n x k or k x n depending on trans.
+  const aRows = effLayoutA === "column-major" ? k : n;
+  const aCols = effLayoutA === "column-major" ? n : k;
+  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.`);
+  if (AIsGpu) {
+    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.");
+  }
+
+  // B: same shape rule as A — netlib's syr2k shares one trans across both operands.
+  const bRows = effLayoutB === "column-major" ? k : n;
+  const bCols = effLayoutB === "column-major" ? n : k;
+  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.`);
+  if (BIsGpu) {
+    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.");
+  }
+
+  // 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.");
+  } else if (C.length < (n - 1) * ldc + n) {
+    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.
+  let effTransA = trans;
+  if (effLayoutA === "column-major")
+    effTransA = effTransA === "no-transpose" ? "transpose" : "no-transpose";
+  let effTransB = trans;
+  if (effLayoutB === "column-major")
+    effTransB = effTransB === "no-transpose" ? "transpose" : "no-transpose";
+
+  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 };
+  }
+
+  // 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 wgCount = useLargeTile
+    ? {
+      x: Math.min(largeWgX, device.limits.maxComputeWorkgroupsPerDimension),
+      y: Math.min(largeWgY, device.limits.maxComputeWorkgroupsPerDimension),
+    }
+    : {
+      x: Math.min(Math.ceil(n / BN_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
+      y: Math.min(Math.ceil(n / BM_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
+    };
+
+  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;
+
+  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",
+    );
+    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 { commandEncoder, querySet } = beginTimedEncoder();
+    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);
+
+    submit(commandEncoder);
+
+    const gpuTimeMs = await extractTimestamp(ts);
+
+    if (CIsGpu) {
+      if (gpuTimeMs !== undefined) return { gpuTimeMs };
+      return {};
+    }
+
+    const result = await extractResult(readBuffer, Float32Array);
+    if (gpuTimeMs !== undefined) return { C: result, gpuTimeMs };
+    return { C: result };
+  } finally {
+    if (!AIsGpu) destroyBuffers(ABuffer);
+    if (!BIsGpu) destroyBuffers(BBuffer);
+    if (!CIsGpu) destroyBuffers(CBuffer);
+    if (paramsBuffer1) destroyBuffers(paramsBuffer1);
+    if (paramsBuffer2) destroyBuffers(paramsBuffer2);
+  }
+}

package/src/ssyrk/ssyrk.d.mts

@@ -0,0 +1,90 @@
+import { GpuMatrix } from "../classes/GpuMatrix.mjs";
+
+/**
+ * Performs the symmetric rank-k update C := uplo(alpha * op(A) * 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.
+ *
+ * - `trans='no-transpose'`: op(A) = A (n×k), computes alpha*A*A^T + beta*C
+ * - `trans='transpose'`:    op(A) = A^T (A stored k×n), computes alpha*A^T*A + beta*C
+ *
+ * No dedicated kernel — this is `sgemmtr`'s exact kernel
+ * (`sgemmtr_small.wgsl`/`sgemmtr_large.wgsl`) with `A` aliased into both
+ * operand slots (`B := A`).
+ *
+ * {@includeCode ../../examples/ssyrk/ssyrk.js}
+ *
+ * **Browser (standalone HTML):**
+ * {@includeCode ../../examples/ssyrk/web/ssyrk.html}
+ *
+ * @param device - GPUDevice from `init()`
+ * @param uplo   - `'lower'` to update only `col <= row`, `'upper'` for `col >= row`
+ * @param trans  - `'no-transpose'` for A*A^T, `'transpose'` for A^T*A
+ * @param n      - order of C (C is n×n); rows of op(A)
+ * @param k      - columns of op(A) — the shared/contracted dimension
+ * @param alpha  - scalar multiplier for op(A)*op(A)^T
+ * @param A      - Float32Array, row-major or column-major (see `layout`)
+ * @param lda    - leading dimension of A as stored
+ * @param beta   - scalar multiplier for C
+ * @param C      - Float32Array input/output matrix, row-major or column-major
+ * @param ldc    - leading dimension of C as stored
+ * @param layout - storage layout shared by A/C when they're Float32Array
+ *   (default: `'row-major'`) — same handling as `sgemmtr`; column-major C
+ *   flips `uplo` internally so it still names the triangle you asked for
+ * @returns updated C as a Float32Array (only the requested triangle changed)
+ * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/ssyrk/ssyrk.mjs#L21">Source code: ssyrk.mjs (L21)</a>
+ * @category BLAS Level 3
+ */
+export declare function ssyrk(
+  device: GPUDevice,
+  uplo: 'lower' | 'upper',
+  trans: 'no-transpose' | 'transpose',
+  n: number,
+  k: number,
+  alpha: number,
+  A: Float32Array,
+  lda: number,
+  beta: number,
+  C: Float32Array,
+  ldc: number,
+  layout?: 'row-major' | 'column-major',
+): Promise<{ C: Float32Array; gpuTimeMs?: number }>;
+
+/**
+ * Performs the symmetric rank-k update C := uplo(alpha * op(A) * 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
+ * `layout` argument here. A must be a GpuMatrix whenever C is, and vice
+ * versa — mixing a GpuMatrix with a plain Float32Array is not supported.
+ *
+ * {@includeCode ../../examples/ssyrk/gpu.ssyrk.js}
+ *
+ * @param device - GPUDevice from `init()`
+ * @param uplo   - `'lower'` to update only `col <= row`, `'upper'` for `col >= row`
+ * @param trans  - `'no-transpose'` for A*A^T, `'transpose'` for A^T*A
+ * @param n      - order of C (C is n×n); rows of op(A)
+ * @param k      - columns of op(A) — the shared/contracted dimension
+ * @param alpha  - scalar multiplier for op(A)*op(A)^T
+ * @param A      - GpuMatrix
+ * @param lda    - leading dimension of A (must equal A.lda)
+ * @param beta   - scalar multiplier for C
+ * @param C      - GpuMatrix (mutated in place; only the requested triangle changes)
+ * @param ldc    - leading dimension of C (must equal C.lda)
+ * @returns no C — it stays GPU-resident; call `C.read()` yourself for a CPU readback (see the example)
+ * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/ssyrk/ssyrk.mjs#L21">Source code: ssyrk.mjs (L21)</a>
+ * @category BLAS Level 3
+ */
+export declare function ssyrk(
+  device: GPUDevice,
+  uplo: 'lower' | 'upper',
+  trans: 'no-transpose' | 'transpose',
+  n: number,
+  k: number,
+  alpha: number,
+  A: GpuMatrix,
+  lda: number,
+  beta: number,
+  C: GpuMatrix,
+  ldc: number,
+): Promise<{ gpuTimeMs?: number }>;

package/src/ssyrk/ssyrk.mjs

@@ -0,0 +1,176 @@
+import {
+  uploadBuffer,
+  createStorageBuffer,
+  createParamsBuffer,
+  stageReadback,
+  destroyBuffers,
+} from "../util/buffer.mjs";
+import { createBindGroup } from "../util/bindgroup.mjs";
+import { beginTimedEncoder, encodePass, submit } from "../util/compute.mjs";
+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
+
+// 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",
+) {
+  const AIsGpu = A instanceof GpuMatrix;
+  const CIsGpu = C instanceof GpuMatrix;
+
+  if (!(device instanceof GPUDevice))
+    throw new Error("device must be a GPUDevice.");
+  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 (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 (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))
+    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.");
+  if (!CIsGpu && !(C instanceof Float32Array))
+    throw new Error("C must be a Float32Array or GpuMatrix.");
+  if (AIsGpu && !CIsGpu)
+    throw new Error("C must be a GpuMatrix when A is a GpuMatrix.");
+  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 (n === 0) return CIsGpu ? {} : { C };
+
+  // GpuMatrix's own .layout wins over the shared `layout` argument.
+  const effLayoutA = AIsGpu ? A.layout : layout;
+  const effLayoutC = CIsGpu ? C.layout : layout;
+
+  // A: op(A) is n x k; A itself is n x k or k x n depending on trans.
+  const aRows = effLayoutA === "column-major" ? k : n;
+  const aCols = effLayoutA === "column-major" ? n : k;
+  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.`);
+  if (AIsGpu) {
+    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.");
+  }
+
+  // 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.");
+  } else if (C.length < (n - 1) * ldc + n) {
+    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.
+  let transA = trans;
+  if (effLayoutA === "column-major")
+    transA = transA === "no-transpose" ? "transpose" : "no-transpose";
+  // op(B) := op(A)^T (B is A itself) — opposite of transA.
+  let transB = transA === "no-transpose" ? "transpose" : "no-transpose";
+
+  // Column-major C: same swap trick sgemm/sgemmtr use, simplified since A === B and m=n.
+  let uploEff = uplo;
+  if (effLayoutC === "column-major") {
+    [transA, transB] = [
+      transB === "no-transpose" ? "transpose" : "no-transpose",
+      transA === "no-transpose" ? "transpose" : "no-transpose",
+    ];
+    uploEff = uploEff === "lower" ? "upper" : "lower";
+  }
+
+  // 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 ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "ssyrk-A", false);
+  const CBuffer = CIsGpu ? C._buf : uploadBuffer(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(
+    [
+      { 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: lda, type: "u32" }, // gemmtr's lda
+      { value: lda, type: "u32" }, // gemmtr's ldb — B := A, same lda
+      { value: ldc, type: "u32" },
+      { value: transA === "transpose" ? 1 : 0, type: "u32" },
+      { value: transB === "transpose" ? 1 : 0, type: "u32" },
+      { value: uploEff === "upper" ? 1 : 0, type: "u32" },
+    ],
+    "ssyrk-params",
+  );
+
+  try {
+    const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
+      ABuffer,
+      BBuffer,
+      CBuffer,
+      paramsBuffer,
+    ]);
+
+    const wgCount = useLargeTile
+      ? {
+        x: Math.min(largeWgX, device.limits.maxComputeWorkgroupsPerDimension),
+        y: Math.min(largeWgY, device.limits.maxComputeWorkgroupsPerDimension),
+      }
+      : {
+        x: Math.min(Math.ceil(n / BN_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
+        y: Math.min(Math.ceil(n / BM_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
+      };
+    // 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);
+    encodePass(commandEncoder, pipeline, bindGroup, wgCount, passDescriptor);
+    const ts = resolveTimestamp(commandEncoder, querySet);
+    const readBuffer = CIsGpu ? null : stageReadback(commandEncoder, CBuffer);
+
+    submit(commandEncoder);
+
+    const gpuTimeMs = await extractTimestamp(ts);
+
+    if (CIsGpu) {
+      if (gpuTimeMs !== undefined) return { gpuTimeMs };
+      return {};
+    }
+
+    const result = await extractResult(readBuffer, Float32Array);
+    if (gpuTimeMs !== undefined) return { C: result, gpuTimeMs };
+    return { C: result };
+  } finally {
+    if (!AIsGpu) destroyBuffers(ABuffer);
+    destroyBuffers(BBuffer); // always our own buffer, never the caller's
+    if (!CIsGpu) destroyBuffers(CBuffer);
+    destroyBuffers(paramsBuffer);
+  }
+}

package/src/strmm/strmm.d.mts

@@ -0,0 +1,100 @@
+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
+ * `uplo` triangle stored; `B` is a general m×n matrix, overwritten in place.
+ *
+ * - `side='left'`:  `A` is m×m — `A` premultiplies `B`
+ * - `side='right'`: `A` is n×n — `A` postmultiplies `B`
+ *
+ * No dedicated fused kernel — a `triangularize` pass materializes a dense
+ * copy of `op(A)` (zero-filling the unstored triangle, substituting the
+ * implicit diagonal when `diag='unit'`), then a plain `sgemm` pass
+ * (`sgemm_small.wgsl`/`sgemm_large.wgsl`, unmodified) does the actual
+ * multiply, both on one command encoder.
+ *
+ * {@includeCode ../../examples/strmm/strmm.js}
+ *
+ * **Browser (standalone HTML):**
+ * {@includeCode ../../examples/strmm/web/strmm.html}
+ *
+ * @param device - GPUDevice from `init()`
+ * @param side   - `'left'` for op(A)*B, `'right'` for B*op(A)
+ * @param uplo   - `'lower'` if only `A`'s lower triangle is stored, `'upper'` for upper
+ * @param transA - `'no-transpose'` for A, `'transpose'` for A^T
+ * @param diag   - `'unit'` to treat A's diagonal as all-ones (A's diagonal is not read), `'non-unit'` to read it
+ * @param m      - rows of B
+ * @param n      - columns of B
+ * @param alpha  - scalar multiplier for the matrix product
+ * @param A      - Float32Array, triangular, row-major or column-major (see `layout`)
+ * @param lda    - leading dimension of A as stored
+ * @param B      - Float32Array input/output matrix, overwritten with the result, row-major or column-major
+ * @param ldb    - leading dimension of B as stored
+ * @param layout - storage layout shared by A/B when they're Float32Array
+ *   (default: `'row-major'`) — column-major A is a genuine transpose (A isn't
+ *   symmetric like ssymm's), so both `transA` and `uplo` are adjusted
+ *   internally to compensate
+ * @returns updated B as a Float32Array
+ * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/strmm/strmm.mjs#L25">Source code: strmm.mjs (L25)</a>
+ * @category BLAS Level 3
+ */
+export declare function strmm(
+  device: GPUDevice,
+  side: 'left' | 'right',
+  uplo: 'lower' | 'upper',
+  transA: 'no-transpose' | 'transpose',
+  diag: 'unit' | 'non-unit',
+  m: number,
+  n: number,
+  alpha: number,
+  A: Float32Array,
+  lda: number,
+  B: Float32Array,
+  ldb: number,
+  layout?: 'row-major' | 'column-major',
+): Promise<{ B: Float32Array; gpuTimeMs?: number }>;
+
+/**
+ * Performs the triangular matrix-matrix operation
+ * B := alpha * op(A) * B (`side='left'`) or
+ * B := alpha * B * op(A) (`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 —
+ * there is no separate `layout` argument here. A and B must both be
+ * GpuMatrix or both be Float32Array — mixing is not supported.
+ *
+ * {@includeCode ../../examples/strmm/gpu.strmm.js}
+ *
+ * @param device - GPUDevice from `init()`
+ * @param side   - `'left'` for op(A)*B, `'right'` for B*op(A)
+ * @param uplo   - `'lower'` if only `A`'s lower triangle is stored, `'upper'` for upper
+ * @param transA - `'no-transpose'` for A, `'transpose'` for A^T
+ * @param diag   - `'unit'` to treat A's diagonal as all-ones (A's diagonal is not read), `'non-unit'` to read it
+ * @param m      - rows of B
+ * @param n      - columns of B
+ * @param alpha  - scalar multiplier for the matrix product
+ * @param A      - GpuMatrix, triangular
+ * @param lda    - leading dimension of A (must equal A.lda)
+ * @param B      - GpuMatrix (mutated in place)
+ * @param ldb    - leading dimension of B (must equal B.lda)
+ * @returns no B — it stays GPU-resident; call `B.read()` yourself for a CPU readback (see the example)
+ * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/strmm/strmm.mjs#L25">Source code: strmm.mjs (L25)</a>
+ * @category BLAS Level 3
+ */
+export declare function strmm(
+  device: GPUDevice,
+  side: 'left' | 'right',
+  uplo: 'lower' | 'upper',
+  transA: 'no-transpose' | 'transpose',
+  diag: 'unit' | 'non-unit',
+  m: number,
+  n: number,
+  alpha: number,
+  A: GpuMatrix,
+  lda: number,
+  B: GpuMatrix,
+  ldb: number,
+): Promise<{ gpuTimeMs?: number }>;