wgblas 1.2.1 → 2.0.0

package/src/sgemm/sgemm.mjs

@@ -0,0 +1,195 @@
+import {
+  uploadBuffer,
+  createParamsBuffer,
+  stageReadback,
+  destroyBuffers,
+} from "../util/buffer.mjs";
+import { createBindGroup } from "../util/bindgroup.mjs";
+import { runComputePass, submit } from "../util/compute.mjs";
+import { extractResult } from "../util/result.mjs";
+import { extractTimestamp } from "../util/benchmark.mjs";
+import { getPipeline } from "../util/pipeline.mjs";
+import { GpuMatrix } from "../classes/GpuMatrix.mjs";
+
+const BM_SMALL = 32, BN_SMALL = 32; // sgemm_small.wgsl's block tile
+const BM_LARGE = 64, BN_LARGE = 64; // sgemm_large.wgsl's block tile
+const LARGE_TILE_WORKGROUP_THRESHOLD = 36; // large tile needs >= a 6x6 grid of its own tiles to beat the small tile
+
+export async function sgemm(
+  device, transA, transB, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc, layout = "row-major",
+) {
+  let AIsGpu = A instanceof GpuMatrix;
+  let BIsGpu = B instanceof GpuMatrix;
+  const CIsGpu = C instanceof GpuMatrix;
+
+  if (!(device instanceof GPUDevice))
+    throw new Error("device must be a GPUDevice.");
+  if (transA !== "no-transpose" && transA !== "transpose")
+    throw new Error("transA must be 'no-transpose' or 'transpose'.");
+  if (transB !== "no-transpose" && transB !== "transpose")
+    throw new Error("transB must be 'no-transpose' or 'transpose'.");
+  if (layout !== "row-major" && layout !== "column-major")
+    throw new Error("layout must be 'row-major' or 'column-major'.");
+  if (typeof alpha !== "number")
+    throw new Error("alpha must be a number.");
+  if (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(m) ||
+    !Number.isInteger(n) ||
+    !Number.isInteger(k) ||
+    !Number.isInteger(lda) ||
+    !Number.isInteger(ldb) ||
+    !Number.isInteger(ldc)
+  )
+    throw new Error("m, 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 (m < 0 || n < 0 || k < 0) throw new Error("m, n, and k must be non-negative.");
+  if (m === 0 || n === 0) return CIsGpu ? {} : { C };
+
+  // GpuMatrix's own .layout wins over the shared `layout` argument.
+  const effLayoutA = AIsGpu ? A.layout : layout;
+  const effLayoutB = BIsGpu ? B.layout : layout;
+  const effLayoutC = CIsGpu ? C.layout : layout;
+
+  // Shape validation, before any layout-driven swapping below.
+  // A: op(A) is m x k; A itself is m x k or k x m depending on transA.
+  const aRows = effLayoutA === "column-major" ? k : m;
+  const aCols = effLayoutA === "column-major" ? m : k;
+  const aOuter = transA === "no-transpose" ? aRows : aCols; // # of stored chunks
+  const aInner = transA === "no-transpose" ? aCols : aRows; // required chunk length
+  if (lda < aInner)
+    throw new Error(`lda must be >= ${effLayoutA === "column-major" ? "rows" : "cols"} of A as stored.`);
+  if (AIsGpu) {
+    if (lda !== A.lda) throw new Error("lda must match A.lda when A is a GpuMatrix.");
+    const [aLogRows, aLogCols] = transA === "no-transpose" ? [m, k] : [k, m];
+    if (A.rows < aLogRows || A.cols < aLogCols)
+      throw new Error("A is too small for the given m, k, and transA.");
+  } else if (A.length < (aOuter - 1) * lda + aInner) {
+    throw new Error("A does not have enough elements for the given dimensions and lda.");
+  }
+
+  // B: same reasoning as A, with op(B) = k x n.
+  const bRows = effLayoutB === "column-major" ? n : k;
+  const bCols = effLayoutB === "column-major" ? k : n;
+  const bOuter = transB === "no-transpose" ? bRows : bCols;
+  const bInner = transB === "no-transpose" ? bCols : bRows;
+  if (ldb < bInner)
+    throw new Error(`ldb must be >= ${effLayoutB === "column-major" ? "rows" : "cols"} of B as stored.`);
+  if (BIsGpu) {
+    if (ldb !== B.lda) throw new Error("ldb must match B.lda when B is a GpuMatrix.");
+    const [bLogRows, bLogCols] = transB === "no-transpose" ? [k, n] : [n, k];
+    if (B.rows < bLogRows || B.cols < bLogCols)
+      throw new Error("B is too small for the given n, k, and transB.");
+  } else if (B.length < (bOuter - 1) * ldb + bInner) {
+    throw new Error("B does not have enough elements for the given dimensions and ldb.");
+  }
+
+  // C: always m x n (no trans flag) — layout only affects lda/storage order.
+  const cOuter = effLayoutC === "column-major" ? n : m;
+  const cInner = effLayoutC === "column-major" ? m : n;
+  if (ldc < cInner)
+    throw new Error(`ldc must be >= ${effLayoutC === "column-major" ? "rows" : "cols"} of C as stored.`);
+  if (CIsGpu) {
+    if (ldc !== C.lda) throw new Error("ldc must match C.lda when C is a GpuMatrix.");
+    if (C.rows < m || C.cols < n) throw new Error("C is too small for the given m and n.");
+  } else if (C.length < (cOuter - 1) * ldc + cInner) {
+    throw new Error("C does not have enough elements for the given dimensions and ldc.");
+  }
+
+  // Column-major A/B reinterpreted row-major is A^T/B^T — flip the trans flag.
+  if (effLayoutA === "column-major")
+    transA = transA === "no-transpose" ? "transpose" : "no-transpose";
+  if (effLayoutB === "column-major")
+    transB = transB === "no-transpose" ? "transpose" : "no-transpose";
+
+  // Column-major C: compute C^T = op(B)^T * op(A)^T instead (swap A/B, flip trans, swap m<->n).
+  if (effLayoutC === "column-major") {
+    [A, B] = [B, A];
+    [AIsGpu, BIsGpu] = [BIsGpu, AIsGpu];
+    [lda, ldb] = [ldb, lda];
+    [transA, transB] = [
+      transB === "no-transpose" ? "transpose" : "no-transpose",
+      transA === "no-transpose" ? "transpose" : "no-transpose",
+    ];
+    [m, n] = [n, m];
+  }
+
+  // Shape-based auto-select — see sgemm_small.wgsl/sgemm_large.wgsl.
+  const largeWgX = Math.ceil(n / BN_LARGE);
+  const largeWgY = Math.ceil(m / BM_LARGE);
+  const useLargeTile = largeWgX * largeWgY >= LARGE_TILE_WORKGROUP_THRESHOLD;
+
+  const pipeline = await getPipeline(device, useLargeTile ? "sgemm_large" : "sgemm_small");
+
+  const ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "sgemm-A", false);
+  const BBuffer = BIsGpu ? B._buf : uploadBuffer(B, "sgemm-B", false);
+  const CBuffer = CIsGpu ? C._buf : uploadBuffer(C, "sgemm-C", true);
+  const paramsBuffer = createParamsBuffer(
+    [
+      { value: m,   type: "u32" },
+      { value: n,   type: "u32" },
+      { value: k,   type: "u32" },
+      { value: alpha, type: "f32" },
+      { value: beta,  type: "f32" },
+      { value: lda, type: "u32" },
+      { value: ldb, type: "u32" },
+      { value: ldc, type: "u32" },
+      { value: transA === "transpose" ? 1 : 0, type: "u32" },
+      { value: transB === "transpose" ? 1 : 0, type: "u32" },
+    ],
+    "sgemm-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(m / BM_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
+      };
+    const { commandEncoder, ts } = runComputePass(pipeline, bindGroup, wgCount);
+    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);
+    destroyBuffers(paramsBuffer);
+  }
+}

