lib/accy/src/kernel/dsl/test.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const kernel = @import("../root.zig");
3 const namespace = @import("root.zig");
4 const body = namespace.body;
5 const Program = namespace.Program;
6 const Family = namespace.Family;
7
8 fn scale_each(inner: anytype, index: kernel.Index1D, args: anytype) !void {
9 const value = try args.param(.src).load(inner, index);
10 const scaled = try value.mul(inner, args.param(.scale));
11 try args.param(.dst).store(inner, scaled, index);
12 }
13
14 fn scaleFamilyBody(k: anytype, extent: u64, args: anytype) !void {
15 _ = try k.forEach1D("i", extent, 4, args, scale_each);
16 }
17
18 const ScaleFamily = Family(.{
19 .name = "kernel_dsl_scale_family_f32",
20 .parameters = .{
21 .src = kernel.dynamicBuffer(.f32),
22 .dst = kernel.dynamicBuffer(.f32),
23 .scale = kernel.scalar(.f32),
24 },
25 .Instance = u64,
26 .body = scaleFamilyBody,
27 });
28
29 test "kernel Family builds checked graphs from runtime instances" {
30 try std.testing.expectEqual(@as(usize, 0), ScaleFamily.arg(.src));
31 try std.testing.expectEqual(@as(usize, 2), ScaleFamily.arg(.scale));
32 try std.testing.expectEqual(@as(usize, 3), ScaleFamily.schema().len);
33
34 var graph = try ScaleFamily.build(std.testing.allocator, ScaleFamily.Limits.testing, 5);
35 defer graph.deinit();
36 try graph.verify();
37
38 const narrow_launch = try graph.launch();
39 try std.testing.expectEqual(@as(u32, 2), narrow_launch.grid[0]);
40 try std.testing.expectEqual(@as(u32, 4), narrow_launch.block[0]);
41
42 const wide_launch = try ScaleFamily.launch(std.testing.allocator, ScaleFamily.Limits.testing, 9);
43 try std.testing.expectEqual(@as(u32, 3), wide_launch.grid[0]);
44 try std.testing.expectEqual(@as(u32, 4), wide_launch.block[0]);
45
46 var plan = try ScaleFamily.createCheckedPlan(std.testing.allocator, ScaleFamily.Limits.testing, 9, .{});
47 defer plan.deinit();
48 try std.testing.expectEqual(@as(u32, 3), plan.argument_count);
49 try std.testing.expectEqualStrings("kernel_dsl_scale_family_f32", plan.entry_name);
50
51 var input = [_]f32{ 1.0, -2.0, 3.5, 4.0, -0.25 };
52 var output = [_]f32{ 0, 0, 0, 0, 0 };
53 try ScaleFamily.runCpu(std.testing.allocator, ScaleFamily.Limits.testing, 5, &.{
54 kernel.argumentBuffer(f32, input[0..]),
55 kernel.argumentBuffer(f32, output[0..]),
56 kernel.argumentF32(2.0),
57 });
58 try std.testing.expectEqualSlices(f32, &.{ 2.0, -4.0, 7.0, 8.0, -0.5 }, output[0..]);
59
60 var wide_input = [_]f32{ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 };
61 var wide_output = [_]f32{ 0, 0, 0, 0, 0, 0, 0, 0, 0 };
62 try ScaleFamily.runCpu(std.testing.allocator, ScaleFamily.Limits.testing, 9, &.{
63 kernel.argumentBuffer(f32, wide_input[0..]),
64 kernel.argumentBuffer(f32, wide_output[0..]),
65 kernel.argumentF32(-1.0),
66 });
67 try std.testing.expectEqualSlices(f32, &.{ -1.0, -2.0, -3.0, -4.0, -5.0, -6.0, -7.0, -8.0, -9.0 }, wide_output[0..]);
68 }
69
70 fn scaleBody(k: anytype, args: anytype) !void {
71 _ = try k.forEach1D("i", 5, 4, args, scale_each);
72 }
73
74 const Scale = Program(.{
75 .name = "kernel_dsl_scale_f32",
76 .parameters = .{
77 .src = kernel.dynamicBuffer(.f32),
78 .dst = kernel.dynamicBuffer(.f32),
79 .scale = kernel.scalar(.f32),
80 },
81 .body = scaleBody,
82 });
83
84 test "kernel Program builds named typed arguments into a checked graph" {
85 try std.testing.expectEqual(@as(usize, 0), Scale.arg(.src));
86 try std.testing.expectEqual(@as(usize, 1), Scale.arg(.dst));
87 try std.testing.expectEqual(@as(usize, 2), Scale.arg(.scale));
88 try std.testing.expectEqual(@as(usize, 3), Scale.schema().len);
89
90 var graph = try Scale.build(std.testing.allocator, Scale.Limits.testing);
91 defer graph.deinit();
92 try graph.verify();
93
94 var interpreted_graph = try Scale.interpret(std.testing.allocator, Scale.Limits.testing, kernel.interpret.graph());
95 defer interpreted_graph.deinit();
96 try interpreted_graph.verify();
97
98 const launch_value = try graph.launch();
99 const interpreted_launch = try interpreted_graph.launch();
100 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
101 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
102 try std.testing.expectEqual(launch_value.grid, interpreted_launch.grid);
103 try std.testing.expectEqual(launch_value.block, interpreted_launch.block);
104
105 var out = try Scale.createCheckedPlan(std.testing.allocator, Scale.Limits.testing, .{});
106 defer out.deinit();
107 try std.testing.expectEqual(@as(u32, 3), out.argument_count);
108 try std.testing.expectEqualStrings("kernel_dsl_scale_f32", out.entry_name);
109
110 var input = [_]f32{ 1.0, -2.0, 3.5, 4.0, -0.25 };
111 var output = [_]f32{ 0, 0, 0, 0, 0 };
112 try Scale.runCpu(std.testing.allocator, Scale.Limits.testing, &.{
113 kernel.argumentBuffer(f32, input[0..]),
114 kernel.argumentBuffer(f32, output[0..]),
115 kernel.argumentF32(2.0),
116 });
117 try std.testing.expectEqualSlices(f32, &.{ 2.0, -4.0, 7.0, 8.0, -0.5 }, output[0..]);
118 }
119
120 fn interpreted_scale_each(inner: anytype, index: kernel.Index1D, args: anytype) !void {
121 const value = try args.param(.src).load(inner, index);
122 const scaled = try value.mul(inner, args.param(.scale));
123 const shifted = try scaled.add(inner, 1.0);
124 try args.param(.dst).store(inner, shifted, index);
125 }
126
127 fn interpretedScaleBody(k: anytype, args: anytype) !void {
128 _ = try k.forEach1D("i", 5, 4, args, interpreted_scale_each);
129 }
130
131 const InterpretedScale = Program(.{
132 .name = "kernel_dsl_interpreted_scale_f32",
133 .parameters = .{
134 .src = kernel.dynamicBuffer(.f32),
135 .dst = kernel.dynamicBuffer(.f32),
136 .scale = kernel.scalar(.f32),
137 },
138 .body = interpretedScaleBody,
139 });
140
141 const DropAdd = struct {
142 pub fn add(_: *@This(), ctx: anytype) !kernel.Value {
143 return ctx.lhs;
144 }
145 };
146
147 const DropMul = struct {
148 pub fn mul(_: *@This(), ctx: anytype) !kernel.Value {
149 return ctx.lhs;
150 }
151 };
152
153 const MulAsAdd = struct {
154 pub fn mul(_: *@This(), ctx: anytype) !kernel.Value {
155 return ctx.next.add(ctx.lhs, ctx.rhs);
156 }
157 };
158
159 const TransformedScale = InterpretedScale.transform(DropAdd{});
160 const TwiceTransformedScale = TransformedScale.transform(DropMul{});
161 const AdditiveScale = InterpretedScale.transform(kernel.interpret.bindIndexed(MulAsAdd{}));
162
163 fn add_before_mul_each(inner: anytype, index: kernel.Index1D, args: anytype) !void {
164 const value = try args.param(.src).load(inner, index);
165 const shifted = try value.add(inner, 1.0);
166 const scaled = try shifted.mul(inner, args.param(.scale));
167 try args.param(.dst).store(inner, scaled, index);
168 }
169
170 fn addBeforeMulBody(k: anytype, args: anytype) !void {
171 _ = try k.forEach1D("i", 5, 4, args, add_before_mul_each);
172 }
173
174 const AddBeforeMul = Program(.{
175 .name = "kernel_dsl_add_before_mul_f32",
176 .parameters = .{
177 .src = kernel.dynamicBuffer(.f32),
178 .dst = kernel.dynamicBuffer(.f32),
179 .scale = kernel.scalar(.f32),
180 },
181 .body = addBeforeMulBody,
182 });
183
184 fn add_only_each(inner: anytype, index: kernel.Index1D, args: anytype) !void {
185 const value = try args.param(.src).load(inner, index);
186 const shifted = try value.add(inner, args.param(.scale));
187 try args.param(.dst).store(inner, shifted, index);
188 }
189
190 fn addOnlyBody(k: anytype, args: anytype) !void {
191 _ = try k.forEach1D("i", 5, 4, args, add_only_each);
192 }
193
194 const AddOnly = Program(.{
195 .name = "kernel_dsl_add_only_f32",
196 .parameters = .{
197 .src = kernel.dynamicBuffer(.f32),
198 .dst = kernel.dynamicBuffer(.f32),
199 .scale = kernel.scalar(.f32),
200 },
201 .body = addOnlyBody,
202 });
203
204 const MulPresence = struct {
205 has_mul: bool = false,
206
207 pub const Result = struct {
208 graph: kernel.Graph,
209 has_mul: bool,
210
211 pub fn deinit(self: *@This()) void {
212 self.graph.deinit();
213 self.* = undefined;
214 }
215 };
216
217 pub fn mul(self: *@This(), ctx: anytype) !kernel.Value {
218 self.has_mul = true;
219 return ctx.default();
220 }
221
222 pub fn finish(self: *@This(), ctx: anytype) !Result {
223 return .{
224 .graph = try ctx.default(),
225 .has_mul = self.has_mul,
226 };
227 }
228 };
229
230 const DropAddWhenMulPresent = struct {
231 enabled: bool,
232
233 pub fn add(self: *@This(), ctx: anytype) !kernel.Value {
234 if (self.enabled) return ctx.lhs;
235 return ctx.default();
236 }
237 };
238
239 fn dropAddFromMulPresence(analysis: *const MulPresence.Result) DropAddWhenMulPresent {
240 return .{ .enabled = analysis.has_mul };
241 }
242
243 const AddPresence = struct {
244 has_add: bool = false,
245
246 pub const Result = struct {
247 graph: kernel.Graph,
248 has_add: bool,
249
250 pub fn deinit(self: *@This()) void {
251 self.graph.deinit();
252 self.* = undefined;
253 }
254 };
255
256 pub fn add(self: *@This(), ctx: anytype) !kernel.Value {
257 self.has_add = true;
258 return ctx.default();
259 }
260
261 pub fn finish(self: *@This(), ctx: anytype) !Result {
262 return .{
263 .graph = try ctx.default(),
264 .has_add = self.has_add,
265 };
266 }
267 };
268
269 const DropMulWhenAddMissing = struct {
270 enabled: bool,
271
272 pub fn mul(self: *@This(), ctx: anytype) !kernel.Value {
273 if (self.enabled) return ctx.lhs;
274 return ctx.default();
275 }
276 };
277
278 fn dropMulFromAddPresence(analysis: *const AddPresence.Result) DropMulWhenAddMissing {
279 return .{ .enabled = !analysis.has_add };
280 }
281
282 const BorrowedAddFlags = struct {
283 add_count: usize = 0,
284
285 pub const Result = struct {
286 graph: kernel.Graph,
287 allocator: std.mem.Allocator,
288 flags: []const bool,
289
290 pub fn deinit(self: *@This()) void {
291 self.graph.deinit();
292 self.allocator.free(self.flags);
293 self.* = undefined;
294 }
295 };
296
297 pub fn add(self: *@This(), ctx: anytype) !kernel.Value {
298 self.add_count += 1;
299 return ctx.default();
300 }
301
302 pub fn finish(self: *@This(), ctx: anytype) !Result {
303 const flags = try std.testing.allocator.alloc(bool, 1);
304 errdefer std.testing.allocator.free(flags);
305 flags[0] = self.add_count != 0;
306 return .{
307 .graph = try ctx.default(),
308 .allocator = std.testing.allocator,
309 .flags = flags,
310 };
311 }
312 };
313
314 const DropMulFromBorrowedFlags = struct {
315 flags: []const bool,
316
317 pub fn mul(self: *@This(), ctx: anytype) !kernel.Value {
318 if (self.flags[0]) return ctx.lhs;
319 return ctx.default();
320 }
321 };
322
323 fn borrowDropMulFlags(analysis: *const BorrowedAddFlags.Result) DropMulFromBorrowedFlags {
324 return .{ .flags = analysis.flags };
325 }
326
327 const AddBeforeMulAnalysis = AddBeforeMul.analyze(MulPresence{});
328 const AddOnlyAnalysis = AddOnly.analyze(MulPresence{});
329 const DropAddAfterMulAnalysis = AddBeforeMulAnalysis.transform(dropAddFromMulPresence);
330 const KeepAddAfterMulAnalysis = AddOnlyAnalysis.transform(dropAddFromMulPresence);
331 const DropMulAfterAnalyzedDropAdd = DropAddAfterMulAnalysis.transform(DropMul{});
332 const AnalyzeAfterDroppedMul = AddBeforeMul.transform(DropMul{}).analyze(MulPresence{}).transform(dropAddFromMulPresence);
333 const DropMulAfterAnalyzedAddRemoval = DropAddAfterMulAnalysis.analyze(AddPresence{}).transform(dropMulFromAddPresence);
334 const DropMulWithBorrowedAnalysisFlags = AddBeforeMul.analyze(BorrowedAddFlags{}).transform(borrowDropMulFlags);
335
336 const RequireStore = struct {
337 store_count: usize = 0,
338
339 pub const Result: type = kernel.Graph;
340
341 pub fn store(self: *@This(), ctx: anytype) !void {
342 self.store_count += 1;
343 try ctx.default();
344 }
345
346 pub fn finish(self: *@This(), ctx: anytype) !Result {
347 if (self.store_count == 0) return error.KernelProgramRequiresStore;
348 return ctx.default();
349 }
350 };
351
352 const StoreCheckedScale = InterpretedScale.transform(RequireStore{});
353
354 fn no_store_each(inner: anytype, index: kernel.Index1D, args: anytype) !void {
355 const value = try args.param(.src).load(inner, index);
356 _ = try value.add(inner, args.param(.scale));
357 }
358
359 fn noStoreBody(k: anytype, args: anytype) !void {
360 _ = try k.forEach1D("i", 5, 4, args, no_store_each);
361 }
362
363 const NoStore = Program(.{
364 .name = "kernel_dsl_no_store_f32",
365 .parameters = .{
366 .src = kernel.dynamicBuffer(.f32),
367 .scale = kernel.scalar(.f32),
368 },
369 .body = noStoreBody,
370 });
371
372 const StoreCheckedNoStore = NoStore.transform(RequireStore{});
373
374 const ArithmeticCounter = struct {
375 mul_count: usize = 0,
376 add_count: usize = 0,
377 store_count: usize = 0,
378
379 pub const Result = struct {
380 graph: kernel.Graph,
381 mul_count: usize,
382 add_count: usize,
383 store_count: usize,
384
385 pub fn deinit(self: *@This()) void {
386 self.graph.deinit();
387 self.* = undefined;
388 }
389 };
390
391 pub fn mul(self: *@This(), ctx: anytype) !kernel.Value {
392 self.mul_count += 1;
393 return ctx.default();
394 }
395
396 pub fn add(self: *@This(), ctx: anytype) !kernel.Value {
397 self.add_count += 1;
398 return ctx.default();
399 }
400
401 pub fn store(self: *@This(), ctx: anytype) !void {
402 self.store_count += 1;
403 try ctx.default();
404 }
405
406 pub fn finish(self: *@This(), ctx: anytype) !Result {
407 return .{
408 .graph = try ctx.default(),
409 .mul_count = self.mul_count,
410 .add_count = self.add_count,
411 .store_count = self.store_count,
412 };
413 }
414 };
415
416 const WhileCountdownCarry = struct {
417 remaining: kernel.TypedValue(.i32),
418 total: kernel.TypedValue(.i32),
419 };
420
421 fn while_countdown_continue(
422 inner: anytype,
423 carry: WhileCountdownCarry,
424 _: void,
425 ) !kernel.TypedValue(.i1) {
426 return carry.remaining.compare(inner, .gt, 0);
427 }
428
429 fn while_countdown_step(
430 inner: anytype,
431 carry: WhileCountdownCarry,
432 _: void,
433 ) !WhileCountdownCarry {
434 return .{
435 .remaining = try carry.remaining.sub(inner, 1),
436 .total = try carry.total.add(inner, carry.remaining),
437 };
438 }
439
440 fn while_countdown_each(inner: anytype, index: kernel.Index1D, args: anytype) !void {
441 const start = try args.param(.src).load(inner, index);
442 const zero = try inner.constantValue(.i32, 0);
443 const walked = try inner.whileLoop(
444 WhileCountdownCarry{ .remaining = start, .total = zero },
445 {},
446 while_countdown_continue,
447 while_countdown_step,
448 );
449 try args.param(.dst).store(inner, walked.total, index);
450 }
451
452 fn whileCountdownBody(k: anytype, args: anytype) !void {
453 _ = try k.forEach1D("i", 4, 4, args, while_countdown_each);
454 }
455
456 const WhileCountdown = Program(.{
457 .name = "kernel_dsl_while_countdown_i32",
458 .parameters = .{
459 .src = kernel.dynamicBuffer(.i32),
460 .dst = kernel.dynamicBuffer(.i32),
461 },
462 .body = whileCountdownBody,
463 });
464
465 const WhileLoopCounter = struct {
466 while_loop_count: usize = 0,
467
468 pub const Result = struct {
469 graph: kernel.Graph,
470 while_loop_count: usize,
471
472 pub fn deinit(self: *@This()) void {
473 self.graph.deinit();
474 self.* = undefined;
475 }
476 };
477
478 pub fn whileLoop(self: *@This(), ctx: anytype) !@TypeOf(ctx.initial) {
479 self.while_loop_count += 1;
480 return ctx.default();
481 }
482
483 pub fn finish(self: *@This(), ctx: anytype) !Result {
484 return .{
485 .graph = try ctx.default(),
486 .while_loop_count = self.while_loop_count,
487 };
488 }
489 };
490
491 fn tiled_launch_each(inner: anytype, index: kernel.Index2D, args: anytype) !void {
492 const global = try inner.globalIdValue(.i32, .x);
493 const thread = try inner.threadIdValue(.i32, .x);
494 const block = try inner.blockIdValue(.i32, .y);
495 const block_dim = try inner.blockDimValue(.i32, .x);
496 const grid_dim = try inner.gridDimValue(.i32, .y);
497 const global_thread = try global.add(inner, thread);
498 const with_block = try global_thread.add(inner, block);
499 const with_block_dim = try with_block.add(inner, block_dim);
500 const value = try with_block_dim.add(inner, grid_dim);
501 try args.param(.dst).store(inner, value, index);
502 }
503
504 fn tiledLaunchBody(k: anytype, args: anytype) !void {
505 _ = try k.forEach2D(.{
506 .x = kernel.domainAxis("x", 5, 2),
507 .y = kernel.domainAxis("y", 3, 2),
508 }, args, tiled_launch_each);
509 }
510
511 const TiledLaunch = Program(.{
512 .name = "kernel_dsl_tiled_launch_i32",
513 .parameters = .{
514 .dst = kernel.dynamicBuffer(.i32),
515 },
516 .body = tiledLaunchBody,
517 });
518
519 const DomainCounter = struct {
520 index2d_count: usize = 0,
521 global_id_count: usize = 0,
522 thread_id_count: usize = 0,
523 block_id_count: usize = 0,
524 block_dim_count: usize = 0,
525 grid_dim_count: usize = 0,
526 compare_count: usize = 0,
527 constant_index_count: usize = 0,
528 mul_count: usize = 0,
529 add_count: usize = 0,
530 store_count: usize = 0,
531
532 pub const Result = struct {
533 graph: kernel.Graph,
534 index2d_count: usize,
535 global_id_count: usize,
536 thread_id_count: usize,
537 block_id_count: usize,
538 block_dim_count: usize,
539 grid_dim_count: usize,
540 compare_count: usize,
541 constant_index_count: usize,
542 mul_count: usize,
543 add_count: usize,
544 store_count: usize,
545
546 pub fn deinit(self: *@This()) void {
547 self.graph.deinit();
548 self.* = undefined;
549 }
550 };
551
552 pub fn index2D(self: *@This(), ctx: anytype) !kernel.Index2D {
553 self.index2d_count += 1;
554 return ctx.default();
555 }
556
557 pub fn globalId(self: *@This(), ctx: anytype) !kernel.Value {
558 self.global_id_count += 1;
559 return ctx.default();
560 }
561
562 pub fn threadId(self: *@This(), ctx: anytype) !kernel.Value {
563 self.thread_id_count += 1;
564 return ctx.default();
565 }
566
567 pub fn blockId(self: *@This(), ctx: anytype) !kernel.Value {
568 self.block_id_count += 1;
569 return ctx.default();
570 }
571
572 pub fn blockDim(self: *@This(), ctx: anytype) !kernel.Value {
573 self.block_dim_count += 1;
574 return ctx.default();
575 }
576
577 pub fn gridDim(self: *@This(), ctx: anytype) !kernel.Value {
578 self.grid_dim_count += 1;
579 return ctx.default();
580 }
581
582 pub fn compare(self: *@This(), ctx: anytype) !kernel.Value {
583 self.compare_count += 1;
584 return ctx.default();
585 }
586
587 pub fn constantIndex(self: *@This(), ctx: anytype) !kernel.Value {
588 self.constant_index_count += 1;
589 return ctx.default();
590 }
591
592 pub fn mul(self: *@This(), ctx: anytype) !kernel.Value {
593 self.mul_count += 1;
594 return ctx.default();
595 }
596
597 pub fn add(self: *@This(), ctx: anytype) !kernel.Value {
598 self.add_count += 1;
599 return ctx.default();
600 }
601
602 pub fn store(self: *@This(), ctx: anytype) !void {
603 self.store_count += 1;
604 try ctx.default();
605 }
606
607 pub fn finish(self: *@This(), ctx: anytype) !Result {
608 return .{
609 .graph = try ctx.default(),
610 .index2d_count = self.index2d_count,
611 .global_id_count = self.global_id_count,
612 .thread_id_count = self.thread_id_count,
613 .block_id_count = self.block_id_count,
614 .block_dim_count = self.block_dim_count,
615 .grid_dim_count = self.grid_dim_count,
616 .compare_count = self.compare_count,
617 .constant_index_count = self.constant_index_count,
618 .mul_count = self.mul_count,
619 .add_count = self.add_count,
620 .store_count = self.store_count,
621 };
622 }
623 };
624
625 fn matrix_copy_each(inner: anytype, index: kernel.Index2D, args: anytype) !void {
626 const value = try args.param(.src).load(inner, index);
627 try args.param(.dst).store(inner, value, index);
628 }
629
630 fn matrixCopyBody(k: anytype, args: anytype) !void {
631 _ = try k.forEach2D(.{
632 .x = kernel.domainAxis("x", 3, 2),
633 .y = kernel.domainAxis("y", 2, 2),
634 }, args, matrix_copy_each);
635 }
636
637 const MatrixCopy = Program(.{
638 .name = "kernel_dsl_matrix_copy_i32",
639 .parameters = .{
640 .src = kernel.dynamicBuffer(.i32),
641 .dst = kernel.dynamicBuffer(.i32),
642 },
643 .body = matrixCopyBody,
644 });
645
646 const TransposeIndexedStores = struct {
647 pub fn storeIndex(_: *@This(), ctx: anytype) !void {
648 const Index = @TypeOf(ctx.index);
649 if (comptime Index == kernel.Index2D) {
650 const height = try ctx.next.constantIndex(std.math.cast(i64, ctx.index.y.extent) orelse return error.LaunchDimensionOverflow);
651 const column_base = try ctx.next.mul(ctx.index.x.index, height);
652 const transposed_index = try ctx.next.add(column_base, ctx.index.y.index);
653 try ctx.next.store(ctx.value, ctx.memref, transposed_index);
654 return;
655 }
656 try ctx.default();
657 }
658 };
659
660 const TransposedMatrixStores = MatrixCopy.transform(TransposeIndexedStores{});
661
662 fn paired_matrix_stores_each(inner: anytype, index: kernel.Index2D, args: anytype) !void {
663 const value = try args.param(.src).load(inner, index);
664 try args.param(.dst).store(inner, value, index);
665 try args.param(.audit).store(inner, value, index);
666 }
667
668 fn pairedMatrixStoresBody(k: anytype, args: anytype) !void {
669 _ = try k.forEach2D(.{
670 .x = kernel.domainAxis("x", 3, 2),
671 .y = kernel.domainAxis("y", 2, 2),
672 }, args, paired_matrix_stores_each);
673 }
674
675 const PairedMatrixStores = Program(.{
676 .name = "kernel_dsl_paired_matrix_stores_i32",
677 .parameters = .{
678 .src = kernel.dynamicBuffer(.i32),
679 .dst = kernel.dynamicBuffer(.i32),
680 .audit = kernel.dynamicBuffer(.i32),
681 },
682 .body = pairedMatrixStoresBody,
683 });
684
685 const DelayedIndexedStore = struct {
686 value: kernel.Value,
687 memref: kernel.Value,
688 index: kernel.Index2D,
689 };
690
691 const TransposeFirstStoreWhenPaired = struct {
692 stores: [2]DelayedIndexedStore = undefined,
693 store_count: usize = 0,
694 active: bool = false,
695
696 pub fn guardIndex2DDo(self: *@This(), ctx: anytype) !void {
697 const previous_active = self.active;
698 self.active = true;
699 self.store_count = 0;
700 defer {
701 self.active = previous_active;
702 self.store_count = 0;
703 }
704
705 var x_guard = try ctx.guardIndex(ctx.index.x);
706 errdefer x_guard.abort();
707 var y_guard = try ctx.guardIndex(ctx.index.y);
708 errdefer y_guard.abort();
709 try ctx.runBody();
710 try self.flush(ctx.next);
711 try y_guard.leave();
712 try x_guard.leave();
713 }
714
715 pub fn storeIndex(self: *@This(), ctx: anytype) !void {
716 const Index = @TypeOf(ctx.index);
717 if (!self.active or comptime Index != kernel.Index2D) {
718 try ctx.default();
719 return;
720 }
721 if (self.store_count == self.stores.len) return error.DelayedStoreCapacityExceeded;
722 self.stores[self.store_count] = .{
723 .value = ctx.value,
724 .memref = ctx.memref,
725 .index = ctx.index,
726 };
727 self.store_count += 1;
728 }
729
730 fn flush(self: *@This(), next: anytype) !void {
731 for (self.stores[0..self.store_count], 0..) |entry, i| {
732 if (self.store_count >= 2 and i == 0) {
733 const height = try next.constantIndex(std.math.cast(i64, entry.index.y.extent) orelse return error.LaunchDimensionOverflow);
734 const column_base = try next.mul(entry.index.x.index, height);
735 const transposed_index = try next.add(column_base, entry.index.y.index);
736 try next.store(entry.value, entry.memref, transposed_index);
737 } else {
738 try next.store(entry.value, entry.memref, try entry.index.linear(next));
739 }
740 }
741 }
742 };
743
744 const DelayedTransposedMatrixStores = MatrixCopy.transform(TransposeFirstStoreWhenPaired{});
745 const DelayedTransposedPairedStores = PairedMatrixStores.transform(TransposeFirstStoreWhenPaired{});
746
747 fn paired_volume_stores_each(inner: anytype, index: kernel.Index3D, args: anytype) !void {
748 const value = try args.param(.src).load(inner, index);
749 try args.param(.dst).store(inner, value, index);
750 try args.param(.audit).store(inner, value, index);
751 }
752
753 fn pairedVolumeStoresBody(k: anytype, args: anytype) !void {
754 _ = try k.forEach3D(.{
755 .x = kernel.domainAxis("x", 2, 2),
756 .y = kernel.domainAxis("y", 2, 2),
757 .z = kernel.domainAxis("z", 2, 2),
758 }, args, paired_volume_stores_each);
759 }
760
761 const PairedVolumeStores = Program(.{
762 .name = "kernel_dsl_paired_volume_stores_i32",
763 .parameters = .{
764 .src = kernel.dynamicBuffer(.i32),
765 .dst = kernel.dynamicBuffer(.i32),
766 .audit = kernel.dynamicBuffer(.i32),
767 },
768 .body = pairedVolumeStoresBody,
769 });
770
771 const DelayedIndexedStore3D = struct {
772 value: kernel.Value,
773 memref: kernel.Value,
774 index: kernel.Index3D,
775 };
776
777 const SwapXzFirstStoreWhenPaired = struct {
778 stores: [2]DelayedIndexedStore3D = undefined,
779 store_count: usize = 0,
780 active: bool = false,
781
782 pub fn guardIndex3DDo(self: *@This(), ctx: anytype) !void {
783 const previous_active = self.active;
784 self.active = true;
785 self.store_count = 0;
786 defer {
787 self.active = previous_active;
788 self.store_count = 0;
789 }
790
791 var x_guard = try ctx.guardIndex(ctx.index.x);
792 errdefer x_guard.abort();
793 var y_guard = try ctx.guardIndex(ctx.index.y);
794 errdefer y_guard.abort();
795 var z_guard = try ctx.guardIndex(ctx.index.z);
796 errdefer z_guard.abort();
797 try ctx.runBody();
798 try self.flush(ctx.next);
799 try z_guard.leave();
800 try y_guard.leave();
801 try x_guard.leave();
802 }
803
804 pub fn storeIndex(self: *@This(), ctx: anytype) !void {
805 const Index = @TypeOf(ctx.index);
806 if (!self.active or comptime Index != kernel.Index3D) {
807 try ctx.default();
808 return;
809 }
810 if (self.store_count == self.stores.len) return error.DelayedStoreCapacityExceeded;
811 self.stores[self.store_count] = .{
812 .value = ctx.value,
813 .memref = ctx.memref,
814 .index = ctx.index,
815 };
816 self.store_count += 1;
817 }
818
819 fn flush(self: *@This(), next: anytype) !void {
820 for (self.stores[0..self.store_count], 0..) |entry, i| {
821 if (self.store_count >= 2 and i == 0) {
822 const depth = try next.constantIndex(std.math.cast(i64, entry.index.z.extent) orelse return error.LaunchDimensionOverflow);
823 const y_depth = std.math.mul(u64, entry.index.y.extent, entry.index.z.extent) catch return error.LaunchDimensionOverflow;
824 const plane = try next.constantIndex(std.math.cast(i64, y_depth) orelse return error.LaunchDimensionOverflow);
825 const x_base = try next.mul(entry.index.x.index, plane);
826 const y_base = try next.mul(entry.index.y.index, depth);
827 const xy_base = try next.add(x_base, y_base);
828 const swapped_index = try next.add(xy_base, entry.index.z.index);
829 try next.store(entry.value, entry.memref, swapped_index);
830 } else {
831 try next.store(entry.value, entry.memref, try entry.index.linear(next));
832 }
833 }
834 }
835 };
836
837 const DelayedSwappedVolumeStores = PairedVolumeStores.transform(SwapXzFirstStoreWhenPaired{});
838
839 const GuardBodyCounter = struct {
840 guard1d_count: usize = 0,
841
842 pub const Result = struct {
843 graph: kernel.Graph,
844 guard1d_count: usize,
845
846 pub fn deinit(self: *@This()) void {
847 self.graph.deinit();
848 self.* = undefined;
849 }
850 };
851
852 pub fn guardIndexDo(self: *@This(), ctx: anytype) !void {
853 self.guard1d_count += 1;
854 try ctx.default();
855 }
856
857 pub fn finish(self: *@This(), ctx: anytype) !Result {
858 return .{
859 .graph = try ctx.default(),
860 .guard1d_count = self.guard1d_count,
861 };
862 }
863 };
864
865 fn guarded_pair_store(inner: anytype, ctx: anytype) !void {
866 const value = try ctx.args.param(.src).load(inner, ctx.index);
867 try ctx.args.param(.dst).store(inner, value, ctx.index);
868 try ctx.args.param(.audit).store(inner, value, ctx.index);
869 }
870
871 fn guarded_pair_stores_each(inner: anytype, index: kernel.Index1D, args: anytype) !void {
872 const limit = try inner.constantIndex(4);
873 const active = try inner.compare(.lt, index.index, limit);
874 try inner.guardDo(active, .{ .args = args, .index = index }, guarded_pair_store);
875 }
876
877 fn guardedPairStoresBody(k: anytype, args: anytype) !void {
878 _ = try k.forEach1D("i", 6, 4, args, guarded_pair_stores_each);
879 }
880
881 const GuardedPairStores = Program(.{
882 .name = "kernel_dsl_guarded_pair_stores_i32",
883 .parameters = .{
884 .src = kernel.dynamicBuffer(.i32),
885 .dst = kernel.dynamicBuffer(.i32),
886 .audit = kernel.dynamicBuffer(.i32),
887 },
888 .body = guardedPairStoresBody,
889 });
890
891 const DelayedIndexedStore1D = struct {
892 value: kernel.Value,
893 memref: kernel.Value,
894 index: kernel.Index1D,
895 };
896
897 const ShiftFirstGuardStoreWhenPaired = struct {
898 stores: [2]DelayedIndexedStore1D = undefined,
899 store_count: usize = 0,
900 active: bool = false,
901
902 pub fn guardDo(self: *@This(), ctx: anytype) !void {
903 const previous_active = self.active;
904 self.active = true;
905 self.store_count = 0;
906 defer {
907 self.active = previous_active;
908 self.store_count = 0;
909 }
910
911 var active_guard = try ctx.guard(ctx.condition);
912 errdefer active_guard.abort();
913 try ctx.runBody();
914 try self.flush(ctx.next);
915 try active_guard.leave();
916 }
917
918 pub fn storeIndex(self: *@This(), ctx: anytype) !void {
919 const Index = @TypeOf(ctx.index);
920 if (!self.active or comptime Index != kernel.Index1D) {
921 try ctx.default();
922 return;
923 }
924 if (self.store_count == self.stores.len) return error.DelayedStoreCapacityExceeded;
925 self.stores[self.store_count] = .{
926 .value = ctx.value,
927 .memref = ctx.memref,
928 .index = ctx.index,
929 };
930 self.store_count += 1;
931 }
932
933 fn flush(self: *@This(), next: anytype) !void {
934 for (self.stores[0..self.store_count], 0..) |entry, i| {
935 if (self.store_count >= 2 and i == 0) {
936 const offset = try next.constantIndex(2);
937 const shifted_index = try next.add(entry.index.index, offset);
938 try next.store(entry.value, entry.memref, shifted_index);
939 } else {
940 try next.store(entry.value, entry.memref, try entry.index.linear(next));
941 }
942 }
943 }
944 };
945
946 const ShiftedGuardedPairStores = GuardedPairStores.transform(ShiftFirstGuardStoreWhenPaired{});
947
948 const IndexedStoreCounter = struct {
949 store_index_count: usize = 0,
950 store_count: usize = 0,
951 mul_count: usize = 0,
952 add_count: usize = 0,
953
954 pub const Result = struct {
955 graph: kernel.Graph,
956 store_index_count: usize,
957 store_count: usize,
958 mul_count: usize,
959 add_count: usize,
960
961 pub fn deinit(self: *@This()) void {
962 self.graph.deinit();
963 self.* = undefined;
964 }
965 };
966
967 pub fn storeIndex(self: *@This(), ctx: anytype) !void {
968 self.store_index_count += 1;
969 try ctx.default();
970 }
971
972 pub fn store(self: *@This(), ctx: anytype) !void {
973 self.store_count += 1;
974 try ctx.default();
975 }
976
977 pub fn mul(self: *@This(), ctx: anytype) !kernel.Value {
978 self.mul_count += 1;
979 return ctx.default();
980 }
981
982 pub fn add(self: *@This(), ctx: anytype) !kernel.Value {
983 self.add_count += 1;
984 return ctx.default();
985 }
986
987 pub fn finish(self: *@This(), ctx: anytype) !Result {
988 return .{
989 .graph = try ctx.default(),
990 .store_index_count = self.store_index_count,
991 .store_count = self.store_count,
992 .mul_count = self.mul_count,
993 .add_count = self.add_count,
994 };
995 }
996 };
997
998 test "kernel Program interprets generic bodies through user build semantics" {
999 var result = try InterpretedScale.interpret(std.testing.allocator, InterpretedScale.Limits.testing, ArithmeticCounter{});
1000 defer result.deinit();
1001
1002 try std.testing.expectEqual(@as(usize, 1), result.mul_count);
1003 try std.testing.expectEqual(@as(usize, 1), result.add_count);
1004 try std.testing.expectEqual(@as(usize, 1), result.store_count);
1005 try result.graph.verify();
1006
1007 var input = [_]f32{ 1.0, -2.0, 3.5, 4.0, -0.25 };
1008 var output = [_]f32{ 0, 0, 0, 0, 0 };
1009 try result.graph.runCpu(std.testing.allocator, &.{
1010 kernel.argumentBuffer(f32, input[0..]),
1011 kernel.argumentBuffer(f32, output[0..]),
1012 kernel.argumentF32(2.5),
1013 });
1014
1015 try std.testing.expectEqualSlices(f32, &.{ 3.5, -4.0, 9.75, 11.0, 0.375 }, output[0..]);
1016 }
1017
1018 test "kernel Program interprets multidimensional launch programs through user semantics" {
1019 var result = try TiledLaunch.interpret(std.testing.allocator, TiledLaunch.Limits.testing, DomainCounter{});
1020 defer result.deinit();
1021
1022 try std.testing.expectEqual(@as(usize, 1), result.index2d_count);
1023 try std.testing.expectEqual(@as(usize, 1), result.global_id_count);
1024 try std.testing.expectEqual(@as(usize, 1), result.thread_id_count);
1025 try std.testing.expectEqual(@as(usize, 1), result.block_id_count);
1026 try std.testing.expectEqual(@as(usize, 1), result.block_dim_count);
1027 try std.testing.expectEqual(@as(usize, 1), result.grid_dim_count);
1028 try std.testing.expectEqual(@as(usize, 2), result.compare_count);
1029 try std.testing.expectEqual(@as(usize, 1), result.constant_index_count);
1030 try std.testing.expectEqual(@as(usize, 1), result.mul_count);
1031 try std.testing.expectEqual(@as(usize, 5), result.add_count);
1032 try std.testing.expectEqual(@as(usize, 1), result.store_count);
1033 try result.graph.verify();
1034
1035 const launch_value = try result.graph.launch();
1036 try std.testing.expectEqual(@as(u32, 3), launch_value.grid[0]);
1037 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[1]);
1038 try std.testing.expectEqual(@as(u32, 2), launch_value.block[0]);
1039 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
1040
1041 var output = @as([15]i32, @splat(0));
1042 try result.graph.runCpu(std.testing.allocator, &.{
1043 kernel.argumentBuffer(i32, output[0..]),
1044 });
1045
1046 try std.testing.expectEqualSlices(i32, &.{
1047 4, 6, 6, 8, 8,
1048 4, 6, 6, 8, 8,
1049 5, 7, 7, 9, 9,
1050 }, output[0..]);
1051 }
1052
1053 test "kernel Program transform forwards whileLoop through user semantics" {
1054 var result = try WhileCountdown.interpret(std.testing.allocator, WhileCountdown.Limits.testing, WhileLoopCounter{});
1055 defer result.deinit();
1056
1057 try std.testing.expectEqual(@as(usize, 1), result.while_loop_count);
1058 try result.graph.verify();
1059
1060 var input = [_]i32{ 0, 1, 4, 5 };
1061 var output = @as([4]i32, @splat(0));
1062 try result.graph.runCpu(std.testing.allocator, &.{
1063 kernel.argumentBuffer(i32, input[0..]),
1064 kernel.argumentBuffer(i32, output[0..]),
1065 });
1066 try std.testing.expectEqualSlices(i32, &.{ 0, 1, 10, 15 }, output[0..]);
1067 }
1068
1069 test "kernel Program transform rewrites multidimensional indexed stores" {
1070 var normal = try MatrixCopy.interpret(std.testing.allocator, MatrixCopy.Limits.testing, IndexedStoreCounter{});
1071 defer normal.deinit();
1072
1073 try std.testing.expectEqual(@as(usize, 1), normal.store_index_count);
1074 try std.testing.expectEqual(@as(usize, 0), normal.store_count);
1075 try normal.graph.verify();
1076
1077 var input = [_]i32{ 1, 2, 3, 4, 5, 6 };
1078 var output = @as([6]i32, @splat(0));
1079 try normal.graph.runCpu(std.testing.allocator, &.{
1080 kernel.argumentBuffer(i32, input[0..]),
1081 kernel.argumentBuffer(i32, output[0..]),
1082 });
1083 try std.testing.expectEqualSlices(i32, input[0..], output[0..]);
1084
1085 var transformed = try TransposedMatrixStores.interpret(std.testing.allocator, TransposedMatrixStores.Limits.testing, IndexedStoreCounter{});
1086 defer transformed.deinit();
1087
1088 try std.testing.expectEqual(@as(usize, 0), transformed.store_index_count);
1089 try std.testing.expectEqual(@as(usize, 1), transformed.store_count);
1090 try std.testing.expectEqual(@as(usize, 2), transformed.mul_count);
1091 try std.testing.expectEqual(@as(usize, 2), transformed.add_count);
1092 try transformed.graph.verify();
1093
1094 output = @as([6]i32, @splat(0));
1095 try transformed.graph.runCpu(std.testing.allocator, &.{
1096 kernel.argumentBuffer(i32, input[0..]),
1097 kernel.argumentBuffer(i32, output[0..]),
1098 });
1099 try std.testing.expectEqualSlices(i32, &.{ 1, 4, 2, 5, 3, 6 }, output[0..]);
1100 }
1101
1102 test "kernel Program transform can delay indexed stores inside multidimensional guards" {
1103 var single = try DelayedTransposedMatrixStores.interpret(std.testing.allocator, DelayedTransposedMatrixStores.Limits.testing, IndexedStoreCounter{});
1104 defer single.deinit();
1105
1106 try std.testing.expectEqual(@as(usize, 0), single.store_index_count);
1107 try std.testing.expectEqual(@as(usize, 1), single.store_count);
1108 try std.testing.expectEqual(@as(usize, 2), single.mul_count);
1109 try std.testing.expectEqual(@as(usize, 2), single.add_count);
1110 try single.graph.verify();
1111
1112 var input = [_]i32{ 1, 2, 3, 4, 5, 6 };
1113 var output = @as([6]i32, @splat(0));
1114 try single.graph.runCpu(std.testing.allocator, &.{
1115 kernel.argumentBuffer(i32, input[0..]),
1116 kernel.argumentBuffer(i32, output[0..]),
1117 });
1118 try std.testing.expectEqualSlices(i32, input[0..], output[0..]);
1119
1120 var paired = try DelayedTransposedPairedStores.interpret(std.testing.allocator, DelayedTransposedPairedStores.Limits.testing, IndexedStoreCounter{});
1121 defer paired.deinit();
1122
1123 try std.testing.expectEqual(@as(usize, 0), paired.store_index_count);
1124 try std.testing.expectEqual(@as(usize, 2), paired.store_count);
1125 try std.testing.expectEqual(@as(usize, 3), paired.mul_count);
1126 try std.testing.expectEqual(@as(usize, 3), paired.add_count);
1127 try paired.graph.verify();
1128
1129 output = @as([6]i32, @splat(0));
1130 var audit = @as([6]i32, @splat(0));
1131 try paired.graph.runCpu(std.testing.allocator, &.{
1132 kernel.argumentBuffer(i32, input[0..]),
1133 kernel.argumentBuffer(i32, output[0..]),
1134 kernel.argumentBuffer(i32, audit[0..]),
1135 });
1136 try std.testing.expectEqualSlices(i32, &.{ 1, 4, 2, 5, 3, 6 }, output[0..]);
1137 try std.testing.expectEqualSlices(i32, input[0..], audit[0..]);
1138 }
1139
1140 test "kernel Program transform can delay indexed stores inside conditional guards" {
1141 var transformed = try ShiftedGuardedPairStores.interpret(std.testing.allocator, ShiftedGuardedPairStores.Limits.testing, IndexedStoreCounter{});
1142 defer transformed.deinit();
1143
1144 try std.testing.expectEqual(@as(usize, 0), transformed.store_index_count);
1145 try std.testing.expectEqual(@as(usize, 2), transformed.store_count);
1146 try std.testing.expectEqual(@as(usize, 0), transformed.mul_count);
1147 try std.testing.expectEqual(@as(usize, 1), transformed.add_count);
1148 try transformed.graph.verify();
1149
1150 var input = [_]i32{ 1, 2, 3, 4, 5, 6 };
1151 var output = @as([6]i32, @splat(0));
1152 var audit = @as([6]i32, @splat(0));
1153 try transformed.graph.runCpu(std.testing.allocator, &.{
1154 kernel.argumentBuffer(i32, input[0..]),
1155 kernel.argumentBuffer(i32, output[0..]),
1156 kernel.argumentBuffer(i32, audit[0..]),
1157 });
1158 try std.testing.expectEqualSlices(i32, &.{ 0, 0, 1, 2, 3, 4 }, output[0..]);
1159 try std.testing.expectEqualSlices(i32, &.{ 1, 2, 3, 4, 0, 0 }, audit[0..]);
1160 }
1161
1162 test "kernel Program exposes one dimensional guard body semantics" {
1163 var result = try InterpretedScale.interpret(std.testing.allocator, InterpretedScale.Limits.testing, GuardBodyCounter{});
1164 defer result.deinit();
1165
1166 try std.testing.expectEqual(@as(usize, 1), result.guard1d_count);
1167 try result.graph.verify();
1168 }
1169
1170 test "kernel Program transform can delay indexed stores inside three dimensional guards" {
1171 var transformed = try DelayedSwappedVolumeStores.interpret(std.testing.allocator, DelayedSwappedVolumeStores.Limits.testing, IndexedStoreCounter{});
1172 defer transformed.deinit();
1173
1174 try std.testing.expectEqual(@as(usize, 0), transformed.store_index_count);
1175 try std.testing.expectEqual(@as(usize, 2), transformed.store_count);
1176 try std.testing.expectEqual(@as(usize, 6), transformed.mul_count);
1177 try std.testing.expectEqual(@as(usize, 6), transformed.add_count);
1178 try transformed.graph.verify();
1179
1180 var input = [_]i32{ 1, 2, 3, 4, 5, 6, 7, 8 };
1181 var output = @as([8]i32, @splat(0));
1182 var audit = @as([8]i32, @splat(0));
1183 try transformed.graph.runCpu(std.testing.allocator, &.{
1184 kernel.argumentBuffer(i32, input[0..]),
1185 kernel.argumentBuffer(i32, output[0..]),
1186 kernel.argumentBuffer(i32, audit[0..]),
1187 });
1188 try std.testing.expectEqualSlices(i32, &.{ 1, 5, 3, 7, 2, 6, 4, 8 }, output[0..]);
1189 try std.testing.expectEqualSlices(i32, input[0..], audit[0..]);
1190 }
1191
1192 const FinishCounter = struct {
1193 finish_count: usize = 0,
1194
1195 pub const Result = struct {
1196 inner: ArithmeticCounter.Result,
1197 finish_count: usize,
1198
1199 pub fn deinit(self: *@This()) void {
1200 self.inner.deinit();
1201 self.* = undefined;
1202 }
1203 };
1204
1205 pub fn finish(self: *@This(), ctx: anytype) !Result {
1206 self.finish_count += 1;
1207 return .{
1208 .inner = try ctx.default(),
1209 .finish_count = self.finish_count,
1210 };
1211 }
1212 };
1213
1214 const IndexedFinishCounter = struct {
1215 finish_count: usize = 0,
1216
1217 pub const Result = struct {
1218 inner: IndexedStoreCounter.Result,
1219 finish_count: usize,
1220
1221 pub fn deinit(self: *@This()) void {
1222 self.inner.deinit();
1223 self.* = undefined;
1224 }
1225 };
1226
1227 pub fn finish(self: *@This(), ctx: anytype) !Result {
1228 self.finish_count += 1;
1229 return .{
1230 .inner = try ctx.default(),
1231 .finish_count = self.finish_count,
1232 };
1233 }
1234 };
1235
1236 test "kernel Program stacks user build semantics" {
1237 var result = try InterpretedScale.interpret(
1238 std.testing.allocator,
1239 InterpretedScale.Limits.testing,
1240 kernel.interpret.stack(.{
1241 kernel.interpret.with(FinishCounter{}),
1242 kernel.interpret.with(ArithmeticCounter{}),
1243 }),
1244 );
1245 defer result.deinit();
1246
1247 try std.testing.expectEqual(@as(usize, 1), result.finish_count);
1248 try std.testing.expectEqual(@as(usize, 1), result.inner.mul_count);
1249 try std.testing.expectEqual(@as(usize, 1), result.inner.add_count);
1250 try std.testing.expectEqual(@as(usize, 1), result.inner.store_count);
1251 try result.inner.graph.verify();
1252 }
1253
1254 test "kernel Program binds unhandled high level operations through transparent semantics" {
1255 var result = try MatrixCopy.interpret(
1256 std.testing.allocator,
1257 MatrixCopy.Limits.testing,
1258 kernel.interpret.stack(.{
1259 kernel.interpret.bind(IndexedFinishCounter{}),
1260 kernel.interpret.with(IndexedStoreCounter{}),
1261 }),
1262 );
1263 defer result.deinit();
1264
1265 try std.testing.expectEqual(@as(usize, 1), result.finish_count);
1266 try std.testing.expectEqual(@as(usize, 1), result.inner.store_index_count);
1267 try std.testing.expectEqual(@as(usize, 0), result.inner.store_count);
1268 try result.inner.graph.verify();
1269
1270 var input = [_]i32{ 1, 2, 3, 4, 5, 6 };
1271 var output = @as([6]i32, @splat(0));
1272 try result.inner.graph.runCpu(std.testing.allocator, &.{
1273 kernel.argumentBuffer(i32, input[0..]),
1274 kernel.argumentBuffer(i32, output[0..]),
1275 });
1276 try std.testing.expectEqualSlices(i32, input[0..], output[0..]);
1277 }
1278
1279 test "kernel Program transform returns another composable Program" {
1280 var result = try TransformedScale.interpret(std.testing.allocator, TransformedScale.Limits.testing, ArithmeticCounter{});
1281 defer result.deinit();
1282
1283 try std.testing.expectEqual(@as(usize, 1), result.mul_count);
1284 try std.testing.expectEqual(@as(usize, 0), result.add_count);
1285 try std.testing.expectEqual(@as(usize, 1), result.store_count);
1286 try result.graph.verify();
1287
1288 var graph = try TransformedScale.interpret(std.testing.allocator, TransformedScale.Limits.testing, kernel.interpret.graph());
1289 defer graph.deinit();
1290 try graph.verify();
1291
1292 var input = [_]f32{ 1.0, -2.0, 3.5, 4.0, -0.25 };
1293 var output = [_]f32{ 0, 0, 0, 0, 0 };
1294 try TransformedScale.runCpu(std.testing.allocator, TransformedScale.Limits.testing, &.{
1295 kernel.argumentBuffer(f32, input[0..]),
1296 kernel.argumentBuffer(f32, output[0..]),
1297 kernel.argumentF32(2.5),
1298 });
1299
1300 try std.testing.expectEqualSlices(f32, &.{ 2.5, -5.0, 8.75, 10.0, -0.625 }, output[0..]);
1301 }
1302
1303 test "kernel Program transform accepts explicit interpreter specs" {
1304 var result = try AdditiveScale.interpret(std.testing.allocator, AdditiveScale.Limits.testing, ArithmeticCounter{});
1305 defer result.deinit();
1306
1307 try std.testing.expectEqual(@as(usize, 0), result.mul_count);
1308 try std.testing.expectEqual(@as(usize, 2), result.add_count);
1309 try std.testing.expectEqual(@as(usize, 1), result.store_count);
1310 try result.graph.verify();
1311
1312 var input = [_]f32{ 1.0, -2.0, 3.5, 4.0, -0.25 };
1313 var output = [_]f32{ 0, 0, 0, 0, 0 };
1314 try AdditiveScale.runCpu(std.testing.allocator, AdditiveScale.Limits.testing, &.{
1315 kernel.argumentBuffer(f32, input[0..]),
1316 kernel.argumentBuffer(f32, output[0..]),
1317 kernel.argumentF32(2.5),
1318 });
1319
1320 try std.testing.expectEqualSlices(f32, &.{ 4.5, 1.5, 7.0, 7.5, 3.25 }, output[0..]);
1321 }
1322
1323 test "kernel Program transforms compose before later interpretation" {
1324 var result = try TwiceTransformedScale.interpret(std.testing.allocator, TwiceTransformedScale.Limits.testing, ArithmeticCounter{});
1325 defer result.deinit();
1326
1327 try std.testing.expectEqual(@as(usize, 0), result.mul_count);
1328 try std.testing.expectEqual(@as(usize, 0), result.add_count);
1329 try std.testing.expectEqual(@as(usize, 1), result.store_count);
1330 try result.graph.verify();
1331
1332 var input = [_]f32{ 1.0, -2.0, 3.5, 4.0, -0.25 };
1333 var output = [_]f32{ 0, 0, 0, 0, 0 };
1334 try TwiceTransformedScale.runCpu(std.testing.allocator, TwiceTransformedScale.Limits.testing, &.{
1335 kernel.argumentBuffer(f32, input[0..]),
1336 kernel.argumentBuffer(f32, output[0..]),
1337 kernel.argumentF32(2.5),
1338 });
1339
1340 try std.testing.expectEqualSlices(f32, input[0..], output[0..]);
1341 }
1342
1343 test "kernel Program analysis view derives transforms from whole-program semantics" {
1344 var mul_analysis = try AddBeforeMulAnalysis.interpret(std.testing.allocator, AddBeforeMulAnalysis.Limits.testing);
1345 defer mul_analysis.deinit();
1346 try std.testing.expect(mul_analysis.has_mul);
1347 try mul_analysis.graph.verify();
1348
1349 var add_only_analysis = try AddOnlyAnalysis.interpret(std.testing.allocator, AddOnlyAnalysis.Limits.testing);
1350 defer add_only_analysis.deinit();
1351 try std.testing.expect(!add_only_analysis.has_mul);
1352 try add_only_analysis.graph.verify();
1353
1354 var transformed = try DropAddAfterMulAnalysis.interpret(std.testing.allocator, DropAddAfterMulAnalysis.Limits.testing, ArithmeticCounter{});
1355 defer transformed.deinit();
1356
1357 try std.testing.expectEqual(@as(usize, 1), transformed.mul_count);
1358 try std.testing.expectEqual(@as(usize, 0), transformed.add_count);
1359 try std.testing.expectEqual(@as(usize, 1), transformed.store_count);
1360 try transformed.graph.verify();
1361
1362 var input = [_]f32{ 1.0, -2.0, 3.5, 4.0, -0.25 };
1363 var output = [_]f32{ 0, 0, 0, 0, 0 };
1364 try DropAddAfterMulAnalysis.runCpu(std.testing.allocator, DropAddAfterMulAnalysis.Limits.testing, &.{
1365 kernel.argumentBuffer(f32, input[0..]),
1366 kernel.argumentBuffer(f32, output[0..]),
1367 kernel.argumentF32(2.5),
1368 });
1369
1370 try std.testing.expectEqualSlices(f32, &.{ 2.5, -5.0, 8.75, 10.0, -0.625 }, output[0..]);
1371
1372 var kept = try KeepAddAfterMulAnalysis.interpret(std.testing.allocator, KeepAddAfterMulAnalysis.Limits.testing, ArithmeticCounter{});
1373 defer kept.deinit();
1374
1375 try std.testing.expectEqual(@as(usize, 0), kept.mul_count);
1376 try std.testing.expectEqual(@as(usize, 1), kept.add_count);
1377 try std.testing.expectEqual(@as(usize, 1), kept.store_count);
1378 try kept.graph.verify();
1379
1380 output = [_]f32{ 0, 0, 0, 0, 0 };
1381 try KeepAddAfterMulAnalysis.runCpu(std.testing.allocator, KeepAddAfterMulAnalysis.Limits.testing, &.{
1382 kernel.argumentBuffer(f32, input[0..]),
1383 kernel.argumentBuffer(f32, output[0..]),
1384 kernel.argumentF32(2.5),
1385 });
1386
1387 try std.testing.expectEqualSlices(f32, &.{ 3.5, 0.5, 6.0, 6.5, 2.25 }, output[0..]);
1388
1389 var composed = try DropMulAfterAnalyzedDropAdd.interpret(std.testing.allocator, DropMulAfterAnalyzedDropAdd.Limits.testing, ArithmeticCounter{});
1390 defer composed.deinit();
1391
1392 try std.testing.expectEqual(@as(usize, 0), composed.mul_count);
1393 try std.testing.expectEqual(@as(usize, 0), composed.add_count);
1394 try std.testing.expectEqual(@as(usize, 1), composed.store_count);
1395 try composed.graph.verify();
1396
1397 output = [_]f32{ 0, 0, 0, 0, 0 };
1398 try DropMulAfterAnalyzedDropAdd.runCpu(std.testing.allocator, DropMulAfterAnalyzedDropAdd.Limits.testing, &.{
1399 kernel.argumentBuffer(f32, input[0..]),
1400 kernel.argumentBuffer(f32, output[0..]),
1401 kernel.argumentF32(2.5),
1402 });
1403
1404 try std.testing.expectEqualSlices(f32, input[0..], output[0..]);
1405
1406 var analyzed_after_transform = try AnalyzeAfterDroppedMul.interpret(std.testing.allocator, AnalyzeAfterDroppedMul.Limits.testing, ArithmeticCounter{});
1407 defer analyzed_after_transform.deinit();
1408
1409 try std.testing.expectEqual(@as(usize, 0), analyzed_after_transform.mul_count);
1410 try std.testing.expectEqual(@as(usize, 1), analyzed_after_transform.add_count);
1411 try std.testing.expectEqual(@as(usize, 1), analyzed_after_transform.store_count);
1412 try analyzed_after_transform.graph.verify();
1413
1414 output = [_]f32{ 0, 0, 0, 0, 0 };
1415 try AnalyzeAfterDroppedMul.runCpu(std.testing.allocator, AnalyzeAfterDroppedMul.Limits.testing, &.{
1416 kernel.argumentBuffer(f32, input[0..]),
1417 kernel.argumentBuffer(f32, output[0..]),
1418 kernel.argumentF32(2.5),
1419 });
1420
1421 try std.testing.expectEqualSlices(f32, &.{ 2.0, -1.0, 4.5, 5.0, 0.75 }, output[0..]);
1422
1423 var chained = try DropMulAfterAnalyzedAddRemoval.interpret(std.testing.allocator, DropMulAfterAnalyzedAddRemoval.Limits.testing, ArithmeticCounter{});
1424 defer chained.deinit();
1425
1426 try std.testing.expectEqual(@as(usize, 0), chained.mul_count);
1427 try std.testing.expectEqual(@as(usize, 0), chained.add_count);
1428 try std.testing.expectEqual(@as(usize, 1), chained.store_count);
1429 try chained.graph.verify();
1430
1431 output = [_]f32{ 0, 0, 0, 0, 0 };
1432 try DropMulAfterAnalyzedAddRemoval.runCpu(std.testing.allocator, DropMulAfterAnalyzedAddRemoval.Limits.testing, &.{
1433 kernel.argumentBuffer(f32, input[0..]),
1434 kernel.argumentBuffer(f32, output[0..]),
1435 kernel.argumentF32(2.5),
1436 });
1437
1438 try std.testing.expectEqualSlices(f32, input[0..], output[0..]);
1439
1440 var borrowed = try DropMulWithBorrowedAnalysisFlags.interpret(std.testing.allocator, DropMulWithBorrowedAnalysisFlags.Limits.testing, ArithmeticCounter{});
1441 defer borrowed.deinit();
1442
1443 try std.testing.expectEqual(@as(usize, 0), borrowed.mul_count);
1444 try std.testing.expectEqual(@as(usize, 1), borrowed.add_count);
1445 try std.testing.expectEqual(@as(usize, 1), borrowed.store_count);
1446 try borrowed.graph.verify();
1447
1448 output = [_]f32{ 0, 0, 0, 0, 0 };
1449 try DropMulWithBorrowedAnalysisFlags.runCpu(std.testing.allocator, DropMulWithBorrowedAnalysisFlags.Limits.testing, &.{
1450 kernel.argumentBuffer(f32, input[0..]),
1451 kernel.argumentBuffer(f32, output[0..]),
1452 kernel.argumentF32(2.5),
1453 });
1454
1455 try std.testing.expectEqualSlices(f32, &.{ 2.0, -1.0, 4.5, 5.0, 0.75 }, output[0..]);
1456 }
1457
1458 test "kernel Program transform can validate a whole build in finish" {
1459 var graph = try StoreCheckedScale.build(std.testing.allocator, StoreCheckedScale.Limits.testing);
1460 defer graph.deinit();
1461 try graph.verify();
1462
1463 try std.testing.expectError(error.KernelProgramRequiresStore, StoreCheckedNoStore.build(std.testing.allocator, StoreCheckedNoStore.Limits.testing));
1464 }
1465
1466 fn positional_copy_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {
1467 const value = try inner.load(ctx.src, try index.linear(inner));
1468 try inner.store(value, ctx.dst, try index.linear(inner));
1469 }
1470
1471 fn positionalBody(k: anytype, args: []const kernel.Value) !void {
1472 _ = try k.forEach1D(
1473 "i",
1474 4,
1475 4,
1476 .{ .src = args[0], .dst = args[1] },
1477 positional_copy_each,
1478 );
1479 }
1480
1481 const PositionalCopy = Program(.{
1482 .name = "kernel_dsl_positional_copy_i32",
1483 .parameters = .{
1484 kernel.dynamicBuffer(.i32),
1485 kernel.dynamicBuffer(.i32),
1486 },
1487 .body = positionalBody,
1488 });
1489
1490 test "kernel Program still accepts positional parameters" {
1491 try std.testing.expectEqual(@as(usize, 2), PositionalCopy.schema().len);
1492 var graph = try PositionalCopy.build(std.testing.allocator, PositionalCopy.Limits.testing);
1493 defer graph.deinit();
1494 try graph.verify();
1495 }