lib/accy/src/tensor/interpret/step.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const program_mod = @import("../root.zig").program;
 2 
 3 pub fn Step(comptime Value: type) type {
 4     return struct {
 5         source: *const program_mod.Program,
 6         values: []const Value,
 7         op: *const program_mod.Operation,
 8 
 9         pub fn read(self: @This(), id: program_mod.Id) Value {
10             return self.values[id.index];
11         }
12 
13         pub fn typeOf(self: @This(), id: program_mod.Id) program_mod.Type {
14             return self.source.typeOf(id);
15         }
16     };
17 }
18 
19 pub fn arguments(comptime Value: type, op: *const program_mod.Operation, values: []const Value, buffer: *[program_mod.max_operation_operands]Value) []const Value {
20     return switch (op.kind) {
21         .parameter, .constant, .iota => buffer[0..0],
22         .unary => |unary| blk: {
23             buffer[0] = values[unary.input.index];
24             break :blk buffer[0..1];
25         },
26         .binary => |binary| blk: {
27             buffer[0] = values[binary.lhs.index];
28             buffer[1] = values[binary.rhs.index];
29             break :blk buffer[0..2];
30         },
31         .broadcast => |broadcast| blk: {
32             buffer[0] = values[broadcast.input.index];
33             break :blk buffer[0..1];
34         },
35         .broadcast_in_dim => |broadcast| blk: {
36             buffer[0] = values[broadcast.input.index];
37             break :blk buffer[0..1];
38         },
39         .reshape => |reshape| blk: {
40             buffer[0] = values[reshape.input.index];
41             break :blk buffer[0..1];
42         },
43         .transpose => |transpose| blk: {
44             buffer[0] = values[transpose.input.index];
45             break :blk buffer[0..1];
46         },
47         .reduce => |reduce| blk: {
48             buffer[0] = values[reduce.input.index];
49             buffer[1] = values[reduce.init.index];
50             break :blk buffer[0..2];
51         },
52         .gather => |gather| blk: {
53             buffer[0] = values[gather.input.index];
54             buffer[1] = values[gather.indices.index];
55             break :blk buffer[0..2];
56         },
57         .scatter_add => |scatter_add| blk: {
58             buffer[0] = values[scatter_add.input.index];
59             buffer[1] = values[scatter_add.indices.index];
60             buffer[2] = values[scatter_add.updates.index];
61             break :blk buffer[0..3];
62         },
63         .sparse_cross_entropy => |sparse_cross_entropy| blk: {
64             buffer[0] = values[sparse_cross_entropy.logits.index];
65             buffer[1] = values[sparse_cross_entropy.targets.index];
66             break :blk buffer[0..2];
67         },
68         .dot_general => |dot| blk: {
69             buffer[0] = values[dot.lhs.index];
70             buffer[1] = values[dot.rhs.index];
71             break :blk buffer[0..2];
72         },
73         .compare => |compare| blk: {
74             buffer[0] = values[compare.lhs.index];
75             buffer[1] = values[compare.rhs.index];
76             break :blk buffer[0..2];
77         },
78         .select => |select| blk: {
79             buffer[0] = values[select.pred.index];
80             buffer[1] = values[select.on_true.index];
81             buffer[2] = values[select.on_false.index];
82             break :blk buffer[0..3];
83         },
84         .custom_call => |custom| blk: {
85             for (custom.operands, 0..) |operand, index| buffer[index] = values[operand.index];
86             break :blk buffer[0..custom.operands.len];
87         },
88         .scan => |scan| blk: {
89             for (scan.inits, 0..) |init, index| buffer[index] = values[init.index];
90             break :blk buffer[0..scan.inits.len];
91         },
92         .projection => |projection| blk: {
93             buffer[0] = values[projection.source.index];
94             break :blk buffer[0..1];
95         },
96     };
97 }