lib/accy/src/tensor/dsl/test.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const accy = @import("../../root.zig");
4 const tensor = @import("../root.zig");
5 const namespace = @import("root.zig");
6 const batch = tensor.batch;
7 const interpret_mod = tensor.interpret;
8 const program_mod = tensor.program;
9 const trace = tensor.trace;
10 const transform = tensor.transform;
11 const types = tensor.types;
12 const body = namespace.body;
13 const output = namespace.output;
14 const parameter = namespace.parameter;
15 const Program = namespace.Program;
16
17 fn elementwiseLoss(_: *trace.Builder, args: anytype) !trace.Value {
18 const product = try args.param(.x).mul(args.param(.y));
19 return try product.sum(.lane);
20 }
21
22 fn positionalLoss(_: *trace.Builder, args: []const trace.Value) !trace.Value {
23 const product = try args[0].mul(args[1]);
24 return try product.sum(.lane);
25 }
26
27 fn arrayForward(_: *trace.Builder, args: anytype) ![2]trace.Value {
28 const sum = try args.param(.x).add(args.param(.y));
29 return .{ sum, try sum.sum(.lane) };
30 }
31
32 fn addBeforeMulLoss(_: *trace.Builder, args: anytype) !trace.Value {
33 const sum = try args.param(.x).add(args.param(.y));
34 const product = try sum.mul(args.param(.z));
35 return try product.sum(.lane);
36 }
37
38 fn selectUnmapped(_: *trace.Builder, args: anytype) !trace.Value {
39 return args.param(.y);
40 }
41
42 fn selectAfterMul(_: *trace.Builder, args: anytype) !trace.Value {
43 _ = try args.param(.x).mul(args.param(.y));
44 return args.param(.z);
45 }
46
47 const ElementwiseOutputs = struct {
48 product: trace.Value,
49 total: trace.Value,
50 };
51
52 fn elementwiseForward(_: *trace.Builder, args: anytype) !ElementwiseOutputs {
53 const product = try args.param(.x).mul(args.param(.y));
54 return .{
55 .product = product,
56 .total = try product.sum(.lane),
57 };
58 }
59
60 const ElementwiseLoss = Program(.{
61 .name = "dsl_elementwise_loss",
62 .parameters = .{
63 .x = types.spec(.f32, .{ .lane = 4 }),
64 .y = types.spec(.f32, .{ .lane = 4 }),
65 },
66 .body = elementwiseLoss,
67 });
68
69 const ElementwiseGradient = ElementwiseLoss.grad(.{ .wrt = ElementwiseLoss.wrt(.{ .x, .y }) });
70
71 const GeneratedMulAsAdd = struct {
72 pub fn mul(_: *@This(), ctx: anytype) !trace.Value {
73 if (ctx.op.id.index == program_mod.synthetic_id.index) {
74 var op = ctx.op.*;
75 op.kind = .{ .binary = .{
76 .op = .add,
77 .lhs = program_mod.synthetic_id,
78 .rhs = program_mod.synthetic_id,
79 } };
80 return ctx.next.bind(&op, ctx.args);
81 }
82 return ctx.default();
83 }
84 };
85
86 const GeneratedSeedAsZero = struct {
87 pub fn constant(_: *@This(), ctx: anytype) !trace.Value {
88 if (ctx.op.id.index == program_mod.synthetic_id.index) {
89 const payload = try ctx.builderHandle().arena.allocator().alloc(u8, try ctx.op.result.byteCount());
90 @memset(payload, 0);
91 var op = ctx.op.*;
92 op.kind = .{ .constant = .{ .payload = payload } };
93 return ctx.next.bind(&op, &.{});
94 }
95 return ctx.default();
96 }
97 };
98
99 const GeneratedBroadcastAsZero = struct {
100 pub fn broadcastInDim(_: *@This(), ctx: anytype) !trace.Value {
101 if (ctx.op.id.index == program_mod.synthetic_id.index) {
102 const payload = try ctx.builderHandle().arena.allocator().alloc(u8, try ctx.op.result.byteCount());
103 @memset(payload, 0);
104 var op = ctx.op.*;
105 op.kind = .{ .constant = .{ .payload = payload } };
106 return ctx.next.bind(&op, &.{});
107 }
108 return ctx.default();
109 }
110 };
111
112 const GeneratedMulAsAddFromAnalysis = struct {
113 analysis: *const MulPresence.Result,
114
115 pub fn mul(self: *@This(), ctx: anytype) !trace.Value {
116 if (self.analysis.has_mul and ctx.op.id.index == program_mod.synthetic_id.index) {
117 var op = ctx.op.*;
118 op.kind = .{ .binary = .{
119 .op = .add,
120 .lhs = program_mod.synthetic_id,
121 .rhs = program_mod.synthetic_id,
122 } };
123 return ctx.next.bind(&op, ctx.args);
124 }
125 return ctx.default();
126 }
127 };
128
129 const GeneratedBroadcastAsZeroFromAnalysis = struct {
130 analysis: *const MulPresence.Result,
131
132 pub fn broadcastInDim(self: *@This(), ctx: anytype) !trace.Value {
133 if (self.analysis.has_mul and ctx.op.id.index == program_mod.synthetic_id.index) {
134 const payload = try ctx.builderHandle().arena.allocator().alloc(u8, try ctx.op.result.byteCount());
135 @memset(payload, 0);
136 var op = ctx.op.*;
137 op.kind = .{ .constant = .{ .payload = payload } };
138 return ctx.next.bind(&op, &.{});
139 }
140 return ctx.default();
141 }
142 };
143
144 const LinearizeAnalysisHooks = struct {
145 linearize: interpret_mod.With(GeneratedMulAsAddFromAnalysis),
146 };
147
148 const PullbackAnalysisHooks = struct {
149 pullback: interpret_mod.With(GeneratedMulAsAddFromAnalysis),
150 };
151
152 const BatchAnalysisHooks = struct {
153 batch: interpret_mod.With(GeneratedBroadcastAsZeroFromAnalysis),
154 };
155
156 fn linearizeHooksFromMulPresence(analysis: *const MulPresence.Result) LinearizeAnalysisHooks {
157 return .{ .linearize = interpret_mod.bind(GeneratedMulAsAddFromAnalysis{ .analysis = analysis }) };
158 }
159
160 fn pullbackHooksFromMulPresence(analysis: *const MulPresence.Result) PullbackAnalysisHooks {
161 return .{ .pullback = interpret_mod.bind(GeneratedMulAsAddFromAnalysis{ .analysis = analysis }) };
162 }
163
164 fn batchHooksFromMulPresence(analysis: *const MulPresence.Result) BatchAnalysisHooks {
165 return .{ .batch = interpret_mod.bind(GeneratedBroadcastAsZeroFromAnalysis{ .analysis = analysis }) };
166 }
167
168 const ElementwiseGradientWithGeneratedRewrite = ElementwiseLoss.gradWith(
169 .{ .wrt = ElementwiseLoss.wrt(.{ .x, .y }) },
170 .{
171 .pullback = interpret_mod.bind(GeneratedMulAsAdd{}),
172 .seed = interpret_mod.bind(GeneratedSeedAsZero{}),
173 },
174 );
175
176 const ElementwiseMulAnalysis = ElementwiseLoss.analyze(MulPresence{});
177
178 const ElementwiseGradientWithAnalysisRewrite = ElementwiseMulAnalysis.gradWith(
179 .{ .wrt = ElementwiseLoss.wrt(.{ .x, .y }) },
180 pullbackHooksFromMulPresence,
181 );
182
183 const BatchedElementwiseGradient = ElementwiseGradient.vmap(.{
184 .axis_size = 8,
185 .in_axes = ElementwiseGradient.inAxes(.{ .x = batch.mapped(0), .y = batch.mapped(0) }),
186 });
187
188 const BatchedElementwiseGradientWithGeneratedRewrite = ElementwiseGradientWithGeneratedRewrite.vmap(.{
189 .axis_size = 8,
190 .in_axes = ElementwiseGradientWithGeneratedRewrite.inAxes(.{ .x = batch.mapped(0), .y = batch.mapped(0) }),
191 });
192
193 const BatchedElementwiseGradientWithAnalysisRewrite = ElementwiseGradientWithAnalysisRewrite.vmap(.{
194 .axis_size = 8,
195 .in_axes = ElementwiseGradientWithAnalysisRewrite.inAxes(.{ .x = batch.mapped(0), .y = batch.mapped(0) }),
196 });
197
198 const ElementwiseJvp = ElementwiseLoss.jvp(.{ .wrt = ElementwiseLoss.wrt(.{ .x, .y }) });
199
200 const ElementwiseJvpWithGeneratedRewrite = ElementwiseLoss.jvpWith(
201 .{ .wrt = ElementwiseLoss.wrt(.{ .x, .y }) },
202 .{ .linearize = interpret_mod.bind(GeneratedMulAsAdd{}) },
203 );
204
205 const ElementwiseJvpWithAnalysisRewrite = ElementwiseMulAnalysis.jvpWith(
206 .{ .wrt = ElementwiseLoss.wrt(.{ .x, .y }) },
207 linearizeHooksFromMulPresence,
208 );
209
210 const dsl_custom_double_target = "accy.dsl.custom.double";
211
212 fn customDoubleLoss(builder: *trace.Builder, args: anytype) !trace.Value {
213 const input = args.param(.x);
214 const doubled = try builder.customCall(dsl_custom_double_target, 1, &.{input}, input.ty);
215 return try doubled.sum(.lane);
216 }
217
218 const DslDoubleJvpRule = struct {
219 pub fn bind(_: *@This(), ctx: anytype) !tensor.Dual {
220 switch (ctx.op.kind) {
221 .custom_call => |custom| {
222 if (std.mem.eql(u8, custom.target, dsl_custom_double_target)) {
223 const builder = ctx.builderHandle();
224 return .{
225 .primal = try builder.customCall(custom.target, custom.version, &.{ctx.args[0].primal}, ctx.op.result),
226 .tangent = try builder.customCall(custom.target, custom.version, &.{ctx.args[0].tangent}, ctx.op.result),
227 };
228 }
229 },
230 else => {},
231 }
232 return ctx.default();
233 }
234 };
235
236 const CustomDoubleLoss = Program(.{
237 .name = "dsl_custom_double_loss",
238 .parameters = .{
239 .x = types.spec(.f32, .{ .lane = 4 }),
240 },
241 .body = customDoubleLoss,
242 });
243
244 const CustomDoubleJvp = CustomDoubleLoss.jvpWith(
245 .{ .wrt = CustomDoubleLoss.wrt(.x) },
246 .{ .jvp = interpret_mod.bind(DslDoubleJvpRule{}) },
247 );
248
249 const BatchedElementwiseJvp = ElementwiseJvp.vmap(.{
250 .axis_size = 8,
251 .in_axes = ElementwiseJvp.inAxes(.{
252 .primals = .{ .x = batch.mapped(0), .y = batch.mapped(0) },
253 .tangents = .{ .x = batch.mapped(0), .y = batch.mapped(0) },
254 }),
255 });
256
257 const BatchedElementwiseJvpWithGeneratedRewrite = ElementwiseJvpWithGeneratedRewrite.vmap(.{
258 .axis_size = 8,
259 .in_axes = ElementwiseJvpWithGeneratedRewrite.inAxes(.{
260 .primals = .{ .x = batch.mapped(0), .y = batch.mapped(0) },
261 .tangents = .{ .x = batch.mapped(0), .y = batch.mapped(0) },
262 }),
263 });
264
265 const BatchedElementwiseJvpWithAnalysisRewrite = ElementwiseJvpWithAnalysisRewrite.vmap(.{
266 .axis_size = 8,
267 .in_axes = ElementwiseJvpWithAnalysisRewrite.inAxes(.{
268 .primals = .{ .x = batch.mapped(0), .y = batch.mapped(0) },
269 .tangents = .{ .x = batch.mapped(0), .y = batch.mapped(0) },
270 }),
271 });
272
273 const ArrayForward = Program(.{
274 .name = "dsl_array_forward",
275 .parameters = .{
276 .x = types.spec(.f32, .{ .lane = 4 }),
277 .y = types.spec(.f32, .{ .lane = 4 }),
278 },
279 .body = arrayForward,
280 });
281
282 const ArrayForwardCopy = ArrayForward.rewrite(struct {}{});
283
284 const ElementwiseForward = Program(.{
285 .name = "dsl_elementwise_forward",
286 .parameters = .{
287 .x = types.spec(.f32, .{ .lane = 4 }),
288 .y = types.spec(.f32, .{ .lane = 4 }),
289 },
290 .body = elementwiseForward,
291 });
292
293 const ElementwiseForwardJvp = ElementwiseForward.jvp(.{ .wrt = ElementwiseForward.wrt(.{ .x, .y }) });
294
295 const SelectUnmapped = Program(.{
296 .name = "dsl_select_unmapped",
297 .parameters = .{
298 .x = types.spec(.f32, .{ .lane = 4 }),
299 .y = types.spec(.f32, .{ .lane = 4 }),
300 },
301 .body = selectUnmapped,
302 });
303
304 const BatchedSelectUnmapped = SelectUnmapped.vmap(.{
305 .axis_size = 8,
306 .in_axes = SelectUnmapped.inAxes(.{ .x = batch.mapped(0) }),
307 });
308
309 const SelectAfterMul = Program(.{
310 .name = "dsl_select_after_mul",
311 .parameters = .{
312 .x = types.spec(.f32, .{ .lane = 4 }),
313 .y = types.spec(.f32, .{ .lane = 4 }),
314 .z = types.spec(.f32, .{ .lane = 4 }),
315 },
316 .body = selectAfterMul,
317 });
318
319 const BatchedSelectUnmappedWithGeneratedRewrite = SelectUnmapped.vmapWith(
320 .{
321 .axis_size = 8,
322 .in_axes = SelectUnmapped.inAxes(.{ .x = batch.mapped(0) }),
323 },
324 .{ .batch = interpret_mod.bind(GeneratedBroadcastAsZero{}) },
325 );
326
327 const SelectAfterMulAnalysis = SelectAfterMul.analyze(MulPresence{});
328
329 const BatchedSelectAfterMulWithAnalysisRewrite = SelectAfterMulAnalysis.vmapWith(
330 .{
331 .axis_size = 8,
332 .in_axes = SelectAfterMul.inAxes(.{ .x = batch.mapped(0), .y = batch.mapped(0) }),
333 },
334 batchHooksFromMulPresence,
335 );
336
337 const BatchedElementwiseForward = ElementwiseForward.vmap(.{
338 .axis_size = 8,
339 .in_axes = ElementwiseForward.inAxes(.{ .x = batch.mapped(0), .y = batch.mapped(0) }),
340 });
341
342 const AddBeforeMulLoss = Program(.{
343 .name = "dsl_add_before_mul_loss",
344 .parameters = .{
345 .x = types.spec(.f32, .{ .lane = 4 }),
346 .y = types.spec(.f32, .{ .lane = 4 }),
347 .z = types.spec(.f32, .{ .lane = 4 }),
348 },
349 .body = addBeforeMulLoss,
350 });
351
352 const AddCounter = struct {
353 count: *usize,
354
355 pub const Value = usize;
356 pub const Result = usize;
357
358 pub fn bind(self: *@This(), op: *const program_mod.Operation, args: []const usize) !usize {
359 _ = args;
360 switch (op.kind) {
361 .binary => |binary| {
362 if (binary.op == .add) self.count.* += 1;
363 },
364 else => {},
365 }
366 return self.count.*;
367 }
368
369 pub fn finish(self: *@This(), outputs: []const usize) !Result {
370 _ = outputs;
371 return self.count.*;
372 }
373 };
374
375 const BinaryCounts = struct {
376 add_count: usize = 0,
377 mul_count: usize = 0,
378 };
379
380 const BinaryCounter = struct {
381 counts: BinaryCounts = .{},
382
383 pub const Value: type = BinaryCounts;
384 pub const Result: type = BinaryCounts;
385
386 pub fn bind(self: *@This(), op: *const program_mod.Operation, args: []const BinaryCounts) !BinaryCounts {
387 _ = args;
388 switch (op.kind) {
389 .binary => |binary| switch (binary.op) {
390 .add => self.counts.add_count += 1,
391 .mul => self.counts.mul_count += 1,
392 else => {},
393 },
394 else => {},
395 }
396 return self.counts;
397 }
398
399 pub fn finish(self: *@This(), outputs: []const BinaryCounts) !Result {
400 _ = outputs;
401 return self.counts;
402 }
403 };
404
405 const CountedGraph = struct {
406 graph: program_mod.Program,
407 add_count: usize,
408 broadcast_in_dim_count: usize,
409 finish_count: usize,
410 mul_count: usize,
411
412 pub fn deinit(self: *@This()) void {
413 self.graph.deinit();
414 self.* = undefined;
415 }
416 };
417
418 const GraphLayerCounter = struct {
419 add_count: usize = 0,
420 broadcast_in_dim_count: usize = 0,
421 finish_count: usize = 0,
422 mul_count: usize = 0,
423
424 pub const Result: type = CountedGraph;
425
426 pub fn add(self: *@This(), ctx: anytype) !trace.Value {
427 self.add_count += 1;
428 return ctx.default();
429 }
430
431 pub fn broadcastInDim(self: *@This(), ctx: anytype) !trace.Value {
432 self.broadcast_in_dim_count += 1;
433 return ctx.default();
434 }
435
436 pub fn mul(self: *@This(), ctx: anytype) !trace.Value {
437 self.mul_count += 1;
438 return ctx.default();
439 }
440
441 pub fn finish(self: *@This(), ctx: anytype, outputs: []const trace.Value) !CountedGraph {
442 self.finish_count += 1;
443 return .{
444 .graph = try ctx.default(outputs),
445 .add_count = self.add_count,
446 .broadcast_in_dim_count = self.broadcast_in_dim_count,
447 .finish_count = self.finish_count,
448 .mul_count = self.mul_count,
449 };
450 }
451 };
452
453 const MulPresence = struct {
454 has_mul: bool = false,
455
456 pub const Value = bool;
457 pub const Result = struct {
458 has_mul: bool,
459 };
460
461 pub fn default(self: *@This(), ctx: anytype) !bool {
462 _ = ctx;
463 return self.has_mul;
464 }
465
466 pub fn mul(self: *@This(), ctx: anytype) !bool {
467 _ = ctx;
468 self.has_mul = true;
469 return self.has_mul;
470 }
471
472 pub fn finish(self: *@This(), outputs: []const bool) !Result {
473 _ = outputs;
474 return .{ .has_mul = self.has_mul };
475 }
476 };
477
478 const DropAddWhenMulPresent = struct {
479 enabled: bool,
480
481 pub fn add(self: *@This(), ctx: *transform.Context) !?trace.Value {
482 if (self.enabled) return ctx.arg(0);
483 return null;
484 }
485 };
486
487 fn dropAddFromMulPresence(analysis: *const MulPresence.Result) DropAddWhenMulPresent {
488 return .{ .enabled = analysis.has_mul };
489 }
490
491 const AddBeforeMulAnalysis = AddBeforeMulLoss.analyze(MulPresence{});
492
493 const AddDroppedAfterMulAnalysis = AddBeforeMulAnalysis.rewrite(dropAddFromMulPresence);
494
495 const BatchedAddDroppedAfterMulAnalysis = AddDroppedAfterMulAnalysis.vmap(.{
496 .axis_size = 8,
497 .in_axes = AddDroppedAfterMulAnalysis.inAxes(.{
498 .x = batch.mapped(0),
499 .y = batch.mapped(0),
500 .z = batch.mapped(0),
501 }),
502 });
503
504 test "tensor Program keeps positional parameters available" {
505 const PositionalLoss = Program(.{
506 .name = "dsl_positional_loss",
507 .parameters = &.{ types.spec(.f32, .{ .lane = 4 }), types.spec(.f32, .{ .lane = 4 }) },
508 .body = positionalLoss,
509 });
510
511 var built = try PositionalLoss.build(std.testing.allocator);
512 defer built.deinit();
513
514 try std.testing.expectEqual(@as(usize, 2), built.parameters.len);
515 try std.testing.expectEqual(@as(usize, 1), built.outputs.len);
516 try types.expectExtents(&.{}, built.typeOf(built.outputs[0]));
517 try PositionalLoss.verify(std.testing.allocator);
518 }
519
520 test "tensor Program supports array output bodies" {
521 var built = try ArrayForward.build(std.testing.allocator);
522 defer built.deinit();
523
524 try std.testing.expectEqual(@as(usize, 2), built.parameters.len);
525 try std.testing.expectEqual(@as(usize, 2), built.outputs.len);
526 try types.expectExtents(&.{4}, built.typeOf(built.outputs[0]));
527 try types.expectExtents(&.{}, built.typeOf(built.outputs[1]));
528 try ArrayForward.verify(std.testing.allocator);
529 }
530
531 test "tensor Program interprets source and derived programs through user semantics" {
532 var add_count: usize = 0;
533 const counted = try ArrayForward.interpret(std.testing.allocator, AddCounter{ .count = &add_count });
534
535 try std.testing.expectEqual(@as(usize, 1), counted);
536
537 var copied_add_count: usize = 0;
538 const counted_copy = try ArrayForwardCopy.interpret(std.testing.allocator, AddCounter{ .count = &copied_add_count });
539
540 try std.testing.expectEqual(@as(usize, 1), counted_copy);
541 }
542
543 test "tensor Program attaches graph semantics specs" {
544 var copied = try ArrayForward.interpret(std.testing.allocator, interpret_mod.bind(GraphLayerCounter{}));
545 defer copied.deinit();
546
547 try std.testing.expectEqual(@as(usize, 1), copied.add_count);
548 try std.testing.expectEqual(@as(usize, 1), copied.finish_count);
549 try std.testing.expectEqual(@as(usize, 2), copied.graph.parameters.len);
550 try std.testing.expectEqual(@as(usize, 2), copied.graph.outputs.len);
551 try types.expectExtents(&.{4}, copied.graph.typeOf(copied.graph.outputs[0]));
552 try types.expectExtents(&.{}, copied.graph.typeOf(copied.graph.outputs[1]));
553
554 var copied_rewrite = try ArrayForwardCopy.interpret(std.testing.allocator, interpret_mod.bind(GraphLayerCounter{}));
555 defer copied_rewrite.deinit();
556
557 try std.testing.expectEqual(@as(usize, 1), copied_rewrite.add_count);
558 try std.testing.expectEqual(@as(usize, 1), copied_rewrite.finish_count);
559 try std.testing.expectEqual(@as(usize, 2), copied_rewrite.graph.parameters.len);
560 try std.testing.expectEqual(@as(usize, 2), copied_rewrite.graph.outputs.len);
561 try types.expectExtents(&.{4}, copied_rewrite.graph.typeOf(copied_rewrite.graph.outputs[0]));
562 try types.expectExtents(&.{}, copied_rewrite.graph.typeOf(copied_rewrite.graph.outputs[1]));
563
564 var copied_jvp = try ElementwiseJvp.interpret(std.testing.allocator, interpret_mod.bind(GraphLayerCounter{}));
565 defer copied_jvp.deinit();
566
567 try std.testing.expectEqual(@as(usize, 1), copied_jvp.add_count);
568 try std.testing.expectEqual(@as(usize, 1), copied_jvp.finish_count);
569 try std.testing.expectEqual(@as(usize, 3), copied_jvp.mul_count);
570 try std.testing.expectEqual(@as(usize, 4), copied_jvp.graph.parameters.len);
571 try std.testing.expectEqual(@as(usize, 2), copied_jvp.graph.outputs.len);
572 try types.expectExtents(&.{}, copied_jvp.graph.typeOf(copied_jvp.graph.outputs[0]));
573 try types.expectExtents(&.{}, copied_jvp.graph.typeOf(copied_jvp.graph.outputs[1]));
574
575 var copied_jvp_with = try ElementwiseJvpWithGeneratedRewrite.interpret(std.testing.allocator, interpret_mod.bind(GraphLayerCounter{}));
576 defer copied_jvp_with.deinit();
577
578 try std.testing.expectEqual(@as(usize, 3), copied_jvp_with.add_count);
579 try std.testing.expectEqual(@as(usize, 1), copied_jvp_with.finish_count);
580 try std.testing.expectEqual(@as(usize, 1), copied_jvp_with.mul_count);
581 try std.testing.expectEqual(@as(usize, 4), copied_jvp_with.graph.parameters.len);
582 try std.testing.expectEqual(@as(usize, 2), copied_jvp_with.graph.outputs.len);
583 try types.expectExtents(&.{}, copied_jvp_with.graph.typeOf(copied_jvp_with.graph.outputs[0]));
584 try types.expectExtents(&.{}, copied_jvp_with.graph.typeOf(copied_jvp_with.graph.outputs[1]));
585
586 var copied_analysis_jvp = try ElementwiseJvpWithAnalysisRewrite.interpret(std.testing.allocator, interpret_mod.bind(GraphLayerCounter{}));
587 defer copied_analysis_jvp.deinit();
588
589 try std.testing.expectEqual(@as(usize, 3), copied_analysis_jvp.add_count);
590 try std.testing.expectEqual(@as(usize, 1), copied_analysis_jvp.finish_count);
591 try std.testing.expectEqual(@as(usize, 1), copied_analysis_jvp.mul_count);
592 try std.testing.expectEqual(@as(usize, 4), copied_analysis_jvp.graph.parameters.len);
593 try std.testing.expectEqual(@as(usize, 2), copied_analysis_jvp.graph.outputs.len);
594 try types.expectExtents(&.{}, copied_analysis_jvp.graph.typeOf(copied_analysis_jvp.graph.outputs[0]));
595 try types.expectExtents(&.{}, copied_analysis_jvp.graph.typeOf(copied_analysis_jvp.graph.outputs[1]));
596
597 var copied_grad = try ElementwiseGradient.interpret(std.testing.allocator, interpret_mod.bind(GraphLayerCounter{}));
598 defer copied_grad.deinit();
599
600 try std.testing.expectEqual(@as(usize, 0), copied_grad.add_count);
601 try std.testing.expectEqual(@as(usize, 1), copied_grad.broadcast_in_dim_count);
602 try std.testing.expectEqual(@as(usize, 1), copied_grad.finish_count);
603 try std.testing.expectEqual(@as(usize, 2), copied_grad.mul_count);
604 try std.testing.expectEqual(@as(usize, 2), copied_grad.graph.parameters.len);
605 try std.testing.expectEqual(@as(usize, 2), copied_grad.graph.outputs.len);
606 try types.expectExtents(&.{4}, copied_grad.graph.typeOf(copied_grad.graph.outputs[ElementwiseGradient.out(.x)]));
607 try types.expectExtents(&.{4}, copied_grad.graph.typeOf(copied_grad.graph.outputs[ElementwiseGradient.out(.y)]));
608
609 var copied_grad_with = try ElementwiseGradientWithGeneratedRewrite.interpret(std.testing.allocator, interpret_mod.bind(GraphLayerCounter{}));
610 defer copied_grad_with.deinit();
611
612 try std.testing.expectEqual(@as(usize, 2), copied_grad_with.add_count);
613 try std.testing.expectEqual(@as(usize, 1), copied_grad_with.broadcast_in_dim_count);
614 try std.testing.expectEqual(@as(usize, 1), copied_grad_with.finish_count);
615 try std.testing.expectEqual(@as(usize, 0), copied_grad_with.mul_count);
616 try std.testing.expectEqual(@as(usize, 2), copied_grad_with.graph.parameters.len);
617 try std.testing.expectEqual(@as(usize, 2), copied_grad_with.graph.outputs.len);
618 try types.expectExtents(&.{4}, copied_grad_with.graph.typeOf(copied_grad_with.graph.outputs[ElementwiseGradientWithGeneratedRewrite.out(.x)]));
619 try types.expectExtents(&.{4}, copied_grad_with.graph.typeOf(copied_grad_with.graph.outputs[ElementwiseGradientWithGeneratedRewrite.out(.y)]));
620
621 var copied_analysis_grad = try ElementwiseGradientWithAnalysisRewrite.interpret(std.testing.allocator, interpret_mod.bind(GraphLayerCounter{}));
622 defer copied_analysis_grad.deinit();
623
624 try std.testing.expectEqual(@as(usize, 2), copied_analysis_grad.add_count);
625 try std.testing.expectEqual(@as(usize, 1), copied_analysis_grad.broadcast_in_dim_count);
626 try std.testing.expectEqual(@as(usize, 1), copied_analysis_grad.finish_count);
627 try std.testing.expectEqual(@as(usize, 0), copied_analysis_grad.mul_count);
628 try std.testing.expectEqual(@as(usize, 2), copied_analysis_grad.graph.parameters.len);
629 try std.testing.expectEqual(@as(usize, 2), copied_analysis_grad.graph.outputs.len);
630 try types.expectExtents(&.{4}, copied_analysis_grad.graph.typeOf(copied_analysis_grad.graph.outputs[ElementwiseGradientWithAnalysisRewrite.out(.x)]));
631 try types.expectExtents(&.{4}, copied_analysis_grad.graph.typeOf(copied_analysis_grad.graph.outputs[ElementwiseGradientWithAnalysisRewrite.out(.y)]));
632
633 var copied_vmap = try BatchedSelectUnmapped.interpret(std.testing.allocator, interpret_mod.bind(GraphLayerCounter{}));
634 defer copied_vmap.deinit();
635
636 try std.testing.expectEqual(@as(usize, 0), copied_vmap.add_count);
637 try std.testing.expectEqual(@as(usize, 1), copied_vmap.broadcast_in_dim_count);
638 try std.testing.expectEqual(@as(usize, 1), copied_vmap.finish_count);
639 try std.testing.expectEqual(@as(usize, 2), copied_vmap.graph.parameters.len);
640 try std.testing.expectEqual(@as(usize, 1), copied_vmap.graph.outputs.len);
641 try types.expectExtents(&.{ 8, 4 }, copied_vmap.graph.typeOf(copied_vmap.graph.parameters[0]));
642 try types.expectExtents(&.{4}, copied_vmap.graph.typeOf(copied_vmap.graph.parameters[1]));
643 try types.expectExtents(&.{ 8, 4 }, copied_vmap.graph.typeOf(copied_vmap.graph.outputs[0]));
644
645 var copied_vmap_with = try BatchedSelectUnmappedWithGeneratedRewrite.interpret(std.testing.allocator, interpret_mod.bind(GraphLayerCounter{}));
646 defer copied_vmap_with.deinit();
647
648 try std.testing.expectEqual(@as(usize, 0), copied_vmap_with.add_count);
649 try std.testing.expectEqual(@as(usize, 0), copied_vmap_with.broadcast_in_dim_count);
650 try std.testing.expectEqual(@as(usize, 1), copied_vmap_with.finish_count);
651 try std.testing.expectEqual(@as(usize, 2), copied_vmap_with.graph.parameters.len);
652 try std.testing.expectEqual(@as(usize, 1), copied_vmap_with.graph.outputs.len);
653 try types.expectExtents(&.{ 8, 4 }, copied_vmap_with.graph.typeOf(copied_vmap_with.graph.outputs[0]));
654
655 var copied_analysis_vmap = try BatchedSelectAfterMulWithAnalysisRewrite.interpret(std.testing.allocator, interpret_mod.bind(GraphLayerCounter{}));
656 defer copied_analysis_vmap.deinit();
657
658 try std.testing.expectEqual(@as(usize, 0), copied_analysis_vmap.add_count);
659 try std.testing.expectEqual(@as(usize, 0), copied_analysis_vmap.broadcast_in_dim_count);
660 try std.testing.expectEqual(@as(usize, 1), copied_analysis_vmap.finish_count);
661 try std.testing.expectEqual(@as(usize, 1), copied_analysis_vmap.mul_count);
662 try std.testing.expectEqual(@as(usize, 3), copied_analysis_vmap.graph.parameters.len);
663 try std.testing.expectEqual(@as(usize, 1), copied_analysis_vmap.graph.outputs.len);
664 try types.expectExtents(&.{ 8, 4 }, copied_analysis_vmap.graph.typeOf(copied_analysis_vmap.graph.outputs[0]));
665 }
666
667 test "tensor Program derives transforms from user analysis semantics" {
668 const analysis = try AddBeforeMulAnalysis.interpret(std.testing.allocator);
669 try std.testing.expect(analysis.has_mul);
670
671 var source = try AddBeforeMulLoss.build(std.testing.allocator);
672 defer source.deinit();
673 var analyzed = try AddDroppedAfterMulAnalysis.build(std.testing.allocator);
674 defer analyzed.deinit();
675
676 try std.testing.expect(analyzed.operationCount() < source.operationCount());
677 try std.testing.expectEqual(@as(usize, 3), analyzed.parameters.len);
678 try std.testing.expectEqual(@as(usize, 1), analyzed.outputs.len);
679 try types.expectExtents(&.{}, analyzed.typeOf(analyzed.outputs[0]));
680
681 const counts = try AddDroppedAfterMulAnalysis.interpret(std.testing.allocator, BinaryCounter{});
682 try std.testing.expectEqual(@as(usize, 0), counts.add_count);
683 try std.testing.expectEqual(@as(usize, 1), counts.mul_count);
684 try AddDroppedAfterMulAnalysis.verify(std.testing.allocator);
685
686 var copied_analysis_rewrite = try AddDroppedAfterMulAnalysis.interpret(std.testing.allocator, interpret_mod.bind(GraphLayerCounter{}));
687 defer copied_analysis_rewrite.deinit();
688 try std.testing.expectEqual(@as(usize, 0), copied_analysis_rewrite.add_count);
689 try std.testing.expectEqual(@as(usize, 1), copied_analysis_rewrite.finish_count);
690 try std.testing.expectEqual(@as(usize, 3), copied_analysis_rewrite.graph.parameters.len);
691 try std.testing.expectEqual(@as(usize, 1), copied_analysis_rewrite.graph.outputs.len);
692 try types.expectExtents(&.{}, copied_analysis_rewrite.graph.typeOf(copied_analysis_rewrite.graph.outputs[0]));
693
694 var batched = try BatchedAddDroppedAfterMulAnalysis.build(std.testing.allocator);
695 defer batched.deinit();
696 try std.testing.expectEqual(@as(usize, 3), batched.parameters.len);
697 try std.testing.expectEqual(@as(usize, 1), batched.outputs.len);
698 try types.expectExtents(&.{8}, batched.typeOf(batched.outputs[0]));
699 try BatchedAddDroppedAfterMulAnalysis.verify(std.testing.allocator);
700 }
701
702 test "tensor Program lowers source and derived programs directly" {
703 const source_module = try ArrayForward.lower(std.testing.allocator);
704 defer source_module.deinit();
705 try source_module.verify();
706
707 const derived_module = try BatchedElementwiseGradient.lower(std.testing.allocator);
708 defer derived_module.deinit();
709 try derived_module.verify();
710
711 const explicit_source_module = try ArrayForward.interpret(std.testing.allocator, tensor.lower.module(std.testing.allocator));
712 defer explicit_source_module.deinit();
713 try explicit_source_module.verify();
714
715 const explicit_jvp_module = try ElementwiseJvpWithGeneratedRewrite.interpret(std.testing.allocator, tensor.lower.module(std.testing.allocator));
716 defer explicit_jvp_module.deinit();
717 try explicit_jvp_module.verify();
718
719 const explicit_grad_module = try ElementwiseGradientWithAnalysisRewrite.interpret(std.testing.allocator, tensor.lower.module(std.testing.allocator));
720 defer explicit_grad_module.deinit();
721 try explicit_grad_module.verify();
722 }
723
724 fn expectPreparedGeneratedKernels(prepared: *const tensor.BackendPreparedJob) !void {
725 const kernel_count = try prepared.generatedKernelCount();
726 try std.testing.expect(kernel_count > 0);
727 var launch_count: usize = 0;
728 for (0..kernel_count) |kernel_index| {
729 const summary = try prepared.generatedKernelSummary(kernel_index);
730 const program = try prepared.generatedKernelProgram(kernel_index);
731 try std.testing.expect(summary.entry_name.len > 0);
732 if (summary.launch_geometry != null) launch_count += 1;
733 const launch = try program.launch();
734 switch (summary.schedule.kind) {
735 .flat => {
736 try std.testing.expectEqual(@as(u32, 64), summary.schedule.threads.x);
737 try std.testing.expect(launch.block[0] > 0);
738 try std.testing.expect(launch.block[0] <= summary.schedule.threads.x);
739 },
740 .matrix => {
741 try std.testing.expectEqual(@as(u32, 16), summary.schedule.threads.x);
742 try std.testing.expectEqual(@as(u32, 16), summary.schedule.threads.y);
743 try std.testing.expect(launch.block[0] > 0);
744 try std.testing.expect(launch.block[0] <= summary.schedule.threads.x);
745 try std.testing.expect(launch.block[1] > 0);
746 try std.testing.expect(launch.block[1] <= summary.schedule.threads.y);
747 },
748 }
749 }
750 try std.testing.expect(launch_count > 0);
751 try std.testing.expectError(error.InvalidIndex, prepared.generatedKernelSummary(kernel_count));
752 }
753
754 fn expectLoadedFragmentMatchesPrepared(
755 allocator: std.mem.Allocator,
756 prepared: *const tensor.BackendPreparedJob,
757 generated_summaries: *const tensor.GeneratedKernelSummaries,
758 artifact_summaries: *const tensor.ArtifactKernelSummaries,
759 fragment: *tensor.LoadedFragment,
760 artifact_format: gpu.ArtifactFormat,
761 ) !void {
762 const prepared_kernel_count = try prepared.generatedKernelCount();
763 try std.testing.expectEqual(prepared_kernel_count, fragment.kernelCount());
764 var executable_summaries = try fragment.copyKernelSummaries(allocator);
765 defer executable_summaries.deinit();
766 try std.testing.expectEqual(prepared_kernel_count, executable_summaries.len());
767 for (generated_summaries.items) |generated_snapshot| {
768 const summary = try executable_summaries.summaryForWork(generated_snapshot.work_item_id);
769 const executable_summary = try fragment.kernelSummaryForWork(generated_snapshot.work_item_id);
770 const generated = try prepared.generatedKernelSummaryForWork(summary.work_item_id);
771 const copied_generated = try generated_summaries.summaryForWork(summary.work_item_id);
772 const artifact_summary = try artifact_summaries.summaryForWork(summary.work_item_id);
773 const generated_program = try prepared.generatedKernelProgramForWork(summary.work_item_id);
774 try std.testing.expectEqual(tensor.ArtifactKernelSource.tensor, summary.source);
775 try std.testing.expectEqualStrings(executable_summary.entry_name, summary.entry_name);
776 try std.testing.expect(summary.entry_name.len > 0);
777 try std.testing.expectEqualStrings(generated.entry_name, summary.entry_name);
778 try std.testing.expectEqualStrings(copied_generated.entry_name, generated.entry_name);
779 try std.testing.expectEqualStrings(artifact_summary.entry_name, summary.entry_name);
780 try std.testing.expectEqual(artifact_summary.compile_argument_count, summary.compile_argument_count);
781 try std.testing.expectEqual(generated.argument_count, summary.compile_argument_count);
782 try std.testing.expectEqual(generated.body_fingerprint, try generated_program.bodyFingerprint(allocator));
783 try std.testing.expectEqual(artifact_format, summary.artifact_format);
784 }
785 }
786
787 test "tensor Program prepares generated kernel programs directly" {
788 var source_prepared = try ElementwiseForward.prepare(std.testing.allocator);
789 defer source_prepared.deinit();
790 try expectPreparedGeneratedKernels(&source_prepared);
791
792 var derived_prepared = try BatchedElementwiseGradient.prepareWith(std.testing.allocator, .{});
793 defer derived_prepared.deinit();
794 try expectPreparedGeneratedKernels(&derived_prepared);
795 }
796
797 test "tensor Program creates executable fragment through BackendHandle" {
798 const allocator = std.testing.allocator;
799 var state = gpu.recording.BackendState{
800 .allocator = allocator,
801 .kind = .vulkan,
802 .format = .vulkan_spirv,
803 };
804
805 const program_options = tensor.FragmentCompilerOptions{ .authored_kernel_diagnostic_id = "tensor/program-executable" };
806 const program_compiled = try BatchedElementwiseGradient.compileFragment(
807 allocator,
808 state.handle(),
809 program_options,
810 );
811 var fragment = try accy.executable.loadFragment(allocator, state.handle(), program_compiled, program_options);
812 defer fragment.deinit();
813
814 try std.testing.expect(state.create_count > 0);
815 try std.testing.expect(state.load_count > 0);
816 var fragment_summaries = try fragment.copyKernelSummaries(allocator);
817 defer fragment_summaries.deinit();
818 try std.testing.expect(fragment_summaries.len() > 0);
819 try std.testing.expectEqual(fragment.kernelCount(), fragment_summaries.len());
820 for (fragment_summaries.items, 0..) |summary, index| {
821 const indexed_summary = try fragment_summaries.summary(index);
822 const executable_summary = try fragment.kernelSummaryForWork(summary.work_item_id);
823 try std.testing.expectEqualStrings(indexed_summary.entry_name, summary.entry_name);
824 try std.testing.expectEqualStrings(executable_summary.entry_name, summary.entry_name);
825 try std.testing.expectEqual(tensor.ArtifactKernelSource.tensor, summary.source);
826 try std.testing.expect(summary.entry_name.len > 0);
827 try std.testing.expectEqual(gpu.ArtifactFormat.vulkan_spirv, summary.artifact_format);
828 }
829
830 var direct_artifact = try BatchedElementwiseGradient.createArtifactJob(
831 allocator,
832 state.handle(),
833 .{ .authored_kernel_diagnostic_id = "tensor/program-direct-artifact" },
834 );
835 defer direct_artifact.deinit();
836 var direct_artifact_summaries = try direct_artifact.copyKernelSummaries(allocator);
837 defer direct_artifact_summaries.deinit();
838 try std.testing.expect(direct_artifact_summaries.len() > 0);
839 try std.testing.expectEqual(direct_artifact.kernelCount(), direct_artifact_summaries.len());
840 for (direct_artifact_summaries.items, 0..) |summary, index| {
841 const indexed_summary = try direct_artifact_summaries.summary(index);
842 try std.testing.expectEqualStrings(indexed_summary.entry_name, summary.entry_name);
843 try std.testing.expectEqual(tensor.ArtifactKernelSource.tensor, summary.source);
844 try std.testing.expect(summary.entry_name.len > 0);
845 try std.testing.expectEqual(gpu.ArtifactFormat.vulkan_spirv, summary.artifact_format);
846 }
847
848 const direct_options = tensor.FragmentCompilerOptions{ .authored_kernel_diagnostic_id = "tensor/program-direct-artifact" };
849 const direct_compiled = try tensor.compileFragmentFromArtifactJob(allocator, direct_artifact);
850 var direct_artifact_fragment = try accy.executable.loadFragment(allocator, state.handle(), direct_compiled, direct_options);
851 defer direct_artifact_fragment.deinit();
852 try std.testing.expectEqual(direct_artifact.kernelCount(), direct_artifact_fragment.kernelCount());
853 var direct_fragment_summaries = try direct_artifact_fragment.copyKernelSummaries(allocator);
854 defer direct_fragment_summaries.deinit();
855 try std.testing.expectEqual(direct_artifact_summaries.len(), direct_fragment_summaries.len());
856 for (direct_artifact_summaries.items) |artifact_summary| {
857 const executable_summary = try direct_fragment_summaries.summaryForWork(artifact_summary.work_item_id);
858 try std.testing.expectEqualStrings(artifact_summary.entry_name, executable_summary.entry_name);
859 try std.testing.expectEqual(artifact_summary.compile_argument_count, executable_summary.compile_argument_count);
860 }
861
862 var staged_state = gpu.recording.BackendState{
863 .allocator = allocator,
864 .kind = .vulkan,
865 .format = .vulkan_spirv,
866 };
867 var prepared = try BatchedElementwiseGradient.prepareFragment(
868 allocator,
869 staged_state.handle(),
870 .{ .authored_kernel_diagnostic_id = "tensor/program-staged-executable" },
871 );
872 defer prepared.deinit();
873 try expectPreparedGeneratedKernels(&prepared);
874 const prepared_kernel_count = try prepared.generatedKernelCount();
875 var generated_summaries = try prepared.copyGeneratedKernelSummaries(allocator);
876 defer generated_summaries.deinit();
877 try std.testing.expectEqual(prepared_kernel_count, generated_summaries.len());
878 for (generated_summaries.items, 0..) |summary, index| {
879 const indexed_summary = try prepared.generatedKernelSummary(index);
880 try std.testing.expectEqualStrings(indexed_summary.entry_name, summary.entry_name);
881 const work_summary = try prepared.generatedKernelSummaryForWork(summary.work_item_id);
882 try std.testing.expectEqualStrings(summary.entry_name, work_summary.entry_name);
883 }
884
885 var artifact_module = try tensor.createArtifactJobFromPreparedJob(
886 allocator,
887 staged_state.handle(),
888 &prepared,
889 .{ .authored_kernel_diagnostic_id = "tensor/program-staged-executable" },
890 );
891 defer artifact_module.deinit();
892 const artifact_fingerprint = artifact_module.fingerprint();
893 try std.testing.expectEqual(prepared_kernel_count, artifact_module.kernelCount());
894 var artifact_summaries = try artifact_module.copyKernelSummaries(allocator);
895 defer artifact_summaries.deinit();
896 try std.testing.expectEqual(prepared_kernel_count, artifact_summaries.len());
897 for (generated_summaries.items) |generated| {
898 const summary = try artifact_summaries.summaryForWork(generated.work_item_id);
899 try std.testing.expectEqual(tensor.ArtifactKernelSource.tensor, summary.source);
900 try std.testing.expectEqualStrings(generated.entry_name, summary.entry_name);
901 }
902 try std.testing.expectError(error.InvalidArtifact, artifact_summaries.summaryForWork(std.math.maxInt(usize)));
903
904 const staged_options = tensor.FragmentCompilerOptions{ .authored_kernel_diagnostic_id = "tensor/program-staged-executable" };
905 const staged_compiled = try tensor.compileFragmentFromArtifactJob(allocator, artifact_module);
906 var staged_fragment = try accy.executable.loadFragment(allocator, staged_state.handle(), staged_compiled, staged_options);
907 defer staged_fragment.deinit();
908
909 try std.testing.expect(staged_state.create_count > 0);
910 try std.testing.expect(staged_state.load_count > 0);
911 try std.testing.expectEqual(prepared_kernel_count, try prepared.generatedKernelCount());
912 try std.testing.expectEqual(artifact_fingerprint, artifact_module.fingerprint());
913 try expectLoadedFragmentMatchesPrepared(allocator, &prepared, &generated_summaries, &artifact_summaries, staged_fragment, .vulkan_spirv);
914
915 const repeat_options = tensor.FragmentCompilerOptions{ .authored_kernel_diagnostic_id = "tensor/program-staged-executable-repeat" };
916 const repeat_compiled = try tensor.compileFragmentFromArtifactJob(allocator, artifact_module);
917 var second_staged_fragment = try accy.executable.loadFragment(allocator, staged_state.handle(), repeat_compiled, repeat_options);
918 defer second_staged_fragment.deinit();
919
920 try std.testing.expectEqual(prepared_kernel_count, try prepared.generatedKernelCount());
921 try std.testing.expectEqual(artifact_fingerprint, artifact_module.fingerprint());
922 try expectLoadedFragmentMatchesPrepared(allocator, &prepared, &generated_summaries, &artifact_summaries, second_staged_fragment, .vulkan_spirv);
923 }
924
925 test "tensor Program supports named multi-output bodies" {
926 var built = try ElementwiseForward.build(std.testing.allocator);
927 defer built.deinit();
928
929 try std.testing.expectEqual(@as(usize, 2), built.parameters.len);
930 try std.testing.expectEqual(@as(usize, 2), built.outputs.len);
931 try std.testing.expectEqual(@as(usize, 0), ElementwiseForward.out(.product));
932 try std.testing.expectEqual(@as(usize, 1), ElementwiseForward.out(.total));
933 try types.expectExtents(&.{4}, built.typeOf(built.outputs[ElementwiseForward.out(.product)]));
934 try types.expectExtents(&.{}, built.typeOf(built.outputs[ElementwiseForward.out(.total)]));
935 try ElementwiseForward.verify(std.testing.allocator);
936 }
937
938 test "tensor Program vmaps named multi-output bodies" {
939 var built = try BatchedElementwiseForward.build(std.testing.allocator);
940 defer built.deinit();
941
942 try std.testing.expectEqual(@as(usize, 2), built.parameters.len);
943 try std.testing.expectEqual(@as(usize, 2), built.outputs.len);
944 try types.expectExtents(&.{ 8, 4 }, built.typeOf(built.outputs[BatchedElementwiseForward.out(.product)]));
945 try types.expectExtents(&.{8}, built.typeOf(built.outputs[BatchedElementwiseForward.out(.total)]));
946 try BatchedElementwiseForward.verify(std.testing.allocator);
947 }
948
949 test "tensor Program names jvp outputs by primal and tangent source output" {
950 var built = try ElementwiseForwardJvp.build(std.testing.allocator);
951 defer built.deinit();
952
953 try std.testing.expectEqual(@as(usize, 4), built.parameters.len);
954 try std.testing.expectEqual(@as(usize, 4), built.outputs.len);
955 try std.testing.expectEqual(@as(usize, 0), ElementwiseForwardJvp.out(.{ .primals = .product }));
956 try std.testing.expectEqual(@as(usize, 1), ElementwiseForwardJvp.out(.{ .primals = .total }));
957 try std.testing.expectEqual(@as(usize, 2), ElementwiseForwardJvp.out(.{ .tangents = .product }));
958 try std.testing.expectEqual(@as(usize, 3), ElementwiseForwardJvp.out(.{ .tangents = .total }));
959 try types.expectExtents(&.{4}, built.typeOf(built.outputs[ElementwiseForwardJvp.out(.{ .primals = .product })]));
960 try types.expectExtents(&.{}, built.typeOf(built.outputs[ElementwiseForwardJvp.out(.{ .primals = .total })]));
961 try types.expectExtents(&.{4}, built.typeOf(built.outputs[ElementwiseForwardJvp.out(.{ .tangents = .product })]));
962 try types.expectExtents(&.{}, built.typeOf(built.outputs[ElementwiseForwardJvp.out(.{ .tangents = .total })]));
963 try ElementwiseForwardJvp.verify(std.testing.allocator);
964 }
965
966 test "tensor Program composes tracing jvp vmap and lowering" {
967 var batched = try BatchedElementwiseJvp.build(std.testing.allocator);
968 defer batched.deinit();
969
970 try std.testing.expectEqual(@as(usize, 4), batched.parameters.len);
971 try std.testing.expectEqual(@as(usize, 2), batched.outputs.len);
972 try std.testing.expectEqual(@as(usize, 0), BatchedElementwiseJvp.out(.primal));
973 try std.testing.expectEqual(@as(usize, 1), BatchedElementwiseJvp.out(.tangent));
974 try types.expectExtents(&.{8}, batched.typeOf(batched.outputs[BatchedElementwiseJvp.out(.primal)]));
975 try types.expectExtents(&.{8}, batched.typeOf(batched.outputs[BatchedElementwiseJvp.out(.tangent)]));
976 try BatchedElementwiseJvp.verify(std.testing.allocator);
977 }
978
979 test "tensor Program jvpWith composes user semantics inside jvp" {
980 const ordinary_counts = try ElementwiseJvp.interpret(std.testing.allocator, BinaryCounter{});
981 const rewritten_counts = try ElementwiseJvpWithGeneratedRewrite.interpret(std.testing.allocator, BinaryCounter{});
982
983 var rewritten = try ElementwiseJvpWithGeneratedRewrite.build(std.testing.allocator);
984 defer rewritten.deinit();
985
986 try std.testing.expect(rewritten_counts.mul_count < ordinary_counts.mul_count);
987 try std.testing.expect(rewritten_counts.add_count > ordinary_counts.add_count);
988 try std.testing.expectEqual(@as(usize, 4), rewritten.parameters.len);
989 try std.testing.expectEqual(@as(usize, 2), rewritten.outputs.len);
990 try std.testing.expectEqual(@as(usize, 0), ElementwiseJvpWithGeneratedRewrite.out(.primal));
991 try std.testing.expectEqual(@as(usize, 1), ElementwiseJvpWithGeneratedRewrite.out(.tangent));
992 try types.expectExtents(&.{}, rewritten.typeOf(rewritten.outputs[ElementwiseJvpWithGeneratedRewrite.out(.primal)]));
993 try types.expectExtents(&.{}, rewritten.typeOf(rewritten.outputs[ElementwiseJvpWithGeneratedRewrite.out(.tangent)]));
994 try ElementwiseJvpWithGeneratedRewrite.verify(std.testing.allocator);
995
996 var batched = try BatchedElementwiseJvpWithGeneratedRewrite.build(std.testing.allocator);
997 defer batched.deinit();
998
999 try std.testing.expectEqual(@as(usize, 4), batched.parameters.len);
1000 try std.testing.expectEqual(@as(usize, 2), batched.outputs.len);
1001 try types.expectExtents(&.{8}, batched.typeOf(batched.outputs[BatchedElementwiseJvpWithGeneratedRewrite.out(.primal)]));
1002 try types.expectExtents(&.{8}, batched.typeOf(batched.outputs[BatchedElementwiseJvpWithGeneratedRewrite.out(.tangent)]));
1003 try BatchedElementwiseJvpWithGeneratedRewrite.verify(std.testing.allocator);
1004 }
1005
1006 test "tensor Program jvpWith accepts custom-call jvp hooks" {
1007 var built = try CustomDoubleJvp.build(std.testing.allocator);
1008 defer built.deinit();
1009
1010 try std.testing.expectEqual(@as(usize, 2), built.parameters.len);
1011 try std.testing.expectEqual(@as(usize, 2), built.outputs.len);
1012 try std.testing.expectEqual(@as(usize, 0), CustomDoubleJvp.out(.primal));
1013 try std.testing.expectEqual(@as(usize, 1), CustomDoubleJvp.out(.tangent));
1014
1015 var custom_calls: usize = 0;
1016 for (built.operations) |op| {
1017 switch (op.kind) {
1018 .custom_call => |custom| {
1019 try std.testing.expectEqualStrings(dsl_custom_double_target, custom.target);
1020 custom_calls += 1;
1021 },
1022 else => {},
1023 }
1024 }
1025 try std.testing.expectEqual(@as(usize, 2), custom_calls);
1026 }
1027
1028 test "tensor Program vmapWith composes user semantics inside batch" {
1029 var rewritten = try BatchedSelectUnmappedWithGeneratedRewrite.build(std.testing.allocator);
1030 defer rewritten.deinit();
1031
1032 try std.testing.expectEqual(@as(usize, 2), rewritten.parameters.len);
1033 try std.testing.expectEqual(@as(usize, 1), rewritten.outputs.len);
1034 try std.testing.expect(rewritten.isZeroConstant(rewritten.outputs[0]));
1035 try types.expectExtents(&.{ 8, 4 }, rewritten.typeOf(rewritten.outputs[0]));
1036 try BatchedSelectUnmappedWithGeneratedRewrite.verify(std.testing.allocator);
1037 }
1038
1039 test "tensor Program derives jvp grad and vmap phase hooks from analysis semantics" {
1040 const ordinary_jvp_counts = try ElementwiseJvp.interpret(std.testing.allocator, BinaryCounter{});
1041 const analyzed_jvp_counts = try ElementwiseJvpWithAnalysisRewrite.interpret(std.testing.allocator, BinaryCounter{});
1042
1043 var analyzed_jvp = try ElementwiseJvpWithAnalysisRewrite.build(std.testing.allocator);
1044 defer analyzed_jvp.deinit();
1045
1046 try std.testing.expect(analyzed_jvp_counts.mul_count < ordinary_jvp_counts.mul_count);
1047 try std.testing.expect(analyzed_jvp_counts.add_count > ordinary_jvp_counts.add_count);
1048 try std.testing.expectEqual(@as(usize, 4), analyzed_jvp.parameters.len);
1049 try std.testing.expectEqual(@as(usize, 2), analyzed_jvp.outputs.len);
1050 try ElementwiseJvpWithAnalysisRewrite.verify(std.testing.allocator);
1051
1052 var batched_jvp = try BatchedElementwiseJvpWithAnalysisRewrite.build(std.testing.allocator);
1053 defer batched_jvp.deinit();
1054
1055 try std.testing.expectEqual(@as(usize, 4), batched_jvp.parameters.len);
1056 try std.testing.expectEqual(@as(usize, 2), batched_jvp.outputs.len);
1057 try types.expectExtents(&.{8}, batched_jvp.typeOf(batched_jvp.outputs[BatchedElementwiseJvpWithAnalysisRewrite.out(.primal)]));
1058 try types.expectExtents(&.{8}, batched_jvp.typeOf(batched_jvp.outputs[BatchedElementwiseJvpWithAnalysisRewrite.out(.tangent)]));
1059 try BatchedElementwiseJvpWithAnalysisRewrite.verify(std.testing.allocator);
1060
1061 const ordinary_gradient_counts = try ElementwiseGradient.interpret(std.testing.allocator, BinaryCounter{});
1062 const analyzed_gradient_counts = try ElementwiseGradientWithAnalysisRewrite.interpret(std.testing.allocator, BinaryCounter{});
1063
1064 var analyzed_gradient = try ElementwiseGradientWithAnalysisRewrite.build(std.testing.allocator);
1065 defer analyzed_gradient.deinit();
1066
1067 try std.testing.expect(analyzed_gradient_counts.mul_count < ordinary_gradient_counts.mul_count);
1068 try std.testing.expectEqual(@as(usize, 2), analyzed_gradient.parameters.len);
1069 try std.testing.expectEqual(@as(usize, 2), analyzed_gradient.outputs.len);
1070 try ElementwiseGradientWithAnalysisRewrite.verify(std.testing.allocator);
1071
1072 var batched_gradient = try BatchedElementwiseGradientWithAnalysisRewrite.build(std.testing.allocator);
1073 defer batched_gradient.deinit();
1074
1075 try std.testing.expectEqual(@as(usize, 2), batched_gradient.parameters.len);
1076 try std.testing.expectEqual(@as(usize, 2), batched_gradient.outputs.len);
1077 try types.expectExtents(&.{ 8, 4 }, batched_gradient.typeOf(batched_gradient.outputs[BatchedElementwiseGradientWithAnalysisRewrite.out(.x)]));
1078 try types.expectExtents(&.{ 8, 4 }, batched_gradient.typeOf(batched_gradient.outputs[BatchedElementwiseGradientWithAnalysisRewrite.out(.y)]));
1079 try BatchedElementwiseGradientWithAnalysisRewrite.verify(std.testing.allocator);
1080
1081 var analyzed_batch = try BatchedSelectAfterMulWithAnalysisRewrite.build(std.testing.allocator);
1082 defer analyzed_batch.deinit();
1083
1084 try std.testing.expectEqual(@as(usize, 3), analyzed_batch.parameters.len);
1085 try std.testing.expectEqual(@as(usize, 1), analyzed_batch.outputs.len);
1086 try std.testing.expect(analyzed_batch.isZeroConstant(analyzed_batch.outputs[0]));
1087 try types.expectExtents(&.{ 8, 4 }, analyzed_batch.typeOf(analyzed_batch.outputs[0]));
1088 try BatchedSelectAfterMulWithAnalysisRewrite.verify(std.testing.allocator);
1089 }
1090
1091 test "tensor Program composes tracing grad vmap and lowering" {
1092 var batched = try BatchedElementwiseGradient.build(std.testing.allocator);
1093 defer batched.deinit();
1094
1095 try std.testing.expectEqual(@as(usize, 2), batched.parameters.len);
1096 try std.testing.expectEqual(@as(usize, 2), batched.outputs.len);
1097 try std.testing.expectEqual(@as(usize, 0), BatchedElementwiseGradient.out(.x));
1098 try std.testing.expectEqual(@as(usize, 1), BatchedElementwiseGradient.out(.y));
1099 try types.expectExtents(&.{ 8, 4 }, batched.typeOf(batched.outputs[BatchedElementwiseGradient.out(.x)]));
1100 try types.expectExtents(&.{ 8, 4 }, batched.typeOf(batched.outputs[BatchedElementwiseGradient.out(.y)]));
1101 try BatchedElementwiseGradient.verify(std.testing.allocator);
1102 }
1103
1104 test "tensor Program gradWith composes user semantics inside grad" {
1105 const ordinary_counts = try ElementwiseGradient.interpret(std.testing.allocator, BinaryCounter{});
1106 const rewritten_counts = try ElementwiseGradientWithGeneratedRewrite.interpret(std.testing.allocator, BinaryCounter{});
1107
1108 var rewritten = try ElementwiseGradientWithGeneratedRewrite.build(std.testing.allocator);
1109 defer rewritten.deinit();
1110
1111 try std.testing.expect(rewritten_counts.mul_count < ordinary_counts.mul_count);
1112 try std.testing.expectEqual(@as(usize, 2), rewritten.parameters.len);
1113 try std.testing.expectEqual(@as(usize, 2), rewritten.outputs.len);
1114 try std.testing.expectEqual(@as(usize, 0), ElementwiseGradientWithGeneratedRewrite.out(.x));
1115 try std.testing.expectEqual(@as(usize, 1), ElementwiseGradientWithGeneratedRewrite.out(.y));
1116 try types.expectExtents(&.{4}, rewritten.typeOf(rewritten.outputs[ElementwiseGradientWithGeneratedRewrite.out(.x)]));
1117 try types.expectExtents(&.{4}, rewritten.typeOf(rewritten.outputs[ElementwiseGradientWithGeneratedRewrite.out(.y)]));
1118 try ElementwiseGradientWithGeneratedRewrite.verify(std.testing.allocator);
1119
1120 var batched = try BatchedElementwiseGradientWithGeneratedRewrite.build(std.testing.allocator);
1121 defer batched.deinit();
1122
1123 try std.testing.expectEqual(@as(usize, 2), batched.parameters.len);
1124 try std.testing.expectEqual(@as(usize, 2), batched.outputs.len);
1125 try types.expectExtents(&.{ 8, 4 }, batched.typeOf(batched.outputs[BatchedElementwiseGradientWithGeneratedRewrite.out(.x)]));
1126 try types.expectExtents(&.{ 8, 4 }, batched.typeOf(batched.outputs[BatchedElementwiseGradientWithGeneratedRewrite.out(.y)]));
1127 try BatchedElementwiseGradientWithGeneratedRewrite.verify(std.testing.allocator);
1128 }