wgblas 1.2.0 → 2.0.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (70) hide show
  1. package/LICENSE +1 -1
  2. package/README.md +36 -54
  3. package/dist/wgblas.browser.js +779 -34
  4. package/index.d.mts +13 -0
  5. package/index.mjs +8 -0
  6. package/package.json +56 -3
  7. package/src/dasum/dasum.d.mts +2 -2
  8. package/src/dasum/dasum.mjs +6 -4
  9. package/src/idamax/idamax.d.mts +51 -0
  10. package/src/idamax/idamax.mjs +128 -0
  11. package/src/init.mjs +9 -1
  12. package/src/isamax/isamax.d.mts +1 -1
  13. package/src/sasum/sasum.d.mts +1 -1
  14. package/src/saxpy/saxpy.d.mts +1 -1
  15. package/src/scopy/scopy.d.mts +1 -1
  16. package/src/sdot/sdot.d.mts +1 -1
  17. package/src/sgemm/sgemm.d.mts +102 -0
  18. package/src/sgemm/sgemm.mjs +195 -0
  19. package/src/sgemmtr/sgemmtr.d.mts +104 -0
  20. package/src/sgemmtr/sgemmtr.mjs +203 -0
  21. package/src/sgemv/sgemv.d.mts +1 -38
  22. package/src/sgemv/sgemv.mjs +4 -0
  23. package/src/sger/sger.d.mts +1 -34
  24. package/src/sger/sger.mjs +2 -0
  25. package/src/shaders/block_transfer.wgsl +42 -0
  26. package/src/shaders/browser-shaders.mjs +26 -0
  27. package/src/shaders/dasum.wgsl +3 -2
  28. package/src/shaders/f64/dekker.wgsl +4 -85
  29. package/src/shaders/f64/utils/abs.wgsl +10 -0
  30. package/src/shaders/f64/utils/add.wgsl +77 -0
  31. package/src/shaders/f64/utils/equal.wgsl +7 -0
  32. package/src/shaders/f64/utils/greater.wgsl +12 -0
  33. package/src/shaders/f64/utils/multiply.wgsl +81 -0
  34. package/src/shaders/idamax.wgsl +96 -0
  35. package/src/shaders/reduction/argmaxF64.wgsl +50 -0
  36. package/src/shaders/reduction/sumF64.wgsl +2 -2
  37. package/src/shaders/sgemm_large.wgsl +117 -0
  38. package/src/shaders/sgemm_small.wgsl +112 -0
  39. package/src/shaders/sgemmtr_large.wgsl +117 -0
  40. package/src/shaders/sgemmtr_small.wgsl +110 -0
  41. package/src/shaders/symmetrize.wgsl +31 -0
  42. package/src/shaders/triangularize.wgsl +44 -0
  43. package/src/snrm2/snrm2.d.mts +1 -1
  44. package/src/srot/srot.d.mts +1 -1
  45. package/src/srotm/srotm.d.mts +1 -1
  46. package/src/sscal/sscal.d.mts +1 -1
  47. package/src/sswap/sswap.d.mts +1 -1
  48. package/src/ssymm/ssymm.d.mts +103 -0
  49. package/src/ssymm/ssymm.mjs +209 -0
  50. package/src/ssymv/ssymv.d.mts +1 -36
  51. package/src/ssymv/ssymv.mjs +2 -0
  52. package/src/ssyr/ssyr.d.mts +1 -30
  53. package/src/ssyr/ssyr.mjs +2 -0
  54. package/src/ssyr2/ssyr2.d.mts +1 -34
  55. package/src/ssyr2/ssyr2.mjs +2 -0
  56. package/src/ssyr2k/ssyr2k.d.mts +100 -0
  57. package/src/ssyr2k/ssyr2k.mjs +201 -0
  58. package/src/ssyrk/ssyrk.d.mts +90 -0
  59. package/src/ssyrk/ssyrk.mjs +176 -0
  60. package/src/strmm/strmm.d.mts +100 -0
  61. package/src/strmm/strmm.mjs +211 -0
  62. package/src/strmv/strmv.d.mts +1 -36
  63. package/src/strmv/strmv.mjs +2 -0
  64. package/src/strsm/strsm.d.mts +99 -0
  65. package/src/strsm/strsm.mjs +342 -0
  66. package/src/strsv/strsv.d.mts +1 -32
  67. package/src/strsv/strsv.mjs +2 -0
  68. package/src/util/buffer.mjs +4 -2
  69. package/src/util/compute.mjs +6 -3
  70. package/src/util/f64.mjs +3 -3
