lib/accy/src/kernel/library/linalg.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4
5 const artifact_product = @import("../../artifact/model/root.zig");
6 const shape = @import("../../choir/shape/root.zig");
7 const entry = @import("entry.zig");
8 const extent_mod = @import("extent.zig");
9 const geometry_mod = @import("geometry.zig");
10 const kernel = @import("../root.zig");
11 const tuning = @import("tuning.zig");
12
13 const DType = choir_abi.DType;
14 const indexExtent = extent_mod.indexExtent;
15 const indexProduct = extent_mod.indexProduct;
16 const runtimeExtentArgument = extent_mod.runtimeExtentArgument;
17
18 pub const MatrixProduct = struct {
19 m: u64,
20 n: u64,
21 k: u64,
22 dtype: DType = .f32,
23 accumulation_dtype: DType = .f32,
24 threads: entry.Threads2D = .{},
25 row_axis: []const u8 = "m",
26 col_axis: []const u8 = "n",
27 reduction_axis: []const u8 = "k",
28 };
29
30 pub const BatchedMatrixProduct = struct {
31 batch: u64,
32 m: u64,
33 n: u64,
34 k: u64,
35 threads: entry.Threads3D = .{},
36 batch_axis: []const u8 = "b",
37 row_axis: []const u8 = "m",
38 col_axis: []const u8 = "n",
39 reduction_axis: []const u8 = "k",
40 };
41
42 pub const MatrixVectorProduct = struct {
43 m: u64,
44 k: u64,
45 threads: u32 = 8,
46 row_axis: []const u8 = "m",
47 reduction_axis: []const u8 = "k",
48 };
49
50 pub const OuterProduct = struct {
51 m: u64,
52 n: u64,
53 threads: entry.Threads2D = .{},
54 lhs_axis: []const u8 = "m",
55 rhs_axis: []const u8 = "n",
56 };
57
58 pub fn matrixProductAccumulationDType(dtype: DType) ?DType {
59 return switch (dtype) {
60 .f32, .f16 => .f32,
61 else => null,
62 };
63 }
64
65 fn matrixProductAccumulationZero(inner: anytype, spec: MatrixProduct) !kernel.Value {
66 return switch (spec.accumulation_dtype) {
67 .f32 => inner.constantFloat(.f32, 0.0),
68 .f16 => inner.constantFloat(.f16, 0.0),
69 else => error.UnsupportedDType,
70 };
71 }
72
73 fn matrixProductAccumulationValue(inner: anytype, spec: MatrixProduct, value: anytype) !kernel.Value {
74 return switch (spec.accumulation_dtype) {
75 .f32 => if (comptime @TypeOf(value).scalar_dtype == .f32) value.raw() else (try value.cast(inner, .f32)).raw(),
76 .f16 => if (comptime @TypeOf(value).scalar_dtype == .f16) value.raw() else (try value.cast(inner, .f16)).raw(),
77 else => error.UnsupportedDType,
78 };
79 }
80
81 fn matrixProductOutputValue(inner: anytype, spec: MatrixProduct, value: kernel.Value) !kernel.Value {
82 if (spec.dtype == spec.accumulation_dtype) return value;
83 return switch (spec.dtype) {
84 .f32 => inner.cast(value, .f32),
85 .f16 => inner.cast(value, .f16),
86 else => error.UnsupportedDType,
87 };
88 }
89
90 fn matrixProductSpecialization(comptime spec: MatrixProduct) entry.Specialization {
91 return .{
92 .dtype = spec.dtype,
93 .accumulation_dtype = spec.accumulation_dtype,
94 .operation = .{ .linalg = .matrix_product },
95 .equation = "mk,kn->mn",
96 .inputs = &.{
97 entry.shape2D(spec.row_axis, spec.m, spec.reduction_axis, spec.k),
98 entry.shape2D(spec.reduction_axis, spec.k, spec.col_axis, spec.n),
99 },
100 .outputs = &.{entry.shape2D(spec.row_axis, spec.m, spec.col_axis, spec.n)},
101 .reductions = &.{entry.reduction("dot", .dot_product, entry.shape1D(spec.reduction_axis, spec.k))},
102 .launch = entry.launch2D(spec.n, spec.m, spec.threads.x, spec.threads.y),
103 .schedule = entry.threadBlocks2D(spec.col_axis, spec.n, spec.row_axis, spec.m, spec.threads.x, spec.threads.y),
104 };
105 }
106
107 fn batchedMatrixProductSpecialization(comptime spec: BatchedMatrixProduct) entry.Specialization {
108 return .{
109 .dtype = .f32,
110 .operation = .{ .linalg = .batched_matrix_product },
111 .equation = "bmk,bkn->bmn",
112 .inputs = &.{
113 entry.shape3D(spec.batch_axis, spec.batch, spec.row_axis, spec.m, spec.reduction_axis, spec.k),
114 entry.shape3D(spec.batch_axis, spec.batch, spec.reduction_axis, spec.k, spec.col_axis, spec.n),
115 },
116 .outputs = &.{entry.shape3D(spec.batch_axis, spec.batch, spec.row_axis, spec.m, spec.col_axis, spec.n)},
117 .reductions = &.{entry.reduction("dot", .dot_product, entry.shape1D(spec.reduction_axis, spec.k))},
118 .launch = entry.launch3D(spec.n, spec.m, spec.batch, spec.threads.x, spec.threads.y, spec.threads.z),
119 .schedule = entry.threadBlocks3D(spec.col_axis, spec.n, spec.row_axis, spec.m, spec.batch_axis, spec.batch, spec.threads.x, spec.threads.y, spec.threads.z),
120 };
121 }
122
123 fn matrixVectorProductSpecialization(comptime spec: MatrixVectorProduct) entry.Specialization {
124 return .{
125 .dtype = .f32,
126 .operation = .{ .linalg = .matrix_vector_product },
127 .equation = "mk,k->m",
128 .inputs = &.{
129 entry.shape2D(spec.row_axis, spec.m, spec.reduction_axis, spec.k),
130 entry.shape1D(spec.reduction_axis, spec.k),
131 },
132 .outputs = &.{entry.shape1D(spec.row_axis, spec.m)},
133 .reductions = &.{entry.reduction("dot", .dot_product, entry.shape1D(spec.reduction_axis, spec.k))},
134 .launch = entry.launch1D(spec.m, spec.threads),
135 .schedule = entry.threadBlocks1D(spec.row_axis, spec.m, spec.threads),
136 };
137 }
138
139 fn outerProductSpecialization(comptime spec: OuterProduct) entry.Specialization {
140 return .{
141 .dtype = .f32,
142 .operation = .{ .linalg = .outer_product },
143 .equation = "m,n->mn",
144 .inputs = &.{
145 entry.shape1D(spec.lhs_axis, spec.m),
146 entry.shape1D(spec.rhs_axis, spec.n),
147 },
148 .outputs = &.{entry.shape2D(spec.lhs_axis, spec.m, spec.rhs_axis, spec.n)},
149 .launch = entry.launch2D(spec.n, spec.m, spec.threads.x, spec.threads.y),
150 .schedule = entry.threadBlocks2D(spec.rhs_axis, spec.n, spec.lhs_axis, spec.m, spec.threads.x, spec.threads.y),
151 };
152 }
153
154 fn matrix_product_cell_sum_accumulate(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
155 const k_stride = try fold_inner.constantIndex(ctx.k_extent);
156 const n_stride = try fold_inner.constantIndex(ctx.n_extent);
157 const lhs_row_offset = try fold_inner.mul(ctx.row, k_stride);
158 const lhs_index = try fold_inner.add(lhs_row_offset, offset);
159 const rhs_row_offset = try fold_inner.mul(offset, n_stride);
160 const rhs_index = try fold_inner.add(rhs_row_offset, ctx.col);
161 const lhs_value = try ctx.lhs.load(fold_inner, lhs_index);
162 const rhs_value = try ctx.rhs.load(fold_inner, rhs_index);
163 const lhs_acc = try matrixProductAccumulationValue(fold_inner, ctx.spec, lhs_value);
164 const rhs_acc = try matrixProductAccumulationValue(fold_inner, ctx.spec, rhs_value);
165 const product = try fold_inner.mul(lhs_acc, rhs_acc);
166 return fold_inner.add(acc, product);
167 }
168
169 pub fn matrixProductCellSum(
170 inner: anytype,
171 spec: MatrixProduct,
172 lhs: anytype,
173 rhs: anytype,
174 row: kernel.Value,
175 col: kernel.Value,
176 ) !kernel.Value {
177 const zero = try matrixProductAccumulationZero(inner, spec);
178 return inner.foldRange(0, try indexExtent(spec.k), 1, zero, .{
179 .spec = spec,
180 .lhs = lhs,
181 .rhs = rhs,
182 .row = row,
183 .col = col,
184 .k_extent = try indexExtent(spec.k),
185 .n_extent = try indexExtent(spec.n),
186 }, matrix_product_cell_sum_accumulate);
187 }
188
189 pub fn outerProductCell(
190 inner: anytype,
191 lhs: anytype,
192 rhs: anytype,
193 row: kernel.Value,
194 col: kernel.Value,
195 ) !kernel.Value {
196 const lhs_value = try lhs.load(inner, row);
197 const rhs_value = try rhs.load(inner, col);
198 const product = try lhs_value.mul(inner, rhs_value);
199 return product.raw();
200 }
201
202 pub fn matrixProductOutputIndex(inner: anytype, spec: MatrixProduct, row: kernel.Value, col: kernel.Value) !kernel.Value {
203 const n_stride = try inner.constantIndex(try indexExtent(spec.n));
204 const out_row_offset = try inner.mul(row, n_stride);
205 return inner.add(out_row_offset, col);
206 }
207
208 pub fn outerProductOutputIndex(inner: anytype, spec: OuterProduct, row: kernel.Value, col: kernel.Value) !kernel.Value {
209 const n_stride = try inner.constantIndex(try indexExtent(spec.n));
210 const out_row_offset = try inner.mul(row, n_stride);
211 return inner.add(out_row_offset, col);
212 }
213
214 fn outerProductRuntimeOutputIndex(inner: anytype, row: kernel.Value, col: kernel.Value, n_extent: kernel.Value) !kernel.Value {
215 const out_row_offset = try inner.mul(row, n_extent);
216 return inner.add(out_row_offset, col);
217 }
218
219 fn batched_matrix_product_cell_sum_accumulate(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
220 const lhs_batch_stride = try fold_inner.constantIndex(ctx.lhs_batch_extent);
221 const lhs_row_stride = try fold_inner.constantIndex(ctx.k_extent);
222 const rhs_batch_stride = try fold_inner.constantIndex(ctx.rhs_batch_extent);
223 const rhs_row_stride = try fold_inner.constantIndex(ctx.n_extent);
224 const lhs_batch_offset = try fold_inner.mul(ctx.batch, lhs_batch_stride);
225 const lhs_row_offset = try fold_inner.mul(ctx.row, lhs_row_stride);
226 const lhs_batch_row_offset = try fold_inner.add(lhs_batch_offset, lhs_row_offset);
227 const lhs_index = try fold_inner.add(lhs_batch_row_offset, offset);
228 const rhs_batch_offset = try fold_inner.mul(ctx.batch, rhs_batch_stride);
229 const rhs_row_offset = try fold_inner.mul(offset, rhs_row_stride);
230 const rhs_batch_row_offset = try fold_inner.add(rhs_batch_offset, rhs_row_offset);
231 const rhs_index = try fold_inner.add(rhs_batch_row_offset, ctx.col);
232 const lhs_value = try ctx.lhs.load(fold_inner, lhs_index);
233 const rhs_value = try ctx.rhs.load(fold_inner, rhs_index);
234 const product = try lhs_value.mul(fold_inner, rhs_value);
235 return fold_inner.add(acc, product.raw());
236 }
237
238 pub fn batchedMatrixProductCellSum(
239 inner: anytype,
240 spec: BatchedMatrixProduct,
241 lhs: anytype,
242 rhs: anytype,
243 batch: kernel.Value,
244 row: kernel.Value,
245 col: kernel.Value,
246 ) !kernel.Value {
247 const zero = try inner.constantFloat(.f32, 0.0);
248 return inner.foldRange(0, try indexExtent(spec.k), 1, zero, .{
249 .lhs = lhs,
250 .rhs = rhs,
251 .batch = batch,
252 .row = row,
253 .col = col,
254 .lhs_batch_extent = try indexProduct(spec.m, spec.k),
255 .k_extent = try indexExtent(spec.k),
256 .rhs_batch_extent = try indexProduct(spec.k, spec.n),
257 .n_extent = try indexExtent(spec.n),
258 }, batched_matrix_product_cell_sum_accumulate);
259 }
260
261 pub fn batchedMatrixProductOutputIndex(
262 inner: anytype,
263 spec: BatchedMatrixProduct,
264 batch: kernel.Value,
265 row: kernel.Value,
266 col: kernel.Value,
267 ) !kernel.Value {
268 const batch_stride = try inner.constantIndex(try indexProduct(spec.m, spec.n));
269 const row_stride = try inner.constantIndex(try indexExtent(spec.n));
270 const batch_offset = try inner.mul(batch, batch_stride);
271 const row_offset = try inner.mul(row, row_stride);
272 const batch_row_offset = try inner.add(batch_offset, row_offset);
273 return inner.add(batch_row_offset, col);
274 }
275
276 fn batchedMatrixProductRuntimeOutputIndex(
277 inner: anytype,
278 batch: kernel.Value,
279 row: kernel.Value,
280 col: kernel.Value,
281 m_extent: kernel.Value,
282 n_extent: kernel.Value,
283 ) !kernel.Value {
284 const batch_stride = try inner.mul(m_extent, n_extent);
285 const batch_offset = try inner.mul(batch, batch_stride);
286 const row_offset = try inner.mul(row, n_extent);
287 const batch_row_offset = try inner.add(batch_offset, row_offset);
288 return inner.add(batch_row_offset, col);
289 }
290
291 fn matrix_vector_product_row_sum_accumulate(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
292 const k_stride = try fold_inner.constantIndex(ctx.k_extent);
293 const matrix_row_offset = try fold_inner.mul(ctx.row, k_stride);
294 const matrix_index = try fold_inner.add(matrix_row_offset, offset);
295 const matrix_value = try ctx.matrix.load(fold_inner, matrix_index);
296 const vector_value = try ctx.vector.load(fold_inner, offset);
297 const product = try matrix_value.mul(fold_inner, vector_value);
298 return fold_inner.add(acc, product.raw());
299 }
300
301 pub fn matrixVectorProductRowSum(
302 inner: anytype,
303 spec: MatrixVectorProduct,
304 matrix: anytype,
305 vector: anytype,
306 row: kernel.Value,
307 ) !kernel.Value {
308 const zero = try inner.constantFloat(.f32, 0.0);
309 return inner.foldRange(0, try indexExtent(spec.k), 1, zero, .{
310 .matrix = matrix,
311 .vector = vector,
312 .row = row,
313 .k_extent = try indexExtent(spec.k),
314 }, matrix_vector_product_row_sum_accumulate);
315 }
316
317 fn matrix_vector_product_body_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {
318 const sum = try matrixVectorProductRowSum(inner, ctx.spec, ctx.args.param(.matrix), ctx.args.param(.vector), index.index);
319 try ctx.args.param(.dst).store(inner, sum, index);
320 }
321
322 fn matrixVectorProductBody(k: anytype, spec: MatrixVectorProduct, args: anytype) !void {
323 _ = try k.forEach1D(spec.row_axis, spec.m, .{ .spec = spec, .args = args }, matrix_vector_product_body_each);
324 }
325
326 fn outer_product_body_each(inner: anytype, index: kernel.Index2D, ctx: anytype) !void {
327 const value = try outerProductCell(inner, ctx.args.param(.lhs), ctx.args.param(.rhs), index.y.index, index.x.index);
328 const out_index = try outerProductOutputIndex(inner, ctx.spec, index.y.index, index.x.index);
329 try ctx.args.param(.dst).store(inner, value, out_index);
330 }
331
332 fn outerProductBody(k: anytype, spec: OuterProduct, args: anytype) !void {
333 _ = try k.forEach2D(.{
334 .x = kernel.logical.axis(spec.rhs_axis, spec.n),
335 .y = kernel.logical.axis(spec.lhs_axis, spec.m),
336 }, .{ .spec = spec, .args = args }, outer_product_body_each);
337 }
338
339 fn batched_matrix_product_body_each(inner: anytype, index: kernel.Index3D, ctx: anytype) !void {
340 const sum = try batchedMatrixProductCellSum(inner, ctx.spec, ctx.args.param(.lhs), ctx.args.param(.rhs), index.z.index, index.y.index, index.x.index);
341 const out_index = try batchedMatrixProductOutputIndex(inner, ctx.spec, index.z.index, index.y.index, index.x.index);
342 try ctx.args.param(.dst).store(inner, sum, out_index);
343 }
344
345 fn batchedMatrixProductBody(k: anytype, spec: BatchedMatrixProduct, args: anytype) !void {
346 _ = try k.forEach3D(.{
347 .x = kernel.logical.axis(spec.col_axis, spec.n),
348 .y = kernel.logical.axis(spec.row_axis, spec.m),
349 .z = kernel.logical.axis(spec.batch_axis, spec.batch),
350 }, .{ .spec = spec, .args = args }, batched_matrix_product_body_each);
351 }
352
353 fn matrix_vector_product_runtime_row_sum_accumulate(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
354 const matrix_row_offset = try fold_inner.mul(ctx.row, ctx.k_extent);
355 const matrix_index = try fold_inner.add(matrix_row_offset, offset);
356 const matrix_value = try ctx.matrix.load(fold_inner, matrix_index);
357 const vector_value = try ctx.vector.load(fold_inner, offset);
358 const product = try matrix_value.mul(fold_inner, vector_value);
359 return fold_inner.add(acc, product.raw());
360 }
361
362 fn matrixVectorProductRuntimeRowSum(
363 inner: anytype,
364 matrix: anytype,
365 vector: anytype,
366 row: kernel.Value,
367 k_extent: kernel.Value,
368 ) !kernel.Value {
369 const lower = try inner.constantIndex(0);
370 const step = try inner.constantIndex(1);
371 const zero = try inner.constantFloat(.f32, 0.0);
372 return inner.fold(lower, k_extent, step, zero, .{
373 .matrix = matrix,
374 .vector = vector,
375 .row = row,
376 .k_extent = k_extent,
377 }, matrix_vector_product_runtime_row_sum_accumulate);
378 }
379
380 fn matrix_product_body_each(inner: anytype, index: kernel.Index2D, ctx: anytype) !void {
381 const sum = try matrixProductCellSum(inner, ctx.spec, ctx.args.param(.lhs), ctx.args.param(.rhs), index.y.index, index.x.index);
382 const out_index = try matrixProductOutputIndex(inner, ctx.spec, index.y.index, index.x.index);
383 try ctx.args.param(.dst).store(inner, try matrixProductOutputValue(inner, ctx.spec, sum), out_index);
384 }
385
386 fn matrixProductBody(k: anytype, spec: MatrixProduct, args: anytype) !void {
387 _ = try k.forEach2D(.{
388 .x = kernel.logical.axis(spec.col_axis, spec.n),
389 .y = kernel.logical.axis(spec.row_axis, spec.m),
390 }, .{ .spec = spec, .args = args }, matrix_product_body_each);
391 }
392
393 fn matrix_product_runtime_cell_sum_accumulate(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
394 const lhs_row_offset = try fold_inner.mul(ctx.row, ctx.k_extent);
395 const lhs_index = try fold_inner.add(lhs_row_offset, offset);
396 const rhs_row_offset = try fold_inner.mul(offset, ctx.n_extent);
397 const rhs_index = try fold_inner.add(rhs_row_offset, ctx.col);
398 const lhs_value = try ctx.lhs.load(fold_inner, lhs_index);
399 const rhs_value = try ctx.rhs.load(fold_inner, rhs_index);
400 const lhs_acc = try matrixProductAccumulationValue(fold_inner, ctx.spec, lhs_value);
401 const rhs_acc = try matrixProductAccumulationValue(fold_inner, ctx.spec, rhs_value);
402 const product = try fold_inner.mul(lhs_acc, rhs_acc);
403 return fold_inner.add(acc, product);
404 }
405
406 fn matrixProductRuntimeCellSum(
407 inner: anytype,
408 spec: MatrixProduct,
409 lhs: anytype,
410 rhs: anytype,
411 row: kernel.Value,
412 col: kernel.Value,
413 n_extent: kernel.Value,
414 k_extent: kernel.Value,
415 ) !kernel.Value {
416 const lower = try inner.constantIndex(0);
417 const step = try inner.constantIndex(1);
418 const zero = try matrixProductAccumulationZero(inner, spec);
419 return inner.fold(lower, k_extent, step, zero, .{
420 .spec = spec,
421 .lhs = lhs,
422 .rhs = rhs,
423 .row = row,
424 .col = col,
425 .n_extent = n_extent,
426 .k_extent = k_extent,
427 }, matrix_product_runtime_cell_sum_accumulate);
428 }
429
430 fn matrixProductRuntimeOutputIndex(inner: anytype, row: kernel.Value, col: kernel.Value, n_extent: kernel.Value) !kernel.Value {
431 const out_row_offset = try inner.mul(row, n_extent);
432 return inner.add(out_row_offset, col);
433 }
434
435 fn batched_matrix_product_runtime_cell_sum_accumulate(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
436 const lhs_batch_stride = try fold_inner.mul(ctx.m_extent, ctx.k_extent);
437 const lhs_batch_offset = try fold_inner.mul(ctx.batch, lhs_batch_stride);
438 const lhs_row_offset = try fold_inner.mul(ctx.row, ctx.k_extent);
439 const lhs_batch_row_offset = try fold_inner.add(lhs_batch_offset, lhs_row_offset);
440 const lhs_index = try fold_inner.add(lhs_batch_row_offset, offset);
441 const rhs_batch_stride = try fold_inner.mul(ctx.k_extent, ctx.n_extent);
442 const rhs_batch_offset = try fold_inner.mul(ctx.batch, rhs_batch_stride);
443 const rhs_row_offset = try fold_inner.mul(offset, ctx.n_extent);
444 const rhs_batch_row_offset = try fold_inner.add(rhs_batch_offset, rhs_row_offset);
445 const rhs_index = try fold_inner.add(rhs_batch_row_offset, ctx.col);
446 const lhs_value = try ctx.lhs.load(fold_inner, lhs_index);
447 const rhs_value = try ctx.rhs.load(fold_inner, rhs_index);
448 const product = try lhs_value.mul(fold_inner, rhs_value);
449 return fold_inner.add(acc, product.raw());
450 }
451
452 fn batchedMatrixProductRuntimeCellSum(
453 inner: anytype,
454 lhs: anytype,
455 rhs: anytype,
456 batch: kernel.Value,
457 row: kernel.Value,
458 col: kernel.Value,
459 m_extent: kernel.Value,
460 n_extent: kernel.Value,
461 k_extent: kernel.Value,
462 ) !kernel.Value {
463 const lower = try inner.constantIndex(0);
464 const step = try inner.constantIndex(1);
465 const zero = try inner.constantFloat(.f32, 0.0);
466 return inner.fold(lower, k_extent, step, zero, .{
467 .lhs = lhs,
468 .rhs = rhs,
469 .batch = batch,
470 .row = row,
471 .col = col,
472 .m_extent = m_extent,
473 .n_extent = n_extent,
474 .k_extent = k_extent,
475 }, batched_matrix_product_runtime_cell_sum_accumulate);
476 }
477
478 fn matrix_product_runtime_body_row_active(inner: anytype, ctx: anytype) !void {
479 const col_active = try inner.compare(.lt, ctx.col, ctx.n_extent);
480 try inner.guardDo(col_active, ctx, matrix_product_runtime_body_col_active);
481 }
482
483 fn matrix_product_runtime_body_col_active(active_inner: anytype, active_ctx: anytype) !void {
484 const sum = try matrixProductRuntimeCellSum(
485 active_inner,
486 active_ctx.spec,
487 active_ctx.args.param(.lhs),
488 active_ctx.args.param(.rhs),
489 active_ctx.row,
490 active_ctx.col,
491 active_ctx.n_extent,
492 active_ctx.k_extent,
493 );
494 const out_index = try matrixProductRuntimeOutputIndex(active_inner, active_ctx.row, active_ctx.col, active_ctx.n_extent);
495 try active_ctx.args.param(.dst).store(active_inner, try matrixProductOutputValue(active_inner, active_ctx.spec, sum), out_index);
496 }
497
498 fn matrixProductRuntimeBody(k: anytype, spec: MatrixProduct, args: anytype) !void {
499 const row = try k.globalId(.y);
500 const col = try k.globalId(.x);
501 const m_extent = try k.castIndex(args.param(.m).raw());
502 const n_extent = try k.castIndex(args.param(.n).raw());
503 const k_extent = try k.castIndex(args.param(.k).raw());
504 const row_active = try k.compare(.lt, row, m_extent);
505 try k.guardDo(row_active, .{
506 .args = args,
507 .spec = spec,
508 .row = row,
509 .col = col,
510 .n_extent = n_extent,
511 .k_extent = k_extent,
512 }, matrix_product_runtime_body_row_active);
513 }
514
515 fn batched_matrix_product_runtime_body_batch_active(inner: anytype, ctx: anytype) !void {
516 const row_active = try inner.compare(.lt, ctx.row, ctx.m_extent);
517 try inner.guardDo(row_active, ctx, batched_matrix_product_runtime_body_row_active);
518 }
519
520 fn batched_matrix_product_runtime_body_row_active(row_inner: anytype, row_ctx: anytype) !void {
521 const col_active = try row_inner.compare(.lt, row_ctx.col, row_ctx.n_extent);
522 try row_inner.guardDo(col_active, row_ctx, batched_matrix_product_runtime_body_col_active);
523 }
524
525 fn batched_matrix_product_runtime_body_col_active(active_inner: anytype, active_ctx: anytype) !void {
526 const sum = try batchedMatrixProductRuntimeCellSum(
527 active_inner,
528 active_ctx.args.param(.lhs),
529 active_ctx.args.param(.rhs),
530 active_ctx.batch,
531 active_ctx.row,
532 active_ctx.col,
533 active_ctx.m_extent,
534 active_ctx.n_extent,
535 active_ctx.k_extent,
536 );
537 const out_index = try batchedMatrixProductRuntimeOutputIndex(
538 active_inner,
539 active_ctx.batch,
540 active_ctx.row,
541 active_ctx.col,
542 active_ctx.m_extent,
543 active_ctx.n_extent,
544 );
545 try active_ctx.args.param(.dst).store(active_inner, sum, out_index);
546 }
547
548 fn batchedMatrixProductRuntimeBody(k: anytype, spec: BatchedMatrixProduct, args: anytype) !void {
549 _ = spec;
550 const batch = try k.globalId(.z);
551 const row = try k.globalId(.y);
552 const col = try k.globalId(.x);
553 const batch_extent = try k.castIndex(args.param(.batch).raw());
554 const m_extent = try k.castIndex(args.param(.m).raw());
555 const n_extent = try k.castIndex(args.param(.n).raw());
556 const k_extent = try k.castIndex(args.param(.k).raw());
557 const batch_active = try k.compare(.lt, batch, batch_extent);
558 try k.guardDo(batch_active, .{
559 .args = args,
560 .batch = batch,
561 .row = row,
562 .col = col,
563 .m_extent = m_extent,
564 .n_extent = n_extent,
565 .k_extent = k_extent,
566 }, batched_matrix_product_runtime_body_batch_active);
567 }
568
569 fn matrix_vector_product_runtime_body_row_active(inner: anytype, ctx: anytype) !void {
570 const sum = try matrixVectorProductRuntimeRowSum(
571 inner,
572 ctx.args.param(.matrix),
573 ctx.args.param(.vector),
574 ctx.row,
575 ctx.k_extent,
576 );
577 try ctx.args.param(.dst).store(inner, sum, ctx.row);
578 }
579
580 fn matrixVectorProductRuntimeBody(k: anytype, spec: MatrixVectorProduct, args: anytype) !void {
581 _ = spec;
582 const row = try k.globalId(.x);
583 const m_extent = try k.castIndex(args.param(.m).raw());
584 const k_extent = try k.castIndex(args.param(.k).raw());
585 const row_active = try k.compare(.lt, row, m_extent);
586 try k.guardDo(row_active, .{
587 .args = args,
588 .row = row,
589 .k_extent = k_extent,
590 }, matrix_vector_product_runtime_body_row_active);
591 }
592
593 fn outer_product_runtime_body_row_active(inner: anytype, ctx: anytype) !void {
594 const col_active = try inner.compare(.lt, ctx.col, ctx.n_extent);
595 try inner.guardDo(col_active, ctx, outer_product_runtime_body_col_active);
596 }
597
598 fn outer_product_runtime_body_col_active(active_inner: anytype, active_ctx: anytype) !void {
599 const value = try outerProductCell(
600 active_inner,
601 active_ctx.args.param(.lhs),
602 active_ctx.args.param(.rhs),
603 active_ctx.row,
604 active_ctx.col,
605 );
606 const out_index = try outerProductRuntimeOutputIndex(active_inner, active_ctx.row, active_ctx.col, active_ctx.n_extent);
607 try active_ctx.args.param(.dst).store(active_inner, value, out_index);
608 }
609
610 fn outerProductRuntimeBody(k: anytype, spec: OuterProduct, args: anytype) !void {
611 _ = spec;
612 const row = try k.globalId(.y);
613 const col = try k.globalId(.x);
614 const m_extent = try k.castIndex(args.param(.m).raw());
615 const n_extent = try k.castIndex(args.param(.n).raw());
616 const row_active = try k.compare(.lt, row, m_extent);
617 try k.guardDo(row_active, .{
618 .args = args,
619 .row = row,
620 .col = col,
621 .n_extent = n_extent,
622 }, outer_product_runtime_body_row_active);
623 }
624
625 fn matrixProductFamilySchedule(instance: MatrixProduct) kernel.logical.schedule.ThreadBlocks {
626 return kernel.logical.schedule.threadBlocks(.{
627 .x = instance.threads.x,
628 .y = instance.threads.y,
629 });
630 }
631
632 fn matrixVectorProductFamilySchedule(instance: MatrixVectorProduct) kernel.logical.schedule.ThreadBlocks {
633 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
634 }
635
636 fn batchedMatrixProductFamilySchedule(instance: BatchedMatrixProduct) kernel.logical.schedule.ThreadBlocks {
637 return kernel.logical.schedule.threadBlocks(.{
638 .x = instance.threads.x,
639 .y = instance.threads.y,
640 .z = instance.threads.z,
641 });
642 }
643
644 fn outerProductFamilySchedule(instance: OuterProduct) kernel.logical.schedule.ThreadBlocks {
645 return kernel.logical.schedule.threadBlocks(.{
646 .x = instance.threads.x,
647 .y = instance.threads.y,
648 });
649 }
650
651 fn matrixProductFamily(comptime dtype: DType) type {
652 return kernel.logical.Family(.{
653 .name = std.fmt.comptimePrint("accy_kernel_linalg_matmul_{s}", .{dtype.name()}),
654 .parameters = .{
655 .dst = kernel.dynamicBuffer(dtype),
656 .lhs = kernel.dynamicBuffer(dtype),
657 .rhs = kernel.dynamicBuffer(dtype),
658 },
659 .Instance = MatrixProduct,
660 .schedule = matrixProductFamilySchedule,
661 .body = matrixProductBody,
662 });
663 }
664
665 fn matrixProductRuntimeFamily(comptime dtype: DType) type {
666 return kernel.logical.Family(.{
667 .name = std.fmt.comptimePrint("accy_kernel_linalg_matmul_runtime_{s}", .{dtype.name()}),
668 .parameters = .{
669 .dst = kernel.dynamicBuffer(dtype),
670 .lhs = kernel.dynamicBuffer(dtype),
671 .rhs = kernel.dynamicBuffer(dtype),
672 .m = kernel.scalar(.i32),
673 .n = kernel.scalar(.i32),
674 .k = kernel.scalar(.i32),
675 },
676 .Instance = MatrixProduct,
677 .schedule = matrixProductFamilySchedule,
678 .body = matrixProductRuntimeBody,
679 });
680 }
681
682 fn matrixVectorProductFamily(comptime dtype: DType) type {
683 return kernel.logical.Family(.{
684 .name = std.fmt.comptimePrint("accy_kernel_linalg_matvec_{s}", .{dtype.name()}),
685 .parameters = .{
686 .dst = kernel.dynamicBuffer(dtype),
687 .matrix = kernel.dynamicBuffer(dtype),
688 .vector = kernel.dynamicBuffer(dtype),
689 },
690 .Instance = MatrixVectorProduct,
691 .schedule = matrixVectorProductFamilySchedule,
692 .body = matrixVectorProductBody,
693 });
694 }
695
696 fn matrixVectorProductRuntimeFamily(comptime dtype: DType) type {
697 return kernel.logical.Family(.{
698 .name = std.fmt.comptimePrint("accy_kernel_linalg_matvec_runtime_{s}", .{dtype.name()}),
699 .parameters = .{
700 .dst = kernel.dynamicBuffer(dtype),
701 .matrix = kernel.dynamicBuffer(dtype),
702 .vector = kernel.dynamicBuffer(dtype),
703 .m = kernel.scalar(.i32),
704 .k = kernel.scalar(.i32),
705 },
706 .Instance = MatrixVectorProduct,
707 .schedule = matrixVectorProductFamilySchedule,
708 .body = matrixVectorProductRuntimeBody,
709 });
710 }
711
712 fn batchedMatrixProductFamily(comptime dtype: DType) type {
713 return kernel.logical.Family(.{
714 .name = std.fmt.comptimePrint("accy_kernel_linalg_batched_matmul_{s}", .{dtype.name()}),
715 .parameters = .{
716 .dst = kernel.dynamicBuffer(dtype),
717 .lhs = kernel.dynamicBuffer(dtype),
718 .rhs = kernel.dynamicBuffer(dtype),
719 },
720 .Instance = BatchedMatrixProduct,
721 .schedule = batchedMatrixProductFamilySchedule,
722 .body = batchedMatrixProductBody,
723 });
724 }
725
726 fn batchedMatrixProductRuntimeFamily(comptime dtype: DType) type {
727 return kernel.logical.Family(.{
728 .name = std.fmt.comptimePrint("accy_kernel_linalg_batched_matmul_runtime_{s}", .{dtype.name()}),
729 .parameters = .{
730 .dst = kernel.dynamicBuffer(dtype),
731 .lhs = kernel.dynamicBuffer(dtype),
732 .rhs = kernel.dynamicBuffer(dtype),
733 .batch = kernel.scalar(.i32),
734 .m = kernel.scalar(.i32),
735 .n = kernel.scalar(.i32),
736 .k = kernel.scalar(.i32),
737 },
738 .Instance = BatchedMatrixProduct,
739 .schedule = batchedMatrixProductFamilySchedule,
740 .body = batchedMatrixProductRuntimeBody,
741 });
742 }
743
744 fn outerProductFamily(comptime dtype: DType) type {
745 return kernel.logical.Family(.{
746 .name = std.fmt.comptimePrint("accy_kernel_linalg_outer_{s}", .{dtype.name()}),
747 .parameters = .{
748 .dst = kernel.dynamicBuffer(dtype),
749 .lhs = kernel.dynamicBuffer(dtype),
750 .rhs = kernel.dynamicBuffer(dtype),
751 },
752 .Instance = OuterProduct,
753 .schedule = outerProductFamilySchedule,
754 .body = outerProductBody,
755 });
756 }
757
758 fn outerProductRuntimeFamily(comptime dtype: DType) type {
759 return kernel.logical.Family(.{
760 .name = std.fmt.comptimePrint("accy_kernel_linalg_outer_runtime_{s}", .{dtype.name()}),
761 .parameters = .{
762 .dst = kernel.dynamicBuffer(dtype),
763 .lhs = kernel.dynamicBuffer(dtype),
764 .rhs = kernel.dynamicBuffer(dtype),
765 .m = kernel.scalar(.i32),
766 .n = kernel.scalar(.i32),
767 },
768 .Instance = OuterProduct,
769 .schedule = outerProductFamilySchedule,
770 .body = outerProductRuntimeBody,
771 });
772 }
773
774 pub const MatrixProductFamilyF32 = matrixProductFamily(.f32);
775 pub const MatrixProductFamilyF16 = matrixProductFamily(.f16);
776 pub const MatrixProductRuntimeFamilyF32 = matrixProductRuntimeFamily(.f32);
777 pub const MatrixProductRuntimeFamilyF16 = matrixProductRuntimeFamily(.f16);
778 pub const MatrixVectorProductFamilyF32 = matrixVectorProductFamily(.f32);
779 pub const MatrixVectorProductRuntimeFamilyF32 = matrixVectorProductRuntimeFamily(.f32);
780 pub const BatchedMatrixProductFamilyF32 = batchedMatrixProductFamily(.f32);
781 pub const BatchedMatrixProductRuntimeFamilyF32 = batchedMatrixProductRuntimeFamily(.f32);
782 pub const OuterProductFamilyF32 = outerProductFamily(.f32);
783 pub const OuterProductRuntimeFamilyF32 = outerProductRuntimeFamily(.f32);
784
785 pub const matrix_product_family_version: u32 = 1;
786 pub const matrix_vector_product_family_version: u32 = 1;
787 pub const batched_matrix_product_family_version: u32 = 1;
788 pub const outer_product_family_version: u32 = 1;
789 const matrix_product_thread_caps = geometry_mod.ThreadCaps{
790 .budget = 256,
791 .x_max = 64,
792 .y_max = 16,
793 };
794 const batched_matrix_product_thread_budget: u32 = 256;
795 const batched_matrix_product_z_max: u32 = 8;
796 const matrix_vector_product_thread_caps = geometry_mod.ThreadCaps1D{};
797 const outer_product_thread_caps = geometry_mod.ThreadCaps{
798 .budget = 256,
799 .x_max = 64,
800 .y_max = 16,
801 };
802
803 pub fn matrixProductThreadsForExtents(m: u64, n: u64) entry.Threads2D {
804 return geometry_mod.threadsForGrid(.{ .rows = m, .cols = n }, matrix_product_thread_caps);
805 }
806
807 pub fn matrixProductThreadCandidatesForExtents(m: u64, n: u64) geometry_mod.ThreadCandidates {
808 return geometry_mod.threadCandidatesForGrid(.{ .rows = m, .cols = n }, matrix_product_thread_caps);
809 }
810
811 pub fn batchedMatrixProductThreadsForExtents(batch: u64, m: u64, n: u64) entry.Threads3D {
812 const matrix_threads = matrixProductThreadsForExtents(m, n);
813 const xy_threads = matrix_threads.x * matrix_threads.y;
814 const z_budget = @max(@as(u32, 1), batched_matrix_product_thread_budget / xy_threads);
815 const z_extent: u32 = @intCast(@max(@as(u64, 1), @min(batch, @min(@as(u64, batched_matrix_product_z_max), @as(u64, z_budget)))));
816 return .{ .x = matrix_threads.x, .y = matrix_threads.y, .z = z_extent };
817 }
818
819 pub fn matrixVectorProductThreadsForExtents(m: u64) u32 {
820 return geometry_mod.threadsForExtent(m, matrix_vector_product_thread_caps);
821 }
822
823 pub fn matrixVectorProductThreadCandidatesForExtents(m: u64) geometry_mod.Thread1DCandidates {
824 return geometry_mod.threadCandidatesForExtent(m, matrix_vector_product_thread_caps);
825 }
826
827 pub fn outerProductThreadsForExtents(m: u64, n: u64) entry.Threads2D {
828 return geometry_mod.threadsForGrid(.{ .rows = m, .cols = n }, outer_product_thread_caps);
829 }
830
831 pub fn matrixProductInstanceTarget(allocator: std.mem.Allocator, instance: MatrixProduct) ![]u8 {
832 return std.fmt.allocPrint(
833 allocator,
834 "accy.kernel.linalg.matmul{d}x{d}x{d}_{d}x{d}_{s}",
835 .{ instance.m, instance.n, instance.k, instance.threads.x, instance.threads.y, instance.dtype.name() },
836 );
837 }
838
839 pub fn matrixProductInstanceEntryName(allocator: std.mem.Allocator, instance: MatrixProduct) ![]u8 {
840 return std.fmt.allocPrint(
841 allocator,
842 "accy_kernel_linalg_matmul{d}x{d}x{d}_{d}x{d}_{s}",
843 .{ instance.m, instance.n, instance.k, instance.threads.x, instance.threads.y, instance.dtype.name() },
844 );
845 }
846
847 pub fn matrixProductFamilyTarget(allocator: std.mem.Allocator, instance: MatrixProduct) ![]u8 {
848 return std.fmt.allocPrint(
849 allocator,
850 "accy.kernel.linalg.matmul_family_{d}x{d}_{s}",
851 .{ instance.threads.x, instance.threads.y, instance.dtype.name() },
852 );
853 }
854
855 pub fn matrixProductFamilyEntryName(allocator: std.mem.Allocator, instance: MatrixProduct) ![]u8 {
856 return std.fmt.allocPrint(
857 allocator,
858 "accy_kernel_linalg_matmul_family_{d}x{d}_{s}",
859 .{ instance.threads.x, instance.threads.y, instance.dtype.name() },
860 );
861 }
862
863 pub fn batchedMatrixProductInstanceTarget(allocator: std.mem.Allocator, instance: BatchedMatrixProduct) ![]u8 {
864 return std.fmt.allocPrint(
865 allocator,
866 "accy.kernel.linalg.batched_matmul{d}x{d}x{d}x{d}_{d}x{d}x{d}_f32",
867 .{ instance.batch, instance.m, instance.n, instance.k, instance.threads.x, instance.threads.y, instance.threads.z },
868 );
869 }
870
871 pub fn batchedMatrixProductInstanceEntryName(allocator: std.mem.Allocator, instance: BatchedMatrixProduct) ![]u8 {
872 return std.fmt.allocPrint(
873 allocator,
874 "accy_kernel_linalg_batched_matmul{d}x{d}x{d}x{d}_{d}x{d}x{d}_f32",
875 .{ instance.batch, instance.m, instance.n, instance.k, instance.threads.x, instance.threads.y, instance.threads.z },
876 );
877 }
878
879 pub fn batchedMatrixProductFamilyTarget(allocator: std.mem.Allocator, instance: BatchedMatrixProduct) ![]u8 {
880 return std.fmt.allocPrint(
881 allocator,
882 "accy.kernel.linalg.batched_matmul_family_{d}x{d}x{d}_f32",
883 .{ instance.threads.x, instance.threads.y, instance.threads.z },
884 );
885 }
886
887 pub fn batchedMatrixProductFamilyEntryName(allocator: std.mem.Allocator, instance: BatchedMatrixProduct) ![]u8 {
888 return std.fmt.allocPrint(
889 allocator,
890 "accy_kernel_linalg_batched_matmul_family_{d}x{d}x{d}_f32",
891 .{ instance.threads.x, instance.threads.y, instance.threads.z },
892 );
893 }
894
895 pub fn matrixVectorProductInstanceTarget(allocator: std.mem.Allocator, instance: MatrixVectorProduct) ![]u8 {
896 return std.fmt.allocPrint(
897 allocator,
898 "accy.kernel.linalg.matvec{d}x{d}_{d}x_f32",
899 .{ instance.m, instance.k, instance.threads },
900 );
901 }
902
903 pub fn matrixVectorProductInstanceEntryName(allocator: std.mem.Allocator, instance: MatrixVectorProduct) ![]u8 {
904 return std.fmt.allocPrint(
905 allocator,
906 "accy_kernel_linalg_matvec{d}x{d}_{d}x_f32",
907 .{ instance.m, instance.k, instance.threads },
908 );
909 }
910
911 pub fn matrixVectorProductFamilyTarget(allocator: std.mem.Allocator, instance: MatrixVectorProduct) ![]u8 {
912 return std.fmt.allocPrint(
913 allocator,
914 "accy.kernel.linalg.matvec_family_{d}x_f32",
915 .{instance.threads},
916 );
917 }
918
919 pub fn matrixVectorProductFamilyEntryName(allocator: std.mem.Allocator, instance: MatrixVectorProduct) ![]u8 {
920 return std.fmt.allocPrint(
921 allocator,
922 "accy_kernel_linalg_matvec_family_{d}x_f32",
923 .{instance.threads},
924 );
925 }
926
927 pub fn outerProductInstanceTarget(allocator: std.mem.Allocator, instance: OuterProduct) ![]u8 {
928 return std.fmt.allocPrint(
929 allocator,
930 "accy.kernel.linalg.outer{d}x{d}_{d}x{d}_f32",
931 .{ instance.m, instance.n, instance.threads.x, instance.threads.y },
932 );
933 }
934
935 pub fn outerProductInstanceEntryName(allocator: std.mem.Allocator, instance: OuterProduct) ![]u8 {
936 return std.fmt.allocPrint(
937 allocator,
938 "accy_kernel_linalg_outer{d}x{d}_{d}x{d}_f32",
939 .{ instance.m, instance.n, instance.threads.x, instance.threads.y },
940 );
941 }
942
943 pub fn outerProductFamilyTarget(allocator: std.mem.Allocator, instance: OuterProduct) ![]u8 {
944 return std.fmt.allocPrint(
945 allocator,
946 "accy.kernel.linalg.outer_family_{d}x{d}_f32",
947 .{ instance.threads.x, instance.threads.y },
948 );
949 }
950
951 pub fn outerProductFamilyEntryName(allocator: std.mem.Allocator, instance: OuterProduct) ![]u8 {
952 return std.fmt.allocPrint(
953 allocator,
954 "accy_kernel_linalg_outer_family_{d}x{d}_f32",
955 .{ instance.threads.x, instance.threads.y },
956 );
957 }
958
959 pub fn matrixProductTuningExtents(instance: MatrixProduct) [3]u64 {
960 return .{ instance.m, instance.n, instance.k };
961 }
962
963 pub fn matrixProductTuningOperation(instance: MatrixProduct) entry.Operation {
964 _ = instance;
965 return .{ .linalg = .matrix_product };
966 }
967
968 pub fn matrixProductFamilyTuningKey(
969 backing_allocator: std.mem.Allocator,
970 device_fingerprint: u64,
971 instance: MatrixProduct,
972 ) !tuning.FamilyTuningKey {
973 const family_fingerprint = try matrixProductFamilyFingerprint(backing_allocator, instance);
974 const extents = matrixProductTuningExtents(instance);
975 return tuning.FamilyTuningKey.init(
976 device_fingerprint,
977 family_fingerprint,
978 entry.operationFingerprint(matrixProductTuningOperation(instance)),
979 instance.dtype,
980 matrix_product_family_version,
981 extents[0..],
982 ) orelse unreachable;
983 }
984
985 pub fn resolveMatrixProductSchedule(
986 backing_allocator: std.mem.Allocator,
987 reader: tuning.FamilyTuningReader,
988 instance: MatrixProduct,
989 ) !?entry.Threads2D {
990 const key = try matrixProductFamilyTuningKey(backing_allocator, reader.device_fingerprint, instance);
991 const record = reader.table.find(key) orelse return null;
992 const thread_candidates = matrixProductThreadCandidatesForExtents(instance.m, instance.n);
993 for (thread_candidates.slice()) |threads| {
994 var candidate = instance;
995 candidate.threads = threads;
996 const target = try matrixProductFamilyTarget(backing_allocator, candidate);
997 defer backing_allocator.free(target);
998 if (std.mem.eql(u8, target, record.target)) return threads;
999 }
1000 return null;
1001 }
1002
1003 /// A caller builds this reader from tuning records, each a measured winner pairing a device and
1004 /// problem key with the fastest thread shape, so matrix-product kernels use the fastest known
1005 /// thread shape on one device. The reader holds a device identity, a code format, and a list of
1006 /// tuning records. `resolve` returns the recorded thread shape for a matrix product, or null when
1007 /// the product has fewer than two possible thread shapes or no record matches, and null leaves the
1008 /// default schedule in place. `resolve` returns `error.InvalidArtifact` when a matching record
1009 /// names a thread shape outside the product's possible shapes. The reader borrows the records and
1010 /// the device name strings, so the caller keeps them alive while the reader is used, and a stage
1011 /// recipe that stores the reader copies their values.
1012 pub const MatrixProductScheduleReader = struct {
1013 device: gpu.DeviceIdentity,
1014 format: gpu.ArtifactFormat,
1015 records: []const tuning.MatrixProductFamilyScheduleTuningRecord = &.{},
1016
1017 pub fn resolve(
1018 self: MatrixProductScheduleReader,
1019 instance: MatrixProduct,
1020 ) gpu.BackendError!?entry.Threads2D {
1021 const threads = matrixProductThreadCandidatesForExtents(instance.m, instance.n);
1022 const candidates = threads.slice();
1023 if (candidates.len < 2) return null;
1024 const capacity = tuning.matrix_product_family_schedule_tuning_max_candidates;
1025 std.debug.assert(candidates.len <= capacity);
1026 var values: [capacity]tuning.MatrixProductFamilyScheduleThreads = undefined;
1027 for (candidates, 0..) |candidate, index| {
1028 values[index] = .{ .x = candidate.x, .y = candidate.y };
1029 }
1030 const key = try tuning.MatrixProductFamilyScheduleTuningKey.init(self.device, .{
1031 .format = self.format,
1032 .m = instance.m,
1033 .n = instance.n,
1034 .k = instance.k,
1035 .dtype = instance.dtype,
1036 .accumulation_dtype = instance.accumulation_dtype,
1037 .family_version = matrix_product_family_version,
1038 .candidates = values[0..candidates.len],
1039 });
1040 for (self.records) |record| {
1041 if (!record.key.eql(key)) continue;
1042 for (candidates) |candidate| {
1043 if (candidate.x == record.selection.threads.x and
1044 candidate.y == record.selection.threads.y) return candidate;
1045 }
1046 return error.InvalidArtifact;
1047 }
1048 return null;
1049 }
1050
1051 pub fn eql(self: MatrixProductScheduleReader, other: MatrixProductScheduleReader) bool {
1052 if (self.format != other.format or !sameDevice(self.device, other.device)) return false;
1053 if (self.records.len != other.records.len) return false;
1054 for (self.records, other.records) |lhs, rhs| {
1055 if (lhs.version != rhs.version or !lhs.key.eql(rhs.key)) return false;
1056 if (!std.meta.eql(lhs.selection, rhs.selection)) return false;
1057 }
1058 return true;
1059 }
1060
1061 fn sameDevice(lhs: gpu.DeviceIdentity, rhs: gpu.DeviceIdentity) bool {
1062 if (lhs.backend != rhs.backend or lhs.family != rhs.family or
1063 lhs.vendor_id != rhs.vendor_id or lhs.device_id != rhs.device_id) return false;
1064 if (!std.mem.eql(u8, lhs.name, rhs.name)) return false;
1065 if (lhs.driver_version == null or rhs.driver_version == null) {
1066 return lhs.driver_version == null and rhs.driver_version == null;
1067 }
1068 return std.mem.eql(u8, lhs.driver_version.?, rhs.driver_version.?);
1069 }
1070 };
1071
1072 pub fn matrixProductRuntimeArguments(instance: MatrixProduct) ![3]choir_abi.ScalarArgument {
1073 return .{
1074 .{ .u32 = try runtimeExtentArgument(instance.m) },
1075 .{ .u32 = try runtimeExtentArgument(instance.n) },
1076 .{ .u32 = try runtimeExtentArgument(instance.k) },
1077 };
1078 }
1079
1080 pub fn matrixProductShapeProfileDimensions(instance: MatrixProduct) [3]artifact_product.KernelCallShapeProfileDimension {
1081 const bounds = matrixProductRuntimeExtentBounds();
1082 return .{
1083 .{
1084 .name = instance.row_axis,
1085 .runtime_scalar_argument_index = 0,
1086 .bounds = bounds,
1087 },
1088 .{
1089 .name = instance.col_axis,
1090 .runtime_scalar_argument_index = 1,
1091 .bounds = bounds,
1092 },
1093 .{
1094 .name = instance.reduction_axis,
1095 .runtime_scalar_argument_index = 2,
1096 .bounds = bounds,
1097 },
1098 };
1099 }
1100
1101 pub fn batchedMatrixProductRuntimeArguments(instance: BatchedMatrixProduct) ![4]choir_abi.ScalarArgument {
1102 return .{
1103 .{ .u32 = try runtimeExtentArgument(instance.batch) },
1104 .{ .u32 = try runtimeExtentArgument(instance.m) },
1105 .{ .u32 = try runtimeExtentArgument(instance.n) },
1106 .{ .u32 = try runtimeExtentArgument(instance.k) },
1107 };
1108 }
1109
1110 pub fn batchedMatrixProductShapeProfileDimensions(instance: BatchedMatrixProduct) [4]artifact_product.KernelCallShapeProfileDimension {
1111 const bounds = batchedMatrixProductRuntimeExtentBounds();
1112 return .{
1113 .{
1114 .name = instance.batch_axis,
1115 .runtime_scalar_argument_index = 0,
1116 .bounds = bounds,
1117 },
1118 .{
1119 .name = instance.row_axis,
1120 .runtime_scalar_argument_index = 1,
1121 .bounds = bounds,
1122 },
1123 .{
1124 .name = instance.col_axis,
1125 .runtime_scalar_argument_index = 2,
1126 .bounds = bounds,
1127 },
1128 .{
1129 .name = instance.reduction_axis,
1130 .runtime_scalar_argument_index = 3,
1131 .bounds = bounds,
1132 },
1133 };
1134 }
1135
1136 pub fn matrixVectorProductRuntimeArguments(instance: MatrixVectorProduct) ![2]choir_abi.ScalarArgument {
1137 return .{
1138 .{ .u32 = try runtimeExtentArgument(instance.m) },
1139 .{ .u32 = try runtimeExtentArgument(instance.k) },
1140 };
1141 }
1142
1143 pub fn matrixVectorProductShapeProfileDimensions(instance: MatrixVectorProduct) [2]artifact_product.KernelCallShapeProfileDimension {
1144 const bounds = matrixVectorProductRuntimeExtentBounds();
1145 return .{
1146 .{
1147 .name = instance.row_axis,
1148 .runtime_scalar_argument_index = 0,
1149 .bounds = bounds,
1150 },
1151 .{
1152 .name = instance.reduction_axis,
1153 .runtime_scalar_argument_index = 1,
1154 .bounds = bounds,
1155 },
1156 };
1157 }
1158
1159 pub fn outerProductRuntimeArguments(instance: OuterProduct) ![2]choir_abi.ScalarArgument {
1160 return .{
1161 .{ .u32 = try runtimeExtentArgument(instance.m) },
1162 .{ .u32 = try runtimeExtentArgument(instance.n) },
1163 };
1164 }
1165
1166 pub fn outerProductShapeProfileDimensions(instance: OuterProduct) [2]artifact_product.KernelCallShapeProfileDimension {
1167 const bounds = outerProductRuntimeExtentBounds();
1168 return .{
1169 .{
1170 .name = instance.lhs_axis,
1171 .runtime_scalar_argument_index = 0,
1172 .bounds = bounds,
1173 },
1174 .{
1175 .name = instance.rhs_axis,
1176 .runtime_scalar_argument_index = 1,
1177 .bounds = bounds,
1178 },
1179 };
1180 }
1181
1182 fn matrixProductDerivedLaunch(instance: MatrixProduct) !artifact_product.KernelCallLaunch {
1183 if (instance.threads.x == 0 or instance.threads.y == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
1184 return .{ .derived = .{
1185 .grid = .{
1186 .{ .runtime_u32_ceil_div = .{ .argument_index = 1, .divisor = instance.threads.x } },
1187 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = instance.threads.y } },
1188 .{ .fixed = 1 },
1189 },
1190 .threadgroup = .{ instance.threads.x, instance.threads.y, 1 },
1191 } };
1192 }
1193
1194 fn batchedMatrixProductDerivedLaunch(instance: BatchedMatrixProduct) !artifact_product.KernelCallLaunch {
1195 if (instance.threads.x == 0 or instance.threads.y == 0 or instance.threads.z == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
1196 return .{ .derived = .{
1197 .grid = .{
1198 .{ .runtime_u32_ceil_div = .{ .argument_index = 2, .divisor = instance.threads.x } },
1199 .{ .runtime_u32_ceil_div = .{ .argument_index = 1, .divisor = instance.threads.y } },
1200 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = instance.threads.z } },
1201 },
1202 .threadgroup = .{ instance.threads.x, instance.threads.y, instance.threads.z },
1203 } };
1204 }
1205
1206 fn matrixVectorProductDerivedLaunch(instance: MatrixVectorProduct) !artifact_product.KernelCallLaunch {
1207 if (instance.threads == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
1208 return .{ .derived = .{
1209 .grid = .{
1210 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = instance.threads } },
1211 .{ .fixed = 1 },
1212 .{ .fixed = 1 },
1213 },
1214 .threadgroup = .{ instance.threads, 1, 1 },
1215 } };
1216 }
1217
1218 fn outerProductDerivedLaunch(instance: OuterProduct) !artifact_product.KernelCallLaunch {
1219 if (instance.threads.x == 0 or instance.threads.y == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
1220 return .{ .derived = .{
1221 .grid = .{
1222 .{ .runtime_u32_ceil_div = .{ .argument_index = 1, .divisor = instance.threads.x } },
1223 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = instance.threads.y } },
1224 .{ .fixed = 1 },
1225 },
1226 .threadgroup = .{ instance.threads.x, instance.threads.y, 1 },
1227 } };
1228 }
1229
1230 pub fn createMatrixProductFamilyArtifact(
1231 allocator: std.mem.Allocator,
1232 handle: kernel.BackendHandle,
1233 instance: MatrixProduct,
1234 options: entry.ArtifactOptions,
1235 ) !kernel.OwnedKernelCallArtifact {
1236 const target = try matrixProductFamilyTarget(allocator, instance);
1237 defer allocator.free(target);
1238 const entry_name = try matrixProductFamilyEntryName(allocator, instance);
1239 defer allocator.free(entry_name);
1240 const family_fingerprint = options.shape_family_fingerprint orelse try matrixProductFamilyFingerprint(allocator, instance);
1241 const shape_profile_dimensions = matrixProductShapeProfileDimensions(instance);
1242 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1243 .name = "matrix_product",
1244 .fingerprint = family_fingerprint,
1245 .dimensions = shape_profile_dimensions[0..],
1246 };
1247
1248 var graph = switch (instance.dtype) {
1249 .f32 => try MatrixProductRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
1250 .f16 => try MatrixProductRuntimeFamilyF16.buildNamed(allocator, options.limits, entry_name, instance),
1251 else => return error.UnsupportedDType,
1252 };
1253 defer graph.deinit();
1254 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1255 .target = target,
1256 .version = matrix_product_family_version,
1257 .format = options.format,
1258 .kernel_plan = options.kernel_plan,
1259 .element_count_argument = options.element_count_argument,
1260 .shape_family_fingerprint = family_fingerprint,
1261 .shape_profile = shape_profile,
1262 .launch = options.launch orelse try matrixProductDerivedLaunch(instance),
1263 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 3 else options.runtime_scalar_argument_count,
1264 .static_arguments = options.static_arguments,
1265 });
1266 }
1267
1268 pub fn createBatchedMatrixProductFamilyArtifact(
1269 allocator: std.mem.Allocator,
1270 handle: kernel.BackendHandle,
1271 instance: BatchedMatrixProduct,
1272 options: entry.ArtifactOptions,
1273 ) !kernel.OwnedKernelCallArtifact {
1274 const target = try batchedMatrixProductFamilyTarget(allocator, instance);
1275 defer allocator.free(target);
1276 const entry_name = try batchedMatrixProductFamilyEntryName(allocator, instance);
1277 defer allocator.free(entry_name);
1278 const family_fingerprint = options.shape_family_fingerprint orelse try batchedMatrixProductFamilyFingerprint(allocator, instance);
1279 const shape_profile_dimensions = batchedMatrixProductShapeProfileDimensions(instance);
1280 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1281 .name = "batched_matrix_product",
1282 .fingerprint = family_fingerprint,
1283 .dimensions = shape_profile_dimensions[0..],
1284 };
1285
1286 var graph = try BatchedMatrixProductRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance);
1287 defer graph.deinit();
1288 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1289 .target = target,
1290 .version = batched_matrix_product_family_version,
1291 .format = options.format,
1292 .kernel_plan = options.kernel_plan,
1293 .element_count_argument = options.element_count_argument,
1294 .shape_family_fingerprint = family_fingerprint,
1295 .shape_profile = shape_profile,
1296 .launch = options.launch orelse try batchedMatrixProductDerivedLaunch(instance),
1297 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 4 else options.runtime_scalar_argument_count,
1298 .static_arguments = options.static_arguments,
1299 });
1300 }
1301
1302 pub fn createMatrixVectorProductFamilyArtifact(
1303 allocator: std.mem.Allocator,
1304 handle: kernel.BackendHandle,
1305 instance: MatrixVectorProduct,
1306 options: entry.ArtifactOptions,
1307 ) !kernel.OwnedKernelCallArtifact {
1308 const target = try matrixVectorProductFamilyTarget(allocator, instance);
1309 defer allocator.free(target);
1310 const entry_name = try matrixVectorProductFamilyEntryName(allocator, instance);
1311 defer allocator.free(entry_name);
1312 const family_fingerprint = options.shape_family_fingerprint orelse try matrixVectorProductFamilyFingerprint(allocator, instance);
1313 const shape_profile_dimensions = matrixVectorProductShapeProfileDimensions(instance);
1314 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1315 .name = "matrix_vector_product",
1316 .fingerprint = family_fingerprint,
1317 .dimensions = shape_profile_dimensions[0..],
1318 };
1319
1320 var graph = try MatrixVectorProductRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance);
1321 defer graph.deinit();
1322 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1323 .target = target,
1324 .version = matrix_vector_product_family_version,
1325 .format = options.format,
1326 .kernel_plan = options.kernel_plan,
1327 .element_count_argument = options.element_count_argument,
1328 .shape_family_fingerprint = family_fingerprint,
1329 .shape_profile = shape_profile,
1330 .launch = options.launch orelse try matrixVectorProductDerivedLaunch(instance),
1331 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 2 else options.runtime_scalar_argument_count,
1332 .static_arguments = options.static_arguments,
1333 });
1334 }
1335
1336 pub fn createOuterProductFamilyArtifact(
1337 allocator: std.mem.Allocator,
1338 handle: kernel.BackendHandle,
1339 instance: OuterProduct,
1340 options: entry.ArtifactOptions,
1341 ) !kernel.OwnedKernelCallArtifact {
1342 const target = try outerProductFamilyTarget(allocator, instance);
1343 defer allocator.free(target);
1344 const entry_name = try outerProductFamilyEntryName(allocator, instance);
1345 defer allocator.free(entry_name);
1346 const family_fingerprint = options.shape_family_fingerprint orelse try outerProductFamilyFingerprint(allocator, instance);
1347 const shape_profile_dimensions = outerProductShapeProfileDimensions(instance);
1348 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1349 .name = "outer_product",
1350 .fingerprint = family_fingerprint,
1351 .dimensions = shape_profile_dimensions[0..],
1352 };
1353
1354 var graph = try OuterProductRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance);
1355 defer graph.deinit();
1356 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1357 .target = target,
1358 .version = outer_product_family_version,
1359 .format = options.format,
1360 .kernel_plan = options.kernel_plan,
1361 .element_count_argument = options.element_count_argument,
1362 .shape_family_fingerprint = family_fingerprint,
1363 .shape_profile = shape_profile,
1364 .launch = options.launch orelse try outerProductDerivedLaunch(instance),
1365 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 2 else options.runtime_scalar_argument_count,
1366 .static_arguments = options.static_arguments,
1367 });
1368 }
1369
1370 pub fn matrixProductFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: MatrixProduct) !u64 {
1371 var family = try matrixProductShapeFamily(backing_allocator, instance);
1372 defer family.deinit();
1373 return shape.fingerprint(family);
1374 }
1375
1376 pub fn matrixProductShapeFamily(backing_allocator: std.mem.Allocator, instance: MatrixProduct) !shape.Family {
1377 var builder = try shape.Builder.init(backing_allocator, "matrix_product");
1378 errdefer builder.deinit();
1379
1380 const m = try builder.symbol(instance.row_axis);
1381 const n = try builder.symbol(instance.col_axis);
1382 const k = try builder.symbol(instance.reduction_axis);
1383
1384 const m_expr = try builder.symbolExpression(m);
1385 const n_expr = try builder.symbolExpression(n);
1386 const k_expr = try builder.symbolExpression(k);
1387
1388 _ = try builder.tensor("lhs", &.{ m_expr, k_expr });
1389 _ = try builder.tensor("rhs", &.{ k_expr, n_expr });
1390 _ = try builder.tensor("out", &.{ m_expr, n_expr });
1391 try builder.assumeBounds(m_expr, matrixProductRuntimeExtentBounds());
1392 try builder.assumeBounds(n_expr, matrixProductRuntimeExtentBounds());
1393 try builder.assumeBounds(k_expr, matrixProductRuntimeExtentBounds());
1394
1395 return builder.finish();
1396 }
1397
1398 fn matrixProductRuntimeExtentBounds() shape.Bounds {
1399 return .{ .min = 1, .max = extent_mod.runtime_extent_max };
1400 }
1401
1402 pub fn batchedMatrixProductFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: BatchedMatrixProduct) !u64 {
1403 var family = try batchedMatrixProductShapeFamily(backing_allocator, instance);
1404 defer family.deinit();
1405 return shape.fingerprint(family);
1406 }
1407
1408 pub fn batchedMatrixProductShapeFamily(backing_allocator: std.mem.Allocator, instance: BatchedMatrixProduct) !shape.Family {
1409 var builder = try shape.Builder.init(backing_allocator, "batched_matrix_product");
1410 errdefer builder.deinit();
1411
1412 const batch = try builder.symbol(instance.batch_axis);
1413 const m = try builder.symbol(instance.row_axis);
1414 const n = try builder.symbol(instance.col_axis);
1415 const k = try builder.symbol(instance.reduction_axis);
1416
1417 const batch_expr = try builder.symbolExpression(batch);
1418 const m_expr = try builder.symbolExpression(m);
1419 const n_expr = try builder.symbolExpression(n);
1420 const k_expr = try builder.symbolExpression(k);
1421
1422 _ = try builder.tensor("lhs", &.{ batch_expr, m_expr, k_expr });
1423 _ = try builder.tensor("rhs", &.{ batch_expr, k_expr, n_expr });
1424 _ = try builder.tensor("out", &.{ batch_expr, m_expr, n_expr });
1425 try builder.assumeBounds(batch_expr, batchedMatrixProductRuntimeExtentBounds());
1426 try builder.assumeBounds(m_expr, batchedMatrixProductRuntimeExtentBounds());
1427 try builder.assumeBounds(n_expr, batchedMatrixProductRuntimeExtentBounds());
1428 try builder.assumeBounds(k_expr, batchedMatrixProductRuntimeExtentBounds());
1429
1430 return builder.finish();
1431 }
1432
1433 fn batchedMatrixProductRuntimeExtentBounds() shape.Bounds {
1434 return .{ .min = 1, .max = extent_mod.runtime_extent_max };
1435 }
1436
1437 pub fn matrixVectorProductFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: MatrixVectorProduct) !u64 {
1438 var family = try matrixVectorProductShapeFamily(backing_allocator, instance);
1439 defer family.deinit();
1440 return shape.fingerprint(family);
1441 }
1442
1443 pub fn matrixVectorProductShapeFamily(backing_allocator: std.mem.Allocator, instance: MatrixVectorProduct) !shape.Family {
1444 var builder = try shape.Builder.init(backing_allocator, "matrix_vector_product");
1445 errdefer builder.deinit();
1446
1447 const m = try builder.symbol(instance.row_axis);
1448 const k = try builder.symbol(instance.reduction_axis);
1449
1450 const m_expr = try builder.symbolExpression(m);
1451 const k_expr = try builder.symbolExpression(k);
1452
1453 _ = try builder.tensor("matrix", &.{ m_expr, k_expr });
1454 _ = try builder.tensor("vector", &.{k_expr});
1455 _ = try builder.tensor("out", &.{m_expr});
1456 try builder.assumeBounds(m_expr, matrixVectorProductRuntimeExtentBounds());
1457 try builder.assumeBounds(k_expr, matrixVectorProductRuntimeExtentBounds());
1458
1459 return builder.finish();
1460 }
1461
1462 fn matrixVectorProductRuntimeExtentBounds() shape.Bounds {
1463 return .{ .min = 1, .max = extent_mod.runtime_extent_max };
1464 }
1465
1466 pub fn outerProductFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: OuterProduct) !u64 {
1467 var family = try outerProductShapeFamily(backing_allocator, instance);
1468 defer family.deinit();
1469 return shape.fingerprint(family);
1470 }
1471
1472 pub fn outerProductShapeFamily(backing_allocator: std.mem.Allocator, instance: OuterProduct) !shape.Family {
1473 var builder = try shape.Builder.init(backing_allocator, "outer_product");
1474 errdefer builder.deinit();
1475
1476 const m = try builder.symbol(instance.lhs_axis);
1477 const n = try builder.symbol(instance.rhs_axis);
1478
1479 const m_expr = try builder.symbolExpression(m);
1480 const n_expr = try builder.symbolExpression(n);
1481
1482 _ = try builder.tensor("lhs", &.{m_expr});
1483 _ = try builder.tensor("rhs", &.{n_expr});
1484 _ = try builder.tensor("out", &.{ m_expr, n_expr });
1485 try builder.assumeBounds(m_expr, outerProductRuntimeExtentBounds());
1486 try builder.assumeBounds(n_expr, outerProductRuntimeExtentBounds());
1487
1488 return builder.finish();
1489 }
1490
1491 fn outerProductRuntimeExtentBounds() shape.Bounds {
1492 return .{ .min = 1, .max = extent_mod.runtime_extent_max };
1493 }
1494
1495 pub fn matrixProductFamilySpecialization(backing_allocator: std.mem.Allocator, instance: MatrixProduct) !entry.OwnedSpecialization {
1496 var owned = entry.OwnedSpecialization.init(backing_allocator);
1497 errdefer owned.deinit();
1498 const lifetime_allocator = owned.allocator();
1499
1500 const inputs = try lifetime_allocator.alloc(entry.Shape, 2);
1501 inputs[0] = try entry.runtimeShape2D(lifetime_allocator, instance.row_axis, instance.m, instance.reduction_axis, instance.k);
1502 inputs[1] = try entry.runtimeShape2D(lifetime_allocator, instance.reduction_axis, instance.k, instance.col_axis, instance.n);
1503
1504 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
1505 outputs[0] = try entry.runtimeShape2D(lifetime_allocator, instance.row_axis, instance.m, instance.col_axis, instance.n);
1506
1507 const reductions = try lifetime_allocator.alloc(entry.Reduction, 1);
1508 reductions[0] = try entry.runtimeReduction(
1509 lifetime_allocator,
1510 "dot",
1511 .dot_product,
1512 try entry.runtimeShape1D(lifetime_allocator, instance.reduction_axis, instance.k),
1513 );
1514
1515 owned.value = .{
1516 .dtype = instance.dtype,
1517 .accumulation_dtype = instance.accumulation_dtype,
1518 .operation = .{ .linalg = .matrix_product },
1519 .equation = "mk,kn->mn",
1520 .inputs = inputs,
1521 .outputs = outputs,
1522 .reductions = reductions,
1523 .schedule = try entry.runtimeThreadBlocks2D(lifetime_allocator, instance.col_axis, instance.n, instance.row_axis, instance.m, instance.threads.x, instance.threads.y),
1524 };
1525 owned.value.launch = owned.value.schedule.?.launch();
1526 var family = try matrixProductShapeFamily(backing_allocator, instance);
1527 errdefer family.deinit();
1528 try owned.takeShapeFamily(&family);
1529 return owned;
1530 }
1531
1532 pub fn matrixProductInstanceFromSpecialization(specialization: entry.Specialization) ?MatrixProduct {
1533 if (!specialization.scheduleMatchesLaunch()) return null;
1534 if (!specialization.operationIs(.{ .linalg = .matrix_product })) return null;
1535 const dtype = specialization.dtype orelse return null;
1536 const accumulation_dtype = specialization.accumulation_dtype orelse return null;
1537 if (matrixProductAccumulationDType(dtype) != accumulation_dtype) return null;
1538 const equation = specialization.equation orelse return null;
1539 if (!std.mem.eql(u8, equation, "mk,kn->mn")) return null;
1540 if (specialization.inputs.len != 2 or specialization.outputs.len != 1 or specialization.reductions.len != 1) return null;
1541 const lhs = specialization.inputs[0];
1542 const rhs = specialization.inputs[1];
1543 const output = specialization.outputs[0];
1544 const reduction = specialization.reductions[0];
1545 if (lhs.axes.len != 2 or rhs.axes.len != 2 or output.axes.len != 2) return null;
1546 if (reduction.shape.axes.len != 1) return null;
1547 const m = lhs.axes[0].extent;
1548 const k = lhs.axes[1].extent;
1549 const n = rhs.axes[1].extent;
1550 if (!std.mem.eql(u8, lhs.axes[1].name, rhs.axes[0].name)) return null;
1551 if (!std.mem.eql(u8, lhs.axes[0].name, output.axes[0].name)) return null;
1552 if (!std.mem.eql(u8, rhs.axes[1].name, output.axes[1].name)) return null;
1553 if (!std.mem.eql(u8, reduction.shape.axes[0].name, lhs.axes[1].name)) return null;
1554 if (rhs.axes[0].extent != k) return null;
1555 if (output.axes[0].extent != m or output.axes[1].extent != n) return null;
1556 if (reduction.shape.axes[0].extent != k) return null;
1557 const launch = specialization.launch orelse return null;
1558 if (launch.threadgroup[0] == 0 or launch.threadgroup[1] == 0) return null;
1559 return .{
1560 .m = m,
1561 .n = n,
1562 .k = k,
1563 .dtype = dtype,
1564 .accumulation_dtype = accumulation_dtype,
1565 .threads = .{ .x = launch.threadgroup[0], .y = launch.threadgroup[1] },
1566 .row_axis = lhs.axes[0].name,
1567 .col_axis = rhs.axes[1].name,
1568 .reduction_axis = lhs.axes[1].name,
1569 };
1570 }
1571
1572 pub fn batchedMatrixProductFamilySpecialization(backing_allocator: std.mem.Allocator, instance: BatchedMatrixProduct) !entry.OwnedSpecialization {
1573 var owned = entry.OwnedSpecialization.init(backing_allocator);
1574 errdefer owned.deinit();
1575 const lifetime_allocator = owned.allocator();
1576
1577 const inputs = try lifetime_allocator.alloc(entry.Shape, 2);
1578 inputs[0] = try entry.runtimeShape3D(lifetime_allocator, instance.batch_axis, instance.batch, instance.row_axis, instance.m, instance.reduction_axis, instance.k);
1579 inputs[1] = try entry.runtimeShape3D(lifetime_allocator, instance.batch_axis, instance.batch, instance.reduction_axis, instance.k, instance.col_axis, instance.n);
1580
1581 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
1582 outputs[0] = try entry.runtimeShape3D(lifetime_allocator, instance.batch_axis, instance.batch, instance.row_axis, instance.m, instance.col_axis, instance.n);
1583
1584 const reductions = try lifetime_allocator.alloc(entry.Reduction, 1);
1585 reductions[0] = try entry.runtimeReduction(
1586 lifetime_allocator,
1587 "dot",
1588 .dot_product,
1589 try entry.runtimeShape1D(lifetime_allocator, instance.reduction_axis, instance.k),
1590 );
1591
1592 owned.value = .{
1593 .dtype = .f32,
1594 .operation = .{ .linalg = .batched_matrix_product },
1595 .equation = "bmk,bkn->bmn",
1596 .inputs = inputs,
1597 .outputs = outputs,
1598 .reductions = reductions,
1599 .schedule = try entry.runtimeThreadBlocks3D(
1600 lifetime_allocator,
1601 instance.col_axis,
1602 instance.n,
1603 instance.row_axis,
1604 instance.m,
1605 instance.batch_axis,
1606 instance.batch,
1607 instance.threads.x,
1608 instance.threads.y,
1609 instance.threads.z,
1610 ),
1611 };
1612 owned.value.launch = owned.value.schedule.?.launch();
1613 var family = try batchedMatrixProductShapeFamily(backing_allocator, instance);
1614 errdefer family.deinit();
1615 try owned.takeShapeFamily(&family);
1616 return owned;
1617 }
1618
1619 pub fn batchedMatrixProductInstanceFromSpecialization(specialization: entry.Specialization) ?BatchedMatrixProduct {
1620 if (!specialization.scheduleMatchesLaunch()) return null;
1621 if (!specialization.operationIs(.{ .linalg = .batched_matrix_product })) return null;
1622 const dtype = specialization.dtype orelse return null;
1623 if (dtype != .f32) return null;
1624 const equation = specialization.equation orelse return null;
1625 if (!std.mem.eql(u8, equation, "bmk,bkn->bmn")) return null;
1626 if (specialization.inputs.len != 2 or specialization.outputs.len != 1 or specialization.reductions.len != 1) return null;
1627 const lhs = specialization.inputs[0];
1628 const rhs = specialization.inputs[1];
1629 const output = specialization.outputs[0];
1630 const reduction = specialization.reductions[0];
1631 if (lhs.axes.len != 3 or rhs.axes.len != 3 or output.axes.len != 3) return null;
1632 if (reduction.shape.axes.len != 1) return null;
1633 const batch = lhs.axes[0].extent;
1634 const m = lhs.axes[1].extent;
1635 const k = lhs.axes[2].extent;
1636 const n = rhs.axes[2].extent;
1637 if (!std.mem.eql(u8, lhs.axes[0].name, rhs.axes[0].name)) return null;
1638 if (!std.mem.eql(u8, lhs.axes[0].name, output.axes[0].name)) return null;
1639 if (!std.mem.eql(u8, lhs.axes[1].name, output.axes[1].name)) return null;
1640 if (!std.mem.eql(u8, lhs.axes[2].name, rhs.axes[1].name)) return null;
1641 if (!std.mem.eql(u8, lhs.axes[2].name, reduction.shape.axes[0].name)) return null;
1642 if (!std.mem.eql(u8, rhs.axes[2].name, output.axes[2].name)) return null;
1643 if (rhs.axes[0].extent != batch or output.axes[0].extent != batch) return null;
1644 if (output.axes[1].extent != m or rhs.axes[1].extent != k or reduction.shape.axes[0].extent != k or output.axes[2].extent != n) return null;
1645 const launch = specialization.launch orelse return null;
1646 if (launch.threadgroup[0] == 0 or launch.threadgroup[1] == 0 or launch.threadgroup[2] == 0) return null;
1647 return .{
1648 .batch = batch,
1649 .m = m,
1650 .n = n,
1651 .k = k,
1652 .threads = .{ .x = launch.threadgroup[0], .y = launch.threadgroup[1], .z = launch.threadgroup[2] },
1653 .batch_axis = lhs.axes[0].name,
1654 .row_axis = lhs.axes[1].name,
1655 .col_axis = rhs.axes[2].name,
1656 .reduction_axis = lhs.axes[2].name,
1657 };
1658 }
1659
1660 pub fn matrixVectorProductFamilySpecialization(backing_allocator: std.mem.Allocator, instance: MatrixVectorProduct) !entry.OwnedSpecialization {
1661 var owned = entry.OwnedSpecialization.init(backing_allocator);
1662 errdefer owned.deinit();
1663 const lifetime_allocator = owned.allocator();
1664
1665 const inputs = try lifetime_allocator.alloc(entry.Shape, 2);
1666 inputs[0] = try entry.runtimeShape2D(lifetime_allocator, instance.row_axis, instance.m, instance.reduction_axis, instance.k);
1667 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, instance.reduction_axis, instance.k);
1668
1669 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
1670 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.row_axis, instance.m);
1671
1672 const reductions = try lifetime_allocator.alloc(entry.Reduction, 1);
1673 reductions[0] = try entry.runtimeReduction(
1674 lifetime_allocator,
1675 "dot",
1676 .dot_product,
1677 try entry.runtimeShape1D(lifetime_allocator, instance.reduction_axis, instance.k),
1678 );
1679
1680 owned.value = .{
1681 .dtype = .f32,
1682 .operation = .{ .linalg = .matrix_vector_product },
1683 .equation = "mk,k->m",
1684 .inputs = inputs,
1685 .outputs = outputs,
1686 .reductions = reductions,
1687 .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, instance.row_axis, instance.m, instance.threads),
1688 };
1689 owned.value.launch = owned.value.schedule.?.launch();
1690 var family = try matrixVectorProductShapeFamily(backing_allocator, instance);
1691 errdefer family.deinit();
1692 try owned.takeShapeFamily(&family);
1693 return owned;
1694 }
1695
1696 pub fn matrixVectorProductInstanceFromSpecialization(specialization: entry.Specialization) ?MatrixVectorProduct {
1697 if (!specialization.scheduleMatchesLaunch()) return null;
1698 if (!specialization.operationIs(.{ .linalg = .matrix_vector_product })) return null;
1699 const dtype = specialization.dtype orelse return null;
1700 if (dtype != .f32) return null;
1701 const equation = specialization.equation orelse return null;
1702 if (!std.mem.eql(u8, equation, "mk,k->m")) return null;
1703 if (specialization.inputs.len != 2 or specialization.outputs.len != 1 or specialization.reductions.len != 1) return null;
1704 const matrix = specialization.inputs[0];
1705 const vector = specialization.inputs[1];
1706 const output = specialization.outputs[0];
1707 const reduction = specialization.reductions[0];
1708 if (matrix.axes.len != 2 or vector.axes.len != 1 or output.axes.len != 1) return null;
1709 if (reduction.shape.axes.len != 1) return null;
1710 const m = matrix.axes[0].extent;
1711 const k = matrix.axes[1].extent;
1712 if (!std.mem.eql(u8, matrix.axes[1].name, vector.axes[0].name)) return null;
1713 if (!std.mem.eql(u8, matrix.axes[0].name, output.axes[0].name)) return null;
1714 if (!std.mem.eql(u8, reduction.shape.axes[0].name, matrix.axes[1].name)) return null;
1715 if (vector.axes[0].extent != k) return null;
1716 if (output.axes[0].extent != m) return null;
1717 if (reduction.shape.axes[0].extent != k) return null;
1718 const launch = specialization.launch orelse return null;
1719 if (launch.threadgroup[0] == 0) return null;
1720 return .{
1721 .m = m,
1722 .k = k,
1723 .threads = launch.threadgroup[0],
1724 .row_axis = matrix.axes[0].name,
1725 .reduction_axis = matrix.axes[1].name,
1726 };
1727 }
1728
1729 pub fn outerProductFamilySpecialization(backing_allocator: std.mem.Allocator, instance: OuterProduct) !entry.OwnedSpecialization {
1730 var owned = entry.OwnedSpecialization.init(backing_allocator);
1731 errdefer owned.deinit();
1732 const lifetime_allocator = owned.allocator();
1733
1734 const inputs = try lifetime_allocator.alloc(entry.Shape, 2);
1735 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.lhs_axis, instance.m);
1736 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, instance.rhs_axis, instance.n);
1737
1738 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
1739 outputs[0] = try entry.runtimeShape2D(lifetime_allocator, instance.lhs_axis, instance.m, instance.rhs_axis, instance.n);
1740
1741 owned.value = .{
1742 .dtype = .f32,
1743 .operation = .{ .linalg = .outer_product },
1744 .equation = "m,n->mn",
1745 .inputs = inputs,
1746 .outputs = outputs,
1747 .schedule = try entry.runtimeThreadBlocks2D(lifetime_allocator, instance.rhs_axis, instance.n, instance.lhs_axis, instance.m, instance.threads.x, instance.threads.y),
1748 };
1749 owned.value.launch = owned.value.schedule.?.launch();
1750 var family = try outerProductShapeFamily(backing_allocator, instance);
1751 errdefer family.deinit();
1752 try owned.takeShapeFamily(&family);
1753 return owned;
1754 }
1755
1756 pub fn outerProductInstanceFromSpecialization(specialization: entry.Specialization) ?OuterProduct {
1757 if (!specialization.scheduleMatchesLaunch()) return null;
1758 if (!specialization.operationIs(.{ .linalg = .outer_product })) return null;
1759 const dtype = specialization.dtype orelse return null;
1760 if (dtype != .f32) return null;
1761 const equation = specialization.equation orelse return null;
1762 if (!std.mem.eql(u8, equation, "m,n->mn")) return null;
1763 if (specialization.inputs.len != 2 or specialization.outputs.len != 1 or specialization.reductions.len != 0) return null;
1764 const lhs = specialization.inputs[0];
1765 const rhs = specialization.inputs[1];
1766 const output = specialization.outputs[0];
1767 if (lhs.axes.len != 1 or rhs.axes.len != 1 or output.axes.len != 2) return null;
1768 const m = lhs.axes[0].extent;
1769 const n = rhs.axes[0].extent;
1770 if (!std.mem.eql(u8, lhs.axes[0].name, output.axes[0].name)) return null;
1771 if (!std.mem.eql(u8, rhs.axes[0].name, output.axes[1].name)) return null;
1772 if (output.axes[0].extent != m or output.axes[1].extent != n) return null;
1773 const launch = specialization.launch orelse return null;
1774 if (launch.threadgroup[0] == 0 or launch.threadgroup[1] == 0) return null;
1775 return .{
1776 .m = m,
1777 .n = n,
1778 .threads = .{ .x = launch.threadgroup[0], .y = launch.threadgroup[1] },
1779 .lhs_axis = lhs.axes[0].name,
1780 .rhs_axis = rhs.axes[0].name,
1781 };
1782 }
1783
1784 fn batchedMatrixProductProgram(comptime spec: BatchedMatrixProduct) type {
1785 const Body = struct {
1786 fn run(k: anytype, args: anytype) !void {
1787 try batchedMatrixProductBody(k, spec, args);
1788 }
1789 };
1790
1791 return kernel.logical.Program(.{
1792 .name = std.fmt.comptimePrint(
1793 "accy_kernel_linalg_batched_matmul{}x{}x{}x{}_{}x{}x{}_f32",
1794 .{ spec.batch, spec.m, spec.n, spec.k, spec.threads.x, spec.threads.y, spec.threads.z },
1795 ),
1796 .parameters = .{
1797 .dst = kernel.dynamicBuffer(.f32),
1798 .lhs = kernel.dynamicBuffer(.f32),
1799 .rhs = kernel.dynamicBuffer(.f32),
1800 },
1801 .body = Body.run,
1802 }).withSchedule(kernel.logical.schedule.threadBlocks(.{
1803 .x = spec.threads.x,
1804 .y = spec.threads.y,
1805 .z = spec.threads.z,
1806 }));
1807 }
1808
1809 fn matrixProductProgram(comptime spec: MatrixProduct) type {
1810 const Body = struct {
1811 fn run(k: anytype, args: anytype) !void {
1812 try matrixProductBody(k, spec, args);
1813 }
1814 };
1815
1816 return kernel.logical.Program(.{
1817 .name = std.fmt.comptimePrint(
1818 "accy_kernel_linalg_matmul{}x{}x{}_{}x{}_{s}",
1819 .{ spec.m, spec.n, spec.k, spec.threads.x, spec.threads.y, spec.dtype.name() },
1820 ),
1821 .parameters = .{
1822 .dst = kernel.dynamicBuffer(spec.dtype),
1823 .lhs = kernel.dynamicBuffer(spec.dtype),
1824 .rhs = kernel.dynamicBuffer(spec.dtype),
1825 },
1826 .body = Body.run,
1827 }).withSchedule(kernel.logical.schedule.threadBlocks(.{
1828 .x = spec.threads.x,
1829 .y = spec.threads.y,
1830 }));
1831 }
1832
1833 fn matrixVectorProductProgram(comptime spec: MatrixVectorProduct) type {
1834 const Body = struct {
1835 fn run(k: anytype, args: anytype) !void {
1836 try matrixVectorProductBody(k, spec, args);
1837 }
1838 };
1839
1840 return kernel.logical.Program(.{
1841 .name = std.fmt.comptimePrint(
1842 "accy_kernel_linalg_matvec{}x{}_{}x_f32",
1843 .{ spec.m, spec.k, spec.threads },
1844 ),
1845 .parameters = .{
1846 .dst = kernel.dynamicBuffer(.f32),
1847 .matrix = kernel.dynamicBuffer(.f32),
1848 .vector = kernel.dynamicBuffer(.f32),
1849 },
1850 .body = Body.run,
1851 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
1852 }
1853
1854 fn outerProductProgram(comptime spec: OuterProduct) type {
1855 const Body = struct {
1856 fn run(k: anytype, args: anytype) !void {
1857 try outerProductBody(k, spec, args);
1858 }
1859 };
1860
1861 return kernel.logical.Program(.{
1862 .name = std.fmt.comptimePrint(
1863 "accy_kernel_linalg_outer{}x{}_{}x{}_f32",
1864 .{ spec.m, spec.n, spec.threads.x, spec.threads.y },
1865 ),
1866 .parameters = .{
1867 .dst = kernel.dynamicBuffer(.f32),
1868 .lhs = kernel.dynamicBuffer(.f32),
1869 .rhs = kernel.dynamicBuffer(.f32),
1870 },
1871 .body = Body.run,
1872 }).withSchedule(kernel.logical.schedule.threadBlocks(.{
1873 .x = spec.threads.x,
1874 .y = spec.threads.y,
1875 }));
1876 }
1877
1878 pub fn batchedMatrixProductF32(comptime spec: BatchedMatrixProduct) type {
1879 return entry.Entry(batchedMatrixProductProgram(spec), .{
1880 .target = std.fmt.comptimePrint(
1881 "accy.kernel.linalg.batched_matmul{}x{}x{}x{}_{}x{}x{}_f32",
1882 .{ spec.batch, spec.m, spec.n, spec.k, spec.threads.x, spec.threads.y, spec.threads.z },
1883 ),
1884 .layer = .logical,
1885 .category = .linalg,
1886 .specialization = batchedMatrixProductSpecialization(spec),
1887 });
1888 }
1889
1890 pub fn matrixProductF32(comptime spec: MatrixProduct) type {
1891 return entry.Entry(matrixProductProgram(spec), .{
1892 .target = std.fmt.comptimePrint(
1893 "accy.kernel.linalg.matmul{}x{}x{}_{}x{}_{s}",
1894 .{ spec.m, spec.n, spec.k, spec.threads.x, spec.threads.y, spec.dtype.name() },
1895 ),
1896 .layer = .logical,
1897 .category = .linalg,
1898 .specialization = matrixProductSpecialization(spec),
1899 });
1900 }
1901
1902 pub fn matrixVectorProductF32(comptime spec: MatrixVectorProduct) type {
1903 return entry.Entry(matrixVectorProductProgram(spec), .{
1904 .target = std.fmt.comptimePrint(
1905 "accy.kernel.linalg.matvec{}x{}_{}x_f32",
1906 .{ spec.m, spec.k, spec.threads },
1907 ),
1908 .layer = .logical,
1909 .category = .linalg,
1910 .specialization = matrixVectorProductSpecialization(spec),
1911 });
1912 }
1913
1914 pub fn outerProductF32(comptime spec: OuterProduct) type {
1915 return entry.Entry(outerProductProgram(spec), .{
1916 .target = std.fmt.comptimePrint(
1917 "accy.kernel.linalg.outer{}x{}_{}x{}_f32",
1918 .{ spec.m, spec.n, spec.threads.x, spec.threads.y },
1919 ),
1920 .layer = .logical,
1921 .category = .linalg,
1922 .specialization = outerProductSpecialization(spec),
1923 });
1924 }
1925
1926 pub const BatchedMatrixProduct2x2x3x4F32 = batchedMatrixProductF32(.{
1927 .batch = 2,
1928 .m = 2,
1929 .n = 3,
1930 .k = 4,
1931 .threads = .{ .x = 3, .y = 2, .z = 2 },
1932 });
1933
1934 pub const MatrixProduct2x3x4F32 = matrixProductF32(.{
1935 .m = 2,
1936 .n = 3,
1937 .k = 4,
1938 .threads = .{ .x = 2, .y = 2 },
1939 });
1940
1941 pub const MatrixProduct4x16x8F32 = matrixProductF32(.{
1942 .m = 4,
1943 .n = 16,
1944 .k = 8,
1945 .threads = .{ .x = 8, .y = 4 },
1946 });
1947
1948 pub const MatrixProduct4x16x8ThreadBlocks4x2F32 = matrixProductF32(.{
1949 .m = 4,
1950 .n = 16,
1951 .k = 8,
1952 .threads = .{ .x = 4, .y = 2 },
1953 });
1954
1955 pub const MatrixProduct8x12x16F32 = matrixProductF32(.{
1956 .m = 8,
1957 .n = 12,
1958 .k = 16,
1959 .threads = .{ .x = 4, .y = 4 },
1960 });
1961
1962 pub const MatrixVectorProduct4x8F32 = matrixVectorProductF32(.{
1963 .m = 4,
1964 .k = 8,
1965 .threads = 4,
1966 });
1967
1968 pub const OuterProduct4x3F32 = outerProductF32(.{
1969 .m = 4,
1970 .n = 3,
1971 .threads = .{ .x = 3, .y = 2 },
1972 });
1973
1974 test "linalg batched matrix product entry runs on CPU and records schedule" {
1975 var lhs = [_]f32{
1976 1.0, 2.0, 3.0, 4.0,
1977 5.0, 6.0, 7.0, 8.0,
1978 2.0, 0.0, -2.0, 1.0,
1979 1.0, 3.0, 5.0, 7.0,
1980 };
1981 var rhs = [_]f32{
1982 1.0, 0.0, 2.0,
1983 0.0, 1.0, 3.0,
1984 1.0, 1.0, 0.0,
1985 2.0, 0.0, 1.0,
1986 -1.0, 2.0, 0.0,
1987 3.0, 1.0, -2.0,
1988 0.0, 4.0, 1.0,
1989 2.0, -1.0, 3.0,
1990 };
1991 var dst = @as([12]f32, @splat(0.0));
1992
1993 try BatchedMatrixProduct2x2x3x4F32.runCpu(std.testing.allocator, BatchedMatrixProduct2x2x3x4F32.Limits.testing, &.{
1994 kernel.argumentBuffer(f32, dst[0..]),
1995 kernel.argumentBuffer(f32, lhs[0..]),
1996 kernel.argumentBuffer(f32, rhs[0..]),
1997 });
1998 try std.testing.expectEqualSlices(f32, &.{ 12.0, 5.0, 12.0, 28.0, 13.0, 36.0, 0.0, -5.0, 1.0, 22.0, 18.0, 20.0 }, dst[0..]);
1999
2000 const launch_value = try BatchedMatrixProduct2x2x3x4F32.launch(std.testing.allocator, BatchedMatrixProduct2x2x3x4F32.Limits.testing);
2001 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
2002 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);
2003 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[2]);
2004 try std.testing.expectEqual(@as(u32, 3), launch_value.block[0]);
2005 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
2006 try std.testing.expectEqual(@as(u32, 2), launch_value.block[2]);
2007 }
2008
2009 test "linalg batched matrix product entry carries einsum specialization metadata" {
2010 const BatchedMatrixProduct3x2x4x5F32 = batchedMatrixProductF32(.{
2011 .batch = 3,
2012 .m = 2,
2013 .n = 4,
2014 .k = 5,
2015 .threads = .{ .x = 2, .y = 2, .z = 1 },
2016 });
2017
2018 try std.testing.expect(BatchedMatrixProduct3x2x4x5F32.specialization.operationIs(.{ .linalg = .batched_matrix_product }));
2019 try std.testing.expectEqualStrings("bmk,bkn->bmn", BatchedMatrixProduct3x2x4x5F32.specialization.equation.?);
2020 try std.testing.expectEqualStrings("accy.kernel.linalg.batched_matmul3x2x4x5_2x2x1_f32", BatchedMatrixProduct3x2x4x5F32.target);
2021 try std.testing.expectEqual(@as(usize, 2), BatchedMatrixProduct3x2x4x5F32.specialization.inputs.len);
2022 try std.testing.expectEqual(@as(u64, 30), BatchedMatrixProduct3x2x4x5F32.specialization.inputs[0].elementCount().?);
2023 try std.testing.expectEqual(@as(u64, 60), BatchedMatrixProduct3x2x4x5F32.specialization.inputs[1].elementCount().?);
2024 try std.testing.expectEqual(@as(u64, 24), BatchedMatrixProduct3x2x4x5F32.specialization.outputs[0].elementCount().?);
2025 try std.testing.expectEqualStrings("dot", BatchedMatrixProduct3x2x4x5F32.specialization.reductions[0].name);
2026 try std.testing.expectEqual(entry.ReductionOperator.dot_product, BatchedMatrixProduct3x2x4x5F32.specialization.reductions[0].operator);
2027 try std.testing.expectEqual(@as(u64, 5), BatchedMatrixProduct3x2x4x5F32.specialization.reductions[0].shape.elementCount().?);
2028 try std.testing.expectEqual(@as(u32, 2), BatchedMatrixProduct3x2x4x5F32.specialization.launch.?.grid[0]);
2029 try std.testing.expectEqual(@as(u32, 1), BatchedMatrixProduct3x2x4x5F32.specialization.launch.?.grid[1]);
2030 try std.testing.expectEqual(@as(u32, 3), BatchedMatrixProduct3x2x4x5F32.specialization.launch.?.grid[2]);
2031 try std.testing.expectEqualDeep(BatchedMatrixProduct3x2x4x5F32.specialization.launch.?, BatchedMatrixProduct3x2x4x5F32.specialization.schedule.?.launch());
2032 try std.testing.expectEqual(@as(usize, 5), BatchedMatrixProduct3x2x4x5F32.specialization.schedule.?.bindings.len);
2033 try std.testing.expectEqualStrings("n_tile", BatchedMatrixProduct3x2x4x5F32.specialization.schedule.?.bindings[0].axis);
2034 try std.testing.expectEqual(kernel.BindTarget.block_x, BatchedMatrixProduct3x2x4x5F32.specialization.schedule.?.bindings[0].target);
2035 try std.testing.expectEqualStrings("b_lane", BatchedMatrixProduct3x2x4x5F32.specialization.schedule.?.bindings[4].axis);
2036 try std.testing.expectEqual(kernel.BindTarget.thread_z, BatchedMatrixProduct3x2x4x5F32.specialization.schedule.?.bindings[4].target);
2037
2038 var snapshot = try BatchedMatrixProduct3x2x4x5F32.scheduleSnapshot(std.testing.allocator, BatchedMatrixProduct3x2x4x5F32.Limits.testing);
2039 defer snapshot.deinit(std.testing.allocator);
2040 try std.testing.expect(BatchedMatrixProduct3x2x4x5F32.specialization.schedule.?.matchesSnapshot(&snapshot));
2041 }
2042
2043 test "linalg batched matrix product entry creates registry-ready artifact" {
2044 const allocator = std.testing.allocator;
2045 var state = gpu.recording.BackendState{
2046 .allocator = allocator,
2047 .kind = .cuda,
2048 .format = .cuda_ptx,
2049 };
2050
2051 var call_artifact = try BatchedMatrixProduct2x2x3x4F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = BatchedMatrixProduct2x2x3x4F32.Limits.testing });
2052 defer call_artifact.deinit();
2053
2054 const artifact = call_artifact.registry().find(BatchedMatrixProduct2x2x3x4F32.target, BatchedMatrixProduct2x2x3x4F32.version, .cuda_ptx) orelse {
2055 return error.TestExpectedKernelCallArtifact;
2056 };
2057 try std.testing.expectEqualStrings(BatchedMatrixProduct2x2x3x4F32.name, artifact.entry_name);
2058 try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);
2059 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
2060 switch (artifact.launch) {
2061 .fixed => |geometry| {
2062 try std.testing.expectEqual(BatchedMatrixProduct2x2x3x4F32.specialization.launch.?.grid[0], geometry.grid[0]);
2063 try std.testing.expectEqual(BatchedMatrixProduct2x2x3x4F32.specialization.launch.?.grid[2], geometry.grid[2]);
2064 try std.testing.expectEqual(BatchedMatrixProduct2x2x3x4F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
2065 try std.testing.expectEqual(BatchedMatrixProduct2x2x3x4F32.specialization.launch.?.threadgroup[2], geometry.threadgroup[2]);
2066 },
2067 else => return error.TestExpectedFixedLaunch,
2068 }
2069 }
2070
2071 test "linalg batched matrix product family matches the fixed entry at its extents" {
2072 const instance = BatchedMatrixProduct{
2073 .batch = 2,
2074 .m = 2,
2075 .n = 3,
2076 .k = 4,
2077 .threads = .{ .x = 3, .y = 2, .z = 2 },
2078 };
2079
2080 const family_launch = try BatchedMatrixProductFamilyF32.launch(std.testing.allocator, BatchedMatrixProductFamilyF32.Limits.testing, instance);
2081 const fixed_launch = try BatchedMatrixProduct2x2x3x4F32.launch(std.testing.allocator, BatchedMatrixProduct2x2x3x4F32.Limits.testing);
2082 try std.testing.expectEqualDeep(fixed_launch, family_launch);
2083
2084 var family_snapshot = try BatchedMatrixProductFamilyF32.scheduleSnapshot(std.testing.allocator, BatchedMatrixProductFamilyF32.Limits.testing, instance);
2085 defer family_snapshot.deinit(std.testing.allocator);
2086 var fixed_snapshot = try BatchedMatrixProduct2x2x3x4F32.scheduleSnapshot(std.testing.allocator, BatchedMatrixProduct2x2x3x4F32.Limits.testing);
2087 defer fixed_snapshot.deinit(std.testing.allocator);
2088 try std.testing.expectEqual(fixed_snapshot.fingerprint(), family_snapshot.fingerprint());
2089
2090 var lhs = [_]f32{
2091 1.0, 2.0, 3.0, 4.0,
2092 5.0, 6.0, 7.0, 8.0,
2093 2.0, 0.0, -2.0, 1.0,
2094 1.0, 3.0, 5.0, 7.0,
2095 };
2096 var rhs = [_]f32{
2097 1.0, 0.0, 2.0,
2098 0.0, 1.0, 3.0,
2099 1.0, 1.0, 0.0,
2100 2.0, 0.0, 1.0,
2101 -1.0, 2.0, 0.0,
2102 3.0, 1.0, -2.0,
2103 0.0, 4.0, 1.0,
2104 2.0, -1.0, 3.0,
2105 };
2106 var family_dst = @as([12]f32, @splat(0.0));
2107 var fixed_dst = @as([12]f32, @splat(0.0));
2108
2109 try BatchedMatrixProductFamilyF32.runCpu(std.testing.allocator, BatchedMatrixProductFamilyF32.Limits.testing, instance, &.{
2110 kernel.argumentBuffer(f32, family_dst[0..]),
2111 kernel.argumentBuffer(f32, lhs[0..]),
2112 kernel.argumentBuffer(f32, rhs[0..]),
2113 });
2114 try BatchedMatrixProduct2x2x3x4F32.runCpu(std.testing.allocator, BatchedMatrixProduct2x2x3x4F32.Limits.testing, &.{
2115 kernel.argumentBuffer(f32, fixed_dst[0..]),
2116 kernel.argumentBuffer(f32, lhs[0..]),
2117 kernel.argumentBuffer(f32, rhs[0..]),
2118 });
2119 try std.testing.expectEqualSlices(f32, fixed_dst[0..], family_dst[0..]);
2120 }
2121
2122 test "linalg batched matrix product family executes fresh runtime extents" {
2123 const instance = BatchedMatrixProduct{
2124 .batch = 3,
2125 .m = 5,
2126 .n = 6,
2127 .k = 4,
2128 .threads = .{ .x = 4, .y = 2, .z = 2 },
2129 };
2130
2131 var lhs: [3 * 5 * 4]f32 = undefined;
2132 for (&lhs, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index % 11)) * 0.25 - 1.0;
2133 var rhs: [3 * 4 * 6]f32 = undefined;
2134 for (&rhs, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index % 13)) * 0.125 - 0.5;
2135
2136 var expected: [3 * 5 * 6]f32 = undefined;
2137 for (0..3) |batch| {
2138 for (0..5) |row| {
2139 for (0..6) |col| {
2140 var sum: f32 = 0.0;
2141 for (0..4) |offset| {
2142 const lhs_index = batch * 5 * 4 + row * 4 + offset;
2143 const rhs_index = batch * 4 * 6 + offset * 6 + col;
2144 sum += lhs[lhs_index] * rhs[rhs_index];
2145 }
2146 expected[batch * 5 * 6 + row * 6 + col] = sum;
2147 }
2148 }
2149 }
2150
2151 var dst = @as([(3 * 5 * 6)]f32, @splat(0.0));
2152 try BatchedMatrixProductFamilyF32.runCpu(std.testing.allocator, BatchedMatrixProductFamilyF32.Limits.testing, instance, &.{
2153 kernel.argumentBuffer(f32, dst[0..]),
2154 kernel.argumentBuffer(f32, lhs[0..]),
2155 kernel.argumentBuffer(f32, rhs[0..]),
2156 });
2157 try std.testing.expectEqualSlices(f32, expected[0..], dst[0..]);
2158
2159 const launch_value = try BatchedMatrixProductFamilyF32.launch(std.testing.allocator, BatchedMatrixProductFamilyF32.Limits.testing, instance);
2160 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
2161 try std.testing.expectEqual(@as(u32, 3), launch_value.grid[1]);
2162 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[2]);
2163 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
2164 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
2165 try std.testing.expectEqual(@as(u32, 2), launch_value.block[2]);
2166 }
2167
2168 test "linalg batched matrix product family identity and metadata" {
2169 const instance = BatchedMatrixProduct{
2170 .batch = 3,
2171 .m = 5,
2172 .n = 6,
2173 .k = 4,
2174 .threads = .{ .x = 4, .y = 2, .z = 2 },
2175 .batch_axis = "batches",
2176 .row_axis = "rows",
2177 .col_axis = "cols",
2178 .reduction_axis = "depth",
2179 };
2180 var owned = try batchedMatrixProductFamilySpecialization(std.testing.allocator, instance);
2181 defer owned.deinit();
2182 const specialization = owned.value;
2183
2184 const fixed_target = try batchedMatrixProductInstanceTarget(std.testing.allocator, .{
2185 .batch = 2,
2186 .m = 2,
2187 .n = 3,
2188 .k = 4,
2189 .threads = .{ .x = 3, .y = 2, .z = 2 },
2190 });
2191 defer std.testing.allocator.free(fixed_target);
2192 try std.testing.expectEqualStrings(BatchedMatrixProduct2x2x3x4F32.target, fixed_target);
2193
2194 const fixed_entry = try batchedMatrixProductInstanceEntryName(std.testing.allocator, .{
2195 .batch = 2,
2196 .m = 2,
2197 .n = 3,
2198 .k = 4,
2199 .threads = .{ .x = 3, .y = 2, .z = 2 },
2200 });
2201 defer std.testing.allocator.free(fixed_entry);
2202 try std.testing.expectEqualStrings(BatchedMatrixProduct2x2x3x4F32.name, fixed_entry);
2203
2204 const family_target = try batchedMatrixProductFamilyTarget(std.testing.allocator, instance);
2205 defer std.testing.allocator.free(family_target);
2206 try std.testing.expectEqualStrings("accy.kernel.linalg.batched_matmul_family_4x2x2_f32", family_target);
2207
2208 const family_entry = try batchedMatrixProductFamilyEntryName(std.testing.allocator, instance);
2209 defer std.testing.allocator.free(family_entry);
2210 try std.testing.expectEqualStrings("accy_kernel_linalg_batched_matmul_family_4x2x2_f32", family_entry);
2211 try std.testing.expectEqual(BatchedMatrixProduct2x2x3x4F32.version, batched_matrix_product_family_version);
2212
2213 try std.testing.expect(specialization.operationIs(.{ .linalg = .batched_matrix_product }));
2214 try std.testing.expectEqualStrings("bmk,bkn->bmn", specialization.equation.?);
2215 try std.testing.expectEqual(@as(usize, 2), specialization.inputs.len);
2216 try std.testing.expectEqual(@as(u64, 60), specialization.inputs[0].elementCount().?);
2217 try std.testing.expectEqual(@as(u64, 72), specialization.inputs[1].elementCount().?);
2218 try std.testing.expectEqual(@as(u64, 90), specialization.outputs[0].elementCount().?);
2219 try std.testing.expect(specialization.reductionMatches(0, .{ .name = "dot", .operator = .dot_product, .extents = &.{4} }));
2220 try std.testing.expectEqualStrings("batches", specialization.inputs[0].axes[0].name);
2221 try std.testing.expectEqualStrings("rows", specialization.inputs[0].axes[1].name);
2222 try std.testing.expectEqualStrings("depth", specialization.inputs[0].axes[2].name);
2223 try std.testing.expectEqualStrings("cols", specialization.inputs[1].axes[2].name);
2224 try std.testing.expect(specialization.scheduleMatchesLaunch());
2225 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.grid[0]);
2226 try std.testing.expectEqual(@as(u32, 3), specialization.launch.?.grid[1]);
2227 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.grid[2]);
2228 try std.testing.expectEqual(@as(u32, 4), specialization.launch.?.threadgroup[0]);
2229 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.threadgroup[1]);
2230 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.threadgroup[2]);
2231 try std.testing.expectEqual(@as(usize, 6), specialization.schedule.?.bindings.len);
2232 try std.testing.expectEqualStrings("cols_tile", specialization.schedule.?.bindings[0].axis);
2233 try std.testing.expectEqual(kernel.BindTarget.block_x, specialization.schedule.?.bindings[0].target);
2234 try std.testing.expectEqualStrings("batches_lane", specialization.schedule.?.bindings[5].axis);
2235 try std.testing.expectEqual(kernel.BindTarget.thread_z, specialization.schedule.?.bindings[5].target);
2236 try std.testing.expect(specialization.shape_family != null);
2237 try std.testing.expectEqual(try batchedMatrixProductFamilyFingerprint(std.testing.allocator, instance), specialization.shapeFamilyFingerprint().?);
2238 }
2239
2240 test "linalg batched matrix product reconstructs family instance from specialization" {
2241 const instance = BatchedMatrixProduct{
2242 .batch = 3,
2243 .m = 5,
2244 .n = 6,
2245 .k = 4,
2246 .threads = .{ .x = 4, .y = 2, .z = 2 },
2247 .batch_axis = "batches",
2248 .row_axis = "rows",
2249 .col_axis = "cols",
2250 .reduction_axis = "depth",
2251 };
2252 var owned = try batchedMatrixProductFamilySpecialization(std.testing.allocator, instance);
2253 defer owned.deinit();
2254 const reconstructed = batchedMatrixProductInstanceFromSpecialization(owned.value) orelse return error.TestExpectedBatchedMatrixProductInstance;
2255
2256 try std.testing.expectEqual(instance.batch, reconstructed.batch);
2257 try std.testing.expectEqual(instance.m, reconstructed.m);
2258 try std.testing.expectEqual(instance.n, reconstructed.n);
2259 try std.testing.expectEqual(instance.k, reconstructed.k);
2260 try std.testing.expectEqual(instance.threads.x, reconstructed.threads.x);
2261 try std.testing.expectEqual(instance.threads.y, reconstructed.threads.y);
2262 try std.testing.expectEqual(instance.threads.z, reconstructed.threads.z);
2263 try std.testing.expectEqualStrings(instance.batch_axis, reconstructed.batch_axis);
2264 try std.testing.expectEqualStrings(instance.row_axis, reconstructed.row_axis);
2265 try std.testing.expectEqualStrings(instance.col_axis, reconstructed.col_axis);
2266 try std.testing.expectEqualStrings(instance.reduction_axis, reconstructed.reduction_axis);
2267 }
2268
2269 test "linalg batched matrix product runtime family executes explicit runtime extents" {
2270 const allocator = std.testing.allocator;
2271 const compiled = BatchedMatrixProduct{
2272 .batch = 1,
2273 .m = 1,
2274 .n = 1,
2275 .k = 1,
2276 .threads = .{ .x = 4, .y = 2, .z = 2 },
2277 };
2278 const runtime = BatchedMatrixProduct{
2279 .batch = 2,
2280 .m = 3,
2281 .n = 4,
2282 .k = 5,
2283 .threads = compiled.threads,
2284 };
2285
2286 var graph = try BatchedMatrixProductRuntimeFamilyF32.build(allocator, BatchedMatrixProductRuntimeFamilyF32.Limits.testing, compiled);
2287 defer graph.deinit();
2288
2289 var lhs: [2 * 3 * 5]f32 = undefined;
2290 for (&lhs, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index % 9)) * 0.5 - 1.0;
2291 var rhs: [2 * 5 * 4]f32 = undefined;
2292 for (&rhs, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index % 7)) * 0.25 - 0.5;
2293
2294 var expected: [2 * 3 * 4]f32 = undefined;
2295 for (0..2) |batch| {
2296 for (0..3) |row| {
2297 for (0..4) |col| {
2298 var sum: f32 = 0.0;
2299 for (0..5) |offset| {
2300 const lhs_index = batch * 3 * 5 + row * 5 + offset;
2301 const rhs_index = batch * 5 * 4 + offset * 4 + col;
2302 sum += lhs[lhs_index] * rhs[rhs_index];
2303 }
2304 expected[batch * 3 * 4 + row * 4 + col] = sum;
2305 }
2306 }
2307 }
2308
2309 var dst = @as([(2 * 3 * 4)]f32, @splat(0.0));
2310 const launch_value = try entry.runtimeLaunch3D(runtime.n, runtime.m, runtime.batch, runtime.threads.x, runtime.threads.y, runtime.threads.z);
2311 try graph.runCpuWithLaunch(allocator, &.{
2312 kernel.argumentBuffer(f32, dst[0..]),
2313 kernel.argumentBuffer(f32, lhs[0..]),
2314 kernel.argumentBuffer(f32, rhs[0..]),
2315 kernel.argumentI32(@intCast(runtime.batch)),
2316 kernel.argumentI32(@intCast(runtime.m)),
2317 kernel.argumentI32(@intCast(runtime.n)),
2318 kernel.argumentI32(@intCast(runtime.k)),
2319 }, .{
2320 .grid = launch_value.grid,
2321 .block = launch_value.threadgroup,
2322 });
2323 try std.testing.expectEqualSlices(f32, expected[0..], dst[0..]);
2324 }
2325
2326 test "linalg batched matrix product family artifact carries runtime launch contract" {
2327 const allocator = std.testing.allocator;
2328 var state = gpu.recording.BackendState{
2329 .allocator = allocator,
2330 .kind = .cuda,
2331 .format = .cuda_ptx,
2332 };
2333 const instance = BatchedMatrixProduct{
2334 .batch = 3,
2335 .m = 5,
2336 .n = 6,
2337 .k = 4,
2338 .threads = .{ .x = 4, .y = 2, .z = 2 },
2339 };
2340
2341 var family_artifact = try createBatchedMatrixProductFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
2342 defer family_artifact.deinit();
2343 var fixed_artifact = try BatchedMatrixProduct2x2x3x4F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = BatchedMatrixProduct2x2x3x4F32.Limits.testing });
2344 defer fixed_artifact.deinit();
2345
2346 const family_entry = family_artifact.entry();
2347 const fixed_entry = fixed_artifact.entry();
2348 try std.testing.expect(!std.mem.eql(u8, fixed_entry.target, family_entry.target));
2349 try std.testing.expectEqualStrings("accy.kernel.linalg.batched_matmul_family_4x2x2_f32", family_entry.target);
2350 try std.testing.expectEqualStrings("accy_kernel_linalg_batched_matmul_family_4x2x2_f32", family_entry.entry_name);
2351 try std.testing.expectEqual(@as(u32, 7), family_entry.argument_count);
2352 try std.testing.expectEqual(@as(u32, 4), family_entry.runtime_scalar_argument_count);
2353 try std.testing.expect(family_entry.required_dtypes.contains(.f32));
2354 try std.testing.expect(family_entry.required_dtypes.contains(.i32));
2355 try std.testing.expect(fixed_entry.shape_family_fingerprint == null);
2356 try std.testing.expect(family_entry.shape_family_fingerprint != null);
2357 const profile = family_entry.shape_profile orelse return error.TestExpectedShapeProfile;
2358 try std.testing.expectEqualStrings("batched_matrix_product", profile.name);
2359 try std.testing.expectEqual(family_entry.shape_family_fingerprint.?, profile.fingerprint);
2360 try std.testing.expectEqual(@as(usize, 4), profile.dimensions.len);
2361 const batch_dimension = profile.runtimeScalarDimension(0) orelse return error.TestExpectedShapeProfile;
2362 try std.testing.expectEqualStrings("b", batch_dimension.name);
2363 const m_dimension = profile.runtimeScalarDimension(1) orelse return error.TestExpectedShapeProfile;
2364 try std.testing.expectEqualStrings("m", m_dimension.name);
2365 const n_dimension = profile.runtimeScalarDimension(2) orelse return error.TestExpectedShapeProfile;
2366 try std.testing.expectEqualStrings("n", n_dimension.name);
2367 const k_dimension = profile.runtimeScalarDimension(3) orelse return error.TestExpectedShapeProfile;
2368 try std.testing.expectEqualStrings("k", k_dimension.name);
2369 switch (family_entry.launch) {
2370 .derived => |launch| {
2371 try std.testing.expectEqual(@as(u32, 4), launch.threadgroup[0]);
2372 try std.testing.expectEqual(@as(u32, 2), launch.threadgroup[1]);
2373 try std.testing.expectEqual(@as(u32, 2), launch.threadgroup[2]);
2374 switch (launch.grid[0]) {
2375 .runtime_u32_ceil_div => |axis| {
2376 try std.testing.expectEqual(@as(u32, 2), axis.argument_index);
2377 try std.testing.expectEqual(@as(u32, 4), axis.divisor);
2378 },
2379 else => return error.TestExpectedDerivedLaunch,
2380 }
2381 switch (launch.grid[1]) {
2382 .runtime_u32_ceil_div => |axis| {
2383 try std.testing.expectEqual(@as(u32, 1), axis.argument_index);
2384 try std.testing.expectEqual(@as(u32, 2), axis.divisor);
2385 },
2386 else => return error.TestExpectedDerivedLaunch,
2387 }
2388 switch (launch.grid[2]) {
2389 .runtime_u32_ceil_div => |axis| {
2390 try std.testing.expectEqual(@as(u32, 0), axis.argument_index);
2391 try std.testing.expectEqual(@as(u32, 2), axis.divisor);
2392 },
2393 else => return error.TestExpectedDerivedLaunch,
2394 }
2395 const args = try batchedMatrixProductRuntimeArguments(instance);
2396 const geometry = try launch.geometry(args[0..]);
2397 try std.testing.expectEqual(@as(u32, 2), geometry.grid[0]);
2398 try std.testing.expectEqual(@as(u32, 3), geometry.grid[1]);
2399 try std.testing.expectEqual(@as(u32, 2), geometry.grid[2]);
2400 try std.testing.expectEqual(@as(u32, 4), geometry.threadgroup[0]);
2401 try std.testing.expectEqual(@as(u32, 2), geometry.threadgroup[1]);
2402 try std.testing.expectEqual(@as(u32, 2), geometry.threadgroup[2]);
2403 },
2404 .fixed => return error.TestExpectedDerivedLaunch,
2405 }
2406 }
2407
2408 test "linalg matrix product entry runs on CPU and records schedule" {
2409 var lhs = [_]f32{
2410 1.0, 2.0, 3.0, 4.0,
2411 5.0, 6.0, 7.0, 8.0,
2412 };
2413 var rhs = [_]f32{
2414 1.0, 0.0, 2.0,
2415 0.0, 1.0, 3.0,
2416 1.0, 1.0, 0.0,
2417 2.0, 0.0, 1.0,
2418 };
2419 var dst = @as([6]f32, @splat(0.0));
2420
2421 try MatrixProduct2x3x4F32.runCpu(std.testing.allocator, MatrixProduct2x3x4F32.Limits.testing, &.{
2422 kernel.argumentBuffer(f32, dst[0..]),
2423 kernel.argumentBuffer(f32, lhs[0..]),
2424 kernel.argumentBuffer(f32, rhs[0..]),
2425 });
2426 try std.testing.expectEqualSlices(f32, &.{ 12.0, 5.0, 12.0, 28.0, 13.0, 36.0 }, dst[0..]);
2427
2428 const launch_value = try MatrixProduct2x3x4F32.launch(std.testing.allocator, MatrixProduct2x3x4F32.Limits.testing);
2429 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
2430 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);
2431 try std.testing.expectEqual(@as(u32, 2), launch_value.block[0]);
2432 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
2433 }
2434
2435 test "linalg matrix product family matches the fixed entry at its extents" {
2436 const instance = MatrixProduct{ .m = 2, .n = 3, .k = 4, .threads = .{ .x = 2, .y = 2 } };
2437
2438 const family_launch = try MatrixProductFamilyF32.launch(std.testing.allocator, MatrixProductFamilyF32.Limits.testing, instance);
2439 const fixed_launch = try MatrixProduct2x3x4F32.launch(std.testing.allocator, MatrixProduct2x3x4F32.Limits.testing);
2440 try std.testing.expectEqualDeep(fixed_launch, family_launch);
2441
2442 var family_snapshot = try MatrixProductFamilyF32.scheduleSnapshot(std.testing.allocator, MatrixProductFamilyF32.Limits.testing, instance);
2443 defer family_snapshot.deinit(std.testing.allocator);
2444 var fixed_snapshot = try MatrixProduct2x3x4F32.scheduleSnapshot(std.testing.allocator, MatrixProduct2x3x4F32.Limits.testing);
2445 defer fixed_snapshot.deinit(std.testing.allocator);
2446 try std.testing.expectEqual(fixed_snapshot.fingerprint(), family_snapshot.fingerprint());
2447
2448 var lhs = [_]f32{
2449 1.0, 2.0, 3.0, 4.0,
2450 5.0, 6.0, 7.0, 8.0,
2451 };
2452 var rhs = [_]f32{
2453 1.0, 0.0, 2.0,
2454 0.0, 1.0, 3.0,
2455 1.0, 1.0, 0.0,
2456 2.0, 0.0, 1.0,
2457 };
2458 var family_dst = @as([6]f32, @splat(0.0));
2459 var fixed_dst = @as([6]f32, @splat(0.0));
2460
2461 try MatrixProductFamilyF32.runCpu(std.testing.allocator, MatrixProductFamilyF32.Limits.testing, instance, &.{
2462 kernel.argumentBuffer(f32, family_dst[0..]),
2463 kernel.argumentBuffer(f32, lhs[0..]),
2464 kernel.argumentBuffer(f32, rhs[0..]),
2465 });
2466 try MatrixProduct2x3x4F32.runCpu(std.testing.allocator, MatrixProduct2x3x4F32.Limits.testing, &.{
2467 kernel.argumentBuffer(f32, fixed_dst[0..]),
2468 kernel.argumentBuffer(f32, lhs[0..]),
2469 kernel.argumentBuffer(f32, rhs[0..]),
2470 });
2471 try std.testing.expectEqualSlices(f32, fixed_dst[0..], family_dst[0..]);
2472 }
2473
2474 fn expectApproxF16Slices(expected: []const f16, actual: []const f16) !void {
2475 try std.testing.expectEqual(expected.len, actual.len);
2476 for (expected, actual) |expected_value, actual_value| {
2477 try std.testing.expectApproxEqAbs(
2478 @as(f32, @floatCast(expected_value)),
2479 @as(f32, @floatCast(actual_value)),
2480 0.001,
2481 );
2482 }
2483 }
2484
2485 test "linalg matrix product f16 family accumulates in f32" {
2486 const instance = MatrixProduct{
2487 .m = 2,
2488 .n = 3,
2489 .k = 4,
2490 .dtype = .f16,
2491 .accumulation_dtype = .f32,
2492 .threads = .{ .x = 2, .y = 2 },
2493 };
2494
2495 var lhs = [_]f16{
2496 0.5, -1.0, 2.25, 0.125,
2497 3.0, -0.5, 1.5, -2.0,
2498 };
2499 var rhs = [_]f16{
2500 1.0, -0.5, 2.0,
2501 0.25, 1.5, -1.0,
2502 -2.0, 0.75, 0.5,
2503 3.0, -4.0, 0.25,
2504 };
2505 var expected: [6]f16 = undefined;
2506 for (0..2) |row| {
2507 for (0..3) |col| {
2508 var sum: f32 = 0.0;
2509 for (0..4) |offset| {
2510 sum += @as(f32, @floatCast(lhs[row * 4 + offset])) * @as(f32, @floatCast(rhs[offset * 3 + col]));
2511 }
2512 expected[row * 3 + col] = @floatCast(sum);
2513 }
2514 }
2515
2516 var dst = @as([6]f16, @splat(0.0));
2517 try MatrixProductFamilyF16.runCpu(std.testing.allocator, MatrixProductFamilyF16.Limits.testing, instance, &.{
2518 kernel.argumentBuffer(f16, dst[0..]),
2519 kernel.argumentBuffer(f16, lhs[0..]),
2520 kernel.argumentBuffer(f16, rhs[0..]),
2521 });
2522 try expectApproxF16Slices(expected[0..], dst[0..]);
2523
2524 var owned = try matrixProductFamilySpecialization(std.testing.allocator, instance);
2525 defer owned.deinit();
2526 try std.testing.expectEqual(@as(?DType, .f16), owned.value.dtype);
2527 try std.testing.expectEqual(@as(?DType, .f32), owned.value.accumulation_dtype);
2528
2529 var plan = try MatrixProductFamilyF16.createCheckedPlan(std.testing.allocator, MatrixProductFamilyF16.Limits.testing, instance, .{});
2530 defer plan.deinit();
2531 try std.testing.expectEqual(@as(u32, 3), plan.argument_count);
2532 try std.testing.expectEqualStrings("accy_kernel_linalg_matmul_f16", plan.entry_name);
2533 }
2534
2535 fn matrixProductOccupancy(m: u64, n: u64, threads: entry.Threads2D) f64 {
2536 const grid_x = (n + threads.x - 1) / threads.x;
2537 const grid_y = (m + threads.y - 1) / threads.y;
2538 const launched = grid_x * grid_y * threads.x * threads.y;
2539 return @as(f64, @floatFromInt(m * n)) / @as(f64, @floatFromInt(launched));
2540 }
2541
2542 test "linalg matrix product thread heuristic keeps occupancy high" {
2543 const skinny = matrixProductThreadsForExtents(1, 1000);
2544 try std.testing.expectEqual(@as(u32, 1), skinny.y);
2545 try std.testing.expectEqual(@as(u32, 64), skinny.x);
2546 try std.testing.expect(matrixProductOccupancy(1, 1000, skinny) >= 0.9);
2547 try std.testing.expect(matrixProductOccupancy(1, 1000, .{ .x = 16, .y = 16 }) < 0.07);
2548
2549 const tall = matrixProductThreadsForExtents(1000, 2);
2550 try std.testing.expectEqual(@as(u32, 16), tall.y);
2551 try std.testing.expectEqual(@as(u32, 2), tall.x);
2552 try std.testing.expect(matrixProductOccupancy(1000, 2, tall) >= 0.9);
2553
2554 const tiny = matrixProductThreadsForExtents(5, 7);
2555 try std.testing.expectEqual(@as(u32, 5), tiny.y);
2556 try std.testing.expectEqual(@as(u32, 7), tiny.x);
2557 try std.testing.expect(matrixProductOccupancy(5, 7, tiny) == 1.0);
2558
2559 const dense = matrixProductThreadsForExtents(1024, 1024);
2560 try std.testing.expectEqual(@as(u32, 16), dense.y);
2561 try std.testing.expectEqual(@as(u32, 16), dense.x);
2562 try std.testing.expect(matrixProductOccupancy(1024, 1024, dense) == 1.0);
2563
2564 const just_over_square = matrixProductThreadsForExtents(17, 17);
2565 try std.testing.expectEqual(@as(u32, 9), just_over_square.y);
2566 try std.testing.expectEqual(@as(u32, 17), just_over_square.x);
2567 try std.testing.expect(matrixProductOccupancy(17, 17, just_over_square) >= 0.9);
2568
2569 const just_over_rect = matrixProductThreadsForExtents(17, 9);
2570 try std.testing.expectEqual(@as(u32, 9), just_over_rect.y);
2571 try std.testing.expectEqual(@as(u32, 9), just_over_rect.x);
2572 try std.testing.expect(matrixProductOccupancy(17, 9, just_over_rect) >= 0.9);
2573
2574 const larger_square = matrixProductThreadsForExtents(33, 33);
2575 try std.testing.expectEqual(@as(u32, 7), larger_square.y);
2576 try std.testing.expectEqual(@as(u32, 33), larger_square.x);
2577 try std.testing.expect(matrixProductOccupancy(33, 33, larger_square) >= 0.9);
2578
2579 const larger_rect = matrixProductThreadsForExtents(33, 17);
2580 try std.testing.expectEqual(@as(u32, 11), larger_rect.y);
2581 try std.testing.expectEqual(@as(u32, 17), larger_rect.x);
2582 try std.testing.expect(matrixProductOccupancy(33, 17, larger_rect) == 1.0);
2583 }
2584
2585 fn expectMatrixProductThreadCandidatesLegal(
2586 candidates: geometry_mod.ThreadCandidates,
2587 m: u64,
2588 n: u64,
2589 ) !void {
2590 try std.testing.expect(candidates.count != 0);
2591 for (candidates.slice(), 0..) |candidate, index| {
2592 try std.testing.expect(candidate.x != 0);
2593 try std.testing.expect(candidate.y != 0);
2594 try std.testing.expect(candidate.x <= @min(@max(n, 1), matrix_product_thread_caps.x_max));
2595 try std.testing.expect(candidate.y <= @min(@max(m, 1), matrix_product_thread_caps.y_max));
2596 try std.testing.expect(candidate.x * candidate.y <= matrix_product_thread_caps.budget);
2597 for (candidates.slice()[0..index]) |previous| {
2598 try std.testing.expect(!geometry_mod.threadCandidatesEqual(previous, candidate));
2599 }
2600 }
2601 }
2602
2603 fn expectMatrixProductThreadCandidatesContain(
2604 candidates: geometry_mod.ThreadCandidates,
2605 expected: entry.Threads2D,
2606 ) !void {
2607 for (candidates.slice()) |candidate| {
2608 if (geometry_mod.threadCandidatesEqual(candidate, expected)) return;
2609 }
2610 return error.TestExpectedMatrixProductThreadCandidate;
2611 }
2612
2613 test "linalg matrix product thread candidates expose stable family variants" {
2614 const near_square = matrixProductThreadCandidatesForExtents(17, 17);
2615 try expectMatrixProductThreadCandidatesLegal(near_square, 17, 17);
2616 try std.testing.expect(near_square.count > 2);
2617 try std.testing.expectEqual(@as(u32, 17), near_square.items[0].x);
2618 try std.testing.expectEqual(@as(u32, 9), near_square.items[0].y);
2619 try expectMatrixProductThreadCandidatesContain(near_square, .{ .x = 16, .y = 16 });
2620
2621 const skinny = matrixProductThreadCandidatesForExtents(1, 1000);
2622 try expectMatrixProductThreadCandidatesLegal(skinny, 1, 1000);
2623 try std.testing.expect(skinny.count > 1);
2624 try std.testing.expectEqual(@as(u32, 64), skinny.items[0].x);
2625 try std.testing.expectEqual(@as(u32, 1), skinny.items[0].y);
2626
2627 const first = MatrixProduct{ .m = 17, .n = 17, .k = 13, .threads = near_square.items[0] };
2628 const second = MatrixProduct{ .m = 17, .n = 17, .k = 13, .threads = near_square.items[1] };
2629 const first_target = try matrixProductFamilyTarget(std.testing.allocator, first);
2630 defer std.testing.allocator.free(first_target);
2631 const second_target = try matrixProductFamilyTarget(std.testing.allocator, second);
2632 defer std.testing.allocator.free(second_target);
2633 try std.testing.expect(!std.mem.eql(u8, first_target, second_target));
2634 }
2635
2636 test "linalg matrix product family instance identity matches fixed entry strings" {
2637 const instance = MatrixProduct{ .m = 2, .n = 3, .k = 4, .threads = .{ .x = 2, .y = 2 } };
2638
2639 const target = try matrixProductInstanceTarget(std.testing.allocator, instance);
2640 defer std.testing.allocator.free(target);
2641 try std.testing.expectEqualStrings(MatrixProduct2x3x4F32.target, target);
2642
2643 const entry_name = try matrixProductInstanceEntryName(std.testing.allocator, instance);
2644 defer std.testing.allocator.free(entry_name);
2645 try std.testing.expectEqualStrings(MatrixProduct2x3x4F32.name, entry_name);
2646
2647 try std.testing.expectEqual(MatrixProduct2x3x4F32.version, matrix_product_family_version);
2648
2649 const fresh = MatrixProduct{ .m = 5, .n = 7, .k = 3, .threads = .{ .x = 4, .y = 2 } };
2650 const fresh_target = try matrixProductInstanceTarget(std.testing.allocator, fresh);
2651 defer std.testing.allocator.free(fresh_target);
2652 try std.testing.expectEqualStrings("accy.kernel.linalg.matmul5x7x3_4x2_f32", fresh_target);
2653
2654 const family_target = try matrixProductFamilyTarget(std.testing.allocator, fresh);
2655 defer std.testing.allocator.free(family_target);
2656 try std.testing.expectEqualStrings("accy.kernel.linalg.matmul_family_4x2_f32", family_target);
2657
2658 const family_entry = try matrixProductFamilyEntryName(std.testing.allocator, fresh);
2659 defer std.testing.allocator.free(family_entry);
2660 try std.testing.expectEqualStrings("accy_kernel_linalg_matmul_family_4x2_f32", family_entry);
2661
2662 const fresh_f16 = MatrixProduct{ .m = 5, .n = 7, .k = 3, .dtype = .f16, .threads = .{ .x = 4, .y = 2 } };
2663 const fresh_f16_target = try matrixProductInstanceTarget(std.testing.allocator, fresh_f16);
2664 defer std.testing.allocator.free(fresh_f16_target);
2665 try std.testing.expectEqualStrings("accy.kernel.linalg.matmul5x7x3_4x2_f16", fresh_f16_target);
2666
2667 const family_f16_target = try matrixProductFamilyTarget(std.testing.allocator, fresh_f16);
2668 defer std.testing.allocator.free(family_f16_target);
2669 try std.testing.expectEqualStrings("accy.kernel.linalg.matmul_family_4x2_f16", family_f16_target);
2670 }
2671
2672 fn linalgFamilyTuningTestCapabilities() gpu.BackendCapabilities {
2673 return .{ .identity = .{
2674 .backend = .cuda,
2675 .family = .nvidia_cuda,
2676 .name = "linalg-family-tuning-test-device",
2677 .vendor_id = 0x10de,
2678 .device_id = 0x2684,
2679 } };
2680 }
2681
2682 test "linalg matrix product family tuning keys discriminate dtype and device" {
2683 const allocator = std.testing.allocator;
2684 const caps = linalgFamilyTuningTestCapabilities();
2685 const device = tuning.deviceFingerprint(caps);
2686
2687 const single = try matrixProductFamilyTuningKey(allocator, device, .{
2688 .m = 64,
2689 .n = 64,
2690 .k = 32,
2691 .dtype = .f32,
2692 });
2693 const half = try matrixProductFamilyTuningKey(allocator, device, .{
2694 .m = 64,
2695 .n = 64,
2696 .k = 32,
2697 .dtype = .f16,
2698 .accumulation_dtype = .f32,
2699 });
2700 try std.testing.expect(!single.eql(half));
2701 try std.testing.expect(single.family_fingerprint == half.family_fingerprint);
2702 try std.testing.expect(single.operation_fingerprint == half.operation_fingerprint);
2703
2704 const other_device = try matrixProductFamilyTuningKey(
2705 allocator,
2706 tuning.deviceFingerprint(.{ .identity = .{
2707 .backend = .cuda,
2708 .family = .nvidia_cuda,
2709 .name = "other-linalg-family-tuning-test-device",
2710 .vendor_id = 0x10de,
2711 .device_id = 0x1b80,
2712 } }),
2713 .{ .m = 64, .n = 64, .k = 32 },
2714 );
2715 try std.testing.expect(!other_device.eql(single));
2716 try std.testing.expectEqual(single.family_fingerprint, other_device.family_fingerprint);
2717 try std.testing.expectEqual(single.operation_fingerprint, other_device.operation_fingerprint);
2718 }
2719
2720 test "linalg matrix product family tuning resolves schedules" {
2721 const allocator = std.testing.allocator;
2722 const caps = linalgFamilyTuningTestCapabilities();
2723 const device = tuning.deviceFingerprint(caps);
2724
2725 const probe = MatrixProduct{ .m = 64, .n = 64, .k = 32 };
2726 const candidates = matrixProductThreadCandidatesForExtents(probe.m, probe.n);
2727 try std.testing.expect(candidates.slice().len >= 2);
2728 var winner_instance = probe;
2729 winner_instance.threads = candidates.slice()[candidates.slice().len - 1];
2730 const winner_target = try matrixProductFamilyTarget(allocator, winner_instance);
2731 defer allocator.free(winner_target);
2732
2733 const records = [_]tuning.FamilyTuningRecord{.{
2734 .key = try matrixProductFamilyTuningKey(allocator, device, probe),
2735 .target = winner_target,
2736 .winner_median_ns = 800,
2737 .runner_up_median_ns = 1100,
2738 .sample_count = 30,
2739 }};
2740 const reader = tuning.FamilyTuningReader.init(caps, .{ .records = records[0..] });
2741
2742 const resolved = (try resolveMatrixProductSchedule(allocator, reader, probe)) orelse
2743 return error.TestExpectedSchedule;
2744 try std.testing.expectEqual(winner_instance.threads, resolved);
2745
2746 const miss = try resolveMatrixProductSchedule(allocator, reader, .{ .m = 32, .n = 32, .k = 32 });
2747 try std.testing.expectEqual(@as(?entry.Threads2D, null), miss);
2748 }
2749
2750 test "linalg matrix product family artifact carries runtime launch contract" {
2751 const allocator = std.testing.allocator;
2752 var state = gpu.recording.BackendState{
2753 .allocator = allocator,
2754 .kind = .cuda,
2755 .format = .cuda_ptx,
2756 };
2757 const instance = MatrixProduct{ .m = 2, .n = 3, .k = 4, .threads = .{ .x = 2, .y = 2 } };
2758
2759 var family_artifact = try createMatrixProductFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
2760 defer family_artifact.deinit();
2761 var fixed_artifact = try MatrixProduct2x3x4F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = MatrixProduct2x3x4F32.Limits.testing });
2762 defer fixed_artifact.deinit();
2763
2764 const family_entry = family_artifact.entry();
2765 const fixed_entry = fixed_artifact.entry();
2766 try std.testing.expect(!std.mem.eql(u8, fixed_entry.target, family_entry.target));
2767 try std.testing.expectEqual(fixed_entry.version, family_entry.version);
2768 try std.testing.expectEqual(fixed_entry.format, family_entry.format);
2769 try std.testing.expectEqualStrings("accy.kernel.linalg.matmul_family_2x2_f32", family_entry.target);
2770 try std.testing.expectEqualStrings("accy_kernel_linalg_matmul_family_2x2_f32", family_entry.entry_name);
2771 try std.testing.expectEqual(@as(u32, 6), family_entry.argument_count);
2772 try std.testing.expectEqual(@as(u32, 3), family_entry.runtime_scalar_argument_count);
2773 try std.testing.expect(family_entry.required_dtypes.contains(.i32));
2774 try std.testing.expect(fixed_entry.shape_family_fingerprint == null);
2775 try std.testing.expect(family_entry.shape_family_fingerprint != null);
2776 const profile = family_entry.shape_profile orelse return error.TestExpectedShapeProfile;
2777 try std.testing.expectEqualStrings("matrix_product", profile.name);
2778 try std.testing.expectEqual(family_entry.shape_family_fingerprint.?, profile.fingerprint);
2779 try std.testing.expectEqual(@as(usize, 3), profile.dimensions.len);
2780 const m_dimension = profile.runtimeScalarDimension(0) orelse return error.TestExpectedShapeProfile;
2781 try std.testing.expectEqualStrings("m", m_dimension.name);
2782 try std.testing.expectEqual(@as(?u64, extent_mod.runtime_extent_max), m_dimension.bounds.max);
2783 switch (family_entry.launch) {
2784 .derived => |launch| {
2785 try std.testing.expectEqual(@as(u32, 2), launch.threadgroup[0]);
2786 try std.testing.expectEqual(@as(u32, 2), launch.threadgroup[1]);
2787 switch (launch.grid[0]) {
2788 .runtime_u32_ceil_div => |axis| {
2789 try std.testing.expectEqual(@as(u32, 1), axis.argument_index);
2790 try std.testing.expectEqual(@as(u32, 2), axis.divisor);
2791 },
2792 else => return error.TestExpectedDerivedLaunch,
2793 }
2794 switch (launch.grid[1]) {
2795 .runtime_u32_ceil_div => |axis| {
2796 try std.testing.expectEqual(@as(u32, 0), axis.argument_index);
2797 try std.testing.expectEqual(@as(u32, 2), axis.divisor);
2798 },
2799 else => return error.TestExpectedDerivedLaunch,
2800 }
2801 },
2802 .fixed => return error.TestExpectedDerivedLaunch,
2803 }
2804 }
2805
2806 test "linalg matrix product family artifact resolves fresh extents in registry" {
2807 const allocator = std.testing.allocator;
2808 var state = gpu.recording.BackendState{
2809 .allocator = allocator,
2810 .kind = .cuda,
2811 .format = .cuda_ptx,
2812 };
2813 const instance = MatrixProduct{ .m = 5, .n = 7, .k = 3, .threads = .{ .x = 4, .y = 2 } };
2814
2815 var call_artifact = try createMatrixProductFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
2816 defer call_artifact.deinit();
2817
2818 const sibling = MatrixProduct{ .m = 11, .n = 13, .k = 17, .threads = .{ .x = 4, .y = 2 } };
2819 const sibling_target = try matrixProductFamilyTarget(allocator, sibling);
2820 defer allocator.free(sibling_target);
2821 try std.testing.expectEqualStrings("accy.kernel.linalg.matmul_family_4x2_f32", sibling_target);
2822
2823 const artifact = call_artifact.registry().find(
2824 "accy.kernel.linalg.matmul_family_4x2_f32",
2825 matrix_product_family_version,
2826 .cuda_ptx,
2827 ) orelse return error.TestExpectedKernelCallArtifact;
2828 try std.testing.expectEqualStrings("accy_kernel_linalg_matmul_family_4x2_f32", artifact.entry_name);
2829 try std.testing.expectEqual(@as(u32, 6), artifact.argument_count);
2830 try std.testing.expectEqual(@as(u32, 3), artifact.runtime_scalar_argument_count);
2831 try std.testing.expectEqual(try matrixProductFamilyFingerprint(allocator, instance), artifact.shape_family_fingerprint.?);
2832 const profile = artifact.shape_profile orelse return error.TestExpectedShapeProfile;
2833 try std.testing.expectEqualStrings("matrix_product", profile.name);
2834 try std.testing.expectEqual(artifact.shape_family_fingerprint.?, profile.fingerprint);
2835 const k_dimension = profile.dimension("k") orelse return error.TestExpectedShapeProfile;
2836 try std.testing.expectEqual(@as(u32, 2), k_dimension.runtime_scalar_argument_index);
2837 switch (artifact.launch) {
2838 .derived => |launch| {
2839 const args = try matrixProductRuntimeArguments(instance);
2840 const geometry = try launch.geometry(args[0..]);
2841 try std.testing.expectEqual(@as(u32, 2), geometry.grid[0]);
2842 try std.testing.expectEqual(@as(u32, 3), geometry.grid[1]);
2843 try std.testing.expectEqual(@as(u32, 4), geometry.threadgroup[0]);
2844 try std.testing.expectEqual(@as(u32, 2), geometry.threadgroup[1]);
2845 },
2846 .fixed => return error.TestExpectedDerivedLaunch,
2847 }
2848 }
2849
2850 test "linalg matrix product f16 family artifact carries dtype contract" {
2851 const allocator = std.testing.allocator;
2852 var state = gpu.recording.BackendState{
2853 .allocator = allocator,
2854 .kind = .cuda,
2855 .format = .cuda_ptx,
2856 };
2857 const instance = MatrixProduct{
2858 .m = 5,
2859 .n = 7,
2860 .k = 3,
2861 .dtype = .f16,
2862 .accumulation_dtype = .f32,
2863 .threads = .{ .x = 4, .y = 2 },
2864 };
2865
2866 var call_artifact = try createMatrixProductFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
2867 defer call_artifact.deinit();
2868
2869 const artifact = call_artifact.registry().find(
2870 "accy.kernel.linalg.matmul_family_4x2_f16",
2871 matrix_product_family_version,
2872 .cuda_ptx,
2873 ) orelse return error.TestExpectedKernelCallArtifact;
2874 try std.testing.expectEqualStrings("accy_kernel_linalg_matmul_family_4x2_f16", artifact.entry_name);
2875 try std.testing.expectEqual(@as(u32, 6), artifact.argument_count);
2876 try std.testing.expect(artifact.required_dtypes.contains(.f16));
2877 try std.testing.expect(artifact.required_dtypes.contains(.i32));
2878 }
2879
2880 test "linalg matrix product runtime family executes explicit runtime extents" {
2881 const allocator = std.testing.allocator;
2882 const compiled = MatrixProduct{ .m = 1, .n = 1, .k = 1, .threads = .{ .x = 4, .y = 2 } };
2883 const runtime = MatrixProduct{ .m = 3, .n = 4, .k = 2, .threads = compiled.threads };
2884
2885 var graph = try MatrixProductRuntimeFamilyF32.build(allocator, MatrixProductRuntimeFamilyF32.Limits.testing, compiled);
2886 defer graph.deinit();
2887
2888 var lhs: [6]f32 = undefined;
2889 for (&lhs, 0..) |*value, index| value.* = @floatFromInt(index + 1);
2890 var rhs: [8]f32 = undefined;
2891 for (&rhs, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index)) * 0.25 + 1.0;
2892
2893 var expected: [12]f32 = undefined;
2894 for (0..3) |row| {
2895 for (0..4) |col| {
2896 var sum: f32 = 0.0;
2897 for (0..2) |offset| {
2898 sum += lhs[row * 2 + offset] * rhs[offset * 4 + col];
2899 }
2900 expected[row * 4 + col] = sum;
2901 }
2902 }
2903
2904 var dst = @as([12]f32, @splat(0.0));
2905 const launch_value = try entry.runtimeLaunch2D(runtime.n, runtime.m, runtime.threads.x, runtime.threads.y);
2906 try graph.runCpuWithLaunch(allocator, &.{
2907 kernel.argumentBuffer(f32, dst[0..]),
2908 kernel.argumentBuffer(f32, lhs[0..]),
2909 kernel.argumentBuffer(f32, rhs[0..]),
2910 kernel.argumentI32(@intCast(runtime.m)),
2911 kernel.argumentI32(@intCast(runtime.n)),
2912 kernel.argumentI32(@intCast(runtime.k)),
2913 }, .{
2914 .grid = launch_value.grid,
2915 .block = launch_value.threadgroup,
2916 });
2917 try std.testing.expectEqualSlices(f32, expected[0..], dst[0..]);
2918 }
2919
2920 test "linalg matrix product family records fixed-entry specialization metadata" {
2921 const instance = MatrixProduct{ .m = 2, .n = 3, .k = 4, .threads = .{ .x = 2, .y = 2 } };
2922 var owned = try matrixProductFamilySpecialization(std.testing.allocator, instance);
2923 defer owned.deinit();
2924 const specialization = owned.value;
2925
2926 try std.testing.expect(specialization.operationIs(.{ .linalg = .matrix_product }));
2927 try std.testing.expectEqualStrings(MatrixProduct2x3x4F32.specialization.equation.?, specialization.equation.?);
2928 try std.testing.expectEqual(MatrixProduct2x3x4F32.specialization.dtype, specialization.dtype);
2929 try std.testing.expectEqual(@as(usize, 2), specialization.inputs.len);
2930 try std.testing.expectEqual(MatrixProduct2x3x4F32.specialization.inputs[0].elementCount().?, specialization.inputs[0].elementCount().?);
2931 try std.testing.expectEqual(MatrixProduct2x3x4F32.specialization.inputs[1].elementCount().?, specialization.inputs[1].elementCount().?);
2932 try std.testing.expectEqual(MatrixProduct2x3x4F32.specialization.outputs[0].elementCount().?, specialization.outputs[0].elementCount().?);
2933 try std.testing.expectEqualStrings(MatrixProduct2x3x4F32.specialization.inputs[0].axes[0].name, specialization.inputs[0].axes[0].name);
2934 try std.testing.expectEqualStrings(MatrixProduct2x3x4F32.specialization.inputs[1].axes[1].name, specialization.inputs[1].axes[1].name);
2935 try std.testing.expectEqualStrings("dot", specialization.reductions[0].name);
2936 try std.testing.expect(specialization.reductionMatches(0, .{ .name = "dot", .operator = .dot_product, .extents = &.{4} }));
2937 try std.testing.expect(specialization.reductionDependenciesAreValid());
2938 try std.testing.expectEqualDeep(MatrixProduct2x3x4F32.specialization.launch.?, specialization.launch.?);
2939 try std.testing.expectEqualDeep(MatrixProduct2x3x4F32.specialization.schedule.?.launch(), specialization.schedule.?.launch());
2940 try std.testing.expect(specialization.shape_family != null);
2941 try std.testing.expectEqual(try matrixProductFamilyFingerprint(std.testing.allocator, instance), specialization.shapeFamilyFingerprint().?);
2942
2943 var snapshot = try MatrixProductFamilyF32.scheduleSnapshot(std.testing.allocator, MatrixProductFamilyF32.Limits.testing, instance);
2944 defer snapshot.deinit(std.testing.allocator);
2945 try std.testing.expect(specialization.schedule.?.matchesSnapshot(&snapshot));
2946 }
2947
2948 test "linalg matrix product family executes fresh runtime extents" {
2949 const instance = MatrixProduct{ .m = 5, .n = 7, .k = 3, .threads = .{ .x = 4, .y = 2 } };
2950
2951 var lhs: [15]f32 = undefined;
2952 for (&lhs, 0..) |*value, index| value.* = @floatFromInt(index + 1);
2953 var rhs: [21]f32 = undefined;
2954 for (&rhs, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index)) * 0.5 - 2.0;
2955
2956 var expected: [35]f32 = undefined;
2957 for (0..5) |row| {
2958 for (0..7) |col| {
2959 var sum: f32 = 0.0;
2960 for (0..3) |offset| {
2961 sum += lhs[row * 3 + offset] * rhs[offset * 7 + col];
2962 }
2963 expected[row * 7 + col] = sum;
2964 }
2965 }
2966
2967 var dst = @as([35]f32, @splat(0.0));
2968 try MatrixProductFamilyF32.runCpu(std.testing.allocator, MatrixProductFamilyF32.Limits.testing, instance, &.{
2969 kernel.argumentBuffer(f32, dst[0..]),
2970 kernel.argumentBuffer(f32, lhs[0..]),
2971 kernel.argumentBuffer(f32, rhs[0..]),
2972 });
2973 try std.testing.expectEqualSlices(f32, expected[0..], dst[0..]);
2974
2975 const launch_value = try MatrixProductFamilyF32.launch(std.testing.allocator, MatrixProductFamilyF32.Limits.testing, instance);
2976 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
2977 try std.testing.expectEqual(@as(u32, 3), launch_value.grid[1]);
2978 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
2979 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
2980
2981 var plan = try MatrixProductFamilyF32.createCheckedPlan(std.testing.allocator, MatrixProductFamilyF32.Limits.testing, instance, .{});
2982 defer plan.deinit();
2983 try std.testing.expectEqual(@as(u32, 3), plan.argument_count);
2984 try std.testing.expectEqualStrings("accy_kernel_linalg_matmul_f32", plan.entry_name);
2985 }
2986
2987 test "linalg matrix product family records fresh runtime specialization metadata" {
2988 const instance = MatrixProduct{
2989 .m = 5,
2990 .n = 7,
2991 .k = 3,
2992 .threads = .{ .x = 4, .y = 2 },
2993 .row_axis = "rows",
2994 .col_axis = "columns",
2995 .reduction_axis = "depth",
2996 };
2997 var owned = try matrixProductFamilySpecialization(std.testing.allocator, instance);
2998 defer owned.deinit();
2999 const specialization = owned.value;
3000
3001 try std.testing.expect(specialization.operationIs(.{ .linalg = .matrix_product }));
3002 try std.testing.expectEqualStrings("mk,kn->mn", specialization.equation.?);
3003 try std.testing.expectEqual(@as(usize, 2), specialization.inputs.len);
3004 try std.testing.expectEqual(@as(u64, 15), specialization.inputs[0].elementCount().?);
3005 try std.testing.expectEqual(@as(u64, 21), specialization.inputs[1].elementCount().?);
3006 try std.testing.expectEqual(@as(u64, 35), specialization.outputs[0].elementCount().?);
3007 try std.testing.expectEqualStrings("rows", specialization.inputs[0].axes[0].name);
3008 try std.testing.expectEqualStrings("depth", specialization.inputs[0].axes[1].name);
3009 try std.testing.expectEqualStrings("columns", specialization.outputs[0].axes[1].name);
3010 try std.testing.expect(specialization.reductionMatches(0, .{ .name = "dot", .operator = .dot_product, .extents = &.{3} }));
3011 try std.testing.expectEqual(@as(u64, 210), specialization.estimatedElementOps().?);
3012 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.grid[0]);
3013 try std.testing.expectEqual(@as(u32, 3), specialization.launch.?.grid[1]);
3014 try std.testing.expectEqual(@as(u32, 4), specialization.launch.?.threadgroup[0]);
3015 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.threadgroup[1]);
3016 try std.testing.expect(specialization.scheduleMatchesLaunch());
3017 try std.testing.expectEqual(@as(usize, 4), specialization.schedule.?.bindings.len);
3018 try std.testing.expect(specialization.shape_family != null);
3019 try std.testing.expectEqual(try matrixProductFamilyFingerprint(std.testing.allocator, instance), specialization.shapeFamilyFingerprint().?);
3020 try std.testing.expectEqualStrings("columns_tile", specialization.schedule.?.bindings[0].axis);
3021 try std.testing.expectEqual(kernel.BindTarget.block_x, specialization.schedule.?.bindings[0].target);
3022 try std.testing.expectEqualStrings("columns_lane", specialization.schedule.?.bindings[1].axis);
3023 try std.testing.expectEqual(kernel.BindTarget.thread_x, specialization.schedule.?.bindings[1].target);
3024 try std.testing.expectEqualStrings("rows_tile", specialization.schedule.?.bindings[2].axis);
3025 try std.testing.expectEqual(kernel.BindTarget.block_y, specialization.schedule.?.bindings[2].target);
3026 try std.testing.expectEqualStrings("rows_lane", specialization.schedule.?.bindings[3].axis);
3027 try std.testing.expectEqual(kernel.BindTarget.thread_y, specialization.schedule.?.bindings[3].target);
3028
3029 var snapshot = try MatrixProductFamilyF32.scheduleSnapshot(std.testing.allocator, MatrixProductFamilyF32.Limits.testing, instance);
3030 defer snapshot.deinit(std.testing.allocator);
3031 try std.testing.expect(specialization.schedule.?.matchesSnapshot(&snapshot));
3032 }
3033
3034 test "linalg matrix product reconstructs family instance from specialization" {
3035 const instance = MatrixProduct{
3036 .m = 5,
3037 .n = 7,
3038 .k = 3,
3039 .threads = .{ .x = 4, .y = 2 },
3040 .row_axis = "rows",
3041 .col_axis = "columns",
3042 .reduction_axis = "depth",
3043 };
3044 var owned = try matrixProductFamilySpecialization(std.testing.allocator, instance);
3045 defer owned.deinit();
3046 const reconstructed = matrixProductInstanceFromSpecialization(owned.value) orelse return error.TestExpectedMatrixProductInstance;
3047
3048 try std.testing.expectEqual(instance.m, reconstructed.m);
3049 try std.testing.expectEqual(instance.n, reconstructed.n);
3050 try std.testing.expectEqual(instance.k, reconstructed.k);
3051 try std.testing.expectEqual(instance.dtype, reconstructed.dtype);
3052 try std.testing.expectEqual(instance.accumulation_dtype, reconstructed.accumulation_dtype);
3053 try std.testing.expectEqual(instance.threads.x, reconstructed.threads.x);
3054 try std.testing.expectEqual(instance.threads.y, reconstructed.threads.y);
3055 try std.testing.expectEqualStrings(instance.row_axis, reconstructed.row_axis);
3056 try std.testing.expectEqualStrings(instance.col_axis, reconstructed.col_axis);
3057 try std.testing.expectEqualStrings(instance.reduction_axis, reconstructed.reduction_axis);
3058 }
3059
3060 test "linalg matrix product family records untiled runtime launch metadata" {
3061 const instance = MatrixProduct{ .m = 5, .n = 7, .k = 3 };
3062 var owned = try matrixProductFamilySpecialization(std.testing.allocator, instance);
3063 defer owned.deinit();
3064 const specialization = owned.value;
3065 const launch_value = try MatrixProductFamilyF32.launch(std.testing.allocator, MatrixProductFamilyF32.Limits.testing, instance);
3066
3067 try std.testing.expectEqual(launch_value.grid[0], specialization.launch.?.grid[0]);
3068 try std.testing.expectEqual(launch_value.grid[1], specialization.launch.?.grid[1]);
3069 try std.testing.expectEqual(launch_value.block[0], specialization.launch.?.threadgroup[0]);
3070 try std.testing.expectEqual(launch_value.block[1], specialization.launch.?.threadgroup[1]);
3071 try std.testing.expectEqual(@as(u32, 1), specialization.launch.?.grid[0]);
3072 try std.testing.expectEqual(@as(u32, 1), specialization.launch.?.grid[1]);
3073 try std.testing.expectEqual(@as(u32, 7), specialization.launch.?.threadgroup[0]);
3074 try std.testing.expectEqual(@as(u32, 5), specialization.launch.?.threadgroup[1]);
3075 try std.testing.expect(specialization.scheduleMatchesLaunch());
3076 }
3077
3078 test "linalg matrix product family fingerprint names symbolic family not point extents" {
3079 const first = MatrixProduct{ .m = 5, .n = 7, .k = 3, .threads = .{ .x = 4, .y = 2 } };
3080 const second = MatrixProduct{ .m = 11, .n = 13, .k = 17, .threads = .{ .x = 8, .y = 4 } };
3081 const renamed = MatrixProduct{
3082 .m = 5,
3083 .n = 7,
3084 .k = 3,
3085 .row_axis = "row",
3086 .col_axis = "col",
3087 .reduction_axis = "depth",
3088 };
3089
3090 try std.testing.expectEqual(
3091 try matrixProductFamilyFingerprint(std.testing.allocator, first),
3092 try matrixProductFamilyFingerprint(std.testing.allocator, second),
3093 );
3094 try std.testing.expect(
3095 try matrixProductFamilyFingerprint(std.testing.allocator, first) !=
3096 try matrixProductFamilyFingerprint(std.testing.allocator, renamed),
3097 );
3098 }
3099
3100 test "linalg matrix product shape family bounds runtime extents" {
3101 const instance = MatrixProduct{ .m = 5, .n = 7, .k = 3 };
3102 var family = try matrixProductShapeFamily(std.testing.allocator, instance);
3103 defer family.deinit();
3104
3105 var bound_count: usize = 0;
3106 const expected_bounds = matrixProductRuntimeExtentBounds();
3107 for (family.facts) |fact| {
3108 switch (fact.predicate) {
3109 .bound => |bound| {
3110 try std.testing.expectEqual(shape.FactMode.assume, fact.mode);
3111 try std.testing.expectEqual(expected_bounds.min, bound.bounds.min);
3112 try std.testing.expectEqual(expected_bounds.opt, bound.bounds.opt);
3113 try std.testing.expectEqual(expected_bounds.max, bound.bounds.max);
3114 bound_count += 1;
3115 },
3116 else => {},
3117 }
3118 }
3119
3120 try std.testing.expectEqual(@as(usize, 3), bound_count);
3121 }
3122
3123 test "linalg matrix product runtime arguments enforce extent bounds" {
3124 const largest = MatrixProduct{ .m = extent_mod.runtime_extent_max, .n = 1, .k = 1 };
3125 const largest_args = try matrixProductRuntimeArguments(largest);
3126 switch (largest_args[0]) {
3127 .u32 => |value| try std.testing.expectEqual(@as(u32, @intCast(extent_mod.runtime_extent_max)), value),
3128 else => return error.TestExpectedRuntimeExtent,
3129 }
3130
3131 const zero = MatrixProduct{ .m = 0, .n = 1, .k = 1 };
3132 try std.testing.expectError(error.ExtentOverflowsIndexRange, matrixProductRuntimeArguments(zero));
3133
3134 const too_large = MatrixProduct{ .m = extent_mod.runtime_extent_max + 1, .n = 1, .k = 1 };
3135 try std.testing.expectError(error.ExtentOverflowsIndexRange, matrixProductRuntimeArguments(too_large));
3136 }
3137
3138 test "linalg matrix product entry carries einsum specialization metadata" {
3139 const MatrixProduct4x5x6F32 = matrixProductF32(.{
3140 .m = 4,
3141 .n = 5,
3142 .k = 6,
3143 .threads = .{ .x = 4, .y = 2 },
3144 });
3145
3146 try std.testing.expect(MatrixProduct4x5x6F32.specialization.operationIs(.{ .linalg = .matrix_product }));
3147 try std.testing.expectEqualStrings("mk,kn->mn", MatrixProduct4x5x6F32.specialization.equation.?);
3148 try std.testing.expectEqualStrings("accy.kernel.linalg.matmul4x5x6_4x2_f32", MatrixProduct4x5x6F32.target);
3149 try std.testing.expectEqual(@as(usize, 2), MatrixProduct4x5x6F32.specialization.inputs.len);
3150 try std.testing.expectEqual(@as(u64, 24), MatrixProduct4x5x6F32.specialization.inputs[0].elementCount().?);
3151 try std.testing.expectEqual(@as(u64, 30), MatrixProduct4x5x6F32.specialization.inputs[1].elementCount().?);
3152 try std.testing.expectEqual(@as(u64, 20), MatrixProduct4x5x6F32.specialization.outputs[0].elementCount().?);
3153 try std.testing.expectEqualStrings("dot", MatrixProduct4x5x6F32.specialization.reductions[0].name);
3154 try std.testing.expectEqual(entry.ReductionOperator.dot_product, MatrixProduct4x5x6F32.specialization.reductions[0].operator);
3155 try std.testing.expectEqual(@as(u64, 6), MatrixProduct4x5x6F32.specialization.reductions[0].shape.elementCount().?);
3156 try std.testing.expectEqual(@as(u32, 2), MatrixProduct4x5x6F32.specialization.launch.?.grid[0]);
3157 try std.testing.expectEqual(@as(u32, 2), MatrixProduct4x5x6F32.specialization.launch.?.grid[1]);
3158 try std.testing.expectEqualDeep(MatrixProduct4x5x6F32.specialization.launch.?, MatrixProduct4x5x6F32.specialization.schedule.?.launch());
3159 try std.testing.expectEqual(@as(usize, 4), MatrixProduct4x5x6F32.specialization.schedule.?.bindings.len);
3160 try std.testing.expectEqualStrings("n_tile", MatrixProduct4x5x6F32.specialization.schedule.?.bindings[0].axis);
3161 try std.testing.expectEqual(kernel.BindTarget.block_x, MatrixProduct4x5x6F32.specialization.schedule.?.bindings[0].target);
3162 try std.testing.expectEqualStrings("m_lane", MatrixProduct4x5x6F32.specialization.schedule.?.bindings[3].axis);
3163 try std.testing.expectEqual(kernel.BindTarget.thread_y, MatrixProduct4x5x6F32.specialization.schedule.?.bindings[3].target);
3164
3165 var snapshot = try MatrixProduct4x5x6F32.scheduleSnapshot(std.testing.allocator, MatrixProduct4x5x6F32.Limits.testing);
3166 defer snapshot.deinit(std.testing.allocator);
3167 try std.testing.expect(MatrixProduct4x5x6F32.specialization.schedule.?.matchesSnapshot(&snapshot));
3168 }
3169
3170 test "linalg matrix product entry creates registry-ready artifact" {
3171 const allocator = std.testing.allocator;
3172 var state = gpu.recording.BackendState{
3173 .allocator = allocator,
3174 .kind = .cuda,
3175 .format = .cuda_ptx,
3176 };
3177
3178 var call_artifact = try MatrixProduct2x3x4F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = MatrixProduct2x3x4F32.Limits.testing });
3179 defer call_artifact.deinit();
3180
3181 const artifact = call_artifact.registry().find(MatrixProduct2x3x4F32.target, MatrixProduct2x3x4F32.version, .cuda_ptx) orelse {
3182 return error.TestExpectedKernelCallArtifact;
3183 };
3184 try std.testing.expectEqualStrings(MatrixProduct2x3x4F32.name, artifact.entry_name);
3185 try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);
3186 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
3187 switch (artifact.launch) {
3188 .fixed => |geometry| {
3189 try std.testing.expectEqual(MatrixProduct2x3x4F32.specialization.launch.?.grid[0], geometry.grid[0]);
3190 try std.testing.expectEqual(MatrixProduct2x3x4F32.specialization.launch.?.grid[1], geometry.grid[1]);
3191 try std.testing.expectEqual(MatrixProduct2x3x4F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
3192 try std.testing.expectEqual(MatrixProduct2x3x4F32.specialization.launch.?.threadgroup[1], geometry.threadgroup[1]);
3193 },
3194 else => return error.TestExpectedFixedLaunch,
3195 }
3196 }
3197
3198 test "linalg matrix vector product entry runs on CPU and records schedule" {
3199 var matrix = [_]f32{
3200 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,
3201 2.0, 0.0, -2.0, 0.0, 1.0, 0.0, -1.0, 0.0,
3202 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0,
3203 -1.0, -2.0, 3.0, 4.0, -5.0, 6.0, 7.0, -8.0,
3204 };
3205 var vector = [_]f32{ 1.0, 0.5, -1.0, 2.0, 0.25, -0.5, 1.5, -2.0 };
3206 var dst = @as([4]f32, @splat(0.0));
3207
3208 try MatrixVectorProduct4x8F32.runCpu(std.testing.allocator, MatrixVectorProduct4x8F32.Limits.testing, &.{
3209 kernel.argumentBuffer(f32, dst[0..]),
3210 kernel.argumentBuffer(f32, matrix[0..]),
3211 kernel.argumentBuffer(f32, vector[0..]),
3212 });
3213 try std.testing.expectEqualSlices(f32, &.{ -0.25, 2.75, -0.125, 25.25 }, dst[0..]);
3214
3215 const launch_value = try MatrixVectorProduct4x8F32.launch(std.testing.allocator, MatrixVectorProduct4x8F32.Limits.testing);
3216 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
3217 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
3218 }
3219
3220 test "linalg matrix vector product entry carries einsum specialization metadata" {
3221 const MatrixVectorProduct5x6F32 = matrixVectorProductF32(.{
3222 .m = 5,
3223 .k = 6,
3224 .threads = 4,
3225 });
3226
3227 try std.testing.expect(MatrixVectorProduct5x6F32.specialization.operationIs(.{ .linalg = .matrix_vector_product }));
3228 try std.testing.expectEqualStrings("mk,k->m", MatrixVectorProduct5x6F32.specialization.equation.?);
3229 try std.testing.expectEqualStrings("accy.kernel.linalg.matvec5x6_4x_f32", MatrixVectorProduct5x6F32.target);
3230 try std.testing.expectEqual(@as(usize, 2), MatrixVectorProduct5x6F32.specialization.inputs.len);
3231 try std.testing.expectEqual(@as(u64, 30), MatrixVectorProduct5x6F32.specialization.inputs[0].elementCount().?);
3232 try std.testing.expectEqual(@as(u64, 6), MatrixVectorProduct5x6F32.specialization.inputs[1].elementCount().?);
3233 try std.testing.expectEqual(@as(u64, 5), MatrixVectorProduct5x6F32.specialization.outputs[0].elementCount().?);
3234 try std.testing.expectEqualStrings("dot", MatrixVectorProduct5x6F32.specialization.reductions[0].name);
3235 try std.testing.expectEqual(entry.ReductionOperator.dot_product, MatrixVectorProduct5x6F32.specialization.reductions[0].operator);
3236 try std.testing.expectEqual(@as(u64, 6), MatrixVectorProduct5x6F32.specialization.reductions[0].shape.elementCount().?);
3237 try std.testing.expectEqual(@as(u32, 2), MatrixVectorProduct5x6F32.specialization.launch.?.grid[0]);
3238 try std.testing.expectEqual(@as(u32, 4), MatrixVectorProduct5x6F32.specialization.launch.?.threadgroup[0]);
3239 try std.testing.expectEqualDeep(MatrixVectorProduct5x6F32.specialization.launch.?, MatrixVectorProduct5x6F32.specialization.schedule.?.launch());
3240 try std.testing.expectEqual(@as(usize, 2), MatrixVectorProduct5x6F32.specialization.schedule.?.bindings.len);
3241 try std.testing.expectEqualStrings("m_tile", MatrixVectorProduct5x6F32.specialization.schedule.?.bindings[0].axis);
3242 try std.testing.expectEqual(kernel.BindTarget.block_x, MatrixVectorProduct5x6F32.specialization.schedule.?.bindings[0].target);
3243 try std.testing.expectEqualStrings("m_lane", MatrixVectorProduct5x6F32.specialization.schedule.?.bindings[1].axis);
3244 try std.testing.expectEqual(kernel.BindTarget.thread_x, MatrixVectorProduct5x6F32.specialization.schedule.?.bindings[1].target);
3245
3246 var snapshot = try MatrixVectorProduct5x6F32.scheduleSnapshot(std.testing.allocator, MatrixVectorProduct5x6F32.Limits.testing);
3247 defer snapshot.deinit(std.testing.allocator);
3248 try std.testing.expect(MatrixVectorProduct5x6F32.specialization.schedule.?.matchesSnapshot(&snapshot));
3249 }
3250
3251 test "linalg matrix vector product entry creates registry-ready artifact" {
3252 const allocator = std.testing.allocator;
3253 var state = gpu.recording.BackendState{
3254 .allocator = allocator,
3255 .kind = .cuda,
3256 .format = .cuda_ptx,
3257 };
3258
3259 var call_artifact = try MatrixVectorProduct4x8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = MatrixVectorProduct4x8F32.Limits.testing });
3260 defer call_artifact.deinit();
3261
3262 const artifact = call_artifact.registry().find(MatrixVectorProduct4x8F32.target, MatrixVectorProduct4x8F32.version, .cuda_ptx) orelse {
3263 return error.TestExpectedKernelCallArtifact;
3264 };
3265 try std.testing.expectEqualStrings(MatrixVectorProduct4x8F32.name, artifact.entry_name);
3266 try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);
3267 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
3268 switch (artifact.launch) {
3269 .fixed => |geometry| {
3270 try std.testing.expectEqual(MatrixVectorProduct4x8F32.specialization.launch.?.grid[0], geometry.grid[0]);
3271 try std.testing.expectEqual(MatrixVectorProduct4x8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
3272 },
3273 else => return error.TestExpectedFixedLaunch,
3274 }
3275 }
3276
3277 test "linalg matrix vector product family matches the fixed entry at its extents" {
3278 const instance = MatrixVectorProduct{ .m = 4, .k = 8, .threads = 4 };
3279
3280 const family_launch = try MatrixVectorProductFamilyF32.launch(std.testing.allocator, MatrixVectorProductFamilyF32.Limits.testing, instance);
3281 const fixed_launch = try MatrixVectorProduct4x8F32.launch(std.testing.allocator, MatrixVectorProduct4x8F32.Limits.testing);
3282 try std.testing.expectEqualDeep(fixed_launch, family_launch);
3283
3284 var family_snapshot = try MatrixVectorProductFamilyF32.scheduleSnapshot(std.testing.allocator, MatrixVectorProductFamilyF32.Limits.testing, instance);
3285 defer family_snapshot.deinit(std.testing.allocator);
3286 var fixed_snapshot = try MatrixVectorProduct4x8F32.scheduleSnapshot(std.testing.allocator, MatrixVectorProduct4x8F32.Limits.testing);
3287 defer fixed_snapshot.deinit(std.testing.allocator);
3288 try std.testing.expectEqual(fixed_snapshot.fingerprint(), family_snapshot.fingerprint());
3289
3290 var matrix = [_]f32{
3291 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,
3292 2.0, 0.0, -2.0, 0.0, 1.0, 0.0, -1.0, 0.0,
3293 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0,
3294 -1.0, -2.0, 3.0, 4.0, -5.0, 6.0, 7.0, -8.0,
3295 };
3296 var vector = [_]f32{ 1.0, 0.5, -1.0, 2.0, 0.25, -0.5, 1.5, -2.0 };
3297 var family_dst = @as([4]f32, @splat(0.0));
3298 var fixed_dst = @as([4]f32, @splat(0.0));
3299
3300 try MatrixVectorProductFamilyF32.runCpu(std.testing.allocator, MatrixVectorProductFamilyF32.Limits.testing, instance, &.{
3301 kernel.argumentBuffer(f32, family_dst[0..]),
3302 kernel.argumentBuffer(f32, matrix[0..]),
3303 kernel.argumentBuffer(f32, vector[0..]),
3304 });
3305 try MatrixVectorProduct4x8F32.runCpu(std.testing.allocator, MatrixVectorProduct4x8F32.Limits.testing, &.{
3306 kernel.argumentBuffer(f32, fixed_dst[0..]),
3307 kernel.argumentBuffer(f32, matrix[0..]),
3308 kernel.argumentBuffer(f32, vector[0..]),
3309 });
3310 try std.testing.expectEqualSlices(f32, fixed_dst[0..], family_dst[0..]);
3311 }
3312
3313 test "linalg matrix vector product family executes fresh runtime extents" {
3314 const instance = MatrixVectorProduct{ .m = 5, .k = 3, .threads = 4 };
3315
3316 var matrix: [15]f32 = undefined;
3317 for (&matrix, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index)) * 0.5 - 1.0;
3318 var vector = [_]f32{ 2.0, -1.0, 0.25 };
3319
3320 var expected: [5]f32 = undefined;
3321 for (0..5) |row| {
3322 var sum: f32 = 0.0;
3323 for (0..3) |offset| {
3324 sum += matrix[row * 3 + offset] * vector[offset];
3325 }
3326 expected[row] = sum;
3327 }
3328
3329 var dst = @as([5]f32, @splat(0.0));
3330 try MatrixVectorProductFamilyF32.runCpu(std.testing.allocator, MatrixVectorProductFamilyF32.Limits.testing, instance, &.{
3331 kernel.argumentBuffer(f32, dst[0..]),
3332 kernel.argumentBuffer(f32, matrix[0..]),
3333 kernel.argumentBuffer(f32, vector[0..]),
3334 });
3335 try std.testing.expectEqualSlices(f32, expected[0..], dst[0..]);
3336
3337 const launch_value = try MatrixVectorProductFamilyF32.launch(std.testing.allocator, MatrixVectorProductFamilyF32.Limits.testing, instance);
3338 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
3339 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
3340 }
3341
3342 test "linalg matrix vector product family identity and metadata" {
3343 const instance = MatrixVectorProduct{
3344 .m = 5,
3345 .k = 3,
3346 .threads = 4,
3347 .row_axis = "rows",
3348 .reduction_axis = "depth",
3349 };
3350 var owned = try matrixVectorProductFamilySpecialization(std.testing.allocator, instance);
3351 defer owned.deinit();
3352 const specialization = owned.value;
3353
3354 const fixed_target = try matrixVectorProductInstanceTarget(std.testing.allocator, .{ .m = 4, .k = 8, .threads = 4 });
3355 defer std.testing.allocator.free(fixed_target);
3356 try std.testing.expectEqualStrings(MatrixVectorProduct4x8F32.target, fixed_target);
3357
3358 const fixed_entry = try matrixVectorProductInstanceEntryName(std.testing.allocator, .{ .m = 4, .k = 8, .threads = 4 });
3359 defer std.testing.allocator.free(fixed_entry);
3360 try std.testing.expectEqualStrings(MatrixVectorProduct4x8F32.name, fixed_entry);
3361
3362 const family_target = try matrixVectorProductFamilyTarget(std.testing.allocator, instance);
3363 defer std.testing.allocator.free(family_target);
3364 try std.testing.expectEqualStrings("accy.kernel.linalg.matvec_family_4x_f32", family_target);
3365
3366 const family_entry = try matrixVectorProductFamilyEntryName(std.testing.allocator, instance);
3367 defer std.testing.allocator.free(family_entry);
3368 try std.testing.expectEqualStrings("accy_kernel_linalg_matvec_family_4x_f32", family_entry);
3369 try std.testing.expectEqual(MatrixVectorProduct4x8F32.version, matrix_vector_product_family_version);
3370
3371 try std.testing.expect(specialization.operationIs(.{ .linalg = .matrix_vector_product }));
3372 try std.testing.expectEqualStrings("mk,k->m", specialization.equation.?);
3373 try std.testing.expectEqual(@as(usize, 2), specialization.inputs.len);
3374 try std.testing.expectEqual(@as(u64, 15), specialization.inputs[0].elementCount().?);
3375 try std.testing.expectEqual(@as(u64, 3), specialization.inputs[1].elementCount().?);
3376 try std.testing.expectEqual(@as(u64, 5), specialization.outputs[0].elementCount().?);
3377 try std.testing.expectEqualStrings("rows", specialization.inputs[0].axes[0].name);
3378 try std.testing.expectEqualStrings("depth", specialization.inputs[0].axes[1].name);
3379 try std.testing.expect(specialization.reductionMatches(0, .{ .name = "dot", .operator = .dot_product, .extents = &.{3} }));
3380 try std.testing.expect(specialization.scheduleMatchesLaunch());
3381 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.grid[0]);
3382 try std.testing.expectEqual(@as(u32, 4), specialization.launch.?.threadgroup[0]);
3383 try std.testing.expectEqual(@as(usize, 2), specialization.schedule.?.bindings.len);
3384 try std.testing.expectEqualStrings("rows_tile", specialization.schedule.?.bindings[0].axis);
3385 try std.testing.expectEqual(kernel.BindTarget.block_x, specialization.schedule.?.bindings[0].target);
3386 try std.testing.expectEqualStrings("rows_lane", specialization.schedule.?.bindings[1].axis);
3387 try std.testing.expectEqual(kernel.BindTarget.thread_x, specialization.schedule.?.bindings[1].target);
3388 try std.testing.expect(specialization.shape_family != null);
3389 try std.testing.expectEqual(try matrixVectorProductFamilyFingerprint(std.testing.allocator, instance), specialization.shapeFamilyFingerprint().?);
3390 }
3391
3392 test "linalg matrix vector product reconstructs family instance from specialization" {
3393 const instance = MatrixVectorProduct{
3394 .m = 5,
3395 .k = 3,
3396 .threads = 4,
3397 .row_axis = "rows",
3398 .reduction_axis = "depth",
3399 };
3400 var owned = try matrixVectorProductFamilySpecialization(std.testing.allocator, instance);
3401 defer owned.deinit();
3402 const reconstructed = matrixVectorProductInstanceFromSpecialization(owned.value) orelse return error.TestExpectedMatrixVectorProductInstance;
3403
3404 try std.testing.expectEqual(instance.m, reconstructed.m);
3405 try std.testing.expectEqual(instance.k, reconstructed.k);
3406 try std.testing.expectEqual(instance.threads, reconstructed.threads);
3407 try std.testing.expectEqualStrings(instance.row_axis, reconstructed.row_axis);
3408 try std.testing.expectEqualStrings(instance.reduction_axis, reconstructed.reduction_axis);
3409 }
3410
3411 test "linalg matrix vector product runtime family executes explicit runtime extents" {
3412 const allocator = std.testing.allocator;
3413 const compiled = MatrixVectorProduct{ .m = 1, .k = 1, .threads = 4 };
3414 const runtime = MatrixVectorProduct{ .m = 3, .k = 4, .threads = compiled.threads };
3415
3416 var graph = try MatrixVectorProductRuntimeFamilyF32.build(allocator, MatrixVectorProductRuntimeFamilyF32.Limits.testing, compiled);
3417 defer graph.deinit();
3418
3419 var matrix: [12]f32 = undefined;
3420 for (&matrix, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index + 1));
3421 var vector = [_]f32{ 0.5, -1.0, 2.0, 0.25 };
3422
3423 var expected: [3]f32 = undefined;
3424 for (0..3) |row| {
3425 var sum: f32 = 0.0;
3426 for (0..4) |offset| {
3427 sum += matrix[row * 4 + offset] * vector[offset];
3428 }
3429 expected[row] = sum;
3430 }
3431
3432 var dst = @as([3]f32, @splat(0.0));
3433 const launch_value = try entry.runtimeLaunch1D(runtime.m, runtime.threads);
3434 try graph.runCpuWithLaunch(allocator, &.{
3435 kernel.argumentBuffer(f32, dst[0..]),
3436 kernel.argumentBuffer(f32, matrix[0..]),
3437 kernel.argumentBuffer(f32, vector[0..]),
3438 kernel.argumentI32(@intCast(runtime.m)),
3439 kernel.argumentI32(@intCast(runtime.k)),
3440 }, .{
3441 .grid = launch_value.grid,
3442 .block = launch_value.threadgroup,
3443 });
3444 try std.testing.expectEqualSlices(f32, expected[0..], dst[0..]);
3445 }
3446
3447 test "linalg matrix vector product family artifact carries runtime launch contract" {
3448 const allocator = std.testing.allocator;
3449 var state = gpu.recording.BackendState{
3450 .allocator = allocator,
3451 .kind = .cuda,
3452 .format = .cuda_ptx,
3453 };
3454 const instance = MatrixVectorProduct{ .m = 5, .k = 3, .threads = 4 };
3455
3456 var family_artifact = try createMatrixVectorProductFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
3457 defer family_artifact.deinit();
3458 var fixed_artifact = try MatrixVectorProduct4x8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = MatrixVectorProduct4x8F32.Limits.testing });
3459 defer fixed_artifact.deinit();
3460
3461 const family_entry = family_artifact.entry();
3462 const fixed_entry = fixed_artifact.entry();
3463 try std.testing.expect(!std.mem.eql(u8, fixed_entry.target, family_entry.target));
3464 try std.testing.expectEqualStrings("accy.kernel.linalg.matvec_family_4x_f32", family_entry.target);
3465 try std.testing.expectEqualStrings("accy_kernel_linalg_matvec_family_4x_f32", family_entry.entry_name);
3466 try std.testing.expectEqual(@as(u32, 5), family_entry.argument_count);
3467 try std.testing.expectEqual(@as(u32, 2), family_entry.runtime_scalar_argument_count);
3468 try std.testing.expect(family_entry.required_dtypes.contains(.f32));
3469 try std.testing.expect(family_entry.required_dtypes.contains(.i32));
3470 try std.testing.expect(fixed_entry.shape_family_fingerprint == null);
3471 try std.testing.expect(family_entry.shape_family_fingerprint != null);
3472 const profile = family_entry.shape_profile orelse return error.TestExpectedShapeProfile;
3473 try std.testing.expectEqualStrings("matrix_vector_product", profile.name);
3474 try std.testing.expectEqual(family_entry.shape_family_fingerprint.?, profile.fingerprint);
3475 try std.testing.expectEqual(@as(usize, 2), profile.dimensions.len);
3476 const m_dimension = profile.runtimeScalarDimension(0) orelse return error.TestExpectedShapeProfile;
3477 try std.testing.expectEqualStrings("m", m_dimension.name);
3478 const k_dimension = profile.runtimeScalarDimension(1) orelse return error.TestExpectedShapeProfile;
3479 try std.testing.expectEqualStrings("k", k_dimension.name);
3480 switch (family_entry.launch) {
3481 .derived => |launch| {
3482 try std.testing.expectEqual(@as(u32, 4), launch.threadgroup[0]);
3483 switch (launch.grid[0]) {
3484 .runtime_u32_ceil_div => |axis| {
3485 try std.testing.expectEqual(@as(u32, 0), axis.argument_index);
3486 try std.testing.expectEqual(@as(u32, 4), axis.divisor);
3487 },
3488 else => return error.TestExpectedDerivedLaunch,
3489 }
3490 const args = try matrixVectorProductRuntimeArguments(instance);
3491 const geometry = try launch.geometry(args[0..]);
3492 try std.testing.expectEqual(@as(u32, 2), geometry.grid[0]);
3493 try std.testing.expectEqual(@as(u32, 4), geometry.threadgroup[0]);
3494 },
3495 .fixed => return error.TestExpectedDerivedLaunch,
3496 }
3497 }
3498
3499 test "linalg outer product entry runs on CPU and records schedule" {
3500 var lhs = [_]f32{ 1.0, -2.0, 0.5, 3.0 };
3501 var rhs = [_]f32{ 4.0, -1.0, 2.0 };
3502 var dst = @as([12]f32, @splat(0.0));
3503
3504 try OuterProduct4x3F32.runCpu(std.testing.allocator, OuterProduct4x3F32.Limits.testing, &.{
3505 kernel.argumentBuffer(f32, dst[0..]),
3506 kernel.argumentBuffer(f32, lhs[0..]),
3507 kernel.argumentBuffer(f32, rhs[0..]),
3508 });
3509 try std.testing.expectEqualSlices(f32, &.{ 4.0, -1.0, 2.0, -8.0, 2.0, -4.0, 2.0, -0.5, 1.0, 12.0, -3.0, 6.0 }, dst[0..]);
3510
3511 const launch_value = try OuterProduct4x3F32.launch(std.testing.allocator, OuterProduct4x3F32.Limits.testing);
3512 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
3513 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[1]);
3514 try std.testing.expectEqual(@as(u32, 3), launch_value.block[0]);
3515 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
3516 }
3517
3518 test "linalg outer product entry carries einsum specialization metadata" {
3519 const OuterProduct5x6F32 = outerProductF32(.{
3520 .m = 5,
3521 .n = 6,
3522 .threads = .{ .x = 4, .y = 2 },
3523 });
3524
3525 try std.testing.expect(OuterProduct5x6F32.specialization.operationIs(.{ .linalg = .outer_product }));
3526 try std.testing.expectEqualStrings("m,n->mn", OuterProduct5x6F32.specialization.equation.?);
3527 try std.testing.expectEqualStrings("accy.kernel.linalg.outer5x6_4x2_f32", OuterProduct5x6F32.target);
3528 try std.testing.expectEqual(@as(usize, 2), OuterProduct5x6F32.specialization.inputs.len);
3529 try std.testing.expectEqual(@as(u64, 5), OuterProduct5x6F32.specialization.inputs[0].elementCount().?);
3530 try std.testing.expectEqual(@as(u64, 6), OuterProduct5x6F32.specialization.inputs[1].elementCount().?);
3531 try std.testing.expectEqual(@as(u64, 30), OuterProduct5x6F32.specialization.outputs[0].elementCount().?);
3532 try std.testing.expectEqual(@as(usize, 0), OuterProduct5x6F32.specialization.reductions.len);
3533 try std.testing.expectEqual(@as(u32, 2), OuterProduct5x6F32.specialization.launch.?.grid[0]);
3534 try std.testing.expectEqual(@as(u32, 3), OuterProduct5x6F32.specialization.launch.?.grid[1]);
3535 try std.testing.expectEqualDeep(OuterProduct5x6F32.specialization.launch.?, OuterProduct5x6F32.specialization.schedule.?.launch());
3536 try std.testing.expectEqual(@as(usize, 4), OuterProduct5x6F32.specialization.schedule.?.bindings.len);
3537 try std.testing.expectEqualStrings("n_tile", OuterProduct5x6F32.specialization.schedule.?.bindings[0].axis);
3538 try std.testing.expectEqual(kernel.BindTarget.block_x, OuterProduct5x6F32.specialization.schedule.?.bindings[0].target);
3539 try std.testing.expectEqualStrings("m_lane", OuterProduct5x6F32.specialization.schedule.?.bindings[3].axis);
3540 try std.testing.expectEqual(kernel.BindTarget.thread_y, OuterProduct5x6F32.specialization.schedule.?.bindings[3].target);
3541
3542 var snapshot = try OuterProduct5x6F32.scheduleSnapshot(std.testing.allocator, OuterProduct5x6F32.Limits.testing);
3543 defer snapshot.deinit(std.testing.allocator);
3544 try std.testing.expect(OuterProduct5x6F32.specialization.schedule.?.matchesSnapshot(&snapshot));
3545 }
3546
3547 test "linalg outer product entry creates registry-ready artifact" {
3548 const allocator = std.testing.allocator;
3549 var state = gpu.recording.BackendState{
3550 .allocator = allocator,
3551 .kind = .cuda,
3552 .format = .cuda_ptx,
3553 };
3554
3555 var call_artifact = try OuterProduct4x3F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = OuterProduct4x3F32.Limits.testing });
3556 defer call_artifact.deinit();
3557
3558 const artifact = call_artifact.registry().find(OuterProduct4x3F32.target, OuterProduct4x3F32.version, .cuda_ptx) orelse {
3559 return error.TestExpectedKernelCallArtifact;
3560 };
3561 try std.testing.expectEqualStrings(OuterProduct4x3F32.name, artifact.entry_name);
3562 try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);
3563 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
3564 switch (artifact.launch) {
3565 .fixed => |geometry| {
3566 try std.testing.expectEqual(OuterProduct4x3F32.specialization.launch.?.grid[0], geometry.grid[0]);
3567 try std.testing.expectEqual(OuterProduct4x3F32.specialization.launch.?.grid[1], geometry.grid[1]);
3568 try std.testing.expectEqual(OuterProduct4x3F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
3569 try std.testing.expectEqual(OuterProduct4x3F32.specialization.launch.?.threadgroup[1], geometry.threadgroup[1]);
3570 },
3571 else => return error.TestExpectedFixedLaunch,
3572 }
3573 }
3574
3575 test "linalg outer product family matches the fixed entry at its extents" {
3576 const instance = OuterProduct{ .m = 4, .n = 3, .threads = .{ .x = 3, .y = 2 } };
3577
3578 const family_launch = try OuterProductFamilyF32.launch(std.testing.allocator, OuterProductFamilyF32.Limits.testing, instance);
3579 const fixed_launch = try OuterProduct4x3F32.launch(std.testing.allocator, OuterProduct4x3F32.Limits.testing);
3580 try std.testing.expectEqualDeep(fixed_launch, family_launch);
3581
3582 var family_snapshot = try OuterProductFamilyF32.scheduleSnapshot(std.testing.allocator, OuterProductFamilyF32.Limits.testing, instance);
3583 defer family_snapshot.deinit(std.testing.allocator);
3584 var fixed_snapshot = try OuterProduct4x3F32.scheduleSnapshot(std.testing.allocator, OuterProduct4x3F32.Limits.testing);
3585 defer fixed_snapshot.deinit(std.testing.allocator);
3586 try std.testing.expectEqual(fixed_snapshot.fingerprint(), family_snapshot.fingerprint());
3587
3588 var lhs = [_]f32{ 1.0, -2.0, 0.5, 3.0 };
3589 var rhs = [_]f32{ 4.0, -1.0, 2.0 };
3590 var family_dst = @as([12]f32, @splat(0.0));
3591 var fixed_dst = @as([12]f32, @splat(0.0));
3592
3593 try OuterProductFamilyF32.runCpu(std.testing.allocator, OuterProductFamilyF32.Limits.testing, instance, &.{
3594 kernel.argumentBuffer(f32, family_dst[0..]),
3595 kernel.argumentBuffer(f32, lhs[0..]),
3596 kernel.argumentBuffer(f32, rhs[0..]),
3597 });
3598 try OuterProduct4x3F32.runCpu(std.testing.allocator, OuterProduct4x3F32.Limits.testing, &.{
3599 kernel.argumentBuffer(f32, fixed_dst[0..]),
3600 kernel.argumentBuffer(f32, lhs[0..]),
3601 kernel.argumentBuffer(f32, rhs[0..]),
3602 });
3603 try std.testing.expectEqualSlices(f32, fixed_dst[0..], family_dst[0..]);
3604 }
3605
3606 test "linalg outer product family executes fresh runtime extents" {
3607 const instance = OuterProduct{ .m = 5, .n = 6, .threads = .{ .x = 4, .y = 2 } };
3608
3609 var lhs = [_]f32{ 2.0, -1.0, 0.5, 3.0, -4.0 };
3610 var rhs = [_]f32{ 1.5, -2.0, 0.25, 4.0, -0.5, 2.5 };
3611 var expected: [30]f32 = undefined;
3612 for (0..5) |row| {
3613 for (0..6) |col| {
3614 expected[row * 6 + col] = lhs[row] * rhs[col];
3615 }
3616 }
3617
3618 var dst = @as([30]f32, @splat(0.0));
3619 try OuterProductFamilyF32.runCpu(std.testing.allocator, OuterProductFamilyF32.Limits.testing, instance, &.{
3620 kernel.argumentBuffer(f32, dst[0..]),
3621 kernel.argumentBuffer(f32, lhs[0..]),
3622 kernel.argumentBuffer(f32, rhs[0..]),
3623 });
3624 try std.testing.expectEqualSlices(f32, expected[0..], dst[0..]);
3625
3626 const launch_value = try OuterProductFamilyF32.launch(std.testing.allocator, OuterProductFamilyF32.Limits.testing, instance);
3627 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
3628 try std.testing.expectEqual(@as(u32, 3), launch_value.grid[1]);
3629 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
3630 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
3631 }
3632
3633 test "linalg outer product family identity and metadata" {
3634 const instance = OuterProduct{
3635 .m = 5,
3636 .n = 6,
3637 .threads = .{ .x = 4, .y = 2 },
3638 .lhs_axis = "rows",
3639 .rhs_axis = "cols",
3640 };
3641 var owned = try outerProductFamilySpecialization(std.testing.allocator, instance);
3642 defer owned.deinit();
3643 const specialization = owned.value;
3644
3645 const fixed_target = try outerProductInstanceTarget(std.testing.allocator, .{ .m = 4, .n = 3, .threads = .{ .x = 3, .y = 2 } });
3646 defer std.testing.allocator.free(fixed_target);
3647 try std.testing.expectEqualStrings(OuterProduct4x3F32.target, fixed_target);
3648
3649 const fixed_entry = try outerProductInstanceEntryName(std.testing.allocator, .{ .m = 4, .n = 3, .threads = .{ .x = 3, .y = 2 } });
3650 defer std.testing.allocator.free(fixed_entry);
3651 try std.testing.expectEqualStrings(OuterProduct4x3F32.name, fixed_entry);
3652
3653 const family_target = try outerProductFamilyTarget(std.testing.allocator, instance);
3654 defer std.testing.allocator.free(family_target);
3655 try std.testing.expectEqualStrings("accy.kernel.linalg.outer_family_4x2_f32", family_target);
3656
3657 const family_entry = try outerProductFamilyEntryName(std.testing.allocator, instance);
3658 defer std.testing.allocator.free(family_entry);
3659 try std.testing.expectEqualStrings("accy_kernel_linalg_outer_family_4x2_f32", family_entry);
3660 try std.testing.expectEqual(OuterProduct4x3F32.version, outer_product_family_version);
3661
3662 try std.testing.expect(specialization.operationIs(.{ .linalg = .outer_product }));
3663 try std.testing.expectEqualStrings("m,n->mn", specialization.equation.?);
3664 try std.testing.expectEqual(@as(usize, 2), specialization.inputs.len);
3665 try std.testing.expectEqual(@as(u64, 5), specialization.inputs[0].elementCount().?);
3666 try std.testing.expectEqual(@as(u64, 6), specialization.inputs[1].elementCount().?);
3667 try std.testing.expectEqual(@as(u64, 30), specialization.outputs[0].elementCount().?);
3668 try std.testing.expectEqual(@as(usize, 0), specialization.reductions.len);
3669 try std.testing.expectEqualStrings("rows", specialization.inputs[0].axes[0].name);
3670 try std.testing.expectEqualStrings("cols", specialization.inputs[1].axes[0].name);
3671 try std.testing.expect(specialization.scheduleMatchesLaunch());
3672 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.grid[0]);
3673 try std.testing.expectEqual(@as(u32, 3), specialization.launch.?.grid[1]);
3674 try std.testing.expectEqual(@as(u32, 4), specialization.launch.?.threadgroup[0]);
3675 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.threadgroup[1]);
3676 try std.testing.expectEqual(@as(usize, 4), specialization.schedule.?.bindings.len);
3677 try std.testing.expectEqualStrings("cols_tile", specialization.schedule.?.bindings[0].axis);
3678 try std.testing.expectEqual(kernel.BindTarget.block_x, specialization.schedule.?.bindings[0].target);
3679 try std.testing.expectEqualStrings("rows_lane", specialization.schedule.?.bindings[3].axis);
3680 try std.testing.expectEqual(kernel.BindTarget.thread_y, specialization.schedule.?.bindings[3].target);
3681 try std.testing.expect(specialization.shape_family != null);
3682 try std.testing.expectEqual(try outerProductFamilyFingerprint(std.testing.allocator, instance), specialization.shapeFamilyFingerprint().?);
3683 }
3684
3685 test "linalg outer product reconstructs family instance from specialization" {
3686 const instance = OuterProduct{
3687 .m = 5,
3688 .n = 6,
3689 .threads = .{ .x = 4, .y = 2 },
3690 .lhs_axis = "rows",
3691 .rhs_axis = "cols",
3692 };
3693 var owned = try outerProductFamilySpecialization(std.testing.allocator, instance);
3694 defer owned.deinit();
3695 const reconstructed = outerProductInstanceFromSpecialization(owned.value) orelse return error.TestExpectedOuterProductInstance;
3696
3697 try std.testing.expectEqual(instance.m, reconstructed.m);
3698 try std.testing.expectEqual(instance.n, reconstructed.n);
3699 try std.testing.expectEqual(instance.threads.x, reconstructed.threads.x);
3700 try std.testing.expectEqual(instance.threads.y, reconstructed.threads.y);
3701 try std.testing.expectEqualStrings(instance.lhs_axis, reconstructed.lhs_axis);
3702 try std.testing.expectEqualStrings(instance.rhs_axis, reconstructed.rhs_axis);
3703 }
3704
3705 test "linalg outer product runtime family executes explicit runtime extents" {
3706 const allocator = std.testing.allocator;
3707 const compiled = OuterProduct{ .m = 1, .n = 1, .threads = .{ .x = 4, .y = 2 } };
3708 const runtime = OuterProduct{ .m = 3, .n = 4, .threads = compiled.threads };
3709
3710 var graph = try OuterProductRuntimeFamilyF32.build(allocator, OuterProductRuntimeFamilyF32.Limits.testing, compiled);
3711 defer graph.deinit();
3712
3713 var lhs = [_]f32{ 2.0, -1.0, 0.5 };
3714 var rhs = [_]f32{ 1.5, -2.0, 0.25, 4.0 };
3715 var expected: [12]f32 = undefined;
3716 for (0..3) |row| {
3717 for (0..4) |col| {
3718 expected[row * 4 + col] = lhs[row] * rhs[col];
3719 }
3720 }
3721
3722 var dst = @as([12]f32, @splat(0.0));
3723 const launch_value = try entry.runtimeLaunch2D(runtime.n, runtime.m, runtime.threads.x, runtime.threads.y);
3724 try graph.runCpuWithLaunch(allocator, &.{
3725 kernel.argumentBuffer(f32, dst[0..]),
3726 kernel.argumentBuffer(f32, lhs[0..]),
3727 kernel.argumentBuffer(f32, rhs[0..]),
3728 kernel.argumentI32(@intCast(runtime.m)),
3729 kernel.argumentI32(@intCast(runtime.n)),
3730 }, .{
3731 .grid = launch_value.grid,
3732 .block = launch_value.threadgroup,
3733 });
3734 try std.testing.expectEqualSlices(f32, expected[0..], dst[0..]);
3735 }
3736
3737 test "linalg outer product family artifact carries runtime launch contract" {
3738 const allocator = std.testing.allocator;
3739 var state = gpu.recording.BackendState{
3740 .allocator = allocator,
3741 .kind = .cuda,
3742 .format = .cuda_ptx,
3743 };
3744 const instance = OuterProduct{ .m = 5, .n = 6, .threads = .{ .x = 4, .y = 2 } };
3745
3746 var family_artifact = try createOuterProductFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
3747 defer family_artifact.deinit();
3748 var fixed_artifact = try OuterProduct4x3F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = OuterProduct4x3F32.Limits.testing });
3749 defer fixed_artifact.deinit();
3750
3751 const family_entry = family_artifact.entry();
3752 const fixed_entry = fixed_artifact.entry();
3753 try std.testing.expect(!std.mem.eql(u8, fixed_entry.target, family_entry.target));
3754 try std.testing.expectEqualStrings("accy.kernel.linalg.outer_family_4x2_f32", family_entry.target);
3755 try std.testing.expectEqualStrings("accy_kernel_linalg_outer_family_4x2_f32", family_entry.entry_name);
3756 try std.testing.expectEqual(@as(u32, 5), family_entry.argument_count);
3757 try std.testing.expectEqual(@as(u32, 2), family_entry.runtime_scalar_argument_count);
3758 try std.testing.expect(family_entry.required_dtypes.contains(.f32));
3759 try std.testing.expect(family_entry.required_dtypes.contains(.i32));
3760 try std.testing.expect(fixed_entry.shape_family_fingerprint == null);
3761 try std.testing.expect(family_entry.shape_family_fingerprint != null);
3762 const profile = family_entry.shape_profile orelse return error.TestExpectedShapeProfile;
3763 try std.testing.expectEqualStrings("outer_product", profile.name);
3764 try std.testing.expectEqual(family_entry.shape_family_fingerprint.?, profile.fingerprint);
3765 try std.testing.expectEqual(@as(usize, 2), profile.dimensions.len);
3766 const m_dimension = profile.runtimeScalarDimension(0) orelse return error.TestExpectedShapeProfile;
3767 try std.testing.expectEqualStrings("m", m_dimension.name);
3768 const n_dimension = profile.runtimeScalarDimension(1) orelse return error.TestExpectedShapeProfile;
3769 try std.testing.expectEqualStrings("n", n_dimension.name);
3770 switch (family_entry.launch) {
3771 .derived => |launch| {
3772 try std.testing.expectEqual(@as(u32, 4), launch.threadgroup[0]);
3773 try std.testing.expectEqual(@as(u32, 2), launch.threadgroup[1]);
3774 switch (launch.grid[0]) {
3775 .runtime_u32_ceil_div => |axis| {
3776 try std.testing.expectEqual(@as(u32, 1), axis.argument_index);
3777 try std.testing.expectEqual(@as(u32, 4), axis.divisor);
3778 },
3779 else => return error.TestExpectedDerivedLaunch,
3780 }
3781 switch (launch.grid[1]) {
3782 .runtime_u32_ceil_div => |axis| {
3783 try std.testing.expectEqual(@as(u32, 0), axis.argument_index);
3784 try std.testing.expectEqual(@as(u32, 2), axis.divisor);
3785 },
3786 else => return error.TestExpectedDerivedLaunch,
3787 }
3788 const args = try outerProductRuntimeArguments(instance);
3789 const geometry = try launch.geometry(args[0..]);
3790 try std.testing.expectEqual(@as(u32, 2), geometry.grid[0]);
3791 try std.testing.expectEqual(@as(u32, 3), geometry.grid[1]);
3792 try std.testing.expectEqual(@as(u32, 4), geometry.threadgroup[0]);
3793 try std.testing.expectEqual(@as(u32, 2), geometry.threadgroup[1]);
3794 },
3795 .fixed => return error.TestExpectedDerivedLaunch,
3796 }
3797 }