lib/accy/src/tensor/wire/encode.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const accy = @import("../../root.zig");
3 const tensor = @import("../root.zig");
4 const format = @import("format.zig");
5
6 const program_mod = tensor.program;
7 const ByteWriter = accy.artifact.wire.ByteWriter;
8 const Dim = tensor.Dim;
9 const Id = program_mod.Id;
10 const Kind = program_mod.Kind;
11 const Operation = program_mod.Operation;
12 const Program = program_mod.Program;
13 const Subgraph = program_mod.Subgraph;
14 const Type = program_mod.Type;
15
16 const Frame = struct {
17 graph: Subgraph,
18 next: usize = 0,
19 };
20
21 /// A caller uses this to send a tensor program to another process, so the function writes `program`
22 /// as version 1 bytes, which the caller owns. The call returns `error.InvalidProgram` unless the
23 /// operation ids run densely from zero and the parameter operations appear in parameter order, and
24 /// every program made by the program builder meets both. The program builder is the tracer that
25 /// appends one operation at a time and checks each result type as it goes. The call returns
26 /// `error.ScanTooDeep` when scans nest deeper than eight levels and `error.InvalidArtifact` when a
27 /// list or string is longer than a 32-bit count can hold.
28 pub fn encode(allocator: std.mem.Allocator, program: *const Program) ![]u8 {
29 var sink = Sink{ .allocator = allocator };
30 errdefer sink.writer.deinit(allocator);
31 try sink.word(format.magic);
32 try sink.word(format.version);
33 try sink.bytes(program.name);
34
35 var frames: [format.max_scan_depth + 1]Frame = undefined;
36 frames[0] = .{ .graph = .{
37 .values = program.values,
38 .operations = program.operations,
39 .parameters = program.parameters,
40 .outputs = program.outputs,
41 } };
42 var count: usize = 1;
43 try sink.graph(frames[0].graph);
44 while (count != 0) {
45 const frame = &frames[count - 1];
46 std.debug.assert(frame.next <= frame.graph.operations.len);
47 if (frame.next == frame.graph.operations.len) {
48 try sink.idList(frame.graph.outputs);
49 count -= 1;
50 continue;
51 }
52 const op = frame.graph.operations[frame.next];
53 frame.next += 1;
54 try sink.operation(op);
55 if (op.kind != .scan) continue;
56 if (count > format.max_scan_depth) return error.ScanTooDeep;
57 frames[count] = .{ .graph = op.kind.scan.body.* };
58 count += 1;
59 std.debug.assert(count <= frames.len);
60 try sink.graph(op.kind.scan.body.*);
61 }
62 return sink.writer.toOwnedSlice(allocator);
63 }
64
65 const Sink = struct {
66 allocator: std.mem.Allocator,
67 writer: ByteWriter = .{},
68
69 fn word(self: *Sink, value: u32) !void {
70 try self.writer.writeU32(self.allocator, value);
71 }
72
73 fn integer(self: *Sink, value: i64) !void {
74 try self.writer.writeU64(self.allocator, @bitCast(value));
75 }
76
77 fn bytes(self: *Sink, value: []const u8) !void {
78 try self.writer.writeLengthPrefixedBytes(self.allocator, value);
79 }
80
81 fn count(self: *Sink, len: usize) !void {
82 try self.word(std.math.cast(u32, len) orelse return error.InvalidArtifact);
83 }
84
85 fn id(self: *Sink, value: Id) !void {
86 std.debug.assert(value.index != program_mod.synthetic_id.index);
87 try self.word(value.index);
88 }
89
90 fn ids(self: *Sink, values: []const Id) !void {
91 std.debug.assert(values.len <= program_mod.max_operation_operands);
92 for (values) |value| try self.id(value);
93 }
94
95 fn idList(self: *Sink, values: []const Id) !void {
96 try self.count(values.len);
97 for (values) |value| try self.id(value);
98 }
99
100 fn integers(self: *Sink, values: []const i64) !void {
101 try self.count(values.len);
102 for (values) |value| try self.integer(value);
103 }
104
105 fn dims(self: *Sink, values: []const Dim) !void {
106 try self.count(values.len);
107 for (values) |dim| {
108 try self.bytes(dim.name);
109 try self.integer(dim.extent);
110 }
111 }
112
113 fn typed(self: *Sink, value: Type) !void {
114 try self.word(format.dtype_codes.code(value.dtype));
115 try self.dims(value.dims);
116 }
117
118 fn graph(self: *Sink, value: Subgraph) !void {
119 if (value.values.len != value.operations.len) return error.InvalidProgram;
120 var parameter_count: usize = 0;
121 for (value.operations, 0..) |op, index| {
122 if (op.id.index != index) return error.InvalidProgram;
123 const parameter = switch (op.kind) {
124 .parameter => |payload| payload,
125 else => continue,
126 };
127 if (parameter.index != parameter_count) return error.InvalidProgram;
128 if (parameter_count == value.parameters.len) return error.InvalidProgram;
129 if (value.parameters[parameter_count].index != index) return error.InvalidProgram;
130 parameter_count += 1;
131 }
132 if (parameter_count != value.parameters.len) return error.InvalidProgram;
133 try self.count(value.operations.len);
134 }
135
136 fn operation(self: *Sink, op: Operation) !void {
137 try self.word(format.kind_codes.code(std.meta.activeTag(op.kind)));
138 switch (op.kind) {
139 .parameter => try self.typed(op.result),
140 .constant => |constant| {
141 try self.typed(op.result);
142 try self.bytes(constant.payload);
143 },
144 .iota => |iota| {
145 try self.typed(op.result);
146 try self.integer(iota.axis);
147 },
148 .custom_call => |custom| {
149 try self.typed(op.result);
150 try self.bytes(custom.target);
151 try self.word(custom.version);
152 try self.idList(custom.operands);
153 },
154 .broadcast => |broadcast| {
155 try self.dims(op.result.dims);
156 try self.id(broadcast.input);
157 },
158 .broadcast_in_dim => |broadcast| {
159 try self.dims(op.result.dims);
160 try self.id(broadcast.input);
161 try self.integers(broadcast.broadcast_dims);
162 },
163 .reshape => |reshape| {
164 try self.dims(op.result.dims);
165 try self.id(reshape.input);
166 },
167 .scan => |scan| {
168 std.debug.assert(scan.body.parameters.len == scan.inits.len);
169 try self.integer(scan.length);
170 try self.idList(scan.inits);
171 },
172 else => try self.derived(op.kind),
173 }
174 }
175
176 fn derived(self: *Sink, kind: Kind) !void {
177 switch (kind) {
178 .unary => |unary| {
179 try self.word(format.unary_codes.code(unary.op));
180 try self.id(unary.input);
181 },
182 .binary => |binary| {
183 try self.word(format.binary_codes.code(binary.op));
184 try self.ids(&.{ binary.lhs, binary.rhs });
185 },
186 .compare => |compare| {
187 try self.word(format.compare_codes.code(compare.direction));
188 try self.ids(&.{ compare.lhs, compare.rhs });
189 },
190 .select => |select| try self.ids(&.{ select.pred, select.on_true, select.on_false }),
191 .projection => |projection| {
192 try self.id(projection.source);
193 try self.count(projection.index);
194 },
195 .transpose => |transpose| {
196 try self.id(transpose.input);
197 try self.integers(transpose.permutation);
198 },
199 .reduce => |reduce| {
200 try self.ids(&.{ reduce.input, reduce.init });
201 try self.word(format.reducer_codes.code(reduce.reducer));
202 try self.integers(reduce.dimensions);
203 },
204 .gather => |gather| {
205 try self.ids(&.{ gather.input, gather.indices });
206 try self.integer(gather.axis);
207 },
208 .scatter_add => |scatter| {
209 try self.ids(&.{ scatter.input, scatter.indices, scatter.updates });
210 try self.integer(scatter.axis);
211 },
212 .sparse_cross_entropy => |loss| {
213 try self.ids(&.{ loss.logits, loss.targets });
214 try self.integer(loss.axis);
215 },
216 .dot_general => |dot| {
217 try self.ids(&.{ dot.lhs, dot.rhs });
218 const axes = [_][]const i64{
219 dot.lhs_contract,
220 dot.rhs_contract,
221 dot.lhs_batch,
222 dot.rhs_batch,
223 };
224 for (axes) |values| try self.integers(values);
225 },
226 .parameter, .constant, .iota, .custom_call => unreachable,
227 .broadcast, .broadcast_in_dim, .reshape, .scan => unreachable,
228 }
229 }
230 };