lib/accy/src/tensor/trace/test.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const tensor = @import("../root.zig");
3 const trace = @import("root.zig");
4
5 const Builder = trace.Builder;
6 const Type = trace.Type;
7 const Value = trace.Value;
8 const define = trace.define;
9 const program_mod = tensor.program;
10
11 test {
12 @import("test_discovery").discover(trace.builder);
13 @import("test_discovery").discover(trace.definition);
14 @import("test_discovery").discover(trace.value);
15 @import("test_discovery").discover(trace);
16 }
17
18 test "tensor trace builder records an ergonomic dense program" {
19 var builder = try Builder.init(std.testing.allocator, "dense");
20 defer builder.deinit();
21
22 const x = try builder.input(.f32, .{ .m = 4, .k = 8 });
23 const w = try builder.input(.f32, .{ .k = 8, .n = 3 });
24 const b = try builder.input(.f32, .{ .n = 3 });
25 const linear = try x.contract(w, .k);
26 const out = try (try linear.add(b)).tanh();
27 var program = try builder.finish(&.{out});
28 defer program.deinit();
29
30 try std.testing.expectEqual(@as(usize, 7), program.valueCount());
31 try std.testing.expectEqual(@as(usize, 3), program.parameters.len);
32 try tensor.types.expectExtents(&.{ 4, 3 }, program.typeOf(program.outputs[0]));
33 try std.testing.expectEqualStrings("m", program.typeOf(program.outputs[0]).dims[0].name);
34 try std.testing.expectEqualStrings("n", program.typeOf(program.outputs[0]).dims[1].name);
35 }
36
37 test "tensor trace builder replays primitive operations" {
38 var source = try define(std.testing.allocator, "source", &.{
39 tensor.spec(.f32, .{ .lane = 4 }),
40 tensor.spec(.f32, .{ .lane = 4 }),
41 }, defineBody);
42 defer source.deinit();
43
44 var builder = try Builder.init(std.testing.allocator, "replay");
45 defer builder.deinit();
46
47 const x = try builder.operation(&source.operations[0], &.{});
48 const y = try builder.operation(&source.operations[1], &.{});
49 const sum = try builder.operation(&source.operations[2], &.{ x, y });
50 const out = try builder.operation(&source.operations[3], &.{sum});
51 var replayed = try builder.finish(&.{out});
52 defer replayed.deinit();
53
54 try std.testing.expectEqual(source.operationCount(), replayed.operationCount());
55 try tensor.types.expectExtents(&.{4}, replayed.typeOf(replayed.outputs[0]));
56 try std.testing.expectEqual(program_mod.Kind.unary, std.meta.activeTag(replayed.operation(replayed.outputs[0]).kind));
57 }
58
59 fn defineBody(_: *Builder, args: []const Value) !Value {
60 const sum = try args[0].add(args[1]);
61 return try sum.tanh();
62 }
63
64 test "tensor define traces a comptime body function" {
65 var program = try define(std.testing.allocator, "body", &.{
66 tensor.spec(.f32, .{ .lane = 4 }),
67 tensor.spec(.f32, .{ .lane = 4 }),
68 }, defineBody);
69 defer program.deinit();
70
71 try std.testing.expectEqual(@as(usize, 4), program.valueCount());
72 try tensor.types.expectExtents(&.{4}, program.typeOf(program.outputs[0]));
73 }
74
75 test "tensor values identify structural zeros" {
76 var builder = try Builder.init(std.testing.allocator, "structural_zero");
77 defer builder.deinit();
78
79 const scalar_zero = try builder.zeros(.{ .dtype = .f32, .dims = &.{} });
80 const input = try builder.input(.f32, .{ .lane = 4 });
81 const zero = try builder.alignTo(scalar_zero, input.ty.dims);
82 const product = try zero.mul(input);
83 const one = try builder.full(.f32, .{ .lane = 4 }, 1.0);
84
85 try std.testing.expect(scalar_zero.isStructuralZero());
86 try std.testing.expect(zero.isStructuralZero());
87 try std.testing.expect(product.isStructuralZero());
88 try std.testing.expect(!input.isStructuralZero());
89 try std.testing.expect(!one.isStructuralZero());
90 }
91
92 test "tensor value reduces sums with canonical inits by axis name" {
93 var builder = try Builder.init(std.testing.allocator, "sum_inferred_zero");
94 defer builder.deinit();
95
96 const input = try builder.input(.f32, .{ .row = 2, .col = 3 });
97 const row_sum = try input.sum(.col);
98 const total = try input.sum(.{ .row, .col });
99 var program = try builder.finish(&.{ row_sum, total });
100 defer program.deinit();
101
102 try tensor.types.expectExtents(&.{2}, program.typeOf(program.outputs[0]));
103 try tensor.types.expectExtents(&.{}, program.typeOf(program.outputs[1]));
104 try std.testing.expectEqualStrings("row", program.typeOf(program.outputs[0]).dims[0].name);
105 }
106
107 test "tensor select validates predicate dtype and accepts scalar masks" {
108 var builder = try Builder.init(std.testing.allocator, "select_predicate");
109 defer builder.deinit();
110
111 const input = try builder.input(.f32, .{ .lane = 4 });
112 const other = try builder.input(.f32, .{ .lane = 4 });
113 try std.testing.expectError(error.DTypeMismatch, builder.select(input, input, other));
114
115 const threshold = try builder.scalar(.f32, 0.0);
116 const one = try builder.scalar(.f32, 1.0);
117 const scalar_pred = try threshold.compare(.lt, one);
118 const selected = try builder.select(scalar_pred, input, other);
119 try tensor.types.expectExtents(&.{4}, selected.ty);
120 }
121
122 test "tensor custom calls validate kernel contract fields at trace construction" {
123 var builder = try Builder.init(std.testing.allocator, "custom_call_contract");
124 defer builder.deinit();
125
126 const input = try builder.input(.f32, .{ .lane = 4 });
127 try std.testing.expectError(
128 error.InvalidCustomCallContract,
129 builder.customCall("", 1, &.{input}, input.ty),
130 );
131 try std.testing.expectError(
132 error.InvalidCustomCallContract,
133 builder.customCall("accy.custom.scale", 0, &.{input}, input.ty),
134 );
135 }
136
137 test "tensor custom calls align arity with kernel calls" {
138 var builder = try Builder.init(std.testing.allocator, "custom_call_arity");
139 defer builder.deinit();
140
141 const result_ty = try Type.init(builder.arena.allocator(), .f32, &.{
142 .{ .name = "lane", .extent = 4 },
143 });
144 const generated = try builder.customCall("accy.custom.generate", 1, &.{}, result_ty);
145 try std.testing.expectEqual(@as(usize, 0), builder.operations.items[generated.id.index].kind.custom_call.operands.len);
146
147 const a = try builder.input(.f32, .{ .lane = 4 });
148 const b = try builder.input(.f32, .{ .lane = 4 });
149 const c = try builder.input(.f32, .{ .lane = 4 });
150 const d = try builder.input(.f32, .{ .lane = 4 });
151 const e = try builder.input(.f32, .{ .lane = 4 });
152 try std.testing.expectError(
153 error.UnsupportedCustomCallArity,
154 builder.customCall("accy.custom.too_many", 1, &.{ a, b, c, d, e }, a.ty),
155 );
156 }
157
158 test "accy tensor trace declaration coverage" {
159 std.testing.refAllDecls(trace.builder);
160 std.testing.refAllDecls(trace.definition);
161 std.testing.refAllDecls(trace.value);
162 std.testing.refAllDecls(trace);
163 }