package/src/sgemmtr/sgemmtr.d.mts

@@ -0,0 +1,104 @@
+import { GpuMatrix } from "../classes/GpuMatrix.mjs";
+
+/**
+ * Performs the matrix-matrix operation C := uplo(alpha * op(A) * op(B) + beta * C) —
+ * `sgemm`'s operation, but only the triangle of C named by `uplo` is read or
+ * written (`'lower'`: `col <= row`, `'upper'`: `col >= row`; C need not be
+ * square — the test applies over the full m×n grid).
+ *
+ * Same kernels as `sgemm` (`sgemmtr_small.wgsl`/`sgemmtr_large.wgsl`,
+ * identical tiling), with the final output write masked to one triangle.
+ *
+ * {@includeCode ../../examples/sgemmtr/sgemmtr.js}
+ *
+ * **Browser (standalone HTML):**
+ * {@includeCode ../../examples/sgemmtr/web/sgemmtr.html}
+ *
+ * @param device - GPUDevice from `init()`
+ * @param uplo   - `'lower'` to update only `col <= row`, `'upper'` for `col >= row`
+ * @param transA - `'no-transpose'` for A, `'transpose'` for A^T
+ * @param transB - `'no-transpose'` for B, `'transpose'` for B^T
+ * @param m      - rows of op(A) and C
+ * @param n      - columns of op(B) and C
+ * @param k      - columns of op(A), rows of op(B)
+ * @param alpha  - scalar multiplier for op(A)*op(B)
+ * @param A      - Float32Array, row-major or column-major (see `layout`)
+ * @param lda    - leading dimension of A as stored
+ * @param B      - Float32Array, row-major or column-major (see `layout`)
+ * @param ldb    - leading dimension of B 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/B/C when they're Float32Array
+ *   (default: `'row-major'`) — same handling as `sgemm`, plus `uplo` is
+ *   flipped internally for column-major C 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/sgemmtr/sgemmtr.mjs#L19">Source code: sgemmtr.mjs (L19)</a>
+ * @category BLAS Level 3
+ */
+export declare function sgemmtr(
+  device: GPUDevice,
+  uplo: 'lower' | 'upper',
+  transA: 'no-transpose' | 'transpose',
+  transB: 'no-transpose' | 'transpose',
+  m: number,
+  n: number,
+  k: number,
+  alpha: number,
+  A: Float32Array,
+  lda: number,
+  B: Float32Array,
+  ldb: number,
+  beta: number,
+  C: Float32Array,
+  ldc: number,
+  layout?: 'row-major' | 'column-major',
+): Promise<{ C: Float32Array; gpuTimeMs?: number }>;
+
+/**
+ * Performs the matrix-matrix operation C := uplo(alpha * op(A) * op(B) + beta * C)
+ *
+ * A, B, and C are all kept GPU-resident. Each matrix's own `layout` (set at
+ * `GpuMatrix.from` time) determines the operation — there is no separate
+ * `layout` argument here. A and B must be GpuMatrix whenever C is, and vice
+ * versa — mixing a GpuMatrix with a plain Float32Array is not supported.
+ *
+ * {@includeCode ../../examples/sgemmtr/gpu.sgemmtr.js}
+ *
+ * @param device - GPUDevice from `init()`
+ * @param uplo   - `'lower'` to update only `col <= row`, `'upper'` for `col >= row`
+ * @param transA - `'no-transpose'` for A, `'transpose'` for A^T
+ * @param transB - `'no-transpose'` for B, `'transpose'` for B^T
+ * @param m      - rows of op(A) and C
+ * @param n      - columns of op(B) and C
+ * @param k      - columns of op(A), rows of op(B)
+ * @param alpha  - scalar multiplier for op(A)*op(B)
+ * @param A      - GpuMatrix
+ * @param lda    - leading dimension of A (must equal A.lda)
+ * @param B      - GpuMatrix
+ * @param ldb    - leading dimension of B (must equal B.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/sgemmtr/sgemmtr.mjs#L19">Source code: sgemmtr.mjs (L19)</a>
+ * @category BLAS Level 3
+ */
+export declare function sgemmtr(
+  device: GPUDevice,
+  uplo: 'lower' | 'upper',
+  transA: 'no-transpose' | 'transpose',
+  transB: 'no-transpose' | 'transpose',
+  m: number,
+  n: number,
+  k: number,
+  alpha: number,
+  A: GpuMatrix,
+  lda: number,
+  B: GpuMatrix,
+  ldb: number,
+  beta: number,
+  C: GpuMatrix,
+  ldc: number,
+): Promise<{ gpuTimeMs?: number }>;

