lib/accy/src/tensor/dsl/surface/root.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const tensor = @import("../../root.zig");
3 const autodiff = tensor.autodiff;
4 const batch = tensor.batch;
5 const body = @import("../root.zig").body;
6 const execute = @import("execute.zig");
7 const gradient = tensor.gradient;
8 const interpret_mod = tensor.interpret;
9 const lower_mod = tensor.lower;
10 const output = @import("../root.zig").output;
11 const parameter = @import("../root.zig").parameter;
12 const program_mod = tensor.program;
13 const trace = tensor.trace;
14 const transform = tensor.transform;
15
16 const Derivation = enum {
17 rewrite,
18 jvp,
19 jvp_with,
20 grad,
21 grad_with,
22 vmap,
23 vmap_with,
24 };
25
26 pub fn Program(comptime definition: anytype) type {
27 return Source(definition);
28 }
29
30 fn Source(comptime definition: anytype) type {
31 return struct {
32 pub const name = definition.name;
33 pub const parameters = definition.parameters;
34 pub const Layout: type = parameter.Standard(parameters);
35 pub const Outputs: type = output.Source(definition.body);
36 const Methods: type = Surface(@This());
37
38 pub fn build(allocator: std.mem.Allocator) !program_mod.Program {
39 var builder = try trace.Builder.init(allocator, definition.name);
40 errdefer builder.deinit();
41
42 const args = try body.buildInputs(&builder, definition.parameters);
43 const outputs = try body.call(definition, &builder, args);
44 return try builder.finish(outputs);
45 }
46
47 pub const wrt = Methods.wrt;
48 pub const inAxes = Methods.inAxes;
49 pub const out = Methods.out;
50 pub const analyze = Methods.analyze;
51 pub const rewrite = Methods.rewrite;
52 pub const jvp = Methods.jvp;
53 pub const jvpWith = Methods.jvpWith;
54 pub const grad = Methods.grad;
55 pub const gradWith = Methods.gradWith;
56 pub const vmap = Methods.vmap;
57 pub const vmapWith = Methods.vmapWith;
58 pub const interpret = Methods.interpret;
59 pub const lower = Methods.lower;
60 pub const prepare = Methods.prepare;
61 pub const prepareWith = Methods.prepareWith;
62 pub const prepareFragment = Methods.prepareFragment;
63 pub const createArtifactJob = Methods.createArtifactJob;
64 pub const compileFragment = Methods.compileFragment;
65 pub const verify = Methods.verify;
66 };
67 }
68
69 fn Derived(comptime SourceProgram: type, comptime derivation: Derivation, comptime options: anytype) type {
70 return struct {
71 pub const name = SourceProgram.name;
72 pub const parameters = SourceProgram.parameters;
73 pub const Layout: type = DerivedLayout(SourceProgram, derivation, options);
74 pub const Outputs: type = DerivedOutputs(SourceProgram, derivation, options);
75 const Methods: type = Surface(@This());
76
77 pub fn build(allocator: std.mem.Allocator) !program_mod.Program {
78 var source = try SourceProgram.build(allocator);
79 defer source.deinit();
80 if (comptime derivation == .rewrite) return transform.apply(allocator, &source, options);
81 if (comptime derivation == .jvp) return autodiff.jvp(allocator, &source, options);
82 if (comptime derivation == .jvp_with) return autodiff.jvpWith(allocator, &source, options.jvp, options.hooks);
83 if (comptime derivation == .grad) return gradient.grad(allocator, &source, options);
84 if (comptime derivation == .grad_with) return gradient.gradWith(allocator, &source, options.grad, options.hooks);
85 if (comptime derivation == .vmap) return batch.vmap(allocator, &source, options);
86 if (comptime derivation == .vmap_with) return batch.vmapWith(allocator, &source, options.vmap, options.hooks);
87 unreachable;
88 }
89
90 pub const wrt = Methods.wrt;
91 pub const inAxes = Methods.inAxes;
92 pub const out = Methods.out;
93 pub const analyze = Methods.analyze;
94 pub const rewrite = Methods.rewrite;
95 pub const jvp = Methods.jvp;
96 pub const jvpWith = Methods.jvpWith;
97 pub const grad = Methods.grad;
98 pub const gradWith = Methods.gradWith;
99 pub const vmap = Methods.vmap;
100 pub const vmapWith = Methods.vmapWith;
101 pub const lower = Methods.lower;
102 pub const prepare = Methods.prepare;
103 pub const prepareWith = Methods.prepareWith;
104 pub const prepareFragment = Methods.prepareFragment;
105 pub const createArtifactJob = Methods.createArtifactJob;
106 pub const compileFragment = Methods.compileFragment;
107 pub const verify = Methods.verify;
108
109 pub fn interpret(allocator: std.mem.Allocator, initial: anytype) !interpret_mod.result(@TypeOf(initial)) {
110 if (comptime derivation == .rewrite) return execute.rewriteProgram(SourceProgram, options, allocator, initial);
111 if (comptime derivation == .jvp) return execute.jvpProgram(SourceProgram, options, allocator, initial);
112 if (comptime derivation == .jvp_with) return execute.jvpWithProgram(SourceProgram, options.jvp, options.hooks, allocator, initial);
113 if (comptime derivation == .grad) return execute.gradProgram(SourceProgram, options, allocator, initial);
114 if (comptime derivation == .grad_with) return execute.gradWithProgram(SourceProgram, options.grad, options.hooks, allocator, initial);
115 if (comptime derivation == .vmap) return execute.vmapProgram(SourceProgram, options, allocator, initial);
116 if (comptime derivation == .vmap_with) return execute.vmapWithProgram(SourceProgram, options.vmap, options.hooks, allocator, initial);
117 return execute.program(@This(), allocator, initial);
118 }
119 };
120 }
121
122 fn Analysis(comptime ProgramType: type, comptime analysis_initial: anytype) type {
123 return struct {
124 pub const Program: type = ProgramType;
125 pub const name = ProgramType.name;
126 pub const parameters = ProgramType.parameters;
127 pub const Layout: type = ProgramType.Layout;
128 pub const Outputs: type = ProgramType.Outputs;
129
130 pub fn wrt(comptime names: anytype) *const [parameter.nameCount(names)]usize {
131 return ProgramType.Layout.indices(names);
132 }
133
134 pub fn inAxes(comptime axes: anytype) *const [ProgramType.Layout.count]batch.Axis {
135 return ProgramType.Layout.axes(axes);
136 }
137
138 pub fn out(comptime spec: anytype) usize {
139 return ProgramType.Outputs.index(spec);
140 }
141
142 pub fn interpret(allocator: std.mem.Allocator) !interpret_mod.result(@TypeOf(analysis_initial)) {
143 return execute.program(ProgramType, allocator, analysis_initial);
144 }
145
146 pub fn rewrite(comptime make_transform: anytype) type {
147 return AnalysisRewrite(ProgramType, analysis_initial, make_transform);
148 }
149
150 pub fn jvpWith(comptime options: autodiff.JvpOptions, comptime make_hooks: anytype) type {
151 return AnalysisDerived(ProgramType, .jvp_with, .{ .jvp = options }, analysis_initial, make_hooks);
152 }
153
154 pub fn gradWith(comptime options: gradient.Options, comptime make_hooks: anytype) type {
155 return AnalysisDerived(ProgramType, .grad_with, .{ .grad = options }, analysis_initial, make_hooks);
156 }
157
158 pub fn vmapWith(comptime options: batch.Options, comptime make_hooks: anytype) type {
159 return AnalysisDerived(ProgramType, .vmap_with, .{ .vmap = options }, analysis_initial, make_hooks);
160 }
161 };
162 }
163
164 fn AnalysisRewrite(comptime ProgramType: type, comptime analysis_initial: anytype, comptime make_transform: anytype) type {
165 return struct {
166 pub const name = ProgramType.name;
167 pub const parameters = ProgramType.parameters;
168 pub const Layout: type = ProgramType.Layout;
169 pub const Outputs: type = ProgramType.Outputs;
170 const Methods: type = Surface(@This());
171
172 pub fn build(allocator: std.mem.Allocator) !program_mod.Program {
173 var source = try ProgramType.build(allocator);
174 defer source.deinit();
175
176 var analysis = try interpret_mod.run(allocator, &source, analysis_initial);
177 defer execute.deinitIfPresent(&analysis);
178 return transform.apply(allocator, &source, make_transform(&analysis));
179 }
180
181 pub const wrt = Methods.wrt;
182 pub const inAxes = Methods.inAxes;
183 pub const out = Methods.out;
184 pub const analyze = Methods.analyze;
185 pub const rewrite = Methods.rewrite;
186 pub const jvp = Methods.jvp;
187 pub const jvpWith = Methods.jvpWith;
188 pub const grad = Methods.grad;
189 pub const gradWith = Methods.gradWith;
190 pub const vmap = Methods.vmap;
191 pub const vmapWith = Methods.vmapWith;
192 pub const lower = Methods.lower;
193 pub const prepare = Methods.prepare;
194 pub const prepareWith = Methods.prepareWith;
195 pub const prepareFragment = Methods.prepareFragment;
196 pub const createArtifactJob = Methods.createArtifactJob;
197 pub const compileFragment = Methods.compileFragment;
198 pub const verify = Methods.verify;
199
200 pub fn interpret(allocator: std.mem.Allocator, initial: anytype) !interpret_mod.result(@TypeOf(initial)) {
201 return execute.analysisRewriteProgram(ProgramType, analysis_initial, make_transform, allocator, initial);
202 }
203 };
204 }
205
206 fn AnalysisDerived(
207 comptime ProgramType: type,
208 comptime derivation: Derivation,
209 comptime options: anytype,
210 comptime analysis_initial: anytype,
211 comptime make_hooks: anytype,
212 ) type {
213 return struct {
214 pub const name = ProgramType.name;
215 pub const parameters = ProgramType.parameters;
216 pub const Layout: type = DerivedLayout(ProgramType, derivation, options);
217 pub const Outputs: type = DerivedOutputs(ProgramType, derivation, options);
218 const Methods: type = Surface(@This());
219
220 pub fn build(allocator: std.mem.Allocator) !program_mod.Program {
221 var source = try ProgramType.build(allocator);
222 defer source.deinit();
223
224 var analysis = try interpret_mod.run(allocator, &source, analysis_initial);
225 defer execute.deinitIfPresent(&analysis);
226
227 const hooks = make_hooks(&analysis);
228 if (comptime derivation == .jvp_with) return autodiff.jvpWith(allocator, &source, options.jvp, hooks);
229 if (comptime derivation == .grad_with) return gradient.gradWith(allocator, &source, options.grad, hooks);
230 if (comptime derivation == .vmap_with) return batch.vmapWith(allocator, &source, options.vmap, hooks);
231 unreachable;
232 }
233
234 pub const wrt = Methods.wrt;
235 pub const inAxes = Methods.inAxes;
236 pub const out = Methods.out;
237 pub const analyze = Methods.analyze;
238 pub const rewrite = Methods.rewrite;
239 pub const jvp = Methods.jvp;
240 pub const jvpWith = Methods.jvpWith;
241 pub const grad = Methods.grad;
242 pub const gradWith = Methods.gradWith;
243 pub const vmap = Methods.vmap;
244 pub const vmapWith = Methods.vmapWith;
245 pub const lower = Methods.lower;
246 pub const prepare = Methods.prepare;
247 pub const prepareWith = Methods.prepareWith;
248 pub const prepareFragment = Methods.prepareFragment;
249 pub const createArtifactJob = Methods.createArtifactJob;
250 pub const compileFragment = Methods.compileFragment;
251 pub const verify = Methods.verify;
252
253 pub fn interpret(allocator: std.mem.Allocator, initial: anytype) !interpret_mod.result(@TypeOf(initial)) {
254 var source = try ProgramType.build(allocator);
255 defer source.deinit();
256
257 var analysis = try interpret_mod.run(allocator, &source, analysis_initial);
258 defer execute.deinitIfPresent(&analysis);
259
260 const hooks = make_hooks(&analysis);
261 if (comptime derivation == .jvp_with) return execute.jvpSource(allocator, &source, options.jvp, hooks, initial);
262 if (comptime derivation == .grad_with) return execute.gradSource(allocator, &source, options.grad, hooks, initial);
263 if (comptime derivation == .vmap_with) return execute.vmapSource(allocator, &source, options.vmap, hooks, initial);
264 return execute.program(@This(), allocator, initial);
265 }
266 };
267 }
268
269 fn Surface(comptime ProgramType: type) type {
270 return struct {
271 pub fn wrt(comptime names: anytype) *const [parameter.nameCount(names)]usize {
272 return ProgramType.Layout.indices(names);
273 }
274
275 pub fn inAxes(comptime axes: anytype) *const [ProgramType.Layout.count]batch.Axis {
276 return ProgramType.Layout.axes(axes);
277 }
278
279 pub fn out(comptime spec: anytype) usize {
280 return ProgramType.Outputs.index(spec);
281 }
282
283 pub fn rewrite(comptime pass: anytype) type {
284 return Derived(ProgramType, .rewrite, pass);
285 }
286
287 pub fn analyze(comptime analysis_initial: anytype) type {
288 return Analysis(ProgramType, analysis_initial);
289 }
290
291 pub fn jvp(comptime options: autodiff.JvpOptions) type {
292 return Derived(ProgramType, .jvp, options);
293 }
294
295 pub fn jvpWith(comptime options: autodiff.JvpOptions, comptime hooks: anytype) type {
296 return Derived(ProgramType, .jvp_with, .{ .jvp = options, .hooks = hooks });
297 }
298
299 pub fn grad(comptime options: gradient.Options) type {
300 return Derived(ProgramType, .grad, options);
301 }
302
303 pub fn gradWith(comptime options: gradient.Options, comptime hooks: anytype) type {
304 return Derived(ProgramType, .grad_with, .{ .grad = options, .hooks = hooks });
305 }
306
307 pub fn vmap(comptime options: batch.Options) type {
308 return Derived(ProgramType, .vmap, options);
309 }
310
311 pub fn vmapWith(comptime options: batch.Options, comptime hooks: anytype) type {
312 return Derived(ProgramType, .vmap_with, .{ .vmap = options, .hooks = hooks });
313 }
314
315 pub fn interpret(allocator: std.mem.Allocator, initial: anytype) !interpret_mod.result(@TypeOf(initial)) {
316 return execute.program(ProgramType, allocator, initial);
317 }
318
319 pub fn lower(allocator: std.mem.Allocator) !*lower_mod.SemanticModule {
320 return execute.lower(ProgramType, allocator);
321 }
322
323 pub fn prepare(allocator: std.mem.Allocator) !lower_mod.BackendPreparedJob {
324 return execute.prepare(ProgramType, allocator);
325 }
326
327 pub fn prepareWith(allocator: std.mem.Allocator, options: lower_mod.BackendPreparationRunOptions) !lower_mod.BackendPreparedJob {
328 return execute.prepareWith(ProgramType, allocator, options);
329 }
330
331 pub fn prepareFragment(
332 allocator: std.mem.Allocator,
333 handle: lower_mod.BackendHandle,
334 options: lower_mod.FragmentCompilerOptions,
335 ) !lower_mod.BackendPreparedJob {
336 return execute.prepareFragment(ProgramType, allocator, handle, options);
337 }
338
339 pub fn createArtifactJob(
340 allocator: std.mem.Allocator,
341 handle: lower_mod.BackendHandle,
342 options: lower_mod.FragmentCompilerOptions,
343 ) !*lower_mod.ArtifactJob {
344 return execute.createArtifactJob(ProgramType, allocator, handle, options);
345 }
346
347 pub fn compileFragment(
348 allocator: std.mem.Allocator,
349 handle: lower_mod.BackendHandle,
350 options: lower_mod.FragmentCompilerOptions,
351 ) !*lower_mod.CompiledFragment {
352 return execute.compileFragment(ProgramType, allocator, handle, options);
353 }
354
355 pub fn verify(allocator: std.mem.Allocator) !void {
356 try execute.verify(ProgramType, allocator);
357 }
358 };
359 }
360
361 fn DerivedLayout(comptime SourceProgram: type, comptime derivation: Derivation, comptime options: anytype) type {
362 if (comptime derivation == .jvp) return parameter.Jvp(SourceProgram.Layout, SourceProgram.parameters, options);
363 if (comptime derivation == .jvp_with) return parameter.Jvp(SourceProgram.Layout, SourceProgram.parameters, options.jvp);
364 return SourceProgram.Layout;
365 }
366
367 fn DerivedOutputs(comptime SourceProgram: type, comptime derivation: Derivation, comptime options: anytype) type {
368 return output.Derived(SourceProgram, derivation, options);
369 }