lib/accy/src/tensor/dsl/surface/execute.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 gradient = tensor.gradient;
  6 const interpret_mod = tensor.interpret;
  7 const lower_mod = tensor.lower;
  8 const program_mod = tensor.program;
  9 const trace = tensor.trace;
 10 const transform = tensor.transform;
 11 
 12 pub fn lower(comptime SourceProgram: type, allocator: std.mem.Allocator) !*lower_mod.SemanticModule {
 13     return SourceProgram.interpret(allocator, lower_mod.module(allocator));
 14 }
 15 
 16 pub fn prepare(
 17     comptime SourceProgram: type,
 18     allocator: std.mem.Allocator,
 19 ) !lower_mod.BackendPreparedJob {
 20     return prepareWith(SourceProgram, allocator, .{});
 21 }
 22 
 23 pub fn prepareWith(
 24     comptime SourceProgram: type,
 25     allocator: std.mem.Allocator,
 26     options: lower_mod.BackendPreparationRunOptions,
 27 ) !lower_mod.BackendPreparedJob {
 28     var built = try SourceProgram.build(allocator);
 29     defer built.deinit();
 30 
 31     return lower_mod.prepareWith(allocator, &built, options);
 32 }
 33 
 34 pub fn prepareFragment(
 35     comptime SourceProgram: type,
 36     allocator: std.mem.Allocator,
 37     handle: lower_mod.BackendHandle,
 38     options: lower_mod.FragmentCompilerOptions,
 39 ) !lower_mod.BackendPreparedJob {
 40     var built = try SourceProgram.build(allocator);
 41     defer built.deinit();
 42 
 43     return lower_mod.prepareFragment(allocator, handle, &built, options);
 44 }
 45 
 46 pub fn createArtifactJob(
 47     comptime SourceProgram: type,
 48     allocator: std.mem.Allocator,
 49     handle: lower_mod.BackendHandle,
 50     options: lower_mod.FragmentCompilerOptions,
 51 ) !*lower_mod.ArtifactJob {
 52     var built = try SourceProgram.build(allocator);
 53     defer built.deinit();
 54 
 55     return lower_mod.createArtifactJob(allocator, handle, &built, options);
 56 }
 57 
 58 pub fn compileFragment(
 59     comptime SourceProgram: type,
 60     allocator: std.mem.Allocator,
 61     handle: lower_mod.BackendHandle,
 62     options: lower_mod.FragmentCompilerOptions,
 63 ) !*lower_mod.CompiledFragment {
 64     var built = try SourceProgram.build(allocator);
 65     defer built.deinit();
 66 
 67     return lower_mod.compileFragment(allocator, handle, &built, options);
 68 }
 69 
 70 pub fn verify(comptime SourceProgram: type, allocator: std.mem.Allocator) !void {
 71     const module = try lower(SourceProgram, allocator);
 72     defer module.deinit();
 73 
 74     try module.verify();
 75 }
 76 
 77 pub fn program(comptime SourceProgram: type, allocator: std.mem.Allocator, initial: anytype) !interpret_mod.result(@TypeOf(initial)) {
 78     var built = try SourceProgram.build(allocator);
 79     defer built.deinit();
 80 
 81     if (comptime @hasDecl(@TypeOf(initial), "attach")) {
 82         var builder = try trace.Builder.init(allocator, SourceProgram.name);
 83         errdefer builder.deinit();
 84 
 85         return interpret_mod.run(allocator, &built, initial.attach(interpret_mod.Graph{ .builder = &builder }));
 86     }
 87 
 88     return interpret_mod.run(allocator, &built, initial);
 89 }
 90 
 91 pub fn rewriteProgram(comptime SourceProgram: type, comptime pass: anytype, allocator: std.mem.Allocator, initial: anytype) !interpret_mod.result(@TypeOf(initial)) {
 92     var source = try SourceProgram.build(allocator);
 93     defer source.deinit();
 94 
 95     return rewriteSource(allocator, &source, pass, initial);
 96 }
 97 
 98 pub fn analysisRewriteProgram(
 99     comptime ProgramType: type,
100     comptime analysis_initial: anytype,
101     comptime make_transform: anytype,
102     allocator: std.mem.Allocator,
103     initial: anytype,
104 ) !interpret_mod.result(@TypeOf(initial)) {
105     var source = try ProgramType.build(allocator);
106     defer source.deinit();
107 
108     var analysis = try interpret_mod.run(allocator, &source, analysis_initial);
109     defer deinitIfPresent(&analysis);
110 
111     return rewriteSource(allocator, &source, make_transform(&analysis), initial);
112 }
113 
114 fn rewriteSource(allocator: std.mem.Allocator, source: *const program_mod.Program, pass: anytype, initial: anytype) !interpret_mod.result(@TypeOf(initial)) {
115     if (comptime @hasDecl(@TypeOf(initial), "attach")) {
116         var builder = try trace.Builder.init(allocator, source.name);
117         errdefer builder.deinit();
118 
119         const graph = interpret_mod.Graph{ .builder = &builder };
120         return interpret_mod.run(allocator, source, transform.semantics(source, initial.attach(graph), pass));
121     }
122 
123     var rewritten = try transform.apply(allocator, source, pass);
124     defer rewritten.deinit();
125 
126     return interpret_mod.run(allocator, &rewritten, initial);
127 }
128 
129 pub fn jvpProgram(comptime SourceProgram: type, comptime options: autodiff.JvpOptions, allocator: std.mem.Allocator, initial: anytype) !interpret_mod.result(@TypeOf(initial)) {
130     var source = try SourceProgram.build(allocator);
131     defer source.deinit();
132 
133     return jvpSource(allocator, &source, options, .{}, initial);
134 }
135 
136 pub fn jvpWithProgram(
137     comptime SourceProgram: type,
138     comptime options: autodiff.JvpOptions,
139     comptime hooks: anytype,
140     allocator: std.mem.Allocator,
141     initial: anytype,
142 ) !interpret_mod.result(@TypeOf(initial)) {
143     var source = try SourceProgram.build(allocator);
144     defer source.deinit();
145 
146     return jvpSource(allocator, &source, options, hooks, initial);
147 }
148 
149 pub fn jvpSource(
150     allocator: std.mem.Allocator,
151     source: *const program_mod.Program,
152     options: autodiff.JvpOptions,
153     hooks: anytype,
154     initial: anytype,
155 ) !interpret_mod.result(@TypeOf(initial)) {
156     if (comptime @hasDecl(@TypeOf(initial), "attach")) {
157         try autodiff.validate(source, options);
158 
159         var builder = try trace.Builder.init(allocator, source.name);
160         errdefer builder.deinit();
161 
162         const graph = interpret_mod.Graph{ .builder = &builder };
163         const attached = initial.attach(graph);
164         return interpret_mod.run(allocator, source, autodiff.jvpSemanticsWith(source, attached, options, hooks));
165     }
166 
167     var owned = try autodiff.jvpWith(allocator, source, options, hooks);
168     defer owned.deinit();
169 
170     return interpret_mod.run(allocator, &owned, initial);
171 }
172 
173 pub fn gradProgram(comptime SourceProgram: type, comptime options: gradient.Options, allocator: std.mem.Allocator, initial: anytype) !interpret_mod.result(@TypeOf(initial)) {
174     var source = try SourceProgram.build(allocator);
175     defer source.deinit();
176 
177     return gradSource(allocator, &source, options, .{}, initial);
178 }
179 
180 pub fn gradWithProgram(
181     comptime SourceProgram: type,
182     comptime options: gradient.Options,
183     comptime hooks: anytype,
184     allocator: std.mem.Allocator,
185     initial: anytype,
186 ) !interpret_mod.result(@TypeOf(initial)) {
187     var source = try SourceProgram.build(allocator);
188     defer source.deinit();
189 
190     return gradSource(allocator, &source, options, hooks, initial);
191 }
192 
193 pub fn gradSource(
194     allocator: std.mem.Allocator,
195     source: *const program_mod.Program,
196     options: gradient.Options,
197     hooks: anytype,
198     initial: anytype,
199 ) !interpret_mod.result(@TypeOf(initial)) {
200     return gradient.interpretWith(allocator, source, options, hooks, initial);
201 }
202 
203 pub fn vmapProgram(comptime SourceProgram: type, comptime options: batch.Options, allocator: std.mem.Allocator, initial: anytype) !interpret_mod.result(@TypeOf(initial)) {
204     var source = try SourceProgram.build(allocator);
205     defer source.deinit();
206 
207     return vmapSource(allocator, &source, options, .{}, initial);
208 }
209 
210 pub fn vmapWithProgram(
211     comptime SourceProgram: type,
212     comptime options: batch.Options,
213     comptime hooks: anytype,
214     allocator: std.mem.Allocator,
215     initial: anytype,
216 ) !interpret_mod.result(@TypeOf(initial)) {
217     var source = try SourceProgram.build(allocator);
218     defer source.deinit();
219 
220     return vmapSource(allocator, &source, options, hooks, initial);
221 }
222 
223 pub fn vmapSource(
224     allocator: std.mem.Allocator,
225     source: *const program_mod.Program,
226     options: batch.Options,
227     hooks: anytype,
228     initial: anytype,
229 ) !interpret_mod.result(@TypeOf(initial)) {
230     if (comptime @hasDecl(@TypeOf(initial), "attach")) {
231         try batch.validate(source, options);
232 
233         var builder = try trace.Builder.init(allocator, source.name);
234         errdefer builder.deinit();
235 
236         const graph = interpret_mod.Graph{ .builder = &builder };
237         const attached = initial.attach(graph);
238         return interpret_mod.run(allocator, source, batch.semanticsWith(attached, options, hooks));
239     }
240 
241     var owned = try batch.vmapWith(allocator, source, options, hooks);
242     defer owned.deinit();
243 
244     return interpret_mod.run(allocator, &owned, initial);
245 }
246 
247 pub fn deinitIfPresent(value: anytype) void {
248     const Value = @typeInfo(@TypeOf(value)).pointer.child;
249     switch (@typeInfo(Value)) {
250         .@"struct", .@"union", .@"enum", .@"opaque" => if (comptime @hasDecl(Value, "deinit")) value.deinit(),
251         else => {},
252     }
253 }