package/src/sgemmtr/sgemmtr.mjs

@@ -0,0 +1,203 @@
+import {
+  uploadBuffer,
+  createParamsBuffer,
+  stageReadback,
+  destroyBuffers,
+} from "../util/buffer.mjs";
+import { createBindGroup } from "../util/bindgroup.mjs";
+import { runComputePass, submit } from "../util/compute.mjs";
+import { extractResult } from "../util/result.mjs";
+import { extractTimestamp } from "../util/benchmark.mjs";
+import { getPipeline } from "../util/pipeline.mjs";
+import { GpuMatrix } from "../classes/GpuMatrix.mjs";
+
+const BM_SMALL = 32, BN_SMALL = 32; // sgemmtr_small.wgsl's block tile
+const BM_LARGE = 64, BN_LARGE = 64; // sgemmtr_large.wgsl's block tile
+const LARGE_TILE_WORKGROUP_THRESHOLD = 36; // same threshold sgemm uses — see sgemm.mjs
+
+export async function sgemmtr(
+  device, uplo, transA, transB, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc, layout = "row-major",
+) {
+  let AIsGpu = A instanceof GpuMatrix;
+  let BIsGpu = B instanceof GpuMatrix;
+  const CIsGpu = C instanceof GpuMatrix;
+
+  if (!(device instanceof GPUDevice))
+    throw new Error("device must be a GPUDevice.");
+  if (uplo !== "lower" && uplo !== "upper")
+    throw new Error("uplo must be 'lower' or 'upper'.");
+  if (transA !== "no-transpose" && transA !== "transpose")
+    throw new Error("transA must be 'no-transpose' or 'transpose'.");
+  if (transB !== "no-transpose" && transB !== "transpose")
+    throw new Error("transB must be 'no-transpose' or 'transpose'.");
+  if (layout !== "row-major" && layout !== "column-major")
+    throw new Error("layout must be 'row-major' or 'column-major'.");
+  if (typeof alpha !== "number")
+    throw new Error("alpha must be a number.");
+  if (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(m) ||
+    !Number.isInteger(n) ||
+    !Number.isInteger(k) ||
+    !Number.isInteger(lda) ||
+    !Number.isInteger(ldb) ||
+    !Number.isInteger(ldc)
+  )
+    throw new Error("m, 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 (m < 0 || n < 0 || k < 0) throw new Error("m, n, and k must be non-negative.");
+  if (m === 0 || n === 0) return CIsGpu ? {} : { C };
+
+  // GpuMatrix's own .layout wins over the shared `layout` argument.
+  const effLayoutA = AIsGpu ? A.layout : layout;
+  const effLayoutB = BIsGpu ? B.layout : layout;
+  const effLayoutC = CIsGpu ? C.layout : layout;
+
+  // Shape validation, before any layout-driven swapping below.
+  // A: op(A) is m x k; A itself is m x k or k x m depending on transA.
+  const aRows = effLayoutA === "column-major" ? k : m;
+  const aCols = effLayoutA === "column-major" ? m : k;
+  const aOuter = transA === "no-transpose" ? aRows : aCols; // # of stored chunks
+  const aInner = transA === "no-transpose" ? aCols : aRows; // required chunk length
+  if (lda < aInner)
+    throw new Error(`lda must be >= ${effLayoutA === "column-major" ? "rows" : "cols"} of A as stored.`);
+  if (AIsGpu) {
+    if (lda !== A.lda) throw new Error("lda must match A.lda when A is a GpuMatrix.");
+    const [aLogRows, aLogCols] = transA === "no-transpose" ? [m, k] : [k, m];
+    if (A.rows < aLogRows || A.cols < aLogCols)
+      throw new Error("A is too small for the given m, k, and transA.");
+  } else if (A.length < (aOuter - 1) * lda + aInner) {
+    throw new Error("A does not have enough elements for the given dimensions and lda.");
+  }
+
+  // B: same reasoning as A, with op(B) = k x n.
+  const bRows = effLayoutB === "column-major" ? n : k;
+  const bCols = effLayoutB === "column-major" ? k : n;
+  const bOuter = transB === "no-transpose" ? bRows : bCols;
+  const bInner = transB === "no-transpose" ? bCols : bRows;
+  if (ldb < bInner)
+    throw new Error(`ldb must be >= ${effLayoutB === "column-major" ? "rows" : "cols"} of B as stored.`);
+  if (BIsGpu) {
+    if (ldb !== B.lda) throw new Error("ldb must match B.lda when B is a GpuMatrix.");
+    const [bLogRows, bLogCols] = transB === "no-transpose" ? [k, n] : [n, k];
+    if (B.rows < bLogRows || B.cols < bLogCols)
+      throw new Error("B is too small for the given n, k, and transB.");
+  } else if (B.length < (bOuter - 1) * ldb + bInner) {
+    throw new Error("B does not have enough elements for the given dimensions and ldb.");
+  }
+
+  // C: always m x n (no trans flag) — layout only affects lda/storage order.
+  const cOuter = effLayoutC === "column-major" ? n : m;
+  const cInner = effLayoutC === "column-major" ? m : n;
+  if (ldc < cInner)
+    throw new Error(`ldc must be >= ${effLayoutC === "column-major" ? "rows" : "cols"} of C as stored.`);
+  if (CIsGpu) {
+    if (ldc !== C.lda) throw new Error("ldc must match C.lda when C is a GpuMatrix.");
+    if (C.rows < m || C.cols < n) throw new Error("C is too small for the given m and n.");
+  } else if (C.length < (cOuter - 1) * ldc + cInner) {
+    throw new Error("C does not have enough elements for the given dimensions and ldc.");
+  }
+
+  // Column-major A/B reinterpreted row-major is A^T/B^T — flip the trans flag.
+  if (effLayoutA === "column-major")
+    transA = transA === "no-transpose" ? "transpose" : "no-transpose";
+  if (effLayoutB === "column-major")
+    transB = transB === "no-transpose" ? "transpose" : "no-transpose";
+
+  // Column-major C: compute C^T = op(B)^T * op(A)^T instead (swap A/B, flip trans, swap m<->n)
+  // — same trick sgemm uses. uplo(C) in row/col terms becomes uplo(C^T) with row/col swapped,
+  // i.e. the opposite triangle, so uplo must flip here too (same reasoning ssyr's isLower flip
+  // uses for column-major A) — nothing else that's uplo-specific needs to change, since the
+  // shader's row/col test is applied to whatever (m, n) grid it's actually given.
+  if (effLayoutC === "column-major") {
+    [A, B] = [B, A];
+    [AIsGpu, BIsGpu] = [BIsGpu, AIsGpu];
+    [lda, ldb] = [ldb, lda];
+    [transA, transB] = [
+      transB === "no-transpose" ? "transpose" : "no-transpose",
+      transA === "no-transpose" ? "transpose" : "no-transpose",
+    ];
+    [m, n] = [n, m];
+    uplo = uplo === "lower" ? "upper" : "lower";
+  }
+
+  // Shape-based auto-select — see sgemmtr_small.wgsl/sgemmtr_large.wgsl.
+  const largeWgX = Math.ceil(n / BN_LARGE);
+  const largeWgY = Math.ceil(m / BM_LARGE);
+  const useLargeTile = largeWgX * largeWgY >= LARGE_TILE_WORKGROUP_THRESHOLD;
+
+  const pipeline = await getPipeline(device, useLargeTile ? "sgemmtr_large" : "sgemmtr_small");
+
+  const ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "sgemmtr-A", false);
+  const BBuffer = BIsGpu ? B._buf : uploadBuffer(B, "sgemmtr-B", false);
+  const CBuffer = CIsGpu ? C._buf : uploadBuffer(C, "sgemmtr-C", true);
+  const paramsBuffer = createParamsBuffer(
+    [
+      { value: m,   type: "u32" },
+      { value: n,   type: "u32" },
+      { value: k,   type: "u32" },
+      { value: alpha, type: "f32" },
+      { value: beta,  type: "f32" },
+      { value: lda, type: "u32" },
+      { value: ldb, type: "u32" },
+      { value: ldc, type: "u32" },
+      { value: transA === "transpose" ? 1 : 0, type: "u32" },
+      { value: transB === "transpose" ? 1 : 0, type: "u32" },
+      { value: uplo === "upper" ? 1 : 0, type: "u32" },
+    ],
+    "sgemmtr-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(m / BM_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
+      };
+    const { commandEncoder, ts } = runComputePass(pipeline, bindGroup, wgCount);
+    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);
+    destroyBuffers(paramsBuffer);
+  }
+}

