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

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const program_mod = @import("../root.zig").program;
 2 const trace = @import("../root.zig").trace;
 3 const context = @import("context.zig");
 4 const dispatch = @import("dispatch.zig");
 5 const step_mod = @import("step.zig");
 6 
 7 fn stateHandle(comptime StateType: type) type {
 8     return switch (@typeInfo(StateType)) {
 9         .pointer => StateType,
10         else => *StateType,
11     };
12 }
13 
14 fn stateResult(comptime StateType: type) type {
15     return switch (@typeInfo(StateType)) {
16         .pointer => |info| info.child.Result,
17         else => StateType.Result,
18     };
19 }
20 
21 fn LayerResult(comptime Next: type, comptime Impl: type) type {
22     if (comptime @hasDecl(Impl, "finish") and @hasDecl(Impl, "Result")) return Impl.Result;
23     return stateResult(Next);
24 }
25 
26 pub fn Layer(comptime BoundValue: type, comptime Next: type, comptime Impl: type) type {
27     return struct {
28         next: Next,
29         impl: Impl,
30 
31         const Self = @This();
32 
33         pub const Value: type = BoundValue;
34         pub const Result: type = LayerResult(Next, Impl);
35 
36         fn nextHandle(self: *Self) stateHandle(Next) {
37             return switch (@typeInfo(Next)) {
38                 .pointer => self.next,
39                 else => &self.next,
40             };
41         }
42 
43         pub fn operation(self: *Self, step: *step_mod.Step(Value)) !Value {
44             var buffer: [program_mod.max_operation_operands]Value = undefined;
45             return self.bind(step.op, step_mod.arguments(Value, step.op, step.values, &buffer));
46         }
47 
48         pub fn bind(self: *Self, op: *const program_mod.Operation, args: []const Value) !Value {
49             var ctx = context.BindContext(Value, stateHandle(Next)){
50                 .next = self.nextHandle(),
51                 .op = op,
52                 .args = args,
53             };
54             if (comptime @hasDecl(Impl, "bind")) {
55                 return self.impl.bind(&ctx);
56             }
57             if (try dispatch.primitive(Value, &self.impl, &ctx)) |value| return value;
58             return ctx.default();
59         }
60 
61         pub fn finish(self: *Self, outputs: []const Value) !Result {
62             var ctx = context.FinishContext(Value, stateHandle(Next)){ .next = self.nextHandle() };
63             if (comptime @hasDecl(Impl, "finish")) {
64                 return self.impl.finish(&ctx, outputs);
65             }
66             return ctx.default(outputs);
67         }
68 
69         pub fn builderHandle(self: *Self) *trace.Builder {
70             return self.nextHandle().builderHandle();
71         }
72     };
73 }