lib/accy/src/kernel/logical/selection/einsum.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3
4 const accy = @import("../../../root.zig");
5 const library = @import("../../library/root.zig");
6
7 const DType = choir_abi.DType;
8
9 pub const EinsumOperand = struct {
10 indices: []const u8,
11 dims: []const i64,
12 };
13
14 pub const EinsumSchedule = union(enum) {
15 batched_matrix_product: library.BatchedMatrixProductSchedule,
16 matrix_product: library.MatrixProductSchedule,
17 matrix_vector_product: library.MatrixVectorProductSchedule,
18 outer_product: library.OuterProductSchedule,
19 };
20
21 pub const EinsumSelectionRequest = struct {
22 dtype: DType,
23 inputs: []const EinsumOperand,
24 output_indices: []const u8,
25 output_dims: []const i64,
26 schedule: ?EinsumSchedule = null,
27 };
28
29 pub const EinsumKernelKind = enum {
30 batched_matrix_product,
31 dot_product,
32 matrix_product,
33 matrix_vector_product,
34 outer_product,
35 sum_reduction,
36 transpose,
37 };
38
39 pub const SelectedEinsumKernel = struct {
40 kind: EinsumKernelKind,
41 descriptor: library.CatalogDescriptor,
42 };
43
44 pub const OwnedSelectedEinsumKernel = struct {
45 kind: EinsumKernelKind,
46 descriptor: library.OwnedCatalogDescriptor,
47
48 pub fn deinit(self: *OwnedSelectedEinsumKernel) void {
49 self.descriptor.deinit();
50 self.* = undefined;
51 }
52 };
53
54 pub fn selectCatalog(request: EinsumSelectionRequest) ?SelectedEinsumKernel {
55 if (selectReduction(request)) |selected| return selected;
56 if (selectBatchedMatrixProduct(request)) |selected| return selected;
57 if (selectMatrixProduct(request)) |selected| return selected;
58 if (selectMatrixVectorProduct(request)) |selected| return selected;
59 if (selectOuterProduct(request)) |selected| return selected;
60 if (selectTranspose(request)) |selected| return selected;
61 return null;
62 }
63
64 pub fn selectOwnedCatalog(backing_allocator: std.mem.Allocator, request: EinsumSelectionRequest) !?OwnedSelectedEinsumKernel {
65 if (selectCatalog(request)) |selected| return ownedSelectedStatic(selected);
66 if (try selectOwnedBatchedMatrixProduct(backing_allocator, request)) |selected| return selected;
67 if (try selectOwnedMatrixProduct(backing_allocator, request)) |selected| return selected;
68 if (try selectOwnedMatrixVectorProduct(backing_allocator, request)) |selected| return selected;
69 if (try selectOwnedOuterProduct(backing_allocator, request)) |selected| return selected;
70 return null;
71 }
72
73 fn ownedSelectedStatic(selected: SelectedEinsumKernel) OwnedSelectedEinsumKernel {
74 return .{
75 .kind = selected.kind,
76 .descriptor = .{ .descriptor = selected.descriptor },
77 };
78 }
79
80 fn selectReduction(request: EinsumSelectionRequest) ?SelectedEinsumKernel {
81 if (request.schedule != null) return null;
82 const kind = reductionKind(request) orelse return null;
83 var inputs: [2]library.ReductionOperand = undefined;
84 for (request.inputs, 0..) |input, index| {
85 inputs[index] = .{
86 .indices = input.indices,
87 .dims = input.dims,
88 };
89 }
90 const descriptor = library.select(.{ .reduction = .{
91 .dtype = request.dtype,
92 .kind = kind,
93 .inputs = inputs[0..request.inputs.len],
94 .output_indices = request.output_indices,
95 .output_dims = request.output_dims,
96 } }) orelse return null;
97 return .{
98 .kind = reductionKernelKind(kind) orelse return null,
99 .descriptor = descriptor,
100 };
101 }
102
103 fn reductionKind(request: EinsumSelectionRequest) ?library.ReductionKind {
104 if (request.output_indices.len != 0 or request.output_dims.len != 0) return null;
105 return switch (request.inputs.len) {
106 1 => if (request.inputs[0].indices.len == 1) .sum else null,
107 2 => if (request.inputs[0].indices.len == 1 and
108 request.inputs[1].indices.len == 1 and
109 request.inputs[0].indices[0] == request.inputs[1].indices[0])
110 .dot_product
111 else
112 null,
113 else => null,
114 };
115 }
116
117 fn reductionKernelKind(kind: library.ReductionKind) ?EinsumKernelKind {
118 return switch (kind) {
119 .sum => .sum_reduction,
120 .dot_product => .dot_product,
121 else => null,
122 };
123 }
124
125 fn selectBatchedMatrixProduct(request: EinsumSelectionRequest) ?SelectedEinsumKernel {
126 if (request.inputs.len != 2) return null;
127 if (!scheduleAllowsBatchedMatrixProduct(request.schedule)) return null;
128 const descriptor = library.select(.{ .batched_matrix_product = .{
129 .dtype = request.dtype,
130 .lhs_indices = request.inputs[0].indices,
131 .rhs_indices = request.inputs[1].indices,
132 .output_indices = request.output_indices,
133 .lhs_dims = request.inputs[0].dims,
134 .rhs_dims = request.inputs[1].dims,
135 .output_dims = request.output_dims,
136 .schedule = batchedMatrixProductSchedule(request.schedule),
137 } }) orelse return null;
138 return .{
139 .kind = .batched_matrix_product,
140 .descriptor = descriptor,
141 };
142 }
143
144 fn selectOwnedBatchedMatrixProduct(backing_allocator: std.mem.Allocator, request: EinsumSelectionRequest) !?OwnedSelectedEinsumKernel {
145 if (request.inputs.len != 2) return null;
146 if (!scheduleAllowsBatchedMatrixProduct(request.schedule)) return null;
147 var descriptor = (try library.selectOwned(backing_allocator, .{ .batched_matrix_product = .{
148 .dtype = request.dtype,
149 .lhs_indices = request.inputs[0].indices,
150 .rhs_indices = request.inputs[1].indices,
151 .output_indices = request.output_indices,
152 .lhs_dims = request.inputs[0].dims,
153 .rhs_dims = request.inputs[1].dims,
154 .output_dims = request.output_dims,
155 .schedule = batchedMatrixProductSchedule(request.schedule),
156 } })) orelse return null;
157 errdefer descriptor.deinit();
158 return .{
159 .kind = .batched_matrix_product,
160 .descriptor = descriptor,
161 };
162 }
163
164 fn selectMatrixProduct(request: EinsumSelectionRequest) ?SelectedEinsumKernel {
165 if (request.inputs.len != 2) return null;
166 if (!scheduleAllowsMatrixProduct(request.schedule)) return null;
167 const descriptor = library.select(.{ .matrix_product = .{
168 .dtype = request.dtype,
169 .lhs_indices = request.inputs[0].indices,
170 .rhs_indices = request.inputs[1].indices,
171 .output_indices = request.output_indices,
172 .lhs_dims = request.inputs[0].dims,
173 .rhs_dims = request.inputs[1].dims,
174 .output_dims = request.output_dims,
175 .schedule = matrixProductSchedule(request.schedule),
176 } }) orelse return null;
177 return .{
178 .kind = .matrix_product,
179 .descriptor = descriptor,
180 };
181 }
182
183 fn selectOwnedMatrixProduct(backing_allocator: std.mem.Allocator, request: EinsumSelectionRequest) !?OwnedSelectedEinsumKernel {
184 if (request.inputs.len != 2) return null;
185 if (!scheduleAllowsMatrixProduct(request.schedule)) return null;
186 var descriptor = (try library.selectOwned(backing_allocator, .{ .matrix_product = .{
187 .dtype = request.dtype,
188 .lhs_indices = request.inputs[0].indices,
189 .rhs_indices = request.inputs[1].indices,
190 .output_indices = request.output_indices,
191 .lhs_dims = request.inputs[0].dims,
192 .rhs_dims = request.inputs[1].dims,
193 .output_dims = request.output_dims,
194 .schedule = matrixProductSchedule(request.schedule),
195 } })) orelse return null;
196 errdefer descriptor.deinit();
197 return .{
198 .kind = .matrix_product,
199 .descriptor = descriptor,
200 };
201 }
202
203 fn selectMatrixVectorProduct(request: EinsumSelectionRequest) ?SelectedEinsumKernel {
204 if (request.inputs.len != 2) return null;
205 if (!scheduleAllowsMatrixVectorProduct(request.schedule)) return null;
206 const descriptor = library.select(.{ .matrix_vector_product = .{
207 .dtype = request.dtype,
208 .matrix_indices = request.inputs[0].indices,
209 .vector_indices = request.inputs[1].indices,
210 .output_indices = request.output_indices,
211 .matrix_dims = request.inputs[0].dims,
212 .vector_dims = request.inputs[1].dims,
213 .output_dims = request.output_dims,
214 .schedule = matrixVectorProductSchedule(request.schedule),
215 } }) orelse return null;
216 return .{
217 .kind = .matrix_vector_product,
218 .descriptor = descriptor,
219 };
220 }
221
222 fn selectOwnedMatrixVectorProduct(backing_allocator: std.mem.Allocator, request: EinsumSelectionRequest) !?OwnedSelectedEinsumKernel {
223 if (request.inputs.len != 2) return null;
224 if (!scheduleAllowsMatrixVectorProduct(request.schedule)) return null;
225 var descriptor = (try library.selectOwned(backing_allocator, .{ .matrix_vector_product = .{
226 .dtype = request.dtype,
227 .matrix_indices = request.inputs[0].indices,
228 .vector_indices = request.inputs[1].indices,
229 .output_indices = request.output_indices,
230 .matrix_dims = request.inputs[0].dims,
231 .vector_dims = request.inputs[1].dims,
232 .output_dims = request.output_dims,
233 .schedule = matrixVectorProductSchedule(request.schedule),
234 } })) orelse return null;
235 errdefer descriptor.deinit();
236 return .{
237 .kind = .matrix_vector_product,
238 .descriptor = descriptor,
239 };
240 }
241
242 fn selectOuterProduct(request: EinsumSelectionRequest) ?SelectedEinsumKernel {
243 if (request.inputs.len != 2) return null;
244 if (!scheduleAllowsOuterProduct(request.schedule)) return null;
245 const descriptor = library.select(.{ .outer_product = .{
246 .dtype = request.dtype,
247 .lhs_indices = request.inputs[0].indices,
248 .rhs_indices = request.inputs[1].indices,
249 .output_indices = request.output_indices,
250 .lhs_dims = request.inputs[0].dims,
251 .rhs_dims = request.inputs[1].dims,
252 .output_dims = request.output_dims,
253 .schedule = outerProductSchedule(request.schedule),
254 } }) orelse return null;
255 return .{
256 .kind = .outer_product,
257 .descriptor = descriptor,
258 };
259 }
260
261 fn selectOwnedOuterProduct(backing_allocator: std.mem.Allocator, request: EinsumSelectionRequest) !?OwnedSelectedEinsumKernel {
262 if (request.inputs.len != 2) return null;
263 if (!scheduleAllowsOuterProduct(request.schedule)) return null;
264 var descriptor = (try library.selectOwned(backing_allocator, .{ .outer_product = .{
265 .dtype = request.dtype,
266 .lhs_indices = request.inputs[0].indices,
267 .rhs_indices = request.inputs[1].indices,
268 .output_indices = request.output_indices,
269 .lhs_dims = request.inputs[0].dims,
270 .rhs_dims = request.inputs[1].dims,
271 .output_dims = request.output_dims,
272 .schedule = outerProductSchedule(request.schedule),
273 } })) orelse return null;
274 errdefer descriptor.deinit();
275 return .{
276 .kind = .outer_product,
277 .descriptor = descriptor,
278 };
279 }
280
281 fn scheduleAllowsMatrixProduct(schedule: ?EinsumSchedule) bool {
282 const requested = schedule orelse return true;
283 return switch (requested) {
284 .batched_matrix_product => false,
285 .matrix_product => true,
286 .matrix_vector_product => false,
287 .outer_product => false,
288 };
289 }
290
291 fn scheduleAllowsMatrixVectorProduct(schedule: ?EinsumSchedule) bool {
292 const requested = schedule orelse return true;
293 return switch (requested) {
294 .batched_matrix_product => false,
295 .matrix_product => false,
296 .matrix_vector_product => true,
297 .outer_product => false,
298 };
299 }
300
301 fn scheduleAllowsBatchedMatrixProduct(schedule: ?EinsumSchedule) bool {
302 const requested = schedule orelse return true;
303 return switch (requested) {
304 .batched_matrix_product => true,
305 .matrix_product => false,
306 .matrix_vector_product => false,
307 .outer_product => false,
308 };
309 }
310
311 fn scheduleAllowsOuterProduct(schedule: ?EinsumSchedule) bool {
312 const requested = schedule orelse return true;
313 return switch (requested) {
314 .batched_matrix_product => false,
315 .matrix_product => false,
316 .matrix_vector_product => false,
317 .outer_product => true,
318 };
319 }
320
321 fn batchedMatrixProductSchedule(schedule: ?EinsumSchedule) ?library.BatchedMatrixProductSchedule {
322 const requested = schedule orelse return null;
323 return switch (requested) {
324 .batched_matrix_product => |batched_matrix_product| batched_matrix_product,
325 .matrix_product => null,
326 .matrix_vector_product => null,
327 .outer_product => null,
328 };
329 }
330
331 fn matrixProductSchedule(schedule: ?EinsumSchedule) ?library.MatrixProductSchedule {
332 const requested = schedule orelse return null;
333 return switch (requested) {
334 .batched_matrix_product => null,
335 .matrix_product => |matrix_product| matrix_product,
336 .matrix_vector_product => null,
337 .outer_product => null,
338 };
339 }
340
341 fn matrixVectorProductSchedule(schedule: ?EinsumSchedule) ?library.MatrixVectorProductSchedule {
342 const requested = schedule orelse return null;
343 return switch (requested) {
344 .batched_matrix_product => null,
345 .matrix_product => null,
346 .matrix_vector_product => |matrix_vector_product| matrix_vector_product,
347 .outer_product => null,
348 };
349 }
350
351 fn outerProductSchedule(schedule: ?EinsumSchedule) ?library.OuterProductSchedule {
352 const requested = schedule orelse return null;
353 return switch (requested) {
354 .batched_matrix_product => null,
355 .matrix_product => null,
356 .matrix_vector_product => null,
357 .outer_product => |outer_product| outer_product,
358 };
359 }
360
361 fn selectTranspose(request: EinsumSelectionRequest) ?SelectedEinsumKernel {
362 if (request.inputs.len != 1) return null;
363 if (request.schedule != null) return null;
364 const descriptor = library.select(.{ .layout = .{
365 .dtype = request.dtype,
366 .kind = .transpose,
367 .input_indices = request.inputs[0].indices,
368 .output_indices = request.output_indices,
369 .input_dims = request.inputs[0].dims,
370 .output_dims = request.output_dims,
371 } }) orelse return null;
372 return .{
373 .kind = .transpose,
374 .descriptor = descriptor,
375 };
376 }
377
378 test "logical einsum selection chooses scalar sum reduction catalog entry" {
379 const input_dims = [_]i64{8};
380 const output_dims = [_]i64{};
381 const inputs = [_]EinsumOperand{
382 .{ .indices = "i", .dims = &input_dims },
383 };
384
385 const selected = selectCatalog(.{
386 .dtype = .f32,
387 .inputs = &inputs,
388 .output_indices = "",
389 .output_dims = &output_dims,
390 }) orelse return error.TestExpectedReductionSelection;
391
392 try std.testing.expectEqual(EinsumKernelKind.sum_reduction, selected.kind);
393 try std.testing.expectEqualStrings(library.reduction.Sum8F32.target, selected.descriptor.metadata.target);
394 }
395
396 test "logical einsum selection chooses scalar dot product catalog entry" {
397 const input_dims = [_]i64{8};
398 const output_dims = [_]i64{};
399 const inputs = [_]EinsumOperand{
400 .{ .indices = "i", .dims = &input_dims },
401 .{ .indices = "i", .dims = &input_dims },
402 };
403
404 const selected = selectCatalog(.{
405 .dtype = .f32,
406 .inputs = &inputs,
407 .output_indices = "",
408 .output_dims = &output_dims,
409 }) orelse return error.TestExpectedDotProductSelection;
410
411 try std.testing.expectEqual(EinsumKernelKind.dot_product, selected.kind);
412 try std.testing.expectEqualStrings(library.reduction.Dot8F32.target, selected.descriptor.metadata.target);
413 }
414
415 test "logical einsum selection chooses matrix product catalog entry" {
416 const lhs_dims = [_]i64{ 4, 8 };
417 const rhs_dims = [_]i64{ 8, 16 };
418 const output_dims = [_]i64{ 4, 16 };
419 const inputs = [_]EinsumOperand{
420 .{ .indices = "ik", .dims = &lhs_dims },
421 .{ .indices = "kj", .dims = &rhs_dims },
422 };
423
424 const selected = selectCatalog(.{
425 .dtype = .f32,
426 .inputs = &inputs,
427 .output_indices = "ij",
428 .output_dims = &output_dims,
429 }) orelse return error.TestExpectedMatrixProductSelection;
430
431 try std.testing.expectEqual(EinsumKernelKind.matrix_product, selected.kind);
432 try std.testing.expectEqualStrings(library.linalg.MatrixProduct4x16x8F32.target, selected.descriptor.metadata.target);
433 }
434
435 test "logical einsum owned selection chooses batched matrix product family at catalog smoke extents" {
436 const lhs_dims = [_]i64{ 2, 2, 4 };
437 const rhs_dims = [_]i64{ 2, 4, 3 };
438 const output_dims = [_]i64{ 2, 2, 3 };
439 const inputs = [_]EinsumOperand{
440 .{ .indices = "bmk", .dims = &lhs_dims },
441 .{ .indices = "bkn", .dims = &rhs_dims },
442 };
443 const request = EinsumSelectionRequest{
444 .dtype = .f32,
445 .inputs = &inputs,
446 .output_indices = "bmn",
447 .output_dims = &output_dims,
448 };
449
450 try std.testing.expect(selectCatalog(request) == null);
451 var selected = (try selectOwnedCatalog(std.testing.allocator, request)) orelse return error.TestExpectedBatchedMatrixProductSelection;
452 defer selected.deinit();
453 const specialization = selected.descriptor.descriptor.metadata.specialization;
454
455 try std.testing.expectEqual(EinsumKernelKind.batched_matrix_product, selected.kind);
456 try std.testing.expect(selected.descriptor.specialization != null);
457 try std.testing.expectEqualStrings("accy.kernel.linalg.batched_matmul_family_3x2x2_f32", selected.descriptor.descriptor.metadata.target);
458 try std.testing.expect(specialization.shape_family != null);
459 }
460
461 test "logical einsum owned selection chooses batched matrix product family schedule at catalog smoke extents" {
462 const lhs_dims = [_]i64{ 2, 2, 4 };
463 const rhs_dims = [_]i64{ 2, 4, 3 };
464 const output_dims = [_]i64{ 2, 2, 3 };
465 const inputs = [_]EinsumOperand{
466 .{ .indices = "bmk", .dims = &lhs_dims },
467 .{ .indices = "bkn", .dims = &rhs_dims },
468 };
469 const request = EinsumSelectionRequest{
470 .dtype = .f32,
471 .inputs = &inputs,
472 .output_indices = "bmn",
473 .output_dims = &output_dims,
474 .schedule = .{ .batched_matrix_product = .{ .thread_blocks = .{ .x = 3, .y = 2, .z = 2 } } },
475 };
476
477 try std.testing.expect(selectCatalog(request) == null);
478 var selected = (try selectOwnedCatalog(std.testing.allocator, request)) orelse return error.TestExpectedBatchedMatrixProductScheduleSelection;
479 defer selected.deinit();
480 const specialization = selected.descriptor.descriptor.metadata.specialization;
481
482 try std.testing.expectEqual(EinsumKernelKind.batched_matrix_product, selected.kind);
483 try std.testing.expect(selected.descriptor.specialization != null);
484 try std.testing.expectEqualStrings("accy.kernel.linalg.batched_matmul_family_3x2x2_f32", selected.descriptor.descriptor.metadata.target);
485 try std.testing.expectEqual(@as(u32, 3), specialization.launch.?.threadgroup[0]);
486 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.threadgroup[1]);
487 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.threadgroup[2]);
488 }
489
490 test "logical einsum owned selection chooses batched matrix product family" {
491 const lhs_dims = [_]i64{ 3, 5, 4 };
492 const rhs_dims = [_]i64{ 3, 4, 6 };
493 const output_dims = [_]i64{ 3, 5, 6 };
494 const inputs = [_]EinsumOperand{
495 .{ .indices = "bmk", .dims = &lhs_dims },
496 .{ .indices = "bkn", .dims = &rhs_dims },
497 };
498
499 var selected = (try selectOwnedCatalog(std.testing.allocator, .{
500 .dtype = .f32,
501 .inputs = &inputs,
502 .output_indices = "bmn",
503 .output_dims = &output_dims,
504 })) orelse return error.TestExpectedBatchedMatrixProductFamilySelection;
505 defer selected.deinit();
506 const specialization = selected.descriptor.descriptor.metadata.specialization;
507
508 try std.testing.expectEqual(EinsumKernelKind.batched_matrix_product, selected.kind);
509 try std.testing.expectEqualStrings("accy.kernel.linalg.batched_matmul_family_6x5x3_f32", selected.descriptor.descriptor.metadata.target);
510 try std.testing.expect(specialization.shape_family != null);
511 try std.testing.expectEqual(@as(u32, 6), specialization.launch.?.threadgroup[0]);
512 try std.testing.expectEqual(@as(u32, 5), specialization.launch.?.threadgroup[1]);
513 try std.testing.expectEqual(@as(u32, 3), specialization.launch.?.threadgroup[2]);
514 }
515
516 test "logical einsum owned selection chooses batched matrix product family schedule" {
517 const lhs_dims = [_]i64{ 3, 5, 4 };
518 const rhs_dims = [_]i64{ 3, 4, 6 };
519 const output_dims = [_]i64{ 3, 5, 6 };
520 const inputs = [_]EinsumOperand{
521 .{ .indices = "bmk", .dims = &lhs_dims },
522 .{ .indices = "bkn", .dims = &rhs_dims },
523 };
524
525 var selected = (try selectOwnedCatalog(std.testing.allocator, .{
526 .dtype = .f32,
527 .inputs = &inputs,
528 .output_indices = "bmn",
529 .output_dims = &output_dims,
530 .schedule = .{ .batched_matrix_product = .{ .thread_blocks = .{ .x = 4, .y = 2, .z = 2 } } },
531 })) orelse return error.TestExpectedBatchedMatrixProductFamilyScheduleSelection;
532 defer selected.deinit();
533 const specialization = selected.descriptor.descriptor.metadata.specialization;
534
535 try std.testing.expectEqual(EinsumKernelKind.batched_matrix_product, selected.kind);
536 try std.testing.expectEqualStrings("accy.kernel.linalg.batched_matmul_family_4x2x2_f32", selected.descriptor.descriptor.metadata.target);
537 try std.testing.expectEqual(@as(u32, 4), specialization.launch.?.threadgroup[0]);
538 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.threadgroup[1]);
539 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.threadgroup[2]);
540 }
541
542 test "logical einsum selection keeps attention-shaped pure einsum generic" {
543 const query_dims = [_]i64{ 2, 2, 2 };
544 const key_dims = [_]i64{ 2, 3, 2 };
545 const value_dims = [_]i64{ 2, 3, 2 };
546 const output_dims = [_]i64{ 2, 2, 2 };
547 const inputs = [_]EinsumOperand{
548 .{ .indices = "bqh", .dims = &query_dims },
549 .{ .indices = "bkh", .dims = &key_dims },
550 .{ .indices = "bkv", .dims = &value_dims },
551 };
552
553 try std.testing.expect(selectCatalog(.{
554 .dtype = .f32,
555 .inputs = &inputs,
556 .output_indices = "bqv",
557 .output_dims = &output_dims,
558 }) == null);
559 }
560
561 test "logical einsum selection chooses matrix product schedule specialization" {
562 const lhs_dims = [_]i64{ 4, 8 };
563 const rhs_dims = [_]i64{ 8, 16 };
564 const output_dims = [_]i64{ 4, 16 };
565 const inputs = [_]EinsumOperand{
566 .{ .indices = "ik", .dims = &lhs_dims },
567 .{ .indices = "kj", .dims = &rhs_dims },
568 };
569
570 const selected = selectCatalog(.{
571 .dtype = .f32,
572 .inputs = &inputs,
573 .output_indices = "ij",
574 .output_dims = &output_dims,
575 .schedule = .{ .matrix_product = .{ .thread_blocks = .{ .x = 4, .y = 2 } } },
576 }) orelse return error.TestExpectedMatrixProductScheduleSelection;
577
578 try std.testing.expectEqual(EinsumKernelKind.matrix_product, selected.kind);
579 try std.testing.expectEqualStrings(library.linalg.MatrixProduct4x16x8ThreadBlocks4x2F32.target, selected.descriptor.metadata.target);
580 }
581
582 test "logical einsum owned selection preserves fixed matrix product descriptor" {
583 const lhs_dims = [_]i64{ 4, 8 };
584 const rhs_dims = [_]i64{ 8, 16 };
585 const output_dims = [_]i64{ 4, 16 };
586 const inputs = [_]EinsumOperand{
587 .{ .indices = "ik", .dims = &lhs_dims },
588 .{ .indices = "kj", .dims = &rhs_dims },
589 };
590
591 var selected = (try selectOwnedCatalog(std.testing.allocator, .{
592 .dtype = .f32,
593 .inputs = &inputs,
594 .output_indices = "ij",
595 .output_dims = &output_dims,
596 })) orelse return error.TestExpectedMatrixProductSelection;
597 defer selected.deinit();
598
599 try std.testing.expectEqual(EinsumKernelKind.matrix_product, selected.kind);
600 try std.testing.expect(selected.descriptor.specialization == null);
601 try std.testing.expectEqualStrings(library.linalg.MatrixProduct4x16x8F32.target, selected.descriptor.descriptor.metadata.target);
602 }
603
604 test "logical einsum owned selection chooses matrix product family descriptor" {
605 const lhs_dims = [_]i64{ 5, 3 };
606 const rhs_dims = [_]i64{ 3, 7 };
607 const output_dims = [_]i64{ 5, 7 };
608 const inputs = [_]EinsumOperand{
609 .{ .indices = "mk", .dims = &lhs_dims },
610 .{ .indices = "kn", .dims = &rhs_dims },
611 };
612 const request = EinsumSelectionRequest{
613 .dtype = .f32,
614 .inputs = &inputs,
615 .output_indices = "mn",
616 .output_dims = &output_dims,
617 };
618
619 try std.testing.expect(selectCatalog(request) == null);
620 var selected = (try selectOwnedCatalog(std.testing.allocator, request)) orelse return error.TestExpectedMatrixProductFamilySelection;
621 defer selected.deinit();
622 const specialization = selected.descriptor.descriptor.metadata.specialization;
623
624 try std.testing.expectEqual(EinsumKernelKind.matrix_product, selected.kind);
625 try std.testing.expect(selected.descriptor.specialization != null);
626 try std.testing.expectEqualStrings("accy.kernel.linalg.matmul_family_7x5_f32", selected.descriptor.descriptor.metadata.target);
627 try std.testing.expect(specialization.scheduleMatchesLaunch());
628 try std.testing.expectEqual(@as(u32, 7), specialization.launch.?.threadgroup[0]);
629 try std.testing.expectEqual(@as(u32, 5), specialization.launch.?.threadgroup[1]);
630 }
631
632 test "logical einsum owned selection chooses f16 matrix product family descriptor" {
633 const lhs_dims = [_]i64{ 5, 3 };
634 const rhs_dims = [_]i64{ 3, 7 };
635 const output_dims = [_]i64{ 5, 7 };
636 const inputs = [_]EinsumOperand{
637 .{ .indices = "mk", .dims = &lhs_dims },
638 .{ .indices = "kn", .dims = &rhs_dims },
639 };
640 const request = EinsumSelectionRequest{
641 .dtype = .f16,
642 .inputs = &inputs,
643 .output_indices = "mn",
644 .output_dims = &output_dims,
645 };
646
647 try std.testing.expect(selectCatalog(request) == null);
648 var selected = (try selectOwnedCatalog(std.testing.allocator, request)) orelse return error.TestExpectedMatrixProductFamilySelection;
649 defer selected.deinit();
650 const specialization = selected.descriptor.descriptor.metadata.specialization;
651
652 try std.testing.expectEqual(EinsumKernelKind.matrix_product, selected.kind);
653 try std.testing.expect(selected.descriptor.specialization != null);
654 try std.testing.expectEqualStrings("accy.kernel.linalg.matmul_family_7x5_f16", selected.descriptor.descriptor.metadata.target);
655 try std.testing.expectEqual(@as(?DType, .f16), specialization.dtype);
656 try std.testing.expectEqual(@as(?DType, .f32), specialization.accumulation_dtype);
657 try std.testing.expect(specialization.scheduleMatchesLaunch());
658 }
659
660 test "logical einsum owned selection chooses matrix product family schedule" {
661 const lhs_dims = [_]i64{ 5, 3 };
662 const rhs_dims = [_]i64{ 3, 7 };
663 const output_dims = [_]i64{ 5, 7 };
664 const inputs = [_]EinsumOperand{
665 .{ .indices = "mk", .dims = &lhs_dims },
666 .{ .indices = "kn", .dims = &rhs_dims },
667 };
668
669 var selected = (try selectOwnedCatalog(std.testing.allocator, .{
670 .dtype = .f32,
671 .inputs = &inputs,
672 .output_indices = "mn",
673 .output_dims = &output_dims,
674 .schedule = .{ .matrix_product = .{ .thread_blocks = .{ .x = 4, .y = 2 } } },
675 })) orelse return error.TestExpectedMatrixProductFamilyScheduleSelection;
676 defer selected.deinit();
677 const specialization = selected.descriptor.descriptor.metadata.specialization;
678
679 try std.testing.expectEqual(EinsumKernelKind.matrix_product, selected.kind);
680 try std.testing.expectEqualStrings("accy.kernel.linalg.matmul_family_4x2_f32", selected.descriptor.descriptor.metadata.target);
681 try std.testing.expectEqual(@as(u32, 4), specialization.launch.?.threadgroup[0]);
682 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.threadgroup[1]);
683 }
684
685 test "logical einsum owned selection chooses matrix vector product family at catalog smoke extents" {
686 const matrix_dims = [_]i64{ 4, 8 };
687 const vector_dims = [_]i64{8};
688 const output_dims = [_]i64{4};
689 const inputs = [_]EinsumOperand{
690 .{ .indices = "mk", .dims = &matrix_dims },
691 .{ .indices = "k", .dims = &vector_dims },
692 };
693 const request = EinsumSelectionRequest{
694 .dtype = .f32,
695 .inputs = &inputs,
696 .output_indices = "m",
697 .output_dims = &output_dims,
698 };
699
700 try std.testing.expect(selectCatalog(request) == null);
701 var selected = (try selectOwnedCatalog(std.testing.allocator, request)) orelse return error.TestExpectedMatrixVectorProductSelection;
702 defer selected.deinit();
703 const specialization = selected.descriptor.descriptor.metadata.specialization;
704
705 try std.testing.expectEqual(EinsumKernelKind.matrix_vector_product, selected.kind);
706 try std.testing.expect(selected.descriptor.specialization != null);
707 try std.testing.expectEqualStrings("accy.kernel.linalg.matvec_family_4x_f32", selected.descriptor.descriptor.metadata.target);
708 try std.testing.expect(specialization.shape_family != null);
709 }
710
711 test "logical einsum owned selection chooses matrix vector product family descriptor" {
712 const matrix_dims = [_]i64{ 5, 3 };
713 const vector_dims = [_]i64{3};
714 const output_dims = [_]i64{5};
715 const inputs = [_]EinsumOperand{
716 .{ .indices = "mk", .dims = &matrix_dims },
717 .{ .indices = "k", .dims = &vector_dims },
718 };
719 const request = EinsumSelectionRequest{
720 .dtype = .f32,
721 .inputs = &inputs,
722 .output_indices = "m",
723 .output_dims = &output_dims,
724 };
725
726 try std.testing.expect(selectCatalog(request) == null);
727 var selected = (try selectOwnedCatalog(std.testing.allocator, request)) orelse return error.TestExpectedMatrixVectorProductFamilySelection;
728 defer selected.deinit();
729 const specialization = selected.descriptor.descriptor.metadata.specialization;
730
731 try std.testing.expectEqual(EinsumKernelKind.matrix_vector_product, selected.kind);
732 try std.testing.expect(selected.descriptor.specialization != null);
733 try std.testing.expectEqualStrings("accy.kernel.linalg.matvec_family_5x_f32", selected.descriptor.descriptor.metadata.target);
734 try std.testing.expect(specialization.operationIs(.{ .linalg = .matrix_vector_product }));
735 try std.testing.expect(specialization.scheduleMatchesLaunch());
736 try std.testing.expectEqual(@as(u32, 5), specialization.launch.?.threadgroup[0]);
737 }
738
739 test "logical einsum owned selection chooses matrix vector product family schedule" {
740 const matrix_dims = [_]i64{ 5, 3 };
741 const vector_dims = [_]i64{3};
742 const output_dims = [_]i64{5};
743 const inputs = [_]EinsumOperand{
744 .{ .indices = "mk", .dims = &matrix_dims },
745 .{ .indices = "k", .dims = &vector_dims },
746 };
747
748 var selected = (try selectOwnedCatalog(std.testing.allocator, .{
749 .dtype = .f32,
750 .inputs = &inputs,
751 .output_indices = "m",
752 .output_dims = &output_dims,
753 .schedule = .{ .matrix_vector_product = .{ .thread_blocks = 4 } },
754 })) orelse return error.TestExpectedMatrixVectorProductFamilyScheduleSelection;
755 defer selected.deinit();
756 const specialization = selected.descriptor.descriptor.metadata.specialization;
757
758 try std.testing.expectEqual(EinsumKernelKind.matrix_vector_product, selected.kind);
759 try std.testing.expectEqualStrings("accy.kernel.linalg.matvec_family_4x_f32", selected.descriptor.descriptor.metadata.target);
760 try std.testing.expectEqual(@as(u32, 4), specialization.launch.?.threadgroup[0]);
761 }
762
763 test "logical einsum owned selection chooses outer product family at catalog smoke extents" {
764 const lhs_dims = [_]i64{4};
765 const rhs_dims = [_]i64{3};
766 const output_dims = [_]i64{ 4, 3 };
767 const inputs = [_]EinsumOperand{
768 .{ .indices = "m", .dims = &lhs_dims },
769 .{ .indices = "n", .dims = &rhs_dims },
770 };
771 const request = EinsumSelectionRequest{
772 .dtype = .f32,
773 .inputs = &inputs,
774 .output_indices = "mn",
775 .output_dims = &output_dims,
776 };
777
778 try std.testing.expect(selectCatalog(request) == null);
779 var selected = (try selectOwnedCatalog(std.testing.allocator, request)) orelse return error.TestExpectedOuterProductSelection;
780 defer selected.deinit();
781 const specialization = selected.descriptor.descriptor.metadata.specialization;
782
783 try std.testing.expectEqual(EinsumKernelKind.outer_product, selected.kind);
784 try std.testing.expect(selected.descriptor.specialization != null);
785 try std.testing.expectEqualStrings("accy.kernel.linalg.outer_family_3x4_f32", selected.descriptor.descriptor.metadata.target);
786 try std.testing.expect(specialization.shape_family != null);
787 }
788
789 test "logical einsum owned selection chooses outer product family schedule at catalog smoke extents" {
790 const lhs_dims = [_]i64{4};
791 const rhs_dims = [_]i64{3};
792 const output_dims = [_]i64{ 4, 3 };
793 const inputs = [_]EinsumOperand{
794 .{ .indices = "m", .dims = &lhs_dims },
795 .{ .indices = "n", .dims = &rhs_dims },
796 };
797 const request = EinsumSelectionRequest{
798 .dtype = .f32,
799 .inputs = &inputs,
800 .output_indices = "mn",
801 .output_dims = &output_dims,
802 .schedule = .{ .outer_product = .{ .thread_blocks = .{ .x = 3, .y = 2 } } },
803 };
804
805 try std.testing.expect(selectCatalog(request) == null);
806 var selected = (try selectOwnedCatalog(std.testing.allocator, request)) orelse return error.TestExpectedOuterProductScheduleSelection;
807 defer selected.deinit();
808 const specialization = selected.descriptor.descriptor.metadata.specialization;
809
810 try std.testing.expectEqual(EinsumKernelKind.outer_product, selected.kind);
811 try std.testing.expect(selected.descriptor.specialization != null);
812 try std.testing.expectEqualStrings("accy.kernel.linalg.outer_family_3x2_f32", selected.descriptor.descriptor.metadata.target);
813 try std.testing.expectEqual(@as(u32, 3), specialization.launch.?.threadgroup[0]);
814 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.threadgroup[1]);
815 }
816
817 test "logical einsum owned selection chooses outer product family" {
818 const lhs_dims = [_]i64{5};
819 const rhs_dims = [_]i64{6};
820 const output_dims = [_]i64{ 5, 6 };
821 const inputs = [_]EinsumOperand{
822 .{ .indices = "m", .dims = &lhs_dims },
823 .{ .indices = "n", .dims = &rhs_dims },
824 };
825
826 var selected = (try selectOwnedCatalog(std.testing.allocator, .{
827 .dtype = .f32,
828 .inputs = &inputs,
829 .output_indices = "mn",
830 .output_dims = &output_dims,
831 })) orelse return error.TestExpectedOuterProductFamilySelection;
832 defer selected.deinit();
833 const specialization = selected.descriptor.descriptor.metadata.specialization;
834
835 try std.testing.expectEqual(EinsumKernelKind.outer_product, selected.kind);
836 try std.testing.expectEqualStrings("accy.kernel.linalg.outer_family_6x5_f32", selected.descriptor.descriptor.metadata.target);
837 try std.testing.expect(specialization.shape_family != null);
838 try std.testing.expectEqual(@as(u32, 6), specialization.launch.?.threadgroup[0]);
839 try std.testing.expectEqual(@as(u32, 5), specialization.launch.?.threadgroup[1]);
840 }
841
842 test "logical einsum owned selection chooses outer product family schedule" {
843 const lhs_dims = [_]i64{5};
844 const rhs_dims = [_]i64{6};
845 const output_dims = [_]i64{ 5, 6 };
846 const inputs = [_]EinsumOperand{
847 .{ .indices = "m", .dims = &lhs_dims },
848 .{ .indices = "n", .dims = &rhs_dims },
849 };
850
851 var selected = (try selectOwnedCatalog(std.testing.allocator, .{
852 .dtype = .f32,
853 .inputs = &inputs,
854 .output_indices = "mn",
855 .output_dims = &output_dims,
856 .schedule = .{ .outer_product = .{ .thread_blocks = .{ .x = 4, .y = 2 } } },
857 })) orelse return error.TestExpectedOuterProductFamilyScheduleSelection;
858 defer selected.deinit();
859 const specialization = selected.descriptor.descriptor.metadata.specialization;
860
861 try std.testing.expectEqual(EinsumKernelKind.outer_product, selected.kind);
862 try std.testing.expectEqualStrings("accy.kernel.linalg.outer_family_4x2_f32", selected.descriptor.descriptor.metadata.target);
863 try std.testing.expectEqual(@as(u32, 4), specialization.launch.?.threadgroup[0]);
864 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.threadgroup[1]);
865 }
866
867 test "logical einsum selection chooses transpose catalog entry" {
868 const input_dims = [_]i64{ 8, 16 };
869 const output_dims = [_]i64{ 16, 8 };
870 const inputs = [_]EinsumOperand{
871 .{ .indices = "ij", .dims = &input_dims },
872 };
873
874 const selected = selectCatalog(.{
875 .dtype = .f32,
876 .inputs = &inputs,
877 .output_indices = "ji",
878 .output_dims = &output_dims,
879 }) orelse return error.TestExpectedTransposeSelection;
880
881 try std.testing.expectEqual(EinsumKernelKind.transpose, selected.kind);
882 try std.testing.expectEqualStrings(library.layout.Transpose8x16F32.target, selected.descriptor.metadata.target);
883 }
884
885 test "logical einsum selection rejects unknown catalog shape" {
886 const lhs_dims = [_]i64{ 5, 8 };
887 const rhs_dims = [_]i64{ 8, 16 };
888 const output_dims = [_]i64{ 5, 16 };
889 const inputs = [_]EinsumOperand{
890 .{ .indices = "ik", .dims = &lhs_dims },
891 .{ .indices = "kj", .dims = &rhs_dims },
892 };
893
894 try std.testing.expect(selectCatalog(.{
895 .dtype = .f32,
896 .inputs = &inputs,
897 .output_indices = "ij",
898 .output_dims = &output_dims,
899 }) == null);
900
901 const scalar_dims = [_]i64{};
902 const mismatch_lhs_dims = [_]i64{8};
903 const mismatch_rhs_dims = [_]i64{8};
904 const mismatched_dot_inputs = [_]EinsumOperand{
905 .{ .indices = "i", .dims = &mismatch_lhs_dims },
906 .{ .indices = "j", .dims = &mismatch_rhs_dims },
907 };
908 try std.testing.expect(selectCatalog(.{
909 .dtype = .f32,
910 .inputs = &mismatched_dot_inputs,
911 .output_indices = "",
912 .output_dims = &scalar_dims,
913 }) == null);
914 try std.testing.expect(selectCatalog(.{
915 .dtype = .f32,
916 .inputs = mismatched_dot_inputs[0..1],
917 .output_indices = "",
918 .output_dims = &scalar_dims,
919 .schedule = .{ .matrix_product = .{ .thread_blocks = .{ .x = 4, .y = 2 } } },
920 }) == null);
921
922 const matrix_product_lhs_dims = [_]i64{ 4, 8 };
923 const matrix_product_rhs_dims = [_]i64{ 8, 16 };
924 const matrix_product_output_dims = [_]i64{ 4, 16 };
925 const matrix_product_inputs = [_]EinsumOperand{
926 .{ .indices = "ik", .dims = &matrix_product_lhs_dims },
927 .{ .indices = "kj", .dims = &matrix_product_rhs_dims },
928 };
929 try std.testing.expect(selectCatalog(.{
930 .dtype = .f32,
931 .inputs = &matrix_product_inputs,
932 .output_indices = "ij",
933 .output_dims = &matrix_product_output_dims,
934 .schedule = .{ .matrix_product = .{ .thread_blocks = .{ .x = 16, .y = 8 } } },
935 }) == null);
936 try std.testing.expect(selectCatalog(.{
937 .dtype = .f32,
938 .inputs = &matrix_product_inputs,
939 .output_indices = "ij",
940 .output_dims = &matrix_product_output_dims,
941 .schedule = .{ .matrix_vector_product = .{ .thread_blocks = 4 } },
942 }) == null);
943
944 const batched_lhs_dims = [_]i64{ 2, 2, 4 };
945 const batched_rhs_dims = [_]i64{ 2, 4, 3 };
946 const batched_output_dims = [_]i64{ 2, 2, 3 };
947 const batched_inputs = [_]EinsumOperand{
948 .{ .indices = "bmk", .dims = &batched_lhs_dims },
949 .{ .indices = "bkn", .dims = &batched_rhs_dims },
950 };
951 try std.testing.expect(selectCatalog(.{
952 .dtype = .f32,
953 .inputs = &batched_inputs,
954 .output_indices = "bmn",
955 .output_dims = &batched_output_dims,
956 .schedule = .{ .batched_matrix_product = .{ .thread_blocks = .{ .x = 4, .y = 2, .z = 2 } } },
957 }) == null);
958 try std.testing.expect(selectCatalog(.{
959 .dtype = .f32,
960 .inputs = &batched_inputs,
961 .output_indices = "bmn",
962 .output_dims = &batched_output_dims,
963 .schedule = .{ .matrix_product = .{ .thread_blocks = .{ .x = 4, .y = 2 } } },
964 }) == null);
965
966 const outer_lhs_dims = [_]i64{4};
967 const outer_rhs_dims = [_]i64{3};
968 const outer_output_dims = [_]i64{ 4, 3 };
969 const outer_inputs = [_]EinsumOperand{
970 .{ .indices = "m", .dims = &outer_lhs_dims },
971 .{ .indices = "n", .dims = &outer_rhs_dims },
972 };
973 try std.testing.expect(selectCatalog(.{
974 .dtype = .f32,
975 .inputs = &outer_inputs,
976 .output_indices = "mn",
977 .output_dims = &outer_output_dims,
978 .schedule = .{ .outer_product = .{ .thread_blocks = .{ .x = 4, .y = 2 } } },
979 }) == null);
980 try std.testing.expect(selectCatalog(.{
981 .dtype = .f32,
982 .inputs = &outer_inputs,
983 .output_indices = "mn",
984 .output_dims = &outer_output_dims,
985 .schedule = .{ .matrix_product = .{ .thread_blocks = .{ .x = 4, .y = 2 } } },
986 }) == null);
987
988 const transpose_input_dims = [_]i64{ 8, 16 };
989 const transpose_output_dims = [_]i64{ 16, 8 };
990 const transpose_inputs = [_]EinsumOperand{
991 .{ .indices = "ij", .dims = &transpose_input_dims },
992 };
993 try std.testing.expect(selectCatalog(.{
994 .dtype = .f32,
995 .inputs = &transpose_inputs,
996 .output_indices = "ji",
997 .output_dims = &transpose_output_dims,
998 .schedule = .{ .matrix_vector_product = .{ .thread_blocks = 4 } },
999 }) == null);
1000
1001 const attention_query_dims = [_]i64{ 2, 2, 2 };
1002 const attention_key_dims = [_]i64{ 2, 3, 2 };
1003 const attention_value_dims = [_]i64{ 2, 3, 2 };
1004 const attention_output_dims = [_]i64{ 2, 2, 2 };
1005 const attention_inputs = [_]EinsumOperand{
1006 .{ .indices = "bqh", .dims = &attention_query_dims },
1007 .{ .indices = "bkh", .dims = &attention_key_dims },
1008 .{ .indices = "bkv", .dims = &attention_value_dims },
1009 };
1010 try std.testing.expect(selectCatalog(.{
1011 .dtype = .f32,
1012 .inputs = &attention_inputs,
1013 .output_indices = "bqv",
1014 .output_dims = &attention_output_dims,
1015 .schedule = .{ .matrix_product = .{ .thread_blocks = .{ .x = 4, .y = 2 } } },
1016 }) == null);
1017 }