package/src/sgemv/sgemv.d.mts

@@ -50,43 +50,6 @@ export declare function sgemv(
   layout?: 'row-major' | 'column-major',
 ): Promise<{ y: Float32Array; gpuTimeMs?: number }>;
 
-/**
- * Performs the matrix-vector operation y = alpha * op(A) * x + beta * y
- *
- * A is kept GPU-resident; x and y are CPU Float32Arrays. `A`'s own `layout`
- * (set at `GpuMatrix.from` time) determines the operation — there is no
- * separate `layout` argument here.
- *
- * @param device - GPUDevice from `init()`
- * @param trans  - `'no-transpose'` for A, `'transpose'` for A^T
- * @param m      - number of rows in A
- * @param n      - number of columns in A
- * @param alpha  - scalar multiplier for op(A)*x
- * @param A      - GpuMatrix, GPU-resident
- * @param lda    - leading dimension of A (must equal A.lda)
- * @param x      - Float32Array input vector
- * @param incx   - stride for x (must be a positive integer)
- * @param beta   - scalar multiplier for y
- * @param y      - Float32Array input/output vector
- * @param incy   - stride for y (must be a positive integer)
- * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/sgemv/sgemv.mjs#L16">Source code: sgemv.mjs (L16)</a>
- * @category BLAS Level 2
- */
-export declare function sgemv(
-  device: GPUDevice,
-  trans: 'no-transpose' | 'transpose',
-  m: number,
-  n: number,
-  alpha: number,
-  A: GpuMatrix,
-  lda: number,
-  x: Float32Array,
-  incx: number,
-  beta: number,
-  y: Float32Array,
-  incy: number,
-): Promise<{ y: Float32Array; gpuTimeMs?: number }>;
-
 /**
  * Performs the matrix-vector operation y = alpha * op(A) * x + beta * y
  *
@@ -94,7 +57,7 @@ export declare function sgemv(
  * `layout` (set at `GpuMatrix.from` time) determines the operation — there is
  * no separate `layout` argument here.
  *
- * {@includeCode ../../examples/sgemv/gpuvec.sgemv.js}
+ * {@includeCode ../../examples/sgemv/gpu.sgemv.js}
  *
  * @param device - GPUDevice from `init()`
  * @param trans  - `'no-transpose'` for A, `'transpose'` for A^T

package/src/sgemv/sgemv.mjs

@@ -56,6 +56,10 @@ export async function sgemv(device, trans, m, n, alpha, A, lda, x, incx, beta, y
     throw new Error(
       "A must be a GpuMatrix when x and y are GpuVectors.",
     );
+  if (AIsGpu && !xIsGpu)
+    throw new Error(
+      "x and y must be GpuVectors when A is a GpuMatrix.",
+    );
   if (xIsGpu && x._buf === y._buf)
     throw new Error("x and y must not reference the same GPU buffer when both are GpuVectors.");
   if (AIsGpu && lda !== A.lda)

package/src/sger/sger.d.mts

@@ -42,39 +42,6 @@ export declare function sger(
   layout?: 'row-major' | 'column-major',
 ): Promise<{ A: Float32Array; gpuTimeMs?: number }>;
 
-/**
- * Performs the rank-1 update A = alpha * x * y^T + A
- *
- * A is kept GPU-resident; x and y are CPU Float32Arrays. `A`'s own `layout`
- * (set at `GpuMatrix.from` time) determines the operation — there is no
- * separate `layout` argument here.
- *
- * @param device - GPUDevice from `init()`
- * @param m      - number of rows in A
- * @param n      - number of columns in A
- * @param alpha  - scalar multiplier for x*y^T
- * @param x      - Float32Array input vector
- * @param incx   - stride for x (must be a positive integer)
- * @param y      - Float32Array input vector
- * @param incy   - stride for y (must be a positive integer)
- * @param A      - GpuMatrix, GPU-resident
- * @param lda    - leading dimension of A (must equal A.lda)
- * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/sger/sger.mjs#L15">Source code: sger.mjs (L15)</a>
- * @category BLAS Level 2
- */
-export declare function sger(
-  device: GPUDevice,
-  m: number,
-  n: number,
-  alpha: number,
-  x: Float32Array,
-  incx: number,
-  y: Float32Array,
-  incy: number,
-  A: GpuMatrix,
-  lda: number,
-): Promise<{ gpuTimeMs?: number }>;
-
 /**
  * Performs the rank-1 update A = alpha * x * y^T + A
  *
@@ -82,7 +49,7 @@ export declare function sger(
  * `GpuMatrix.from` time) determines the operation — there is no separate
  * `layout` argument here.
  *
- * {@includeCode ../../examples/sger/gpuvec.sger.js}
+ * {@includeCode ../../examples/sger/gpu.sger.js}
  *
  * @param device - GPUDevice from `init()`
  * @param m      - number of rows in A

package/src/sger/sger.mjs

@@ -63,6 +63,8 @@ export async function sger(device, m, n, alpha, x, incx, y, incy, A, lda, layout
     );
   if (xIsGpu && !AIsGpu)
     throw new Error("A must be a GpuMatrix when x and y are GpuVectors.");
+  if (AIsGpu && !xIsGpu)
+    throw new Error("x and y must be GpuVectors when A is a GpuMatrix.");
   if (AIsGpu && xIsGpu && A._buf === x._buf)
     throw new Error("A and x must not reference the same GPU buffer.");
   if (AIsGpu && yIsGpu && A._buf === y._buf)

package/src/shaders/block_transfer.wgsl

@@ -0,0 +1,42 @@
+// block_transfer: gather/scatter/scatter-subtract between a tight (blockLen
+// x otherLen) block and a sub-range of a strided (any ld, row/col-major)
+// buffer — needed since block offsets aren't 256-byte-aligned and block
+// rows/cols aren't always one contiguous range for copyBufferToBuffer.
+
+@group(0) @binding(0) var<storage, read_write> block:   array<f32>; // blockLen x otherLen, block[i*otherLen+j]
+@group(0) @binding(1) var<storage, read_write> strided: array<f32>; // B's or A's own buffer
+
+struct Params {
+  blockStart: u32,
+  blockLen:   u32,
+  otherStart: u32,
+  otherLen:   u32,
+  ld:         u32,
+  isColMajor: u32, // 0 = row-major addressing, 1 = column-major (row/col swapped)
+  blockIsRow: u32, // 1 = blockStart indexes strided's rows, 0 = its columns
+  mode:       u32, // 0 = scatter (strided := block), 1 = scatter_sub (strided -= block), 2 = gather (block := strided)
+}
+
+@group(0) @binding(2) var<uniform> params: Params;
+
+@compute @workgroup_size(8, 8)
+fn main(@builtin(global_invocation_id) gid: vec3u) {
+  let i = gid.y; // index along the blocked axis, within the block
+  let j = gid.x; // index along the other axis, within the block
+  if (i >= params.blockLen || j >= params.otherLen) {
+    return;
+  }
+
+  let row = select(params.otherStart + j, params.blockStart + i, params.blockIsRow == 1u);
+  let col = select(params.blockStart + i, params.otherStart + j, params.blockIsRow == 1u);
+  let stridedIdx = select(row * params.ld + col, col * params.ld + row, params.isColMajor == 1u);
+  let blockIdx = i * params.otherLen + j;
+
+  if (params.mode == 2u) {
+    block[blockIdx] = strided[stridedIdx];
+  } else if (params.mode == 1u) {
+    strided[stridedIdx] -= block[blockIdx];
+  } else {
+    strided[stridedIdx] = block[blockIdx];
+  }
+}