lib/accy/src/kernel/library/entry.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const alloc_arena = @import("alloc_arena");
5 const choir = @import("choir");
6
7 const activation_mod = @import("../../choir/root.zig").activation;
8 const artifact = @import("../../artifact/model/root.zig");
9 const shape_mod = @import("../../choir/shape/root.zig");
10 const kernel = @import("../root.zig");
11
12 pub const Layer = enum {
13 logical,
14 authored,
15 };
16
17 pub const Category = enum {
18 attention,
19 elementwise,
20 reduction,
21 normalization,
22 loss,
23 layout,
24 linalg,
25 stencil,
26 indexing,
27 segmented,
28 scan,
29 sort,
30 sparse,
31 spatial,
32 image,
33 random,
34 compaction,
35 fused,
36 model,
37 };
38
39 pub const ElementwiseOperator = enum {
40 add,
41 axpy,
42 mul,
43 scale,
44 };
45
46 pub const ReductionOperator = enum {
47 sum,
48 maximum,
49 dot_product,
50 sum_squares,
51 sum_exp_shifted,
52 sum_squared_difference,
53 weighted_sum,
54 };
55
56 pub const RowNormalizationParameterization = enum {
57 none,
58 scale,
59 scale_bias,
60 };
61
62 pub const RowNormalizationOperator = union(enum) {
63 softmax,
64 log_softmax,
65 rmsnorm: RowNormalizationParameterization,
66 layernorm: RowNormalizationParameterization,
67 };
68
69 pub const LayoutOperator = enum {
70 transpose,
71 };
72
73 pub const LinalgOperator = enum {
74 batched_matrix_product,
75 matrix_product,
76 matrix_vector_product,
77 outer_product,
78 batched_cholesky,
79 batched_cholesky_solve,
80 batched_inverse,
81 };
82
83 pub const StencilOperator = enum {
84 window,
85 };
86
87 pub const IndexingOperator = enum {
88 gather,
89 scatter,
90 scatter_add,
91 histogram,
92 };
93
94 pub const SegmentedOperator = enum {
95 segment_sum,
96 };
97
98 pub const SortOperator = enum {
99 radix_ascending,
100 top_k_smallest,
101 };
102
103 pub const SparseOperator = enum {
104 coo_spmv,
105 csr_spmv,
106 csr_spmm,
107 ell_spmv,
108 sell_spmv,
109 };
110
111 pub const SpatialOperator = enum {
112 grid_cells,
113 grid_count,
114 grid_neighbor_count,
115 };
116
117 pub const ImageOperator = enum {
118 blur_pass,
119 resize_bilinear,
120 };
121
122 pub const ScanOperator = enum {
123 prefix_sum,
124 prefix_sum_exclusive,
125 };
126
127 pub const RandomOperator = union(enum) {
128 philox: u32,
129 threefry: u32,
130 squares,
131 philox_key_split: u32,
132 philox_key_uniform: u32,
133 philox_key_counter_uniform: u32,
134 philox_fold: u32,
135 threefry_fold: u32,
136 squares_fold,
137 };
138
139 pub const CompactionPredicate = enum {
140 nonzero,
141 greater_than,
142 };
143
144 pub const CompactionOperator = union(enum) {
145 blocks: CompactionPredicate,
146 };
147
148 pub const AttentionOperator = enum {
149 scaled_dot_product,
150 };
151
152 pub const LossOperator = enum {
153 row_sparse_cross_entropy,
154 };
155
156 pub const Operation = union(enum) {
157 attention: AttentionOperator,
158 elementwise: ElementwiseOperator,
159 activation: activation_mod.Kind,
160 reduction: ReductionOperator,
161 row_normalization: RowNormalizationOperator,
162 loss: LossOperator,
163 layout: LayoutOperator,
164 linalg: LinalgOperator,
165 stencil: StencilOperator,
166 indexing: IndexingOperator,
167 segmented: SegmentedOperator,
168 scan: ScanOperator,
169 sort: SortOperator,
170 sparse: SparseOperator,
171 spatial: SpatialOperator,
172 image: ImageOperator,
173 random: RandomOperator,
174 compaction: CompactionOperator,
175 };
176
177 pub fn operationFingerprint(operation: Operation) u64 {
178 var builder = choir.product.incremental.FingerprintBuilder{};
179 builder.updateBytes("accy.kernel.library.operation");
180 hashTagged(&builder, operation);
181 return builder.finish();
182 }
183
184 fn hashTagged(builder: *choir.product.incremental.FingerprintBuilder, value: anytype) void {
185 const Value = @TypeOf(value);
186 switch (@typeInfo(Value)) {
187 .@"union" => {
188 builder.updateEnumTag(std.meta.activeTag(value));
189 switch (value) {
190 inline else => |payload| hashTagged(builder, payload),
191 }
192 },
193 .@"enum" => builder.updateEnumTag(value),
194 .int => builder.updateU64(value),
195 .void => {},
196 else => @compileError("operation payload kind has no fingerprint rule: " ++ @typeName(Value)),
197 }
198 }
199
200 pub const Matrix2D = struct {
201 rows: u64,
202 cols: u64,
203 threads: Threads2D = .{},
204 row_axis: []const u8 = "i",
205 col_axis: []const u8 = "j",
206 };
207
208 pub const Threads2D = struct {
209 x: u32 = 16,
210 y: u32 = 16,
211 };
212
213 pub const Threads3D = struct {
214 x: u32 = 8,
215 y: u32 = 4,
216 z: u32 = 1,
217 };
218
219 pub const Axis = struct {
220 name: []const u8,
221 extent: u64,
222 };
223
224 pub const Vector1D = struct {
225 extent: u64,
226 threads: u32,
227 axis: []const u8 = "i",
228 };
229
230 pub const Shape = struct {
231 axes: []const Axis = &.{},
232
233 pub fn rank(self: Shape) usize {
234 return self.axes.len;
235 }
236
237 pub fn axis(self: Shape, name: []const u8) ?Axis {
238 for (self.axes) |candidate| {
239 if (std.mem.eql(u8, candidate.name, name)) return candidate;
240 }
241 return null;
242 }
243
244 pub fn elementCount(self: Shape) ?u64 {
245 var count: u64 = 1;
246 for (self.axes) |axis_value| {
247 count = std.math.mul(u64, count, axis_value.extent) catch return null;
248 }
249 return count;
250 }
251
252 pub fn matchesExtents(self: Shape, extents: []const u64) bool {
253 if (self.axes.len != extents.len) return false;
254 for (extents, 0..) |extent, index| {
255 if (self.axes[index].extent != extent) return false;
256 }
257 return true;
258 }
259 };
260
261 pub const Launch = struct {
262 grid: [3]u32 = .{ 1, 1, 1 },
263 threadgroup: [3]u32 = .{ 1, 1, 1 },
264 };
265
266 pub const ScheduleBinding = struct {
267 axis: []const u8,
268 target: kernel.BindTarget,
269 extent: u64,
270 };
271
272 pub const Schedule = struct {
273 bindings: []const ScheduleBinding = &.{},
274
275 pub fn launch(self: Schedule) Launch {
276 var result = Launch{};
277 for (self.bindings) |binding| {
278 const extent: u32 = @intCast(binding.extent);
279 switch (binding.target) {
280 .block_x => result.grid[0] = extent,
281 .block_y => result.grid[1] = extent,
282 .block_z => result.grid[2] = extent,
283 .thread_x => result.threadgroup[0] = extent,
284 .thread_y => result.threadgroup[1] = extent,
285 .thread_z => result.threadgroup[2] = extent,
286 }
287 }
288 return result;
289 }
290
291 pub fn matchesSnapshot(self: Schedule, snapshot: *const kernel.ScheduleSnapshot) bool {
292 const axes = snapshot.allAxes();
293 if (self.bindings.len != axes.len) return false;
294 for (self.bindings, axes) |binding, axis_value| {
295 if (!std.mem.eql(u8, binding.axis, axis_value.name)) return false;
296 if (axis_value.bind != binding.target) return false;
297 if (axis_value.extent != binding.extent) return false;
298 }
299 return true;
300 }
301 };
302
303 pub const ReductionContract = struct {
304 name: []const u8,
305 operator: ReductionOperator,
306 extents: []const u64,
307 dependencies: []const []const u8 = &.{},
308 };
309
310 pub const Reduction = struct {
311 name: []const u8,
312 operator: ReductionOperator,
313 shape: Shape,
314 dependencies: []const []const u8 = &.{},
315
316 pub fn matches(self: Reduction, contract: ReductionContract) bool {
317 if (!std.mem.eql(u8, self.name, contract.name)) return false;
318 if (self.operator != contract.operator) return false;
319 if (!self.shape.matchesExtents(contract.extents)) return false;
320 if (self.dependencies.len != contract.dependencies.len) return false;
321 for (contract.dependencies, 0..) |dependency, index| {
322 if (!std.mem.eql(u8, self.dependencies[index], dependency)) return false;
323 }
324 return true;
325 }
326 };
327
328 pub const ReductionReuseContract = struct {
329 reduction: []const u8,
330 extents: []const u64,
331 };
332
333 pub const ReductionReuse = struct {
334 reduction: []const u8,
335 shape: Shape,
336
337 pub fn matches(self: ReductionReuse, contract: ReductionReuseContract) bool {
338 if (!std.mem.eql(u8, self.reduction, contract.reduction)) return false;
339 return self.shape.matchesExtents(contract.extents);
340 }
341 };
342
343 pub const StaticParameter = struct {
344 name: []const u8,
345 value: u64,
346
347 pub fn matches(self: StaticParameter, name: []const u8, value: u64) bool {
348 return std.mem.eql(u8, self.name, name) and self.value == value;
349 }
350 };
351
352 pub const EpilogueOperator = union(enum) {
353 bias_add,
354 activation: activation_mod.Kind,
355 };
356
357 pub const EpilogueContract = struct {
358 operator: EpilogueOperator,
359 input_index: ?usize = null,
360 extents: []const u64 = &.{},
361 };
362
363 pub const EpilogueStep = struct {
364 operator: EpilogueOperator,
365 input_index: ?usize = null,
366 shape: ?Shape = null,
367
368 pub fn matches(self: EpilogueStep, contract: EpilogueContract) bool {
369 if (!std.meta.eql(self.operator, contract.operator)) return false;
370 if (self.input_index != contract.input_index) return false;
371 if (contract.extents.len == 0) return true;
372 const shape = self.shape orelse return false;
373 return shape.matchesExtents(contract.extents);
374 }
375 };
376
377 pub const InputTransformOperator = union(enum) {
378 activation: activation_mod.Kind,
379 residual_add,
380 };
381
382 pub const InputTransformContract = struct {
383 operator: InputTransformOperator,
384 input_index: usize,
385 extents: []const u64 = &.{},
386 };
387
388 pub const InputTransform = struct {
389 operator: InputTransformOperator,
390 input_index: usize,
391 shape: ?Shape = null,
392
393 pub fn matches(self: InputTransform, contract: InputTransformContract) bool {
394 if (!std.meta.eql(self.operator, contract.operator)) return false;
395 if (self.input_index != contract.input_index) return false;
396 if (contract.extents.len == 0) return true;
397 const shape = self.shape orelse return false;
398 return shape.matchesExtents(contract.extents);
399 }
400 };
401
402 pub const Specialization = struct {
403 dtype: ?choir_abi.DType = null,
404 accumulation_dtype: ?choir_abi.DType = null,
405 operation: ?Operation = null,
406 equation: ?[]const u8 = null,
407 inputs: []const Shape = &.{},
408 outputs: []const Shape = &.{},
409 reductions: []const Reduction = &.{},
410 reduction_reuse: []const ReductionReuse = &.{},
411 input_transforms: []const InputTransform = &.{},
412 epilogues: []const EpilogueStep = &.{},
413 static_parameters: []const StaticParameter = &.{},
414 launch: ?Launch = null,
415 schedule: ?Schedule = null,
416 structure: ?[]const u8 = null,
417 layout: ?[]const u8 = null,
418 shape_family: ?*const shape_mod.Family = null,
419
420 pub fn shapeFamilyFingerprint(self: Specialization) ?u64 {
421 const family = self.shape_family orelse return null;
422 return shape_mod.fingerprint(family.*);
423 }
424
425 pub fn operationIs(self: Specialization, operation: Operation) bool {
426 const value = self.operation orelse return false;
427 return std.meta.eql(value, operation);
428 }
429
430 pub fn structureIs(self: Specialization, structure: []const u8) bool {
431 const value = self.structure orelse return false;
432 return std.mem.eql(u8, value, structure);
433 }
434
435 pub fn layoutIs(self: Specialization, layout: []const u8) bool {
436 const value = self.layout orelse return false;
437 return std.mem.eql(u8, value, layout);
438 }
439
440 pub fn staticParameterValue(self: Specialization, name: []const u8) ?u64 {
441 for (self.static_parameters) |parameter| {
442 if (std.mem.eql(u8, parameter.name, name)) return parameter.value;
443 }
444 return null;
445 }
446
447 pub fn staticParameterMatches(self: Specialization, name: []const u8, value: u64) bool {
448 const actual = self.staticParameterValue(name) orelse return false;
449 return actual == value;
450 }
451
452 pub fn scheduleMatchesLaunch(self: Specialization) bool {
453 const schedule = self.schedule orelse return false;
454 const launch = self.launch orelse return false;
455 return std.meta.eql(launch, schedule.launch());
456 }
457
458 pub fn inputHasExtents(self: Specialization, index: usize, extents: []const u64) bool {
459 if (index >= self.inputs.len) return false;
460 return self.inputs[index].matchesExtents(extents);
461 }
462
463 pub fn outputHasExtents(self: Specialization, index: usize, extents: []const u64) bool {
464 if (index >= self.outputs.len) return false;
465 return self.outputs[index].matchesExtents(extents);
466 }
467
468 pub fn reductionMatches(self: Specialization, index: usize, contract: ReductionContract) bool {
469 if (index >= self.reductions.len) return false;
470 return self.reductions[index].matches(contract);
471 }
472
473 pub fn reductionReuseMatches(self: Specialization, index: usize, contract: ReductionReuseContract) bool {
474 if (index >= self.reduction_reuse.len) return false;
475 return self.reduction_reuse[index].matches(contract);
476 }
477
478 pub fn reductionDependenciesAreValid(self: Specialization) bool {
479 for (self.reductions, 0..) |reduction_value, index| {
480 if (reduction_value.name.len == 0) return false;
481 if (reductionNameExists(self.reductions[0..index], reduction_value.name)) return false;
482 for (reduction_value.dependencies) |dependency| {
483 if (dependency.len == 0) return false;
484 if (!reductionNameExists(self.reductions[0..index], dependency)) return false;
485 }
486 }
487 return true;
488 }
489
490 pub fn reductionReuseScopesAreValid(self: Specialization) bool {
491 for (self.reduction_reuse, 0..) |reuse, index| {
492 if (reuse.reduction.len == 0) return false;
493 if (!reductionNameExists(self.reductions, reuse.reduction)) return false;
494 if (reductionReuseExists(self.reduction_reuse[0..index], reuse.reduction)) return false;
495 if (reuse.shape.elementCount() == null) return false;
496 }
497 return true;
498 }
499
500 pub fn staticParametersAreValid(self: Specialization) bool {
501 for (self.static_parameters, 0..) |parameter, index| {
502 if (parameter.name.len == 0) return false;
503 if (staticParameterNameExists(self.static_parameters[0..index], parameter.name)) return false;
504 }
505 return true;
506 }
507
508 pub fn inputTransformMatches(self: Specialization, index: usize, contract: InputTransformContract) bool {
509 if (index >= self.input_transforms.len) return false;
510 return self.input_transforms[index].matches(contract);
511 }
512
513 pub fn epilogueMatches(self: Specialization, index: usize, contract: EpilogueContract) bool {
514 if (index >= self.epilogues.len) return false;
515 return self.epilogues[index].matches(contract);
516 }
517
518 pub fn outputElementCount(self: Specialization) ?u64 {
519 if (self.outputs.len == 0) return null;
520 var count: u64 = 0;
521 for (self.outputs) |shape| {
522 count = saturatedAdd(count, shape.elementCount() orelse return null);
523 }
524 return count;
525 }
526
527 pub fn estimatedElementOps(self: Specialization) ?u64 {
528 const output_count = self.outputElementCount() orelse return null;
529 if (self.reductions.len == 0) return output_count;
530 var total: u64 = 0;
531 for (self.reductions) |reduction_value| {
532 const reduction_count = reduction_value.shape.elementCount() orelse return null;
533 const result_count = self.reductionResultCount(reduction_value.name, output_count) orelse return null;
534 total = saturatedAdd(
535 total,
536 saturatedMul(
537 saturatedMul(result_count, reduction_count),
538 reductionOperatorWeight(reduction_value.operator),
539 ),
540 );
541 }
542 return total;
543 }
544
545 fn reductionResultCount(self: Specialization, name: []const u8, default_count: u64) ?u64 {
546 for (self.reduction_reuse) |reuse| {
547 if (std.mem.eql(u8, reuse.reduction, name)) return reuse.shape.elementCount();
548 }
549 return default_count;
550 }
551 };
552
553 fn reductionNameExists(reductions: []const Reduction, name: []const u8) bool {
554 for (reductions) |reduction_value| {
555 if (std.mem.eql(u8, reduction_value.name, name)) return true;
556 }
557 return false;
558 }
559
560 fn reductionReuseExists(reuses: []const ReductionReuse, reduction_name: []const u8) bool {
561 for (reuses) |reuse| {
562 if (std.mem.eql(u8, reuse.reduction, reduction_name)) return true;
563 }
564 return false;
565 }
566
567 fn staticParameterNameExists(parameters: []const StaticParameter, name: []const u8) bool {
568 for (parameters) |parameter| {
569 if (std.mem.eql(u8, parameter.name, name)) return true;
570 }
571 return false;
572 }
573
574 fn reductionOperatorWeight(operator: ReductionOperator) u64 {
575 return switch (operator) {
576 .dot_product => 2,
577 .sum,
578 .maximum,
579 .sum_squares,
580 .sum_exp_shifted,
581 .sum_squared_difference,
582 => 1,
583 .weighted_sum => 2,
584 };
585 }
586
587 fn saturatedAdd(lhs: u64, rhs: u64) u64 {
588 return std.math.add(u64, lhs, rhs) catch std.math.maxInt(u64);
589 }
590
591 fn saturatedMul(lhs: u64, rhs: u64) u64 {
592 return std.math.mul(u64, lhs, rhs) catch std.math.maxInt(u64);
593 }
594
595 pub const Metadata = struct {
596 target: []const u8,
597 version: u32 = 1,
598 layer: Layer,
599 category: Category,
600 specialization: Specialization = .{},
601 };
602
603 pub const OwnedSpecialization = struct {
604 backing_allocator: std.mem.Allocator,
605 arena: alloc_arena.Arena,
606 value: Specialization = .{},
607 shape_family: ?*shape_mod.Family = null,
608
609 pub fn init(backing_allocator: std.mem.Allocator) OwnedSpecialization {
610 return .{
611 .backing_allocator = backing_allocator,
612 .arena = alloc_arena.Arena.init(backing_allocator),
613 };
614 }
615
616 pub fn allocator(self: *OwnedSpecialization) std.mem.Allocator {
617 return self.arena.allocator();
618 }
619
620 pub fn deinit(self: *OwnedSpecialization) void {
621 self.clearShapeFamily();
622 self.arena.deinit();
623 self.* = undefined;
624 }
625
626 pub fn takeShapeFamily(self: *OwnedSpecialization, family: *shape_mod.Family) !void {
627 const owned = try self.backing_allocator.create(shape_mod.Family);
628 owned.* = family.*;
629 family.* = undefined;
630 self.clearShapeFamily();
631 self.shape_family = owned;
632 self.value.shape_family = owned;
633 }
634
635 fn clearShapeFamily(self: *OwnedSpecialization) void {
636 if (self.shape_family) |family| {
637 var mutable = family;
638 mutable.deinit();
639 self.backing_allocator.destroy(mutable);
640 self.shape_family = null;
641 self.value.shape_family = null;
642 }
643 }
644 };
645
646 pub fn shape1D(comptime axis_name: []const u8, comptime extent: u64) Shape {
647 if (extent == 0) @compileError("kernel library 1D shape extent must be nonzero");
648 return .{ .axes = &.{.{ .name = axis_name, .extent = extent }} };
649 }
650
651 pub fn runtimeShape1D(lifetime_allocator: std.mem.Allocator, axis_name: []const u8, extent: u64) !Shape {
652 if (extent == 0) return error.KernelLibraryShapeExtentMustBeNonzero;
653 const axes = try lifetime_allocator.alloc(Axis, 1);
654 axes[0] = try runtimeAxis(lifetime_allocator, axis_name, extent);
655 return .{ .axes = axes };
656 }
657
658 pub fn shapeScalar() Shape {
659 return .{ .axes = &.{} };
660 }
661
662 pub fn runtimeShapeScalar() Shape {
663 return .{ .axes = &.{} };
664 }
665
666 pub fn shape2D(
667 comptime outer_name: []const u8,
668 comptime outer_extent: u64,
669 comptime inner_name: []const u8,
670 comptime inner_extent: u64,
671 ) Shape {
672 if (outer_extent == 0) @compileError("kernel library 2D outer extent must be nonzero");
673 if (inner_extent == 0) @compileError("kernel library 2D inner extent must be nonzero");
674 return .{ .axes = &.{
675 .{ .name = outer_name, .extent = outer_extent },
676 .{ .name = inner_name, .extent = inner_extent },
677 } };
678 }
679
680 pub fn runtimeShape2D(
681 lifetime_allocator: std.mem.Allocator,
682 outer_name: []const u8,
683 outer_extent: u64,
684 inner_name: []const u8,
685 inner_extent: u64,
686 ) !Shape {
687 if (outer_extent == 0 or inner_extent == 0) return error.KernelLibraryShapeExtentMustBeNonzero;
688 const axes = try lifetime_allocator.alloc(Axis, 2);
689 axes[0] = try runtimeAxis(lifetime_allocator, outer_name, outer_extent);
690 axes[1] = try runtimeAxis(lifetime_allocator, inner_name, inner_extent);
691 return .{ .axes = axes };
692 }
693
694 pub fn shape3D(
695 comptime outer_name: []const u8,
696 comptime outer_extent: u64,
697 comptime middle_name: []const u8,
698 comptime middle_extent: u64,
699 comptime inner_name: []const u8,
700 comptime inner_extent: u64,
701 ) Shape {
702 if (outer_extent == 0) @compileError("kernel library 3D outer extent must be nonzero");
703 if (middle_extent == 0) @compileError("kernel library 3D middle extent must be nonzero");
704 if (inner_extent == 0) @compileError("kernel library 3D inner extent must be nonzero");
705 return .{ .axes = &.{
706 .{ .name = outer_name, .extent = outer_extent },
707 .{ .name = middle_name, .extent = middle_extent },
708 .{ .name = inner_name, .extent = inner_extent },
709 } };
710 }
711
712 pub fn runtimeShape3D(
713 lifetime_allocator: std.mem.Allocator,
714 outer_name: []const u8,
715 outer_extent: u64,
716 middle_name: []const u8,
717 middle_extent: u64,
718 inner_name: []const u8,
719 inner_extent: u64,
720 ) !Shape {
721 if (outer_extent == 0 or middle_extent == 0 or inner_extent == 0) return error.KernelLibraryShapeExtentMustBeNonzero;
722 const axes = try lifetime_allocator.alloc(Axis, 3);
723 axes[0] = try runtimeAxis(lifetime_allocator, outer_name, outer_extent);
724 axes[1] = try runtimeAxis(lifetime_allocator, middle_name, middle_extent);
725 axes[2] = try runtimeAxis(lifetime_allocator, inner_name, inner_extent);
726 return .{ .axes = axes };
727 }
728
729 fn runtimeAxis(lifetime_allocator: std.mem.Allocator, name: []const u8, extent: u64) !Axis {
730 if (name.len == 0) return error.KernelLibraryAxisNameMustBeNonempty;
731 return .{
732 .name = try lifetime_allocator.dupe(u8, name),
733 .extent = extent,
734 };
735 }
736
737 pub fn reduction(comptime name: []const u8, comptime operator: ReductionOperator, shape: Shape) Reduction {
738 return .{
739 .name = name,
740 .operator = operator,
741 .shape = shape,
742 };
743 }
744
745 pub fn runtimeReduction(
746 lifetime_allocator: std.mem.Allocator,
747 name: []const u8,
748 operator: ReductionOperator,
749 shape: Shape,
750 ) !Reduction {
751 return runtimeDependentReduction(lifetime_allocator, name, operator, shape, &.{});
752 }
753
754 pub fn dependentReduction(
755 comptime name: []const u8,
756 comptime operator: ReductionOperator,
757 shape: Shape,
758 comptime dependencies: []const []const u8,
759 ) Reduction {
760 return .{
761 .name = name,
762 .operator = operator,
763 .shape = shape,
764 .dependencies = dependencies,
765 };
766 }
767
768 pub fn runtimeDependentReduction(
769 lifetime_allocator: std.mem.Allocator,
770 name: []const u8,
771 operator: ReductionOperator,
772 shape: Shape,
773 dependencies: []const []const u8,
774 ) !Reduction {
775 if (name.len == 0) return error.KernelLibraryReductionNameMustBeNonempty;
776 const owned_dependencies = try lifetime_allocator.alloc([]const u8, dependencies.len);
777 for (dependencies, 0..) |dependency, index| {
778 if (dependency.len == 0) return error.KernelLibraryReductionDependencyMustBeNonempty;
779 owned_dependencies[index] = try lifetime_allocator.dupe(u8, dependency);
780 }
781 return .{
782 .name = try lifetime_allocator.dupe(u8, name),
783 .operator = operator,
784 .shape = shape,
785 .dependencies = owned_dependencies,
786 };
787 }
788
789 pub fn reductionReuse(comptime reduction_name: []const u8, shape: Shape) ReductionReuse {
790 return .{
791 .reduction = reduction_name,
792 .shape = shape,
793 };
794 }
795
796 pub fn runtimeReductionReuse(
797 lifetime_allocator: std.mem.Allocator,
798 reduction_name: []const u8,
799 shape: Shape,
800 ) !ReductionReuse {
801 if (reduction_name.len == 0) return error.KernelLibraryReductionReuseReductionMustBeNonempty;
802 return .{
803 .reduction = try lifetime_allocator.dupe(u8, reduction_name),
804 .shape = shape,
805 };
806 }
807
808 pub fn staticParameter(comptime name: []const u8, comptime value: u64) StaticParameter {
809 if (name.len == 0) @compileError("kernel library static parameter name must be nonempty");
810 return .{ .name = name, .value = value };
811 }
812
813 pub fn runtimeStaticParameter(lifetime_allocator: std.mem.Allocator, name: []const u8, value: u64) !StaticParameter {
814 if (name.len == 0) return error.KernelLibraryStaticParameterNameMustBeNonempty;
815 return .{
816 .name = try lifetime_allocator.dupe(u8, name),
817 .value = value,
818 };
819 }
820
821 pub fn epilogue(comptime operator: EpilogueOperator) EpilogueStep {
822 return .{ .operator = operator };
823 }
824
825 pub fn inputEpilogue(comptime operator: EpilogueOperator, comptime input_index: usize, shape: Shape) EpilogueStep {
826 return .{
827 .operator = operator,
828 .input_index = input_index,
829 .shape = shape,
830 };
831 }
832
833 pub fn inputTransform(comptime operator: InputTransformOperator, comptime input_index: usize, shape: Shape) InputTransform {
834 return .{
835 .operator = operator,
836 .input_index = input_index,
837 .shape = shape,
838 };
839 }
840
841 pub fn grid1D(comptime extent: u64, comptime threads: u32) u32 {
842 if (extent == 0) @compileError("kernel library 1D extent must be nonzero");
843 if (threads == 0) @compileError("kernel library 1D threadgroup must be nonzero");
844 const thread_count: u64 = threads;
845 if (extent > std.math.maxInt(u64) - (thread_count - 1)) {
846 @compileError("kernel library 1D launch extent overflow");
847 }
848 const biased = extent + thread_count - 1;
849 const blocks = biased / thread_count;
850 if (blocks > std.math.maxInt(u32)) {
851 @compileError("kernel library 1D grid overflow");
852 }
853 return @intCast(blocks);
854 }
855
856 pub fn runtimeGrid1D(extent: u64, threads: u32) !u32 {
857 if (extent == 0) return error.KernelLibraryLaunchExtentMustBeNonzero;
858 if (threads == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
859 const thread_count: u64 = threads;
860 const biased = std.math.add(u64, extent, thread_count - 1) catch return error.KernelLibraryLaunchGridOverflow;
861 const blocks = biased / thread_count;
862 return std.math.cast(u32, blocks) orelse error.KernelLibraryLaunchGridOverflow;
863 }
864
865 pub fn launch1D(comptime extent: u64, comptime threads: u32) Launch {
866 return .{
867 .grid = .{ grid1D(extent, threads), 1, 1 },
868 .threadgroup = .{ threads, 1, 1 },
869 };
870 }
871
872 pub fn runtimeLaunch1D(extent: u64, threads: u32) !Launch {
873 return .{
874 .grid = .{ try runtimeGrid1D(extent, threads), 1, 1 },
875 .threadgroup = .{ threads, 1, 1 },
876 };
877 }
878
879 pub fn launch2D(
880 comptime x_extent: u64,
881 comptime y_extent: u64,
882 comptime x_threads: u32,
883 comptime y_threads: u32,
884 ) Launch {
885 return .{
886 .grid = .{ grid1D(x_extent, x_threads), grid1D(y_extent, y_threads), 1 },
887 .threadgroup = .{ x_threads, y_threads, 1 },
888 };
889 }
890
891 pub fn runtimeLaunch2D(
892 x_extent: u64,
893 y_extent: u64,
894 x_threads: u32,
895 y_threads: u32,
896 ) !Launch {
897 return .{
898 .grid = .{ try runtimeGrid1D(x_extent, x_threads), try runtimeGrid1D(y_extent, y_threads), 1 },
899 .threadgroup = .{ x_threads, y_threads, 1 },
900 };
901 }
902
903 pub fn launch3D(
904 comptime x_extent: u64,
905 comptime y_extent: u64,
906 comptime z_extent: u64,
907 comptime x_threads: u32,
908 comptime y_threads: u32,
909 comptime z_threads: u32,
910 ) Launch {
911 return .{
912 .grid = .{ grid1D(x_extent, x_threads), grid1D(y_extent, y_threads), grid1D(z_extent, z_threads) },
913 .threadgroup = .{ x_threads, y_threads, z_threads },
914 };
915 }
916
917 pub fn runtimeLaunch3D(
918 x_extent: u64,
919 y_extent: u64,
920 z_extent: u64,
921 x_threads: u32,
922 y_threads: u32,
923 z_threads: u32,
924 ) !Launch {
925 return .{
926 .grid = .{ try runtimeGrid1D(x_extent, x_threads), try runtimeGrid1D(y_extent, y_threads), try runtimeGrid1D(z_extent, z_threads) },
927 .threadgroup = .{ x_threads, y_threads, z_threads },
928 };
929 }
930
931 pub fn threadBlocks1D(comptime axis_name: []const u8, comptime extent: u64, comptime threads: u32) Schedule {
932 if (extent == 0) @compileError("kernel library 1D schedule extent must be nonzero");
933 if (threads == 0) @compileError("kernel library 1D schedule threadgroup must be nonzero");
934 if (extent <= threads) {
935 return .{ .bindings = &.{.{
936 .axis = axis_name,
937 .target = .thread_x,
938 .extent = extent,
939 }} };
940 }
941 return .{ .bindings = &.{
942 .{
943 .axis = std.fmt.comptimePrint("{s}_tile", .{axis_name}),
944 .target = .block_x,
945 .extent = @intCast(grid1D(extent, threads)),
946 },
947 .{
948 .axis = std.fmt.comptimePrint("{s}_lane", .{axis_name}),
949 .target = .thread_x,
950 .extent = threads,
951 },
952 } };
953 }
954
955 pub fn runtimeThreadBlocks1D(lifetime_allocator: std.mem.Allocator, axis_name: []const u8, extent: u64, threads: u32) !Schedule {
956 if (extent == 0) return error.KernelLibraryScheduleExtentMustBeNonzero;
957 if (threads == 0) return error.KernelLibraryScheduleThreadgroupMustBeNonzero;
958 const bindings = try lifetime_allocator.alloc(ScheduleBinding, if (extent <= threads) 1 else 2);
959 if (extent <= threads) {
960 bindings[0] = try runtimeThreadBlockAxisBinding(lifetime_allocator, axis_name, .thread_x, extent);
961 } else {
962 bindings[0] = try runtimeThreadBlockTileBinding(lifetime_allocator, axis_name, extent, threads, .block_x);
963 bindings[1] = try runtimeThreadBlockLaneBinding(lifetime_allocator, axis_name, threads, .thread_x);
964 }
965 return .{ .bindings = bindings };
966 }
967
968 pub fn threadBlocks2D(
969 comptime x_axis: []const u8,
970 comptime x_extent: u64,
971 comptime y_axis: []const u8,
972 comptime y_extent: u64,
973 comptime x_threads: u32,
974 comptime y_threads: u32,
975 ) Schedule {
976 if (x_extent == 0 or y_extent == 0) @compileError("kernel library 2D schedule extents must be nonzero");
977 if (x_threads == 0 or y_threads == 0) @compileError("kernel library 2D schedule threadgroups must be nonzero");
978
979 const x_tiled = x_extent > x_threads;
980 const y_tiled = y_extent > y_threads;
981 if (x_tiled and y_tiled) {
982 return .{ .bindings = &.{
983 threadBlockTileBinding(x_axis, x_extent, x_threads, .block_x),
984 threadBlockLaneBinding(x_axis, x_threads, .thread_x),
985 threadBlockTileBinding(y_axis, y_extent, y_threads, .block_y),
986 threadBlockLaneBinding(y_axis, y_threads, .thread_y),
987 } };
988 }
989 if (x_tiled) {
990 return .{ .bindings = &.{
991 threadBlockTileBinding(x_axis, x_extent, x_threads, .block_x),
992 threadBlockLaneBinding(x_axis, x_threads, .thread_x),
993 .{ .axis = y_axis, .target = .thread_y, .extent = y_extent },
994 } };
995 }
996 if (y_tiled) {
997 return .{ .bindings = &.{
998 .{ .axis = x_axis, .target = .thread_x, .extent = x_extent },
999 threadBlockTileBinding(y_axis, y_extent, y_threads, .block_y),
1000 threadBlockLaneBinding(y_axis, y_threads, .thread_y),
1001 } };
1002 }
1003 return .{ .bindings = &.{
1004 .{ .axis = x_axis, .target = .thread_x, .extent = x_extent },
1005 .{ .axis = y_axis, .target = .thread_y, .extent = y_extent },
1006 } };
1007 }
1008
1009 pub fn runtimeThreadBlocks2D(
1010 lifetime_allocator: std.mem.Allocator,
1011 x_axis: []const u8,
1012 x_extent: u64,
1013 y_axis: []const u8,
1014 y_extent: u64,
1015 x_threads: u32,
1016 y_threads: u32,
1017 ) !Schedule {
1018 if (x_extent == 0 or y_extent == 0) return error.KernelLibraryScheduleExtentMustBeNonzero;
1019 if (x_threads == 0 or y_threads == 0) return error.KernelLibraryScheduleThreadgroupMustBeNonzero;
1020
1021 const x_tiled = x_extent > x_threads;
1022 const y_tiled = y_extent > y_threads;
1023 const count: usize = @as(usize, if (x_tiled) 2 else 1) + @as(usize, if (y_tiled) 2 else 1);
1024 const bindings = try lifetime_allocator.alloc(ScheduleBinding, count);
1025 var index: usize = 0;
1026 if (x_tiled) {
1027 bindings[index] = try runtimeThreadBlockTileBinding(lifetime_allocator, x_axis, x_extent, x_threads, .block_x);
1028 index += 1;
1029 bindings[index] = try runtimeThreadBlockLaneBinding(lifetime_allocator, x_axis, x_threads, .thread_x);
1030 index += 1;
1031 } else {
1032 bindings[index] = try runtimeThreadBlockAxisBinding(lifetime_allocator, x_axis, .thread_x, x_extent);
1033 index += 1;
1034 }
1035 if (y_tiled) {
1036 bindings[index] = try runtimeThreadBlockTileBinding(lifetime_allocator, y_axis, y_extent, y_threads, .block_y);
1037 index += 1;
1038 bindings[index] = try runtimeThreadBlockLaneBinding(lifetime_allocator, y_axis, y_threads, .thread_y);
1039 } else {
1040 bindings[index] = try runtimeThreadBlockAxisBinding(lifetime_allocator, y_axis, .thread_y, y_extent);
1041 }
1042 return .{ .bindings = bindings };
1043 }
1044
1045 pub fn threadBlocks3D(
1046 comptime x_axis: []const u8,
1047 comptime x_extent: u64,
1048 comptime y_axis: []const u8,
1049 comptime y_extent: u64,
1050 comptime z_axis: []const u8,
1051 comptime z_extent: u64,
1052 comptime x_threads: u32,
1053 comptime y_threads: u32,
1054 comptime z_threads: u32,
1055 ) Schedule {
1056 if (x_extent == 0 or y_extent == 0 or z_extent == 0) @compileError("kernel library 3D schedule extents must be nonzero");
1057 if (x_threads == 0 or y_threads == 0 or z_threads == 0) @compileError("kernel library 3D schedule threadgroups must be nonzero");
1058
1059 const x_tiled = x_extent > x_threads;
1060 const y_tiled = y_extent > y_threads;
1061 const z_tiled = z_extent > z_threads;
1062 if (x_tiled and y_tiled and z_tiled) {
1063 return .{ .bindings = &.{
1064 threadBlockTileBinding(x_axis, x_extent, x_threads, .block_x),
1065 threadBlockLaneBinding(x_axis, x_threads, .thread_x),
1066 threadBlockTileBinding(y_axis, y_extent, y_threads, .block_y),
1067 threadBlockLaneBinding(y_axis, y_threads, .thread_y),
1068 threadBlockTileBinding(z_axis, z_extent, z_threads, .block_z),
1069 threadBlockLaneBinding(z_axis, z_threads, .thread_z),
1070 } };
1071 }
1072 if (x_tiled and y_tiled) {
1073 return .{ .bindings = &.{
1074 threadBlockTileBinding(x_axis, x_extent, x_threads, .block_x),
1075 threadBlockLaneBinding(x_axis, x_threads, .thread_x),
1076 threadBlockTileBinding(y_axis, y_extent, y_threads, .block_y),
1077 threadBlockLaneBinding(y_axis, y_threads, .thread_y),
1078 .{ .axis = z_axis, .target = .thread_z, .extent = z_extent },
1079 } };
1080 }
1081 if (x_tiled and z_tiled) {
1082 return .{ .bindings = &.{
1083 threadBlockTileBinding(x_axis, x_extent, x_threads, .block_x),
1084 threadBlockLaneBinding(x_axis, x_threads, .thread_x),
1085 .{ .axis = y_axis, .target = .thread_y, .extent = y_extent },
1086 threadBlockTileBinding(z_axis, z_extent, z_threads, .block_z),
1087 threadBlockLaneBinding(z_axis, z_threads, .thread_z),
1088 } };
1089 }
1090 if (y_tiled and z_tiled) {
1091 return .{ .bindings = &.{
1092 .{ .axis = x_axis, .target = .thread_x, .extent = x_extent },
1093 threadBlockTileBinding(y_axis, y_extent, y_threads, .block_y),
1094 threadBlockLaneBinding(y_axis, y_threads, .thread_y),
1095 threadBlockTileBinding(z_axis, z_extent, z_threads, .block_z),
1096 threadBlockLaneBinding(z_axis, z_threads, .thread_z),
1097 } };
1098 }
1099 if (x_tiled) {
1100 return .{ .bindings = &.{
1101 threadBlockTileBinding(x_axis, x_extent, x_threads, .block_x),
1102 threadBlockLaneBinding(x_axis, x_threads, .thread_x),
1103 .{ .axis = y_axis, .target = .thread_y, .extent = y_extent },
1104 .{ .axis = z_axis, .target = .thread_z, .extent = z_extent },
1105 } };
1106 }
1107 if (y_tiled) {
1108 return .{ .bindings = &.{
1109 .{ .axis = x_axis, .target = .thread_x, .extent = x_extent },
1110 threadBlockTileBinding(y_axis, y_extent, y_threads, .block_y),
1111 threadBlockLaneBinding(y_axis, y_threads, .thread_y),
1112 .{ .axis = z_axis, .target = .thread_z, .extent = z_extent },
1113 } };
1114 }
1115 if (z_tiled) {
1116 return .{ .bindings = &.{
1117 .{ .axis = x_axis, .target = .thread_x, .extent = x_extent },
1118 .{ .axis = y_axis, .target = .thread_y, .extent = y_extent },
1119 threadBlockTileBinding(z_axis, z_extent, z_threads, .block_z),
1120 threadBlockLaneBinding(z_axis, z_threads, .thread_z),
1121 } };
1122 }
1123 return .{ .bindings = &.{
1124 .{ .axis = x_axis, .target = .thread_x, .extent = x_extent },
1125 .{ .axis = y_axis, .target = .thread_y, .extent = y_extent },
1126 .{ .axis = z_axis, .target = .thread_z, .extent = z_extent },
1127 } };
1128 }
1129
1130 pub fn runtimeThreadBlocks3D(
1131 lifetime_allocator: std.mem.Allocator,
1132 x_axis: []const u8,
1133 x_extent: u64,
1134 y_axis: []const u8,
1135 y_extent: u64,
1136 z_axis: []const u8,
1137 z_extent: u64,
1138 x_threads: u32,
1139 y_threads: u32,
1140 z_threads: u32,
1141 ) !Schedule {
1142 if (x_extent == 0 or y_extent == 0 or z_extent == 0) return error.KernelLibraryScheduleExtentMustBeNonzero;
1143 if (x_threads == 0 or y_threads == 0 or z_threads == 0) return error.KernelLibraryScheduleThreadgroupMustBeNonzero;
1144
1145 const x_tiled = x_extent > x_threads;
1146 const y_tiled = y_extent > y_threads;
1147 const z_tiled = z_extent > z_threads;
1148 const count: usize = @as(usize, if (x_tiled) 2 else 1) + @as(usize, if (y_tiled) 2 else 1) + @as(usize, if (z_tiled) 2 else 1);
1149 const bindings = try lifetime_allocator.alloc(ScheduleBinding, count);
1150 var index: usize = 0;
1151 if (x_tiled) {
1152 bindings[index] = try runtimeThreadBlockTileBinding(lifetime_allocator, x_axis, x_extent, x_threads, .block_x);
1153 index += 1;
1154 bindings[index] = try runtimeThreadBlockLaneBinding(lifetime_allocator, x_axis, x_threads, .thread_x);
1155 index += 1;
1156 } else {
1157 bindings[index] = try runtimeThreadBlockAxisBinding(lifetime_allocator, x_axis, .thread_x, x_extent);
1158 index += 1;
1159 }
1160 if (y_tiled) {
1161 bindings[index] = try runtimeThreadBlockTileBinding(lifetime_allocator, y_axis, y_extent, y_threads, .block_y);
1162 index += 1;
1163 bindings[index] = try runtimeThreadBlockLaneBinding(lifetime_allocator, y_axis, y_threads, .thread_y);
1164 index += 1;
1165 } else {
1166 bindings[index] = try runtimeThreadBlockAxisBinding(lifetime_allocator, y_axis, .thread_y, y_extent);
1167 index += 1;
1168 }
1169 if (z_tiled) {
1170 bindings[index] = try runtimeThreadBlockTileBinding(lifetime_allocator, z_axis, z_extent, z_threads, .block_z);
1171 index += 1;
1172 bindings[index] = try runtimeThreadBlockLaneBinding(lifetime_allocator, z_axis, z_threads, .thread_z);
1173 } else {
1174 bindings[index] = try runtimeThreadBlockAxisBinding(lifetime_allocator, z_axis, .thread_z, z_extent);
1175 }
1176 return .{ .bindings = bindings };
1177 }
1178
1179 fn threadBlockTileBinding(
1180 comptime axis_name: []const u8,
1181 comptime extent: u64,
1182 comptime threads: u32,
1183 comptime target: kernel.BindTarget,
1184 ) ScheduleBinding {
1185 return .{
1186 .axis = std.fmt.comptimePrint("{s}_tile", .{axis_name}),
1187 .target = target,
1188 .extent = @intCast(grid1D(extent, threads)),
1189 };
1190 }
1191
1192 fn runtimeThreadBlockAxisBinding(
1193 lifetime_allocator: std.mem.Allocator,
1194 axis_name: []const u8,
1195 target: kernel.BindTarget,
1196 extent: u64,
1197 ) !ScheduleBinding {
1198 if (axis_name.len == 0) return error.KernelLibraryAxisNameMustBeNonempty;
1199 return .{
1200 .axis = try lifetime_allocator.dupe(u8, axis_name),
1201 .target = target,
1202 .extent = extent,
1203 };
1204 }
1205
1206 fn runtimeThreadBlockTileBinding(
1207 lifetime_allocator: std.mem.Allocator,
1208 axis_name: []const u8,
1209 extent: u64,
1210 threads: u32,
1211 target: kernel.BindTarget,
1212 ) !ScheduleBinding {
1213 if (axis_name.len == 0) return error.KernelLibraryAxisNameMustBeNonempty;
1214 return .{
1215 .axis = try std.fmt.allocPrint(lifetime_allocator, "{s}_tile", .{axis_name}),
1216 .target = target,
1217 .extent = try runtimeGrid1D(extent, threads),
1218 };
1219 }
1220
1221 fn threadBlockLaneBinding(
1222 comptime axis_name: []const u8,
1223 comptime threads: u32,
1224 comptime target: kernel.BindTarget,
1225 ) ScheduleBinding {
1226 return .{
1227 .axis = std.fmt.comptimePrint("{s}_lane", .{axis_name}),
1228 .target = target,
1229 .extent = threads,
1230 };
1231 }
1232
1233 fn runtimeThreadBlockLaneBinding(
1234 lifetime_allocator: std.mem.Allocator,
1235 axis_name: []const u8,
1236 threads: u32,
1237 target: kernel.BindTarget,
1238 ) !ScheduleBinding {
1239 if (axis_name.len == 0) return error.KernelLibraryAxisNameMustBeNonempty;
1240 return .{
1241 .axis = try std.fmt.allocPrint(lifetime_allocator, "{s}_lane", .{axis_name}),
1242 .target = target,
1243 .extent = threads,
1244 };
1245 }
1246
1247 pub const ArtifactOptions = struct {
1248 limits: kernel.Limits,
1249 format: ?gpu.ArtifactFormat = null,
1250 kernel_plan: kernel.PlanOptions = .{},
1251 element_count_argument: artifact.ElementCountArgument = .none,
1252 shape_family_fingerprint: ?u64 = null,
1253 shape_profile: ?artifact.KernelCallShapeProfile = null,
1254 launch: ?artifact.KernelCallLaunch = null,
1255 runtime_scalar_argument_count: u32 = 0,
1256 static_arguments: []const choir_abi.ScalarArgument = &.{},
1257 };
1258
1259 pub fn Entry(comptime ProgramType: type, comptime entry_metadata: Metadata) type {
1260 return struct {
1261 pub const Program: type = ProgramType;
1262 pub const metadata = entry_metadata;
1263 pub const name = ProgramType.name;
1264 pub const target = entry_metadata.target;
1265 pub const version = entry_metadata.version;
1266 pub const layer = entry_metadata.layer;
1267 pub const category = entry_metadata.category;
1268 pub const specialization = entry_metadata.specialization;
1269 pub const Limits: type = ProgramType.Limits;
1270 pub const arg = ProgramType.arg;
1271 pub const schema = ProgramType.schema;
1272 pub const build = ProgramType.build;
1273 pub const interpret = ProgramType.interpret;
1274 pub const launch = ProgramType.launch;
1275 pub const scheduleSnapshot = ProgramType.scheduleSnapshot;
1276 pub const createPlan = ProgramType.createPlan;
1277 pub const createCheckedPlan = ProgramType.createCheckedPlan;
1278 pub const compileFragment = ProgramType.compileFragment;
1279 pub const createKernelArtifact = ProgramType.createKernelArtifact;
1280 pub const runCpu = ProgramType.runCpu;
1281 pub const runCpuWithDiagnostic = ProgramType.runCpuWithDiagnostic;
1282 pub const verify = ProgramType.verify;
1283
1284 pub fn createKernelCallArtifact(
1285 allocator: std.mem.Allocator,
1286 handle: kernel.BackendHandle,
1287 options: ArtifactOptions,
1288 ) !kernel.OwnedKernelCallArtifact {
1289 return ProgramType.createKernelCallArtifact(allocator, options.limits, handle, .{
1290 .target = entry_metadata.target,
1291 .version = entry_metadata.version,
1292 .format = options.format,
1293 .kernel_plan = options.kernel_plan,
1294 .element_count_argument = options.element_count_argument,
1295 .shape_family_fingerprint = options.shape_family_fingerprint orelse entry_metadata.specialization.shapeFamilyFingerprint(),
1296 .shape_profile = options.shape_profile,
1297 .launch = options.launch,
1298 .runtime_scalar_argument_count = options.runtime_scalar_argument_count,
1299 .static_arguments = options.static_arguments,
1300 });
1301 }
1302 };
1303 }
1304
1305 fn entry_test_copy_body_each(inner: anytype, index: kernel.Index1D, each_args: anytype) !void {
1306 const value = try each_args.param(.src).load(inner, index);
1307 try each_args.param(.dst).store(inner, value, index);
1308 }
1309
1310 fn entryTestCopyBody(k: anytype, args: anytype) !void {
1311 _ = try k.forEach1D("i", 2, args, entry_test_copy_body_each);
1312 }
1313
1314 test "kernel library entry carries metadata and creates plans" {
1315 const Program = kernel.logical.Program(.{
1316 .name = "kernel_library_entry_test_copy_i32",
1317 .parameters = .{
1318 .src = kernel.dynamicBuffer(.i32),
1319 .dst = kernel.dynamicBuffer(.i32),
1320 },
1321 .body = entryTestCopyBody,
1322 });
1323
1324 const Copy = Entry(Program, .{
1325 .target = "accy.kernel.test.copy_i32",
1326 .layer = .logical,
1327 .category = .elementwise,
1328 .specialization = .{
1329 .dtype = .i32,
1330 .inputs = &.{shape1D("i", 2)},
1331 .outputs = &.{shape1D("i", 2)},
1332 .launch = .{ .grid = .{ 1, 1, 1 }, .threadgroup = .{ 2, 1, 1 } },
1333 .schedule = threadBlocks1D("i", 2, 2),
1334 },
1335 });
1336
1337 try std.testing.expectEqualStrings("kernel_library_entry_test_copy_i32", Copy.name);
1338 try std.testing.expectEqualStrings("accy.kernel.test.copy_i32", Copy.target);
1339 try std.testing.expectEqual(@as(u32, 1), Copy.version);
1340 try std.testing.expectEqual(Layer.logical, Copy.layer);
1341 try std.testing.expectEqual(Category.elementwise, Copy.category);
1342 try std.testing.expectEqual(@as(usize, 1), Copy.specialization.outputs[0].rank());
1343 try std.testing.expectEqual(@as(u64, 2), Copy.specialization.outputs[0].elementCount().?);
1344 try std.testing.expectEqual(@as(u32, 2), Copy.specialization.launch.?.threadgroup[0]);
1345 try std.testing.expectEqual(@as(usize, 1), Copy.specialization.schedule.?.bindings.len);
1346 try std.testing.expectEqualStrings("i", Copy.specialization.schedule.?.bindings[0].axis);
1347 try std.testing.expectEqual(kernel.BindTarget.thread_x, Copy.specialization.schedule.?.bindings[0].target);
1348 try std.testing.expectEqual(@as(u64, 2), Copy.specialization.schedule.?.bindings[0].extent);
1349 try std.testing.expectEqualDeep(Copy.specialization.launch.?, Copy.specialization.schedule.?.launch());
1350
1351 var snapshot = try Copy.scheduleSnapshot(std.testing.allocator, Copy.Limits.testing);
1352 defer snapshot.deinit(std.testing.allocator);
1353 try std.testing.expect(Copy.specialization.schedule.?.matchesSnapshot(&snapshot));
1354
1355 var plan = try Copy.createCheckedPlan(std.testing.allocator, Copy.Limits.testing, .{});
1356 defer plan.deinit();
1357 try std.testing.expectEqual(@as(u32, 2), plan.argument_count);
1358 }
1359
1360 test "kernel library reductions carry operator shape and dependencies" {
1361 const first = reduction("row_max", .maximum, shape1D("col", 4));
1362 const second = dependentReduction("row_exp_sum", .sum_exp_shifted, shape1D("col", 4), &.{"row_max"});
1363 const reuse = reductionReuse("row_exp_sum", shape1D("row", 2));
1364
1365 try std.testing.expectEqualStrings("row_max", first.name);
1366 try std.testing.expectEqual(ReductionOperator.maximum, first.operator);
1367 try std.testing.expectEqual(@as(u64, 4), first.shape.elementCount().?);
1368 try std.testing.expectEqual(@as(usize, 0), first.dependencies.len);
1369
1370 try std.testing.expectEqualStrings("row_exp_sum", second.name);
1371 try std.testing.expectEqual(ReductionOperator.sum_exp_shifted, second.operator);
1372 try std.testing.expectEqual(@as(u64, 4), second.shape.elementCount().?);
1373 try std.testing.expectEqual(@as(usize, 1), second.dependencies.len);
1374 try std.testing.expectEqualStrings("row_max", second.dependencies[0]);
1375
1376 try std.testing.expectEqualStrings("row_exp_sum", reuse.reduction);
1377 try std.testing.expectEqual(@as(u64, 2), reuse.shape.elementCount().?);
1378 try std.testing.expect(reuse.matches(.{ .reduction = "row_exp_sum", .extents = &.{2} }));
1379 try std.testing.expect(!reuse.matches(.{ .reduction = "row_max", .extents = &.{2} }));
1380 }
1381
1382 test "kernel library runtime metadata constructors build owned specialization facts" {
1383 var owned = OwnedSpecialization.init(std.testing.allocator);
1384 defer owned.deinit();
1385 const lifetime_allocator = owned.allocator();
1386
1387 const inputs = try lifetime_allocator.alloc(Shape, 1);
1388 inputs[0] = try runtimeShape2D(lifetime_allocator, "row", 5, "col", 7);
1389 const outputs = try lifetime_allocator.alloc(Shape, 1);
1390 outputs[0] = try runtimeShape2D(lifetime_allocator, "row", 5, "col", 7);
1391 const reductions = try lifetime_allocator.alloc(Reduction, 2);
1392 reductions[0] = try runtimeReduction(lifetime_allocator, "row_sum", .sum, try runtimeShape1D(lifetime_allocator, "col", 7));
1393 reductions[1] = try runtimeDependentReduction(lifetime_allocator, "row_variance_sum", .sum_squared_difference, try runtimeShape1D(lifetime_allocator, "col", 7), &.{"row_sum"});
1394 const reduction_reuse = try lifetime_allocator.alloc(ReductionReuse, 2);
1395 reduction_reuse[0] = try runtimeReductionReuse(lifetime_allocator, "row_sum", try runtimeShape1D(lifetime_allocator, "row", 5));
1396 reduction_reuse[1] = try runtimeReductionReuse(lifetime_allocator, "row_variance_sum", try runtimeShape1D(lifetime_allocator, "row", 5));
1397
1398 owned.value = .{
1399 .dtype = .f32,
1400 .operation = .{ .row_normalization = .{ .layernorm = .none } },
1401 .inputs = inputs,
1402 .outputs = outputs,
1403 .reductions = reductions,
1404 .reduction_reuse = reduction_reuse,
1405 .launch = try runtimeLaunch2D(7, 5, 4, 2),
1406 .schedule = try runtimeThreadBlocks2D(lifetime_allocator, "col", 7, "row", 5, 4, 2),
1407 };
1408 const specialization = owned.value;
1409
1410 try std.testing.expect(specialization.operationIs(.{ .row_normalization = .{ .layernorm = .none } }));
1411 try std.testing.expectEqual(@as(u64, 35), specialization.inputs[0].elementCount().?);
1412 try std.testing.expectEqualStrings("row", specialization.inputs[0].axes[0].name);
1413 try std.testing.expectEqualStrings("col", specialization.inputs[0].axes[1].name);
1414 try std.testing.expect(specialization.reductionDependenciesAreValid());
1415 try std.testing.expect(specialization.reductionReuseScopesAreValid());
1416 try std.testing.expect(specialization.reductionMatches(1, .{
1417 .name = "row_variance_sum",
1418 .operator = .sum_squared_difference,
1419 .extents = &.{7},
1420 .dependencies = &.{"row_sum"},
1421 }));
1422 try std.testing.expect(specialization.reductionReuseMatches(1, .{
1423 .reduction = "row_variance_sum",
1424 .extents = &.{5},
1425 }));
1426 try std.testing.expectEqual(@as(u32, 2), specialization.launch.?.grid[0]);
1427 try std.testing.expectEqual(@as(u32, 3), specialization.launch.?.grid[1]);
1428 try std.testing.expect(specialization.scheduleMatchesLaunch());
1429 try std.testing.expectEqual(@as(usize, 4), specialization.schedule.?.bindings.len);
1430 try std.testing.expectEqualStrings("col_tile", specialization.schedule.?.bindings[0].axis);
1431 try std.testing.expectEqual(kernel.BindTarget.block_x, specialization.schedule.?.bindings[0].target);
1432 try std.testing.expectEqualStrings("row_lane", specialization.schedule.?.bindings[3].axis);
1433 try std.testing.expectEqual(kernel.BindTarget.thread_y, specialization.schedule.?.bindings[3].target);
1434
1435 try std.testing.expectError(error.KernelLibraryShapeExtentMustBeNonzero, runtimeShape1D(lifetime_allocator, "i", 0));
1436 try std.testing.expectError(error.KernelLibraryAxisNameMustBeNonempty, runtimeShape1D(lifetime_allocator, "", 1));
1437 try std.testing.expectError(error.KernelLibraryLaunchThreadgroupMustBeNonzero, runtimeLaunch1D(4, 0));
1438 try std.testing.expectError(error.KernelLibraryScheduleExtentMustBeNonzero, runtimeThreadBlocks1D(lifetime_allocator, "i", 0, 1));
1439 try std.testing.expectError(error.KernelLibraryReductionReuseReductionMustBeNonempty, runtimeReductionReuse(lifetime_allocator, "", try runtimeShape1D(lifetime_allocator, "row", 5)));
1440 }
1441
1442 test "kernel library reduction dependency graph validates ordered names" {
1443 const column = shape1D("col", 4);
1444 const valid = Specialization{
1445 .reductions = &.{
1446 reduction("score_dot", .dot_product, column),
1447 dependentReduction("score_max", .maximum, column, &.{"score_dot"}),
1448 dependentReduction("score_exp_sum", .sum_exp_shifted, column, &.{"score_max"}),
1449 dependentReduction("value_weighted_sum", .weighted_sum, column, &.{"score_exp_sum"}),
1450 },
1451 };
1452 try std.testing.expect(valid.reductionDependenciesAreValid());
1453 try std.testing.expect((Specialization{}).reductionDependenciesAreValid());
1454
1455 const duplicate_name = Specialization{
1456 .reductions = &.{
1457 reduction("row_max", .maximum, column),
1458 reduction("row_max", .sum, column),
1459 },
1460 };
1461 try std.testing.expect(!duplicate_name.reductionDependenciesAreValid());
1462
1463 const empty_name = Specialization{
1464 .reductions = &.{reduction("", .sum, column)},
1465 };
1466 try std.testing.expect(!empty_name.reductionDependenciesAreValid());
1467
1468 const unknown_dependency = Specialization{
1469 .reductions = &.{dependentReduction("score_max", .maximum, column, &.{"score_dot"})},
1470 };
1471 try std.testing.expect(!unknown_dependency.reductionDependenciesAreValid());
1472
1473 const forward_dependency = Specialization{
1474 .reductions = &.{
1475 dependentReduction("score_max", .maximum, column, &.{"score_dot"}),
1476 reduction("score_dot", .dot_product, column),
1477 },
1478 };
1479 try std.testing.expect(!forward_dependency.reductionDependenciesAreValid());
1480
1481 const empty_dependency = Specialization{
1482 .reductions = &.{
1483 reduction("score_dot", .dot_product, column),
1484 dependentReduction("score_max", .maximum, column, &.{""}),
1485 },
1486 };
1487 try std.testing.expect(!empty_dependency.reductionDependenciesAreValid());
1488
1489 const self_dependency = Specialization{
1490 .reductions = &.{dependentReduction("score_max", .maximum, column, &.{"score_max"})},
1491 };
1492 try std.testing.expect(!self_dependency.reductionDependenciesAreValid());
1493 }
1494
1495 test "kernel library reduction reuse scopes validate named reductions" {
1496 const column = shape1D("col", 4);
1497 const row = shape1D("row", 2);
1498 const valid = Specialization{
1499 .reductions = &.{
1500 reduction("row_max", .maximum, column),
1501 dependentReduction("row_exp_sum", .sum_exp_shifted, column, &.{"row_max"}),
1502 },
1503 .reduction_reuse = &.{
1504 reductionReuse("row_max", row),
1505 reductionReuse("row_exp_sum", row),
1506 },
1507 };
1508 try std.testing.expect(valid.reductionReuseScopesAreValid());
1509 try std.testing.expect((Specialization{}).reductionReuseScopesAreValid());
1510
1511 const unknown_reduction = Specialization{
1512 .reductions = &.{reduction("row_max", .maximum, column)},
1513 .reduction_reuse = &.{reductionReuse("row_exp_sum", row)},
1514 };
1515 try std.testing.expect(!unknown_reduction.reductionReuseScopesAreValid());
1516
1517 const duplicate_reuse = Specialization{
1518 .reductions = &.{reduction("row_max", .maximum, column)},
1519 .reduction_reuse = &.{
1520 reductionReuse("row_max", row),
1521 reductionReuse("row_max", row),
1522 },
1523 };
1524 try std.testing.expect(!duplicate_reuse.reductionReuseScopesAreValid());
1525
1526 const empty_reduction = Specialization{
1527 .reductions = &.{reduction("row_max", .maximum, column)},
1528 .reduction_reuse = &.{reductionReuse("", row)},
1529 };
1530 try std.testing.expect(!empty_reduction.reductionReuseScopesAreValid());
1531 }
1532
1533 test "kernel library specialization carries named static parameters" {
1534 const specialization = Specialization{
1535 .static_parameters = &.{
1536 staticParameter("slice_size", 4),
1537 staticParameter("rounds", 10),
1538 },
1539 };
1540
1541 try std.testing.expect(specialization.staticParametersAreValid());
1542 try std.testing.expect(specialization.staticParameterMatches("slice_size", 4));
1543 try std.testing.expectEqual(@as(?u64, 10), specialization.staticParameterValue("rounds"));
1544 try std.testing.expectEqual(@as(?u64, null), specialization.staticParameterValue("missing"));
1545
1546 const duplicate = Specialization{
1547 .static_parameters = &.{
1548 staticParameter("slice_size", 4),
1549 staticParameter("slice_size", 8),
1550 },
1551 };
1552 try std.testing.expect(!duplicate.staticParametersAreValid());
1553 try std.testing.expectError(
1554 error.KernelLibraryStaticParameterNameMustBeNonempty,
1555 runtimeStaticParameter(std.testing.allocator, "", 4),
1556 );
1557 }
1558
1559 test "kernel library scalar shape has one element and no axes" {
1560 const shape = shapeScalar();
1561
1562 try std.testing.expectEqual(@as(usize, 0), shape.rank());
1563 try std.testing.expectEqual(@as(u64, 1), shape.elementCount().?);
1564 try std.testing.expect(shape.matchesExtents(&.{}));
1565 }
1566
1567 test "kernel library 3D shape and launch metadata" {
1568 const shape = shape3D("batch", 2, "row", 3, "col", 4);
1569 const launch_value = launch3D(4, 3, 2, 2, 3, 1);
1570
1571 try std.testing.expectEqual(@as(usize, 3), shape.rank());
1572 try std.testing.expectEqual(@as(u64, 24), shape.elementCount().?);
1573 try std.testing.expect(shape.matchesExtents(&.{ 2, 3, 4 }));
1574 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
1575 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);
1576 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[2]);
1577 try std.testing.expectEqual(@as(u32, 2), launch_value.threadgroup[0]);
1578 try std.testing.expectEqual(@as(u32, 3), launch_value.threadgroup[1]);
1579 try std.testing.expectEqual(@as(u32, 1), launch_value.threadgroup[2]);
1580 }
1581
1582 test "kernel library 3D schedule records thread block bindings" {
1583 const tiled = comptime threadBlocks3D("n", 8, "m", 4, "b", 2, 4, 4, 1);
1584 const direct = comptime threadBlocks3D("n", 3, "m", 2, "b", 2, 3, 2, 2);
1585
1586 try std.testing.expectEqual(@as(usize, 5), tiled.bindings.len);
1587 try std.testing.expectEqualStrings("n_tile", tiled.bindings[0].axis);
1588 try std.testing.expectEqual(kernel.BindTarget.block_x, tiled.bindings[0].target);
1589 try std.testing.expectEqualStrings("n_lane", tiled.bindings[1].axis);
1590 try std.testing.expectEqual(kernel.BindTarget.thread_x, tiled.bindings[1].target);
1591 try std.testing.expectEqualStrings("m", tiled.bindings[2].axis);
1592 try std.testing.expectEqual(kernel.BindTarget.thread_y, tiled.bindings[2].target);
1593 try std.testing.expectEqualStrings("b_tile", tiled.bindings[3].axis);
1594 try std.testing.expectEqual(kernel.BindTarget.block_z, tiled.bindings[3].target);
1595 try std.testing.expectEqualStrings("b_lane", tiled.bindings[4].axis);
1596 try std.testing.expectEqual(kernel.BindTarget.thread_z, tiled.bindings[4].target);
1597
1598 try std.testing.expectEqual(@as(usize, 3), direct.bindings.len);
1599 try std.testing.expectEqualStrings("n", direct.bindings[0].axis);
1600 try std.testing.expectEqual(kernel.BindTarget.thread_x, direct.bindings[0].target);
1601 try std.testing.expectEqualStrings("m", direct.bindings[1].axis);
1602 try std.testing.expectEqual(kernel.BindTarget.thread_y, direct.bindings[1].target);
1603 try std.testing.expectEqualStrings("b", direct.bindings[2].axis);
1604 try std.testing.expectEqual(kernel.BindTarget.thread_z, direct.bindings[2].target);
1605 }
1606
1607 test "kernel library epilogues carry operator input shape" {
1608 const activation = epilogue(.{ .activation = .gelu });
1609 const relu_activation = epilogue(.{ .activation = .relu });
1610 const silu_activation = epilogue(.{ .activation = .silu });
1611 const bias = inputEpilogue(.bias_add, 2, shape1D("n", 3));
1612
1613 try std.testing.expect(activation.matches(.{ .operator = .{ .activation = .gelu } }));
1614 try std.testing.expect(!activation.matches(.{ .operator = .bias_add }));
1615 try std.testing.expect(!activation.matches(.{
1616 .operator = .{ .activation = .gelu },
1617 .input_index = 2,
1618 }));
1619 try std.testing.expect(relu_activation.matches(.{ .operator = .{ .activation = .relu } }));
1620 try std.testing.expect(!relu_activation.matches(.{ .operator = .{ .activation = .silu } }));
1621 try std.testing.expect(silu_activation.matches(.{ .operator = .{ .activation = .silu } }));
1622 try std.testing.expect(bias.matches(.{
1623 .operator = .bias_add,
1624 .input_index = 2,
1625 .extents = &.{3},
1626 }));
1627 try std.testing.expect(!bias.matches(.{
1628 .operator = .bias_add,
1629 .input_index = 1,
1630 .extents = &.{3},
1631 }));
1632 try std.testing.expect(!bias.matches(.{
1633 .operator = .bias_add,
1634 .input_index = 2,
1635 .extents = &.{4},
1636 }));
1637 }
1638
1639 test "kernel library input transforms carry operator input shape" {
1640 const gate = inputTransform(.{ .activation = .silu }, 0, shape1D("i", 8));
1641 const gelu_gate = inputTransform(.{ .activation = .gelu }, 0, shape1D("i", 8));
1642 const relu_gate = inputTransform(.{ .activation = .relu }, 0, shape1D("i", 8));
1643 const residual = inputTransform(.residual_add, 1, shape2D("row", 2, "col", 4));
1644
1645 try std.testing.expect(gate.matches(.{
1646 .operator = .{ .activation = .silu },
1647 .input_index = 0,
1648 .extents = &.{8},
1649 }));
1650 try std.testing.expect(!gate.matches(.{
1651 .operator = .{ .activation = .silu },
1652 .input_index = 1,
1653 .extents = &.{8},
1654 }));
1655 try std.testing.expect(!gate.matches(.{
1656 .operator = .{ .activation = .silu },
1657 .input_index = 0,
1658 .extents = &.{4},
1659 }));
1660 try std.testing.expect(gelu_gate.matches(.{
1661 .operator = .{ .activation = .gelu },
1662 .input_index = 0,
1663 .extents = &.{8},
1664 }));
1665 try std.testing.expect(!gelu_gate.matches(.{
1666 .operator = .{ .activation = .silu },
1667 .input_index = 0,
1668 .extents = &.{8},
1669 }));
1670 try std.testing.expect(relu_gate.matches(.{
1671 .operator = .{ .activation = .relu },
1672 .input_index = 0,
1673 .extents = &.{8},
1674 }));
1675 try std.testing.expect(!relu_gate.matches(.{
1676 .operator = .{ .activation = .gelu },
1677 .input_index = 0,
1678 .extents = &.{8},
1679 }));
1680 try std.testing.expect(residual.matches(.{
1681 .operator = .residual_add,
1682 .input_index = 1,
1683 .extents = &.{ 2, 4 },
1684 }));
1685 try std.testing.expect(!residual.matches(.{
1686 .operator = .residual_add,
1687 .input_index = 0,
1688 .extents = &.{ 2, 4 },
1689 }));
1690 try std.testing.expect(!residual.matches(.{
1691 .operator = .{ .activation = .relu },
1692 .input_index = 1,
1693 .extents = &.{ 2, 4 },
1694 }));
1695 }
1696
1697 test "kernel library metadata matches specialization contracts" {
1698 const shape = shape2D("row", 2, "col", 4);
1699 try std.testing.expect(shape.matchesExtents(&.{ 2, 4 }));
1700 try std.testing.expect(!shape.matchesExtents(&.{2}));
1701 try std.testing.expect(!shape.matchesExtents(&.{ 4, 2 }));
1702
1703 const first = reduction("row_max", .maximum, shape1D("col", 4));
1704 const second = dependentReduction("row_exp_sum", .sum_exp_shifted, shape1D("col", 4), &.{"row_max"});
1705 try std.testing.expect(first.matches(.{
1706 .name = "row_max",
1707 .operator = .maximum,
1708 .extents = &.{4},
1709 }));
1710 try std.testing.expect(second.matches(.{
1711 .name = "row_exp_sum",
1712 .operator = .sum_exp_shifted,
1713 .extents = &.{4},
1714 .dependencies = &.{"row_max"},
1715 }));
1716 try std.testing.expect(!second.matches(.{
1717 .name = "row_exp_sum",
1718 .operator = .sum_exp_shifted,
1719 .extents = &.{4},
1720 .dependencies = &.{},
1721 }));
1722
1723 const specialization = Specialization{
1724 .operation = .{ .row_normalization = .softmax },
1725 .inputs = &.{shape},
1726 .outputs = &.{shape},
1727 .reductions = &.{ first, second },
1728 .reduction_reuse = &.{
1729 reductionReuse("row_max", shape1D("row", 2)),
1730 reductionReuse("row_exp_sum", shape1D("row", 2)),
1731 },
1732 .input_transforms = &.{
1733 inputTransform(.{ .activation = .silu }, 0, shape),
1734 inputTransform(.residual_add, 1, shape),
1735 },
1736 .epilogues = &.{epilogue(.{ .activation = .gelu })},
1737 .launch = .{ .grid = .{ 1, 1, 1 }, .threadgroup = .{ 4, 2, 1 } },
1738 .schedule = threadBlocks2D("col", 4, "row", 2, 4, 2),
1739 };
1740 try std.testing.expect(specialization.operationIs(.{ .row_normalization = .softmax }));
1741 try std.testing.expect(!specialization.operationIs(.{ .row_normalization = .log_softmax }));
1742 try std.testing.expect(!specialization.operationIs(.{ .row_normalization = .{ .rmsnorm = .scale } }));
1743 try std.testing.expect(specialization.scheduleMatchesLaunch());
1744 try std.testing.expect(specialization.inputHasExtents(0, &.{ 2, 4 }));
1745 try std.testing.expect(specialization.outputHasExtents(0, &.{ 2, 4 }));
1746 try std.testing.expect(specialization.reductionMatches(1, .{
1747 .name = "row_exp_sum",
1748 .operator = .sum_exp_shifted,
1749 .extents = &.{4},
1750 .dependencies = &.{"row_max"},
1751 }));
1752 try std.testing.expect(specialization.reductionReuseMatches(1, .{
1753 .reduction = "row_exp_sum",
1754 .extents = &.{2},
1755 }));
1756 try std.testing.expect(specialization.inputTransformMatches(0, .{
1757 .operator = .{ .activation = .silu },
1758 .input_index = 0,
1759 .extents = &.{ 2, 4 },
1760 }));
1761 try std.testing.expect(specialization.inputTransformMatches(1, .{
1762 .operator = .residual_add,
1763 .input_index = 1,
1764 .extents = &.{ 2, 4 },
1765 }));
1766 try std.testing.expect(specialization.epilogueMatches(0, .{ .operator = .{ .activation = .gelu } }));
1767 try std.testing.expect(!specialization.inputHasExtents(1, &.{ 2, 4 }));
1768 try std.testing.expect(!specialization.reductionMatches(2, .{
1769 .name = "missing",
1770 .operator = .sum,
1771 .extents = &.{4},
1772 }));
1773 try std.testing.expect(!specialization.inputTransformMatches(2, .{
1774 .operator = .{ .activation = .silu },
1775 .input_index = 0,
1776 }));
1777 try std.testing.expect(!specialization.epilogueMatches(1, .{ .operator = .{ .activation = .gelu } }));
1778 try std.testing.expect(!((Specialization{}).scheduleMatchesLaunch()));
1779 }
1780
1781 test "kernel library specializations estimate output and reduction work" {
1782 const elementwise = Specialization{
1783 .outputs = &.{shape1D("i", 8)},
1784 };
1785 try std.testing.expectEqual(@as(u64, 8), elementwise.outputElementCount().?);
1786 try std.testing.expectEqual(@as(u64, 8), elementwise.estimatedElementOps().?);
1787
1788 const matmul = Specialization{
1789 .outputs = &.{shape2D("m", 4, "n", 5)},
1790 .reductions = &.{reduction("dot", .dot_product, shape1D("k", 6))},
1791 };
1792 try std.testing.expectEqual(@as(u64, 20), matmul.outputElementCount().?);
1793 try std.testing.expectEqual(@as(u64, 240), matmul.estimatedElementOps().?);
1794
1795 const row_softmax = Specialization{
1796 .outputs = &.{shape2D("row", 2, "col", 4)},
1797 .reductions = &.{
1798 reduction("row_max", .maximum, shape1D("col", 4)),
1799 dependentReduction("row_exp_sum", .sum_exp_shifted, shape1D("col", 4), &.{"row_max"}),
1800 },
1801 .reduction_reuse = &.{
1802 reductionReuse("row_max", shape1D("row", 2)),
1803 reductionReuse("row_exp_sum", shape1D("row", 2)),
1804 },
1805 };
1806 try std.testing.expectEqual(@as(u64, 8), row_softmax.outputElementCount().?);
1807 try std.testing.expectEqual(@as(u64, 16), row_softmax.estimatedElementOps().?);
1808
1809 const repeated_row_softmax = Specialization{
1810 .outputs = &.{shape2D("row", 2, "col", 4)},
1811 .reductions = &.{
1812 reduction("row_max", .maximum, shape1D("col", 4)),
1813 dependentReduction("row_exp_sum", .sum_exp_shifted, shape1D("col", 4), &.{"row_max"}),
1814 },
1815 };
1816 try std.testing.expectEqual(@as(u64, 64), repeated_row_softmax.estimatedElementOps().?);
1817
1818 try std.testing.expect((Specialization{}).outputElementCount() == null);
1819 try std.testing.expect((Specialization{}).estimatedElementOps() == null);
1820 }
1821
1822 test "operation fingerprints discriminate semantic operations" {
1823 const inclusive = operationFingerprint(.{ .scan = .prefix_sum });
1824 const exclusive = operationFingerprint(.{ .scan = .prefix_sum_exclusive });
1825 try std.testing.expect(inclusive != exclusive);
1826 try std.testing.expectEqual(inclusive, operationFingerprint(.{ .scan = .prefix_sum }));
1827
1828 const philox_ten = operationFingerprint(.{ .random = .{ .philox = 10 } });
1829 const philox_seven = operationFingerprint(.{ .random = .{ .philox = 7 } });
1830 const threefry_ten = operationFingerprint(.{ .random = .{ .threefry = 10 } });
1831 try std.testing.expect(philox_ten != philox_seven);
1832 try std.testing.expect(philox_ten != threefry_ten);
1833
1834 const nonzero = operationFingerprint(.{ .compaction = .{ .blocks = .nonzero } });
1835 const greater = operationFingerprint(.{ .compaction = .{ .blocks = .greater_than } });
1836 try std.testing.expect(nonzero != greater);
1837
1838 const layernorm = operationFingerprint(.{ .row_normalization = .{ .layernorm = .none } });
1839 const rmsnorm = operationFingerprint(.{ .row_normalization = .{ .rmsnorm = .none } });
1840 try std.testing.expect(layernorm != rmsnorm);
1841 try std.testing.expect(operationFingerprint(.{ .linalg = .matrix_product }) !=
1842 operationFingerprint(.{ .indexing = .gather }));
1843 }