@@ -0,0 +1,201 @@
1
+ import {
2
+ uploadBuffer,
3
+ createParamsBuffer,
4
+ stageReadback,
5
+ destroyBuffers,
6
+ } from "../util/buffer.mjs";
7
+ import { createBindGroup } from "../util/bindgroup.mjs";
8
+ import { beginTimedEncoder, encodePass, submit } from "../util/compute.mjs";
9
+ import { extractResult } from "../util/result.mjs";
10
+ import { resolveTimestamp, extractTimestamp } from "../util/benchmark.mjs";
11
+ import { getPipeline } from "../util/pipeline.mjs";
12
+ import { GpuMatrix } from "../classes/GpuMatrix.mjs";
13
+
14
+ const BM_SMALL = 32, BN_SMALL = 32; // sgemmtr_small.wgsl's block tile
15
+ const BM_LARGE = 64, BN_LARGE = 64; // sgemmtr_large.wgsl's block tile
16
+ const LARGE_TILE_WORKGROUP_THRESHOLD = 36; // same threshold sgemm/sgemmtr/ssyrk use
17
+
18
+ // ssyr2k: C := uplo(alpha*op(A)*op(B)^T + alpha*op(B)*op(A)^T + beta*C). No
19
+ // dedicated shader — two sgemmtr passes on one encoder, second with beta=1.
20
+ export async function ssyr2k(
21
+ device, uplo, trans, n, k, alpha, A, lda, B, ldb, beta, C, ldc, layout = "row-major",
22
+ ) {
23
+ const AIsGpu = A instanceof GpuMatrix;
24
+ const BIsGpu = B instanceof GpuMatrix;
25
+ const CIsGpu = C instanceof GpuMatrix;
26
+
27
+ if (!(device instanceof GPUDevice))
28
+ throw new Error("device must be a GPUDevice.");
29
+ if (uplo !== "lower" && uplo !== "upper")
30
+ throw new Error("uplo must be 'lower' or 'upper'.");
31
+ if (trans !== "no-transpose" && trans !== "transpose")
32
+ throw new Error("trans must be 'no-transpose' or 'transpose'.");
33
+ if (layout !== "row-major" && layout !== "column-major")
34
+ throw new Error("layout must be 'row-major' or 'column-major'.");
35
+ if (typeof alpha !== "number")
36
+ throw new Error("alpha must be a number.");
37
+ if (Number.isNaN(alpha)) throw new Error("alpha must not be NaN.");
38
+ if (!Number.isFinite(alpha)) throw new Error("alpha must be finite.");
39
+ if (typeof beta !== "number")
40
+ throw new Error("beta must be a number.");
41
+ if (Number.isNaN(beta)) throw new Error("beta must not be NaN.");
42
+ if (!Number.isFinite(beta)) throw new Error("beta must be finite.");
43
+ if (
44
+ !Number.isInteger(n) || !Number.isInteger(k) ||
45
+ !Number.isInteger(lda) || !Number.isInteger(ldb) || !Number.isInteger(ldc)
46
+ )
47
+ throw new Error("n, k, lda, ldb, and ldc must be integers.");
48
+ if (!AIsGpu && !(A instanceof Float32Array))
49
+ throw new Error("A must be a Float32Array or GpuMatrix.");
50
+ if (!BIsGpu && !(B instanceof Float32Array))
51
+ throw new Error("B must be a Float32Array or GpuMatrix.");
52
+ if (!CIsGpu && !(C instanceof Float32Array))
53
+ throw new Error("C must be a Float32Array or GpuMatrix.");
54
+ if ((AIsGpu || BIsGpu) && !CIsGpu)
55
+ throw new Error("C must be a GpuMatrix when A or B is a GpuMatrix.");
56
+ if (CIsGpu && (!AIsGpu || !BIsGpu))
57
+ throw new Error("A and B must be GpuMatrix when C is a GpuMatrix.");
58
+ if (n < 0 || k < 0) throw new Error("n and k must be non-negative.");
59
+ if (n === 0) return CIsGpu ? {} : { C };
60
+
61
+ const effLayoutA = AIsGpu ? A.layout : layout;
62
+ const effLayoutB = BIsGpu ? B.layout : layout;
63
+ const effLayoutC = CIsGpu ? C.layout : layout;
64
+
65
+ // A: op(A) is n x k; A itself is n x k or k x n depending on trans.
66
+ const aRows = effLayoutA === "column-major" ? k : n;
67
+ const aCols = effLayoutA === "column-major" ? n : k;
68
+ const aOuter = trans === "no-transpose" ? aRows : aCols;
69
+ const aInner = trans === "no-transpose" ? aCols : aRows;
70
+ if (lda < aInner)
71
+ throw new Error(`lda must be >= ${effLayoutA === "column-major" ? "rows" : "cols"} of A as stored.`);
72
+ if (AIsGpu) {
73
+ if (lda !== A.lda) throw new Error("lda must match A.lda when A is a GpuMatrix.");
74
+ const [aLogRows, aLogCols] = trans === "no-transpose" ? [n, k] : [k, n];
75
+ if (A.rows < aLogRows || A.cols < aLogCols)
76
+ throw new Error("A is too small for the given n, k, and trans.");
77
+ } else if (A.length < (aOuter - 1) * lda + aInner) {
78
+ throw new Error("A does not have enough elements for the given dimensions and lda.");
79
+ }
80
+
81
+ // B: same shape rule as A — netlib's syr2k shares one trans across both operands.
82
+ const bRows = effLayoutB === "column-major" ? k : n;
83
+ const bCols = effLayoutB === "column-major" ? n : k;
84
+ const bOuter = trans === "no-transpose" ? bRows : bCols;
85
+ const bInner = trans === "no-transpose" ? bCols : bRows;
86
+ if (ldb < bInner)
87
+ throw new Error(`ldb must be >= ${effLayoutB === "column-major" ? "rows" : "cols"} of B as stored.`);
88
+ if (BIsGpu) {
89
+ if (ldb !== B.lda) throw new Error("ldb must match B.lda when B is a GpuMatrix.");
90
+ const [bLogRows, bLogCols] = trans === "no-transpose" ? [n, k] : [k, n];
91
+ if (B.rows < bLogRows || B.cols < bLogCols)
92
+ throw new Error("B is too small for the given n, k, and trans.");
93
+ } else if (B.length < (bOuter - 1) * ldb + bInner) {
94
+ throw new Error("B does not have enough elements for the given dimensions and ldb.");
95
+ }
96
+
97
+ // C: always n x n symmetric — layout only affects storage order, not size.
98
+ if (ldc < n) throw new Error("ldc must be >= n.");
99
+ if (CIsGpu) {
100
+ if (ldc !== C.lda) throw new Error("ldc must match C.lda when C is a GpuMatrix.");
101
+ if (C.rows < n || C.cols < n) throw new Error("C is too small for the given n.");
102
+ } else if (C.length < (n - 1) * ldc + n) {
103
+ throw new Error("C does not have enough elements for the given dimensions and ldc.");
104
+ }
105
+
106
+ // Column-major A/B reinterpreted row-major is A^T/B^T — flip trans per operand.
107
+ let effTransA = trans;
108
+ if (effLayoutA === "column-major")
109
+ effTransA = effTransA === "no-transpose" ? "transpose" : "no-transpose";
110
+ let effTransB = trans;
111
+ if (effLayoutB === "column-major")
112
+ effTransB = effTransB === "no-transpose" ? "transpose" : "no-transpose";
113
+
114
+ const uploEff = effLayoutC === "column-major"
115
+ ? (uplo === "lower" ? "upper" : "lower")
116
+ : uplo;
117
+ const flip = (t) => (t === "no-transpose" ? "transpose" : "no-transpose");
118
+
119
+ // One pass: op(X) as-is, op(Y) transposed. transOwnY is explicit, not inferred.
120
+ function passShape(transOwnX, X, ldX, transOwnY, Y, ldY) {
121
+ const transX = transOwnX;
122
+ const transY = flip(transOwnY);
123
+ if (effLayoutC !== "column-major") return { transX, X, ldX, transY, Y, ldY };
124
+ return { transX: flip(transY), X: Y, ldX: ldY, transY: flip(transX), Y: X, ldY: ldX };
125
+ }
126
+
127
+ // Shape-based auto-select — see sgemmtr_small.wgsl/sgemmtr_large.wgsl. m=n=n here (square C).
128
+ const largeWgX = Math.ceil(n / BN_LARGE);
129
+ const largeWgY = Math.ceil(n / BM_LARGE);
130
+ const useLargeTile = largeWgX * largeWgY >= LARGE_TILE_WORKGROUP_THRESHOLD;
131
+ const pipeline = await getPipeline(device, useLargeTile ? "sgemmtr_large" : "sgemmtr_small");
132
+ const wgCount = useLargeTile
133
+ ? {
134
+ x: Math.min(largeWgX, device.limits.maxComputeWorkgroupsPerDimension),
135
+ y: Math.min(largeWgY, device.limits.maxComputeWorkgroupsPerDimension),
136
+ }
137
+ : {
138
+ x: Math.min(Math.ceil(n / BN_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
139
+ y: Math.min(Math.ceil(n / BM_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
140
+ };
141
+
142
+ const ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "ssyr2k-A", false);
143
+ const BBuffer = BIsGpu ? B._buf : uploadBuffer(B, "ssyr2k-B", false);
144
+ const CBuffer = CIsGpu ? C._buf : uploadBuffer(C, "ssyr2k-C", true);
145
+ let paramsBuffer1 = null, paramsBuffer2 = null;
146
+
147
+ try {
148
+ const pass1 = passShape(effTransA, ABuffer, lda, effTransB, BBuffer, ldb);
149
+ const pass2 = passShape(effTransB, BBuffer, ldb, effTransA, ABuffer, lda);
150
+
151
+ const makeParams = (p, betaVal) => createParamsBuffer(
152
+ [
153
+ { value: n, type: "u32" },
154
+ { value: n, type: "u32" },
155
+ { value: k, type: "u32" },
156
+ { value: alpha, type: "f32" },
157
+ { value: betaVal, type: "f32" },
158
+ { value: p.ldX, type: "u32" },
159
+ { value: p.ldY, type: "u32" },
160
+ { value: ldc, type: "u32" },
161
+ { value: p.transX === "transpose" ? 1 : 0, type: "u32" },
162
+ { value: p.transY === "transpose" ? 1 : 0, type: "u32" },
163
+ { value: uploEff === "upper" ? 1 : 0, type: "u32" },
164
+ ],
165
+ "ssyr2k-params",
166
+ );
167
+ paramsBuffer1 = makeParams(pass1, beta);
168
+ paramsBuffer2 = makeParams(pass2, 1.0);
169
+
170
+ const bindGroup1 = createBindGroup(pipeline.getBindGroupLayout(0), [pass1.X, pass1.Y, CBuffer, paramsBuffer1]);
171
+ const bindGroup2 = createBindGroup(pipeline.getBindGroupLayout(0), [pass2.X, pass2.Y, CBuffer, paramsBuffer2]);
172
+
173
+ const { commandEncoder, querySet } = beginTimedEncoder();
174
+ const desc1 = querySet ? { timestampWrites: { querySet, beginningOfPassWriteIndex: 0 } } : undefined;
175
+ const desc2 = querySet ? { timestampWrites: { querySet, endOfPassWriteIndex: 1 } } : undefined;
176
+ encodePass(commandEncoder, pipeline, bindGroup1, wgCount, desc1);
177
+ encodePass(commandEncoder, pipeline, bindGroup2, wgCount, desc2);
178
+
179
+ const ts = resolveTimestamp(commandEncoder, querySet);
180
+ const readBuffer = CIsGpu ? null : stageReadback(commandEncoder, CBuffer);
181
+
182
+ submit(commandEncoder);
183
+
184
+ const gpuTimeMs = await extractTimestamp(ts);
185
+
186
+ if (CIsGpu) {
187
+ if (gpuTimeMs !== undefined) return { gpuTimeMs };
188
+ return {};
189
+ }
190
+
191
+ const result = await extractResult(readBuffer, Float32Array);
192
+ if (gpuTimeMs !== undefined) return { C: result, gpuTimeMs };
193
+ return { C: result };
194
+ } finally {
195
+ if (!AIsGpu) destroyBuffers(ABuffer);
196
+ if (!BIsGpu) destroyBuffers(BBuffer);
197
+ if (!CIsGpu) destroyBuffers(CBuffer);
198
+ if (paramsBuffer1) destroyBuffers(paramsBuffer1);
199
+ if (paramsBuffer2) destroyBuffers(paramsBuffer2);
200
+ }
201
+ }
@@ -0,0 +1,90 @@
1
+ import { GpuMatrix } from "../classes/GpuMatrix.mjs";
2
+
3
+ /**
4
+ * Performs the symmetric rank-k update C := uplo(alpha * op(A) * op(A)^T + beta * C) —
5
+ * only the triangle of C named by `uplo` is read or written (`'lower'`: `col <= row`,
6
+ * `'upper'`: `col >= row`). C is always n×n.
7
+ *
8
+ * - `trans='no-transpose'`: op(A) = A (n×k), computes alpha*A*A^T + beta*C
9
+ * - `trans='transpose'`: op(A) = A^T (A stored k×n), computes alpha*A^T*A + beta*C
10
+ *
11
+ * No dedicated kernel — this is `sgemmtr`'s exact kernel
12
+ * (`sgemmtr_small.wgsl`/`sgemmtr_large.wgsl`) with `A` aliased into both
13
+ * operand slots (`B := A`).
14
+ *
15
+ * {@includeCode ../../examples/ssyrk/ssyrk.js}
16
+ *
17
+ * **Browser (standalone HTML):**
18
+ * {@includeCode ../../examples/ssyrk/web/ssyrk.html}
19
+ *
20
+ * @param device - GPUDevice from `init()`
21
+ * @param uplo - `'lower'` to update only `col <= row`, `'upper'` for `col >= row`
22
+ * @param trans - `'no-transpose'` for A*A^T, `'transpose'` for A^T*A
23
+ * @param n - order of C (C is n×n); rows of op(A)
24
+ * @param k - columns of op(A) — the shared/contracted dimension
25
+ * @param alpha - scalar multiplier for op(A)*op(A)^T
26
+ * @param A - Float32Array, row-major or column-major (see `layout`)
27
+ * @param lda - leading dimension of A as stored
28
+ * @param beta - scalar multiplier for C
29
+ * @param C - Float32Array input/output matrix, row-major or column-major
30
+ * @param ldc - leading dimension of C as stored
31
+ * @param layout - storage layout shared by A/C when they're Float32Array
32
+ * (default: `'row-major'`) — same handling as `sgemmtr`; column-major C
33
+ * flips `uplo` internally so it still names the triangle you asked for
34
+ * @returns updated C as a Float32Array (only the requested triangle changed)
35
+ * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/ssyrk/ssyrk.mjs#L21">Source code: ssyrk.mjs (L21)</a>
36
+ * @category BLAS Level 3
37
+ */
38
+ export declare function ssyrk(
39
+ device: GPUDevice,
40
+ uplo: 'lower' | 'upper',
41
+ trans: 'no-transpose' | 'transpose',
42
+ n: number,
43
+ k: number,
44
+ alpha: number,
45
+ A: Float32Array,
46
+ lda: number,
47
+ beta: number,
48
+ C: Float32Array,
49
+ ldc: number,
50
+ layout?: 'row-major' | 'column-major',
51
+ ): Promise<{ C: Float32Array; gpuTimeMs?: number }>;
52
+
53
+ /**
54
+ * Performs the symmetric rank-k update C := uplo(alpha * op(A) * op(A)^T + beta * C)
55
+ *
56
+ * A and C are both kept GPU-resident. Each matrix's own `layout` (set at
57
+ * `GpuMatrix.from` time) determines the operation — there is no separate
58
+ * `layout` argument here. A must be a GpuMatrix whenever C is, and vice
59
+ * versa — mixing a GpuMatrix with a plain Float32Array is not supported.
60
+ *
61
+ * {@includeCode ../../examples/ssyrk/gpu.ssyrk.js}
62
+ *
63
+ * @param device - GPUDevice from `init()`
64
+ * @param uplo - `'lower'` to update only `col <= row`, `'upper'` for `col >= row`
65
+ * @param trans - `'no-transpose'` for A*A^T, `'transpose'` for A^T*A
66
+ * @param n - order of C (C is n×n); rows of op(A)
67
+ * @param k - columns of op(A) — the shared/contracted dimension
68
+ * @param alpha - scalar multiplier for op(A)*op(A)^T
69
+ * @param A - GpuMatrix
70
+ * @param lda - leading dimension of A (must equal A.lda)
71
+ * @param beta - scalar multiplier for C
72
+ * @param C - GpuMatrix (mutated in place; only the requested triangle changes)
73
+ * @param ldc - leading dimension of C (must equal C.lda)
74
+ * @returns no C — it stays GPU-resident; call `C.read()` yourself for a CPU readback (see the example)
75
+ * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/ssyrk/ssyrk.mjs#L21">Source code: ssyrk.mjs (L21)</a>
76
+ * @category BLAS Level 3
77
+ */
78
+ export declare function ssyrk(
79
+ device: GPUDevice,
80
+ uplo: 'lower' | 'upper',
81
+ trans: 'no-transpose' | 'transpose',
82
+ n: number,
83
+ k: number,
84
+ alpha: number,
85
+ A: GpuMatrix,
86
+ lda: number,
87
+ beta: number,
88
+ C: GpuMatrix,
89
+ ldc: number,
90
+ ): Promise<{ gpuTimeMs?: number }>;
@@ -0,0 +1,176 @@
1
+ import {
2
+ uploadBuffer,
3
+ createStorageBuffer,
4
+ createParamsBuffer,
5
+ stageReadback,
6
+ destroyBuffers,
7
+ } from "../util/buffer.mjs";
8
+ import { createBindGroup } from "../util/bindgroup.mjs";
9
+ import { beginTimedEncoder, encodePass, submit } from "../util/compute.mjs";
10
+ import { extractResult } from "../util/result.mjs";
11
+ import { resolveTimestamp, extractTimestamp } from "../util/benchmark.mjs";
12
+ import { getPipeline } from "../util/pipeline.mjs";
13
+ import { GpuMatrix } from "../classes/GpuMatrix.mjs";
14
+
15
+ const BM_SMALL = 32, BN_SMALL = 32; // sgemmtr_small.wgsl's block tile
16
+ const BM_LARGE = 64, BN_LARGE = 64; // sgemmtr_large.wgsl's block tile
17
+ const LARGE_TILE_WORKGROUP_THRESHOLD = 36; // same threshold sgemm/sgemmtr use
18
+
19
+ // ssyrk: C := uplo(alpha*op(A)*op(A)^T + beta*C). No dedicated shader —
20
+ // sgemmtr's kernel with A duplicated into a separate B buffer (B := A).
21
+ export async function ssyrk(
22
+ device, uplo, trans, n, k, alpha, A, lda, beta, C, ldc, layout = "row-major",
23
+ ) {
24
+ const AIsGpu = A instanceof GpuMatrix;
25
+ const CIsGpu = C instanceof GpuMatrix;
26
+
27
+ if (!(device instanceof GPUDevice))
28
+ throw new Error("device must be a GPUDevice.");
29
+ if (uplo !== "lower" && uplo !== "upper")
30
+ throw new Error("uplo must be 'lower' or 'upper'.");
31
+ if (trans !== "no-transpose" && trans !== "transpose")
32
+ throw new Error("trans must be 'no-transpose' or 'transpose'.");
33
+ if (layout !== "row-major" && layout !== "column-major")
34
+ throw new Error("layout must be 'row-major' or 'column-major'.");
35
+ if (typeof alpha !== "number")
36
+ throw new Error("alpha must be a number.");
37
+ if (Number.isNaN(alpha)) throw new Error("alpha must not be NaN.");
38
+ if (!Number.isFinite(alpha)) throw new Error("alpha must be finite.");
39
+ if (typeof beta !== "number")
40
+ throw new Error("beta must be a number.");
41
+ if (Number.isNaN(beta)) throw new Error("beta must not be NaN.");
42
+ if (!Number.isFinite(beta)) throw new Error("beta must be finite.");
43
+ if (!Number.isInteger(n) || !Number.isInteger(k) || !Number.isInteger(lda) || !Number.isInteger(ldc))
44
+ throw new Error("n, k, lda, and ldc must be integers.");
45
+ if (!AIsGpu && !(A instanceof Float32Array))
46
+ throw new Error("A must be a Float32Array or GpuMatrix.");
47
+ if (!CIsGpu && !(C instanceof Float32Array))
48
+ throw new Error("C must be a Float32Array or GpuMatrix.");
49
+ if (AIsGpu && !CIsGpu)
50
+ throw new Error("C must be a GpuMatrix when A is a GpuMatrix.");
51
+ if (CIsGpu && !AIsGpu)
52
+ throw new Error("A must be a GpuMatrix when C is a GpuMatrix.");
53
+ if (n < 0 || k < 0) throw new Error("n and k must be non-negative.");
54
+ if (n === 0) return CIsGpu ? {} : { C };
55
+
56
+ // GpuMatrix's own .layout wins over the shared `layout` argument.
57
+ const effLayoutA = AIsGpu ? A.layout : layout;
58
+ const effLayoutC = CIsGpu ? C.layout : layout;
59
+
60
+ // A: op(A) is n x k; A itself is n x k or k x n depending on trans.
61
+ const aRows = effLayoutA === "column-major" ? k : n;
62
+ const aCols = effLayoutA === "column-major" ? n : k;
63
+ const aOuter = trans === "no-transpose" ? aRows : aCols;
64
+ const aInner = trans === "no-transpose" ? aCols : aRows;
65
+ if (lda < aInner)
66
+ throw new Error(`lda must be >= ${effLayoutA === "column-major" ? "rows" : "cols"} of A as stored.`);
67
+ if (AIsGpu) {
68
+ if (lda !== A.lda) throw new Error("lda must match A.lda when A is a GpuMatrix.");
69
+ const [aLogRows, aLogCols] = trans === "no-transpose" ? [n, k] : [k, n];
70
+ if (A.rows < aLogRows || A.cols < aLogCols)
71
+ throw new Error("A is too small for the given n, k, and trans.");
72
+ } else if (A.length < (aOuter - 1) * lda + aInner) {
73
+ throw new Error("A does not have enough elements for the given dimensions and lda.");
74
+ }
75
+
76
+ // C: always n x n symmetric — layout only affects storage order, not size.
77
+ if (ldc < n) throw new Error("ldc must be >= n.");
78
+ if (CIsGpu) {
79
+ if (ldc !== C.lda) throw new Error("ldc must match C.lda when C is a GpuMatrix.");
80
+ if (C.rows < n || C.cols < n) throw new Error("C is too small for the given n.");
81
+ } else if (C.length < (n - 1) * ldc + n) {
82
+ throw new Error("C does not have enough elements for the given dimensions and ldc.");
83
+ }
84
+
85
+ // Column-major A reinterpreted row-major is A^T — flip trans, same trick sgemm/sgemmtr use.
86
+ let transA = trans;
87
+ if (effLayoutA === "column-major")
88
+ transA = transA === "no-transpose" ? "transpose" : "no-transpose";
89
+ // op(B) := op(A)^T (B is A itself) — opposite of transA.
90
+ let transB = transA === "no-transpose" ? "transpose" : "no-transpose";
91
+
92
+ // Column-major C: same swap trick sgemm/sgemmtr use, simplified since A === B and m=n.
93
+ let uploEff = uplo;
94
+ if (effLayoutC === "column-major") {
95
+ [transA, transB] = [
96
+ transB === "no-transpose" ? "transpose" : "no-transpose",
97
+ transA === "no-transpose" ? "transpose" : "no-transpose",
98
+ ];
99
+ uploEff = uploEff === "lower" ? "upper" : "lower";
100
+ }
101
+
102
+ // Shape-based auto-select — see sgemmtr_small.wgsl/sgemmtr_large.wgsl. m=n=n here (square C).
103
+ const largeWgX = Math.ceil(n / BN_LARGE);
104
+ const largeWgY = Math.ceil(n / BM_LARGE);
105
+ const useLargeTile = largeWgX * largeWgY >= LARGE_TILE_WORKGROUP_THRESHOLD;
106
+
107
+ const pipeline = await getPipeline(device, useLargeTile ? "sgemmtr_large" : "sgemmtr_small");
108
+
109
+ const ABuffer = AIsGpu ? A._buf : uploadBuffer(A, "ssyrk-A", false);
110
+ const CBuffer = CIsGpu ? C._buf : uploadBuffer(C, "ssyrk-C", true);
111
+ // B := A, but as a genuinely separate buffer — re-uploaded for Float32Array,
112
+ // GPU-copied (see below) for GpuMatrix, since the caller owns ABuffer.
113
+ const BBuffer = AIsGpu
114
+ ? createStorageBuffer(ABuffer.size, "ssyrk-B", GPUBufferUsage.COPY_DST)
115
+ : uploadBuffer(A, "ssyrk-B", false);
116
+ const paramsBuffer = createParamsBuffer(
117
+ [
118
+ { value: n, type: "u32" }, // gemmtr's m
119
+ { value: n, type: "u32" }, // gemmtr's n
120
+ { value: k, type: "u32" },
121
+ { value: alpha, type: "f32" },
122
+ { value: beta, type: "f32" },
123
+ { value: lda, type: "u32" }, // gemmtr's lda
124
+ { value: lda, type: "u32" }, // gemmtr's ldb — B := A, same lda
125
+ { value: ldc, type: "u32" },
126
+ { value: transA === "transpose" ? 1 : 0, type: "u32" },
127
+ { value: transB === "transpose" ? 1 : 0, type: "u32" },
128
+ { value: uploEff === "upper" ? 1 : 0, type: "u32" },
129
+ ],
130
+ "ssyrk-params",
131
+ );
132
+
133
+ try {
134
+ const bindGroup = createBindGroup(pipeline.getBindGroupLayout(0), [
135
+ ABuffer,
136
+ BBuffer,
137
+ CBuffer,
138
+ paramsBuffer,
139
+ ]);
140
+
141
+ const wgCount = useLargeTile
142
+ ? {
143
+ x: Math.min(largeWgX, device.limits.maxComputeWorkgroupsPerDimension),
144
+ y: Math.min(largeWgY, device.limits.maxComputeWorkgroupsPerDimension),
145
+ }
146
+ : {
147
+ x: Math.min(Math.ceil(n / BN_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
148
+ y: Math.min(Math.ceil(n / BM_SMALL), device.limits.maxComputeWorkgroupsPerDimension),
149
+ };
150
+ // Manual encoder (not runComputePass) so the A->B duplicate copy lands
151
+ // on the same command encoder, strictly before the compute pass reads B.
152
+ const { commandEncoder, querySet, passDescriptor } = beginTimedEncoder();
153
+ if (AIsGpu) commandEncoder.copyBufferToBuffer(ABuffer, 0, BBuffer, 0, ABuffer.size);
154
+ encodePass(commandEncoder, pipeline, bindGroup, wgCount, passDescriptor);
155
+ const ts = resolveTimestamp(commandEncoder, querySet);
156
+ const readBuffer = CIsGpu ? null : stageReadback(commandEncoder, CBuffer);
157
+
158
+ submit(commandEncoder);
159
+
160
+ const gpuTimeMs = await extractTimestamp(ts);
161
+
162
+ if (CIsGpu) {
163
+ if (gpuTimeMs !== undefined) return { gpuTimeMs };
164
+ return {};
165
+ }
166
+
167
+ const result = await extractResult(readBuffer, Float32Array);
168
+ if (gpuTimeMs !== undefined) return { C: result, gpuTimeMs };
169
+ return { C: result };
170
+ } finally {
171
+ if (!AIsGpu) destroyBuffers(ABuffer);
172
+ destroyBuffers(BBuffer); // always our own buffer, never the caller's
173
+ if (!CIsGpu) destroyBuffers(CBuffer);
174
+ destroyBuffers(paramsBuffer);
175
+ }
176
+ }
@@ -0,0 +1,100 @@
1
+ import { GpuMatrix } from "../classes/GpuMatrix.mjs";
2
+
3
+ /**
4
+ * Performs the triangular matrix-matrix operation
5
+ * B := alpha * op(A) * B (`side='left'`) or
6
+ * B := alpha * B * op(A) (`side='right'`) — `A` is triangular, only its
7
+ * `uplo` triangle stored; `B` is a general m×n matrix, overwritten in place.
8
+ *
9
+ * - `side='left'`: `A` is m×m — `A` premultiplies `B`
10
+ * - `side='right'`: `A` is n×n — `A` postmultiplies `B`
11
+ *
12
+ * No dedicated fused kernel — a `triangularize` pass materializes a dense
13
+ * copy of `op(A)` (zero-filling the unstored triangle, substituting the
14
+ * implicit diagonal when `diag='unit'`), then a plain `sgemm` pass
15
+ * (`sgemm_small.wgsl`/`sgemm_large.wgsl`, unmodified) does the actual
16
+ * multiply, both on one command encoder.
17
+ *
18
+ * {@includeCode ../../examples/strmm/strmm.js}
19
+ *
20
+ * **Browser (standalone HTML):**
21
+ * {@includeCode ../../examples/strmm/web/strmm.html}
22
+ *
23
+ * @param device - GPUDevice from `init()`
24
+ * @param side - `'left'` for op(A)*B, `'right'` for B*op(A)
25
+ * @param uplo - `'lower'` if only `A`'s lower triangle is stored, `'upper'` for upper
26
+ * @param transA - `'no-transpose'` for A, `'transpose'` for A^T
27
+ * @param diag - `'unit'` to treat A's diagonal as all-ones (A's diagonal is not read), `'non-unit'` to read it
28
+ * @param m - rows of B
29
+ * @param n - columns of B
30
+ * @param alpha - scalar multiplier for the matrix product
31
+ * @param A - Float32Array, triangular, row-major or column-major (see `layout`)
32
+ * @param lda - leading dimension of A as stored
33
+ * @param B - Float32Array input/output matrix, overwritten with the result, row-major or column-major
34
+ * @param ldb - leading dimension of B as stored
35
+ * @param layout - storage layout shared by A/B when they're Float32Array
36
+ * (default: `'row-major'`) — column-major A is a genuine transpose (A isn't
37
+ * symmetric like ssymm's), so both `transA` and `uplo` are adjusted
38
+ * internally to compensate
39
+ * @returns updated B as a Float32Array
40
+ * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/strmm/strmm.mjs#L25">Source code: strmm.mjs (L25)</a>
41
+ * @category BLAS Level 3
42
+ */
43
+ export declare function strmm(
44
+ device: GPUDevice,
45
+ side: 'left' | 'right',
46
+ uplo: 'lower' | 'upper',
47
+ transA: 'no-transpose' | 'transpose',
48
+ diag: 'unit' | 'non-unit',
49
+ m: number,
50
+ n: number,
51
+ alpha: number,
52
+ A: Float32Array,
53
+ lda: number,
54
+ B: Float32Array,
55
+ ldb: number,
56
+ layout?: 'row-major' | 'column-major',
57
+ ): Promise<{ B: Float32Array; gpuTimeMs?: number }>;
58
+
59
+ /**
60
+ * Performs the triangular matrix-matrix operation
61
+ * B := alpha * op(A) * B (`side='left'`) or
62
+ * B := alpha * B * op(A) (`side='right'`)
63
+ *
64
+ * A and B are both kept GPU-resident; B is mutated in place. Each matrix's
65
+ * own `layout` (set at `GpuMatrix.from` time) determines the operation —
66
+ * there is no separate `layout` argument here. A and B must both be
67
+ * GpuMatrix or both be Float32Array — mixing is not supported.
68
+ *
69
+ * {@includeCode ../../examples/strmm/gpu.strmm.js}
70
+ *
71
+ * @param device - GPUDevice from `init()`
72
+ * @param side - `'left'` for op(A)*B, `'right'` for B*op(A)
73
+ * @param uplo - `'lower'` if only `A`'s lower triangle is stored, `'upper'` for upper
74
+ * @param transA - `'no-transpose'` for A, `'transpose'` for A^T
75
+ * @param diag - `'unit'` to treat A's diagonal as all-ones (A's diagonal is not read), `'non-unit'` to read it
76
+ * @param m - rows of B
77
+ * @param n - columns of B
78
+ * @param alpha - scalar multiplier for the matrix product
79
+ * @param A - GpuMatrix, triangular
80
+ * @param lda - leading dimension of A (must equal A.lda)
81
+ * @param B - GpuMatrix (mutated in place)
82
+ * @param ldb - leading dimension of B (must equal B.lda)
83
+ * @returns no B — it stays GPU-resident; call `B.read()` yourself for a CPU readback (see the example)
84
+ * @see <a href="https://github.com/manit2004/wgblas/blob/main/src/strmm/strmm.mjs#L25">Source code: strmm.mjs (L25)</a>
85
+ * @category BLAS Level 3
86
+ */
87
+ export declare function strmm(
88
+ device: GPUDevice,
89
+ side: 'left' | 'right',
90
+ uplo: 'lower' | 'upper',
91
+ transA: 'no-transpose' | 'transpose',
92
+ diag: 'unit' | 'non-unit',
93
+ m: number,
94
+ n: number,
95
+ alpha: number,
96
+ A: GpuMatrix,
97
+ lda: number,
98
+ B: GpuMatrix,
99
+ ldb: number,
100
+ ): Promise<{ gpuTimeMs?: number }>;