lib/accy/src/tensor/wire/format.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 //! Layout, version 1, of a tensor program written as bytes, shared by the encoder and the decoder.
2 //! A program written by one process has to be read back by another with every type intact, and a
3 //! reader has to reject bytes it cannot trust, because a reader that took result types from the
4 //! bytes would trust whatever type the sender wrote, correct or not.
5 //!
6 //! A program's bytes (its *wire bytes*) hold a magic number, a version, the program name and its
7 //! root graph. A graph is written as an operation count, the operations, and the ids of its
8 //! outputs. Each operation starts with a number that says which kind of operation it is (its *kind
9 //! code*), then only the fields that `Builder.operation` needs for that kind, so the decoder works
10 //! out every other result type itself. An operation can repeat a body graph a set number of times,
11 //! carrying values from one step to the next (a *scan*), and such an operation also writes its
12 //! length and the ids of its initial values, with its body graph after it. Integers are
13 //! little-endian, and every list, whether of bytes, ids, dimensions or integers, starts with its
14 //! length as a 32-bit count.
15 const std = @import("std");
16 const tensor = @import("../root.zig");
17
18 const program = tensor.program;
19
20 /// A reader checks these four bytes first to tell wire bytes, the versioned byte encoding of a
21 /// tensor program, from other data: the value is the four bytes "ACTP" read as a little-endian
22 /// 32-bit number.
23 pub const magic: u32 = std.mem.readInt(u32, "ACTP", .little);
24 /// A writer and a reader agree on the layout through this number: the current layout is version 1.
25 /// The version is raised before any change to how this format writes anything.
26 pub const version: u32 = 1;
27 /// A caller relies on this bound when building programs for transfer with nested scans that repeat
28 /// a body graph: a program may nest scans inside scan bodies at most eight levels deep. Encoding a
29 /// deeper program returns `error.ScanTooDeep`.
30 pub const max_scan_depth: usize = 8;
31
32 /// A caller matches on these to tell why a program failed to encode or decode, so the set lists the
33 /// wire failures other than bytes that end early or run past the program: a wrong magic number,
34 /// another version, an unknown code, a reference to a missing value, too many operands and scans
35 /// nested too deep. Bytes that end early or run past the program return `error.InvalidArtifact`.
36 pub const Error = error{
37 BadMagic,
38 UnsupportedVersion,
39 UnknownTag,
40 InvalidReference,
41 TooManyOperands,
42 ScanTooDeep,
43 };
44
45 /// The function assigns each operation kind, element type and option a fixed number on the wire: it
46 /// gives each value of `E` the number of its position in `order`. A compile-time check fails the
47 /// build when `order` leaves out a value of `E` or lists one twice. Decoding a number past the end
48 /// of `order` returns `error.UnknownTag`.
49 pub fn Codes(comptime E: type, comptime order: []const E) type {
50 comptime {
51 if (order.len != @typeInfo(E).@"enum".field_names.len) {
52 @compileError("wire codes must cover every value of " ++ @typeName(E));
53 }
54 for (order, 0..) |value, index| {
55 for (order[0..index]) |previous| {
56 if (previous == value) @compileError("duplicate wire code in " ++ @typeName(E));
57 }
58 }
59 }
60 return struct {
61 pub fn code(value: E) u32 {
62 for (order, 0..) |candidate, index| {
63 if (candidate == value) return @intCast(index);
64 }
65 unreachable;
66 }
67
68 pub fn decode(value_code: u32) Error!E {
69 if (value_code >= order.len) return error.UnknownTag;
70 const value = order[value_code];
71 std.debug.assert(code(value) == value_code);
72 return value;
73 }
74 };
75 }
76
77 pub const Tag = std.meta.Tag(program.Kind);
78
79 pub const kind_codes = Codes(Tag, &.{
80 .parameter,
81 .constant,
82 .unary,
83 .binary,
84 .iota,
85 .broadcast,
86 .broadcast_in_dim,
87 .reshape,
88 .transpose,
89 .reduce,
90 .gather,
91 .scatter_add,
92 .sparse_cross_entropy,
93 .dot_general,
94 .compare,
95 .select,
96 .custom_call,
97 .scan,
98 .projection,
99 });
100
101 pub const dtype_codes = Codes(tensor.DType, &.{
102 .i1, .i8, .i16, .i32, .i64, .u8, .u16, .u32, .u64, .f16, .bf16, .f32, .f64, .key,
103 });
104
105 pub const unary_codes = Codes(program.Unary, &.{
106 .neg, .abs, .exp, .log, .sqrt, .tanh, .sin, .cos, .tan,
107 });
108
109 pub const binary_codes = Codes(program.Binary, &.{ .add, .sub, .mul, .div, .max, .min, .pow });
110
111 pub const reducer_codes = Codes(program.Reducer, &.{ .sum, .max, .min });
112
113 pub const compare_codes = Codes(program.CompareDirection, &.{ .lt, .le, .gt, .ge, .eq, .ne });
114
115 test "tensor wire codes pin enum positions and reject unknown codes" {
116 try std.testing.expectEqual(@as(u32, 0), kind_codes.code(.parameter));
117 try std.testing.expectEqual(@as(u32, 18), kind_codes.code(.projection));
118 try std.testing.expectEqual(Tag.scan, try kind_codes.decode(17));
119 try std.testing.expectEqual(program.Unary.tan, try unary_codes.decode(8));
120 try std.testing.expectEqual(tensor.DType.key, try dtype_codes.decode(13));
121 try std.testing.expectError(error.UnknownTag, unary_codes.decode(9));
122 try std.testing.expectError(error.UnknownTag, dtype_codes.decode(14));
123 try std.testing.expectError(error.UnknownTag, kind_codes.decode(std.math.maxInt(u32)));
124 }