lib/accy/src/tensor/dsl/output.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const tensor = @import("../root.zig");
3 const gradient = tensor.gradient;
4 const trace = tensor.trace;
5 const parameter = @import("parameter.zig");
6
7 pub fn Source(comptime body: anytype) type {
8 return OutputLayout(bodyReturnPayload(body));
9 }
10
11 pub fn Derived(comptime SourceProgram: type, comptime derivation: anytype, comptime options: anytype) type {
12 if (comptime derivation == .jvp) return Jvp(SourceProgram.Outputs);
13 if (comptime derivation == .jvp_with) return Jvp(SourceProgram.Outputs);
14 if (comptime derivation == .grad) return Grad(SourceProgram.Layout, options);
15 if (comptime derivation == .grad_with) return Grad(SourceProgram.Layout, options.grad);
16 return SourceProgram.Outputs;
17 }
18
19 fn OutputLayout(comptime Result: type) type {
20 if (comptime Result == trace.Value) return Positional(1);
21
22 return switch (@typeInfo(Result)) {
23 .array => |array| Positional(array.len),
24 .pointer => |pointer| Pointer(pointer),
25 .@"struct" => |info| if (info.is_tuple) Positional(info.field_names.len) else Named(Result),
26 else => @compileError("tensor Program bodies must return a Value, array, slice, or struct of Values"),
27 };
28 }
29
30 fn Pointer(comptime pointer: std.builtin.Type.Pointer) type {
31 return switch (pointer.size) {
32 .slice => Dynamic,
33 .one => switch (@typeInfo(pointer.child)) {
34 .array => |array| Positional(array.len),
35 else => @compileError("tensor Program pointer outputs must point at arrays or slices of tensor Values"),
36 },
37 else => @compileError("tensor Program pointer outputs must point at arrays or slices of tensor Values"),
38 };
39 }
40
41 fn Positional(comptime output_count: usize) type {
42 return struct {
43 pub const count: ?usize = output_count;
44
45 pub fn index(comptime _: anytype) usize {
46 @compileError("tensor Program outputs are positional");
47 }
48 };
49 }
50
51 const Dynamic = struct {
52 pub const count: ?usize = null;
53
54 pub fn index(comptime _: anytype) usize {
55 @compileError("tensor Program slice outputs are dynamic and positional");
56 }
57 };
58
59 fn Named(comptime Result: type) type {
60 const info = @typeInfo(Result).@"struct";
61 return struct {
62 pub const count: ?usize = info.field_names.len;
63
64 pub fn index(comptime spec: anytype) usize {
65 return fieldIndex(info.field_names, name(spec));
66 }
67 };
68 }
69
70 fn Jvp(comptime SourceOutputLayout: type) type {
71 return struct {
72 pub const count: ?usize = if (SourceOutputLayout.count) |source_count| source_count * 2 else null;
73
74 pub fn index(comptime spec: anytype) usize {
75 const source_count = SourceOutputLayout.count orelse @compileError("tensor jvp output lookup requires statically known source outputs");
76 return switch (@typeInfo(@TypeOf(spec))) {
77 .@"struct" => |info| blk: {
78 if (info.is_tuple or info.field_names.len != 1) @compileError("tensor jvp output lookup requires .primals or .tangents");
79 const field_name = info.field_names[0];
80 if (comptime std.mem.eql(u8, field_name, "primals")) {
81 break :blk SourceOutputLayout.index(@field(spec, field_name));
82 }
83 if (comptime std.mem.eql(u8, field_name, "tangents")) {
84 break :blk source_count + SourceOutputLayout.index(@field(spec, field_name));
85 }
86 @compileError("tensor jvp output lookup requires .primals or .tangents");
87 },
88 .enum_literal => blk: {
89 const output_name = @tagName(spec);
90 if (comptime std.mem.eql(u8, output_name, "primal")) {
91 if (source_count != 1) @compileError("tensor jvp .primal output lookup requires one source output");
92 break :blk 0;
93 }
94 if (comptime std.mem.eql(u8, output_name, "tangent")) {
95 if (source_count != 1) @compileError("tensor jvp .tangent output lookup requires one source output");
96 break :blk source_count;
97 }
98 @compileError("tensor jvp output lookup requires .primal, .tangent, .primals, or .tangents");
99 },
100 else => @compileError("tensor jvp output lookup requires .primal, .tangent, .primals, or .tangents"),
101 };
102 }
103 };
104 }
105
106 fn Grad(comptime SourceLayout: type, comptime options: gradient.Options) type {
107 return struct {
108 pub const count: ?usize = options.wrt.len;
109
110 pub fn index(comptime spec: anytype) usize {
111 const parameter_index = comptime SourceLayout.index(spec);
112 return parameter.wrtOffset(options.wrt, parameter_index);
113 }
114 };
115 }
116
117 fn bodyReturnPayload(comptime body: anytype) type {
118 const return_type = @typeInfo(@TypeOf(body)).@"fn".return_type orelse @compileError("tensor Program body must return values");
119 return switch (@typeInfo(return_type)) {
120 .error_union => |info| info.payload,
121 else => return_type,
122 };
123 }
124
125 fn fieldIndex(comptime field_names: []const [:0]const u8, comptime output_name: []const u8) usize {
126 inline for (field_names, 0..) |field_name, index_value| {
127 if (comptime std.mem.eql(u8, field_name, output_name)) return index_value;
128 }
129 @compileError("unknown tensor output: " ++ output_name);
130 }
131
132 fn name(comptime name_value: anytype) []const u8 {
133 return switch (@typeInfo(@TypeOf(name_value))) {
134 .enum_literal => @tagName(name_value),
135 .pointer => |pointer| blk: {
136 if (pointer.child == u8) break :blk name_value[0..];
137 if (@typeInfo(pointer.child) == .array) {
138 const array = @typeInfo(pointer.child).array;
139 if (array.child == u8) break :blk name_value[0..array.len];
140 }
141 @compileError("tensor output name pointers must point at string literals");
142 },
143 .array => |array| blk: {
144 if (array.child == u8) break :blk name_value[0..];
145 @compileError("tensor output name arrays must contain bytes");
146 },
147 else => @compileError("tensor output names must be enum literals or string literals"),
148 };
149 }