lib/accy/src/tensor/wire/test.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("choir");
3 const accy = @import("../../root.zig");
4 const tensor = @import("../root.zig");
5 const fixture = @import("../../fixture/root.zig");
6 const wire = @import("root.zig");
7
8 const Builder = tensor.Builder;
9 const ByteWriter = accy.artifact.wire.ByteWriter;
10 const Fingerprint = choir.product.incremental.Fingerprint;
11 const Operation = tensor.program.Operation;
12 const Program = tensor.program.Program;
13 const ScanScope = tensor.trace.builder.ScanScope;
14 const Tag = wire.format.Tag;
15 const Value = tensor.Value;
16
17 const Fixture = *const fn (std.mem.Allocator) anyerror!Program;
18
19 const fixtures = [_]Fixture{
20 pointwiseProgram,
21 shapeProgram,
22 scanProgram,
23 lossProgram,
24 lossGradientProgram,
25 deepestScanProgram,
26 };
27
28 test "tensor wire round trip preserves program fingerprints and bytes" {
29 const allocator = std.testing.allocator;
30 var covered: std.EnumSet(Tag) = .empty;
31 for (fixtures) |build| {
32 var source = try build(allocator);
33 defer source.deinit();
34 coverKinds(&covered, source.operations);
35
36 const bytes = try wire.encode(allocator, &source);
37 defer allocator.free(bytes);
38 var decoded = try wire.decode(allocator, bytes);
39 defer decoded.deinit();
40 try std.testing.expectEqual(source.fingerprint(), decoded.fingerprint());
41
42 const again = try wire.encode(allocator, &decoded);
43 defer allocator.free(again);
44 try std.testing.expectEqualSlices(u8, bytes, again);
45 }
46 try std.testing.expect(covered.eql(.full));
47 }
48
49 test "tensor wire rejects every truncated prefix" {
50 const allocator = std.testing.allocator;
51 var source = try scanProgram(allocator);
52 defer source.deinit();
53 const bytes = try wire.encode(allocator, &source);
54 defer allocator.free(bytes);
55 for (0..bytes.len) |len| {
56 try std.testing.expectError(error.InvalidArtifact, wire.decode(allocator, bytes[0..len]));
57 }
58 }
59
60 test "tensor wire rejects foreign headers unknown tags and trailing bytes" {
61 const allocator = std.testing.allocator;
62 var source = try pointwiseProgram(allocator);
63 defer source.deinit();
64 const bytes = try wire.encode(allocator, &source);
65 defer allocator.free(bytes);
66 const corrupt = try allocator.dupe(u8, bytes);
67 defer allocator.free(corrupt);
68
69 const first_operation = 3 * @sizeOf(u32) + source.name.len + @sizeOf(u32);
70 try expectPatchedError(corrupt, bytes, 0, wire.magic +% 1, error.BadMagic);
71 try expectPatchedError(corrupt, bytes, 4, wire.version + 1, error.UnsupportedVersion);
72 try expectPatchedError(corrupt, bytes, first_operation, 19, error.UnknownTag);
73 try expectPatchedError(corrupt, bytes, first_operation + 4, 14, error.UnknownTag);
74
75 const trailing = try std.mem.concat(allocator, u8, &.{ bytes, &.{0} });
76 defer allocator.free(trailing);
77 try std.testing.expectError(error.InvalidArtifact, wire.decode(allocator, trailing));
78 }
79
80 test "tensor wire rejects later references and oversized operand lists" {
81 const allocator = std.testing.allocator;
82 var writer = ByteWriter{};
83 defer writer.deinit(allocator);
84
85 try writeHeader(&writer, 1);
86 try writer.writeU32(allocator, wire.format.kind_codes.code(.unary));
87 try writer.writeU32(allocator, wire.format.unary_codes.code(.neg));
88 try writer.writeU32(allocator, 0);
89 try writer.writeU32(allocator, 0);
90 try std.testing.expectError(error.InvalidReference, wire.decode(allocator, writer.bytes.items));
91
92 writer.bytes.clearRetainingCapacity();
93 try writeHeader(&writer, 1);
94 try writer.writeU32(allocator, wire.format.kind_codes.code(.custom_call));
95 try writer.writeU32(allocator, wire.format.dtype_codes.code(.f32));
96 try writer.writeU32(allocator, 0);
97 try writer.writeLengthPrefixedBytes(allocator, "accy.custom.double");
98 try writer.writeU32(allocator, 1);
99 const operand_count = tensor.program.max_custom_call_operands + 1;
100 try writer.writeU32(allocator, operand_count);
101 for (0..operand_count) |_| try writer.writeU32(allocator, 0);
102 try std.testing.expectError(error.TooManyOperands, wire.decode(allocator, writer.bytes.items));
103 }
104
105 test "tensor wire bounds scan nesting when encoding and decoding" {
106 const allocator = std.testing.allocator;
107 var too_deep = try nestedScanProgram(allocator, wire.max_scan_depth + 1);
108 defer too_deep.deinit();
109 try std.testing.expectError(error.ScanTooDeep, wire.encode(allocator, &too_deep));
110
111 const deepest = try nestedScanBytes(allocator, wire.max_scan_depth);
112 defer allocator.free(deepest);
113 var decoded = try wire.decode(allocator, deepest);
114 decoded.deinit();
115
116 const deeper = try nestedScanBytes(allocator, wire.max_scan_depth + 1);
117 defer allocator.free(deeper);
118 try std.testing.expectError(error.ScanTooDeep, wire.decode(allocator, deeper));
119 }
120
121 test "tensor wire encode and decode survive allocation failures" {
122 const allocator = std.testing.allocator;
123 var source = try nestedScanProgram(allocator, 2);
124 defer source.deinit();
125 const bytes = try wire.encode(allocator, &source);
126 defer allocator.free(bytes);
127 try fixture.checkAllAllocationFailures(encodeAll, .{&source});
128 try fixture.checkAllAllocationFailures(decodeAll, .{ bytes, source.fingerprint() });
129 }
130
131 fn encodeAll(allocator: std.mem.Allocator, source: *const Program) !void {
132 const bytes = try wire.encode(allocator, source);
133 allocator.free(bytes);
134 }
135
136 fn decodeAll(allocator: std.mem.Allocator, bytes: []const u8, expected: Fingerprint) !void {
137 var decoded = try wire.decode(allocator, bytes);
138 defer decoded.deinit();
139 try std.testing.expectEqual(expected, decoded.fingerprint());
140 }
141
142 fn expectPatchedError(
143 corrupt: []u8,
144 bytes: []const u8,
145 offset: usize,
146 value: u32,
147 expected: anyerror,
148 ) !void {
149 @memcpy(corrupt, bytes);
150 std.mem.writeInt(u32, corrupt[offset..][0..4], value, .little);
151 try std.testing.expectError(expected, wire.decode(std.testing.allocator, corrupt));
152 }
153
154 fn writeHeader(writer: *ByteWriter, operation_count: u32) !void {
155 const allocator = std.testing.allocator;
156 try writer.writeU32(allocator, wire.magic);
157 try writer.writeU32(allocator, wire.version);
158 try writer.writeLengthPrefixedBytes(allocator, "wire_crafted");
159 try writer.writeU32(allocator, operation_count);
160 }
161
162 fn coverKinds(covered: *std.EnumSet(Tag), operations: []const Operation) void {
163 for (operations) |op| {
164 covered.insert(std.meta.activeTag(op.kind));
165 switch (op.kind) {
166 .scan => |scan| for (scan.body.operations) |body_op| {
167 covered.insert(std.meta.activeTag(body_op.kind));
168 },
169 else => {},
170 }
171 }
172 }
173
174 fn pointwiseProgram(allocator: std.mem.Allocator) anyerror!Program {
175 var builder = try Builder.init(allocator, "wire_pointwise");
176 errdefer builder.deinit();
177 const x = try builder.input(.f32, .{ .lane = 4 });
178 const y = try builder.input(.f32, .{ .lane = 4 });
179 const sum = try builder.binary(.add, x, try builder.unary(.neg, y));
180 const floor = try builder.full(.f32, .{ .lane = 4 }, 0.5);
181 const below = try builder.compare(.lt, sum, floor);
182 const chosen = try builder.select(below, sum, try builder.unary(.exp, x));
183 return builder.finish(&.{ chosen, below });
184 }
185
186 fn shapeProgram(allocator: std.mem.Allocator) anyerror!Program {
187 var builder = try Builder.init(allocator, "wire_shapes");
188 errdefer builder.deinit();
189 const grid = try builder.input(.f32, .{ .row = 2, .col = 3 });
190 const columns = try builder.iota(.i32, .{ .row = 2, .col = 3 }, .col);
191 const flipped = try builder.transposeBy(grid, &.{ 1, 0 });
192 const flat = try builder.reshapeTo(flipped, &.{.{ .name = "cell", .extent = 6 }});
193 const scale = try builder.scalar(.f32, 2.0);
194 const spread = try builder.broadcastOp(scale, &.{.{ .name = "cell", .extent = 6 }});
195 const doubled = try builder.customCall("accy.custom.double", 1, &.{flat}, flat.ty);
196 const product = try builder.binary(.mul, doubled, spread);
197 return builder.finish(&.{ product, columns });
198 }
199
200 fn scanStep(_: *Builder, carry: anytype) !@TypeOf(carry) {
201 const doubled = try carry.x.add(carry.x);
202 return .{ .x = doubled, .acc = try carry.acc.add(doubled) };
203 }
204
205 fn scanProgram(allocator: std.mem.Allocator) anyerror!Program {
206 var builder = try Builder.init(allocator, "wire_scan");
207 errdefer builder.deinit();
208 const x = try builder.input(.f32, .{ .lane = 2 });
209 const acc = try builder.full(.f32, .{ .lane = 2 }, 0.0);
210 const walked = try builder.scan(.{
211 .length = 3,
212 .init = .{ .x = x, .acc = acc },
213 .body = scanStep,
214 });
215 return builder.finish(&.{ walked.x, walked.acc });
216 }
217
218 const loss_specs = [_]tensor.Spec{
219 tensor.spec(.i32, .{ .batch = 2, .token = 3 }),
220 tensor.spec(.i32, .{ .batch = 2, .token = 3 }),
221 tensor.spec(.f32, .{ .vocab = 11, .channel = 4 }),
222 tensor.spec(.f32, .{ .channel = 4, .vocab = 11 }),
223 };
224
225 fn lossBody(_: *Builder, args: []const Value) !Value {
226 const hidden = try tensor.nn.embedding(args[2], args[0], .vocab);
227 const flat_hidden = try hidden.merge(.{ .batch, .token }, .sample);
228 const flat_targets = try args[1].merge(.{ .batch, .token }, .sample);
229 const logits = try flat_hidden.contract(args[3], .channel);
230 return tensor.nn.sparseCrossEntropyMean(logits, flat_targets, .vocab);
231 }
232
233 fn lossProgram(allocator: std.mem.Allocator) anyerror!Program {
234 return tensor.define(allocator, "wire_loss", &loss_specs, lossBody);
235 }
236
237 fn lossGradientProgram(allocator: std.mem.Allocator) anyerror!Program {
238 var loss = try lossProgram(allocator);
239 defer loss.deinit();
240 return tensor.grad(allocator, &loss, .{ .wrt = &.{ 2, 3 } });
241 }
242
243 fn deepestScanProgram(allocator: std.mem.Allocator) anyerror!Program {
244 return nestedScanProgram(allocator, wire.max_scan_depth);
245 }
246
247 fn nestedScanProgram(allocator: std.mem.Allocator, depth: usize) anyerror!Program {
248 std.debug.assert(depth <= wire.max_scan_depth + 1);
249 var builder = try Builder.init(allocator, "wire_nested_scan");
250 errdefer builder.deinit();
251 var scopes: [wire.max_scan_depth + 1]ScanScope = undefined;
252 var opened: usize = 0;
253 errdefer for (scopes[0..opened]) |*scope| scope.abort();
254 var carry = try builder.input(.f32, .{ .lane = 2 });
255 var parent = &builder;
256 while (opened < depth) : (opened += 1) {
257 scopes[opened] = try parent.scanScope(2, &.{carry});
258 parent = scopes[opened].body();
259 carry = scopes[opened].carry(0);
260 }
261 var next = try carry.add(carry);
262 while (opened != 0) : (opened -= 1) {
263 next = try scopes[opened - 1].finish(&.{next});
264 }
265 return builder.finish(&.{next});
266 }
267
268 fn nestedScanBytes(allocator: std.mem.Allocator, depth: usize) ![]u8 {
269 var writer = ByteWriter{};
270 errdefer writer.deinit(allocator);
271 try writer.writeU32(allocator, wire.magic);
272 try writer.writeU32(allocator, wire.version);
273 try writer.writeLengthPrefixedBytes(allocator, "wire_deep_scan");
274 for (0..depth + 1) |level| {
275 try writer.writeU32(allocator, if (level == depth) 1 else 2);
276 try writer.writeU32(allocator, wire.format.kind_codes.code(.parameter));
277 try writer.writeU32(allocator, wire.format.dtype_codes.code(.f32));
278 try writer.writeU32(allocator, 0);
279 if (level == depth) break;
280 try writer.writeU32(allocator, wire.format.kind_codes.code(.scan));
281 try writer.writeU64(allocator, 1);
282 try writer.writeU32(allocator, 1);
283 try writer.writeU32(allocator, 0);
284 }
285 for (0..depth + 1) |level| {
286 try writer.writeU32(allocator, 1);
287 try writer.writeU32(allocator, if (level == 0) 0 else 1);
288 }
289 return writer.toOwnedSlice(allocator);
290 }