lib/accy/src/tensor/program.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("choir");
3 const type_mod = @import("type/root.zig");
4
5 pub const Type = type_mod.Type;
6 pub const product_name = "accy.tensor_program";
7
8 pub const Id = struct {
9 index: u32,
10 };
11
12 pub const synthetic_id = Id{ .index = std.math.maxInt(u32) };
13
14 pub const Unary = enum {
15 neg,
16 abs,
17 exp,
18 log,
19 sqrt,
20 tanh,
21 sin,
22 cos,
23 tan,
24 };
25
26 pub const Binary = enum {
27 add,
28 sub,
29 mul,
30 div,
31 max,
32 min,
33 pow,
34 };
35
36 pub const Reducer = enum {
37 sum,
38 max,
39 min,
40
41 pub fn name(self: Reducer) []const u8 {
42 return switch (self) {
43 .sum => "sum",
44 .max => "max",
45 .min => "min",
46 };
47 }
48 };
49
50 pub const Parameter = struct {
51 index: usize,
52 };
53
54 pub const Constant = struct {
55 payload: []const u8,
56 };
57
58 pub const UnaryOp = struct {
59 op: Unary,
60 input: Id,
61 };
62
63 pub const BinaryOp = struct {
64 op: Binary,
65 lhs: Id,
66 rhs: Id,
67 };
68
69 pub const Iota = struct {
70 axis: i64,
71 };
72
73 pub const Broadcast = struct {
74 input: Id,
75 sizes: []const i64,
76 };
77
78 pub const BroadcastInDim = struct {
79 input: Id,
80 broadcast_dims: []const i64,
81 };
82
83 pub const Reshape = struct {
84 input: Id,
85 new_shape: []const i64,
86 };
87
88 pub const Transpose = struct {
89 input: Id,
90 permutation: []const i64,
91 };
92
93 pub const Reduce = struct {
94 input: Id,
95 init: Id,
96 reducer: Reducer,
97 dimensions: []const i64,
98 };
99
100 pub const Gather = struct {
101 input: Id,
102 indices: Id,
103 axis: i64,
104 };
105
106 pub const ScatterAdd = struct {
107 input: Id,
108 indices: Id,
109 updates: Id,
110 axis: i64,
111 };
112
113 pub const SparseCrossEntropy = struct {
114 logits: Id,
115 targets: Id,
116 axis: i64,
117 };
118
119 pub const CompareDirection = enum {
120 lt,
121 le,
122 gt,
123 ge,
124 eq,
125 ne,
126 };
127
128 pub const Compare = struct {
129 lhs: Id,
130 rhs: Id,
131 direction: CompareDirection,
132 };
133
134 pub const Select = struct {
135 pred: Id,
136 on_true: Id,
137 on_false: Id,
138 };
139
140 pub const max_custom_call_operands: usize = 4;
141 pub const max_operation_operands: usize = 8;
142 pub const max_scan_carries: usize = 8;
143
144 pub const CustomCall = struct {
145 target: []const u8,
146 version: u32 = 1,
147 operands: []const Id,
148 };
149
150 pub const Subgraph = struct {
151 values: []const Type,
152 operations: []const Operation,
153 parameters: []const Id,
154 outputs: []const Id,
155
156 pub fn typeOf(self: Subgraph, id: Id) Type {
157 return self.values[id.index];
158 }
159
160 pub fn operation(self: Subgraph, id: Id) Operation {
161 return self.operations[id.index];
162 }
163 };
164
165 pub const Scan = struct {
166 length: i64,
167 inits: []const Id,
168 body: *const Subgraph,
169 };
170
171 pub const Projection = struct {
172 source: Id,
173 index: usize,
174 };
175
176 pub const DotGeneral = struct {
177 lhs: Id,
178 rhs: Id,
179 lhs_contract: []const i64,
180 rhs_contract: []const i64,
181 lhs_batch: []const i64,
182 rhs_batch: []const i64,
183 };
184
185 pub const Kind = union(enum) {
186 parameter: Parameter,
187 constant: Constant,
188 unary: UnaryOp,
189 binary: BinaryOp,
190 iota: Iota,
191 broadcast: Broadcast,
192 broadcast_in_dim: BroadcastInDim,
193 reshape: Reshape,
194 transpose: Transpose,
195 reduce: Reduce,
196 gather: Gather,
197 scatter_add: ScatterAdd,
198 sparse_cross_entropy: SparseCrossEntropy,
199 dot_general: DotGeneral,
200 compare: Compare,
201 select: Select,
202 custom_call: CustomCall,
203 scan: Scan,
204 projection: Projection,
205 };
206
207 pub const Operation = struct {
208 id: Id,
209 result: Type,
210 kind: Kind,
211 };
212
213 pub const Program = struct {
214 arena: std.heap.ArenaAllocator,
215 name: []const u8,
216 values: []const Type,
217 operations: []const Operation,
218 parameters: []const Id,
219 outputs: []const Id,
220
221 pub fn deinit(self: *Program) void {
222 self.arena.deinit();
223 }
224
225 pub fn fingerprint(self: Program) choir.product.incremental.Fingerprint {
226 var builder = choir.product.incremental.FingerprintBuilder{};
227 builder.updateBytes(product_name);
228 builder.updateBytes(self.name);
229 builder.updateUsize(self.values.len);
230 for (self.values) |value| updateTypeFingerprint(&builder, value);
231 builder.updateUsize(self.operations.len);
232 for (self.operations) |op| updateOperationFingerprint(&builder, op);
233 builder.updateUsize(self.parameters.len);
234 for (self.parameters) |id| updateIdFingerprint(&builder, id);
235 builder.updateUsize(self.outputs.len);
236 for (self.outputs) |id| updateIdFingerprint(&builder, id);
237 return builder.finish();
238 }
239
240 pub fn productStamp(self: Program) choir.product.incremental.ProductStamp {
241 return choir.product.incremental.productStamp(product_name, self.fingerprint());
242 }
243
244 pub fn valueCount(self: Program) usize {
245 return self.values.len;
246 }
247
248 pub fn containsScan(self: Program) bool {
249 for (self.operations) |op| {
250 if (op.kind == .scan) return true;
251 }
252 return false;
253 }
254
255 pub fn operationCount(self: Program) usize {
256 return self.operations.len;
257 }
258
259 pub fn typeOf(self: Program, id: Id) Type {
260 return self.values[id.index];
261 }
262
263 pub fn operation(self: Program, id: Id) Operation {
264 return self.operations[id.index];
265 }
266
267 pub fn isZeroConstant(self: Program, id: Id) bool {
268 const op = self.operation(id);
269 return switch (op.kind) {
270 .constant => |constant| isZeroPayload(constant.payload),
271 else => false,
272 };
273 }
274
275 pub fn constantPayload(self: Program, id: Id) ?[]const u8 {
276 const op = self.operation(id);
277 return switch (op.kind) {
278 .constant => |constant| constant.payload,
279 else => null,
280 };
281 }
282 };
283
284 pub fn isZeroPayload(bytes: []const u8) bool {
285 for (bytes) |byte| {
286 if (byte != 0) return false;
287 }
288 return true;
289 }
290
291 fn updateIdFingerprint(builder: *choir.product.incremental.FingerprintBuilder, id: Id) void {
292 builder.updateU32(id.index);
293 }
294
295 fn updateTypeFingerprint(builder: *choir.product.incremental.FingerprintBuilder, ty: Type) void {
296 builder.updateEnumTag(ty.dtype);
297 builder.updateUsize(ty.dims.len);
298 for (ty.dims) |dim| {
299 builder.updateBytes(dim.name);
300 builder.updateI64(dim.extent);
301 }
302 }
303
304 fn updateOperationFingerprint(builder: *choir.product.incremental.FingerprintBuilder, op: Operation) void {
305 updateIdFingerprint(builder, op.id);
306 updateTypeFingerprint(builder, op.result);
307 builder.updateEnumTag(std.meta.activeTag(op.kind));
308 switch (op.kind) {
309 .parameter => |parameter| builder.updateUsize(parameter.index),
310 .constant => |constant| builder.updateBytes(constant.payload),
311 .unary => |unary| {
312 builder.updateEnumTag(unary.op);
313 updateIdFingerprint(builder, unary.input);
314 },
315 .binary => |binary| {
316 builder.updateEnumTag(binary.op);
317 updateIdFingerprint(builder, binary.lhs);
318 updateIdFingerprint(builder, binary.rhs);
319 },
320 .compare => |compare| {
321 builder.updateEnumTag(compare.direction);
322 updateIdFingerprint(builder, compare.lhs);
323 updateIdFingerprint(builder, compare.rhs);
324 },
325 .select => |select| {
326 updateIdFingerprint(builder, select.pred);
327 updateIdFingerprint(builder, select.on_true);
328 updateIdFingerprint(builder, select.on_false);
329 },
330 .custom_call => |custom| {
331 builder.updateBytes(custom.target);
332 builder.updateUsize(custom.version);
333 for (custom.operands) |operand| updateIdFingerprint(builder, operand);
334 },
335 .iota => |iota| builder.updateI64(iota.axis),
336 .broadcast => |broadcast| {
337 updateIdFingerprint(builder, broadcast.input);
338 builder.updateI64Slice(broadcast.sizes);
339 },
340 .broadcast_in_dim => |broadcast| {
341 updateIdFingerprint(builder, broadcast.input);
342 builder.updateI64Slice(broadcast.broadcast_dims);
343 },
344 .reshape => |reshape| {
345 updateIdFingerprint(builder, reshape.input);
346 builder.updateI64Slice(reshape.new_shape);
347 },
348 .transpose => |transpose| {
349 updateIdFingerprint(builder, transpose.input);
350 builder.updateI64Slice(transpose.permutation);
351 },
352 .reduce => |reduce| {
353 updateIdFingerprint(builder, reduce.input);
354 updateIdFingerprint(builder, reduce.init);
355 builder.updateEnumTag(reduce.reducer);
356 builder.updateI64Slice(reduce.dimensions);
357 },
358 .gather => |gather| {
359 updateIdFingerprint(builder, gather.input);
360 updateIdFingerprint(builder, gather.indices);
361 builder.updateI64(gather.axis);
362 },
363 .scatter_add => |scatter_add| {
364 updateIdFingerprint(builder, scatter_add.input);
365 updateIdFingerprint(builder, scatter_add.indices);
366 updateIdFingerprint(builder, scatter_add.updates);
367 builder.updateI64(scatter_add.axis);
368 },
369 .sparse_cross_entropy => |sparse_cross_entropy| {
370 updateIdFingerprint(builder, sparse_cross_entropy.logits);
371 updateIdFingerprint(builder, sparse_cross_entropy.targets);
372 builder.updateI64(sparse_cross_entropy.axis);
373 },
374 .dot_general => |dot| {
375 updateIdFingerprint(builder, dot.lhs);
376 updateIdFingerprint(builder, dot.rhs);
377 builder.updateI64Slice(dot.lhs_contract);
378 builder.updateI64Slice(dot.rhs_contract);
379 builder.updateI64Slice(dot.lhs_batch);
380 builder.updateI64Slice(dot.rhs_batch);
381 },
382 .scan => |scan| {
383 builder.updateI64(scan.length);
384 builder.updateUsize(scan.inits.len);
385 for (scan.inits) |id| updateIdFingerprint(builder, id);
386 builder.updateUsize(scan.body.values.len);
387 for (scan.body.values) |value| updateTypeFingerprint(builder, value);
388 builder.updateUsize(scan.body.operations.len);
389 for (scan.body.operations) |body_op| updateOperationFingerprint(builder, body_op);
390 builder.updateUsize(scan.body.parameters.len);
391 for (scan.body.parameters) |id| updateIdFingerprint(builder, id);
392 builder.updateUsize(scan.body.outputs.len);
393 for (scan.body.outputs) |id| updateIdFingerprint(builder, id);
394 },
395 .projection => |projection| {
396 updateIdFingerprint(builder, projection.source);
397 builder.updateUsize(projection.index);
398 },
399 }
400 }
401
402 pub const CloneError = std.mem.Allocator.Error || error{ InvalidDimension, DuplicateAxis, AxisNameEmpty };
403
404 pub fn cloneSubgraph(allocator: std.mem.Allocator, body: *const Subgraph) CloneError!*const Subgraph {
405 const clone = try allocator.create(Subgraph);
406 const values = try allocator.alloc(Type, body.values.len);
407 for (body.values, values) |ty, *slot| slot.* = try cloneType(allocator, ty);
408 const operations = try allocator.alloc(Operation, body.operations.len);
409 for (body.operations, operations) |op, *slot| {
410 slot.* = .{
411 .id = op.id,
412 .result = try cloneType(allocator, op.result),
413 .kind = switch (op.kind) {
414 .parameter, .iota, .projection => op.kind,
415 .constant => |constant| .{ .constant = .{ .payload = try allocator.dupe(u8, constant.payload) } },
416 .unary, .binary, .compare, .select => op.kind,
417 .custom_call => |custom| .{ .custom_call = .{
418 .target = try allocator.dupe(u8, custom.target),
419 .version = custom.version,
420 .operands = try allocator.dupe(Id, custom.operands),
421 } },
422 .broadcast => |broadcast| .{ .broadcast = .{
423 .input = broadcast.input,
424 .sizes = try allocator.dupe(i64, broadcast.sizes),
425 } },
426 .broadcast_in_dim => |broadcast| .{ .broadcast_in_dim = .{
427 .input = broadcast.input,
428 .broadcast_dims = try allocator.dupe(i64, broadcast.broadcast_dims),
429 } },
430 .reshape => |reshape| .{ .reshape = .{
431 .input = reshape.input,
432 .new_shape = try allocator.dupe(i64, reshape.new_shape),
433 } },
434 .transpose => |transpose| .{ .transpose = .{
435 .input = transpose.input,
436 .permutation = try allocator.dupe(i64, transpose.permutation),
437 } },
438 .reduce => |reduce| .{ .reduce = .{
439 .input = reduce.input,
440 .init = reduce.init,
441 .reducer = reduce.reducer,
442 .dimensions = try allocator.dupe(i64, reduce.dimensions),
443 } },
444 .gather, .scatter_add, .sparse_cross_entropy => op.kind,
445 .dot_general => |dot| .{ .dot_general = .{
446 .lhs = dot.lhs,
447 .rhs = dot.rhs,
448 .lhs_contract = try allocator.dupe(i64, dot.lhs_contract),
449 .rhs_contract = try allocator.dupe(i64, dot.rhs_contract),
450 .lhs_batch = try allocator.dupe(i64, dot.lhs_batch),
451 .rhs_batch = try allocator.dupe(i64, dot.rhs_batch),
452 } },
453 .scan => |scan| .{ .scan = .{
454 .length = scan.length,
455 .inits = try allocator.dupe(Id, scan.inits),
456 .body = try cloneSubgraph(allocator, scan.body),
457 } },
458 },
459 };
460 }
461 clone.* = .{
462 .values = values,
463 .operations = operations,
464 .parameters = try allocator.dupe(Id, body.parameters),
465 .outputs = try allocator.dupe(Id, body.outputs),
466 };
467 return clone;
468 }
469
470 fn cloneType(allocator: std.mem.Allocator, ty: Type) !Type {
471 return Type.init(allocator, ty.dtype, ty.dims);
472 }
473
474 fn singleConstantProgram(
475 allocator: std.mem.Allocator,
476 name: []const u8,
477 payload: []const u8,
478 ) !Program {
479 var arena = std.heap.ArenaAllocator.init(allocator);
480 errdefer arena.deinit();
481 const arena_allocator = arena.allocator();
482
483 const owned_name = try arena_allocator.dupe(u8, name);
484 const ty = try Type.init(arena_allocator, .f32, &.{});
485 const owned_payload = try arena_allocator.dupe(u8, payload);
486 const values = try arena_allocator.dupe(Type, &.{ty});
487 const operations = try arena_allocator.dupe(Operation, &.{
488 .{
489 .id = .{ .index = 0 },
490 .result = ty,
491 .kind = .{ .constant = .{ .payload = owned_payload } },
492 },
493 });
494 const outputs = try arena_allocator.dupe(Id, &.{.{ .index = 0 }});
495
496 return .{
497 .arena = arena,
498 .name = owned_name,
499 .values = values,
500 .operations = operations,
501 .parameters = &.{},
502 .outputs = outputs,
503 };
504 }
505
506 test "tensor program identifies zero constants" {
507 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
508 defer arena.deinit();
509 const allocator = arena.allocator();
510
511 const ty = try Type.init(allocator, .f32, &.{});
512 const payload = try allocator.alloc(u8, 4);
513 @memset(payload, 0);
514 const values = try allocator.dupe(Type, &.{ty});
515 const operations = try allocator.dupe(Operation, &.{
516 .{
517 .id = .{ .index = 0 },
518 .result = ty,
519 .kind = .{ .constant = .{ .payload = payload } },
520 },
521 });
522 const outputs = try allocator.dupe(Id, &.{.{ .index = 0 }});
523
524 const program = Program{
525 .arena = arena,
526 .name = "zero",
527 .values = values,
528 .operations = operations,
529 .parameters = &.{},
530 .outputs = outputs,
531 };
532
533 try std.testing.expect(program.isZeroConstant(.{ .index = 0 }));
534 }
535
536 test "tensor program fingerprints summarize current content" {
537 const allocator = std.testing.allocator;
538
539 var first = try singleConstantProgram(allocator, "tensor_program_product_identity", &.{ 0, 0, 0, 0 });
540 defer first.deinit();
541 var same = try singleConstantProgram(allocator, "tensor_program_product_identity", &.{ 0, 0, 0, 0 });
542 defer same.deinit();
543 var changed = try singleConstantProgram(allocator, "tensor_program_product_identity_changed", &.{ 1, 0, 0, 0 });
544 defer changed.deinit();
545
546 const first_stamp = first.productStamp();
547 try std.testing.expectEqualStrings(product_name, first_stamp.name);
548 try std.testing.expectEqual(first.fingerprint(), first_stamp.fingerprint);
549 try std.testing.expectEqual(first.fingerprint(), same.fingerprint());
550 try std.testing.expect(first.fingerprint() != changed.fingerprint());
551 }