lib/accy/src/kernel/logical/program.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const executable = @import("../../executable/root.zig");
  3 const kernel = @import("../root.zig");
  4 const kernel_compile = @import("../compile/root.zig");
  5 const logical_builder = @import("../model/logical/root.zig");
  6 const parameter = @import("../program/root.zig").parameter;
  7 const schedule_mod = logical_builder.schedule;
  8 
  9 pub fn Program(comptime definition: anytype) type {
 10     return struct {
 11         const Default: type = Scheduled(definition, schedule_mod.default());
 12 
 13         pub const name = definition.name;
 14         pub const parameters = definition.parameters;
 15         pub const Layout: type = Default.Layout;
 16         pub const Limits: type = Default.Limits;
 17 
 18         pub fn withSchedule(comptime schedule: anytype) type {
 19             return Scheduled(definition, schedule);
 20         }
 21 
 22         pub const arg = Default.arg;
 23         pub const schema = Default.schema;
 24         pub const build = Default.build;
 25         pub const interpret = Default.interpret;
 26         pub const analyze = Default.analyze;
 27         pub const transform = Default.transform;
 28         pub const launch = Default.launch;
 29         pub const scheduleSnapshot = Default.scheduleSnapshot;
 30         pub const createPlan = Default.createPlan;
 31         pub const createCheckedPlan = Default.createCheckedPlan;
 32         pub const compileFragment = Default.compileFragment;
 33         pub const createKernelArtifact = Default.createKernelArtifact;
 34         pub const createKernelCallArtifact = Default.createKernelCallArtifact;
 35         pub const runCpu = Default.runCpu;
 36         pub const runCpuWithDiagnostic = Default.runCpuWithDiagnostic;
 37         pub const verify = Default.verify;
 38     };
 39 }
 40 
 41 fn Scheduled(comptime definition: anytype, comptime schedule: anytype) type {
 42     return struct {
 43         const Self = @This();
 44         const ProgramSource: type = kernel.Program(ScheduledDefinition(definition, schedule));
 45 
 46         pub const name = definition.name;
 47         pub const parameters = definition.parameters;
 48         pub const Layout: type = ProgramSource.Layout;
 49         pub const Limits: type = ProgramSource.Limits;
 50         pub const schedule_policy = schedule;
 51 
 52         pub fn withSchedule(comptime next_schedule: anytype) type {
 53             return Scheduled(definition, next_schedule);
 54         }
 55 
 56         pub const arg = ProgramSource.arg;
 57         pub const schema = ProgramSource.schema;
 58         pub const build = ProgramSource.build;
 59         pub const interpret = ProgramSource.interpret;
 60         pub fn analyze(comptime analysis_impl: anytype) type {
 61             return Analysis(Self, analysis_impl);
 62         }
 63 
 64         pub fn transform(comptime impl: anytype) type {
 65             return Derived(Self, kernel.interpret.asSpec(impl));
 66         }
 67 
 68         pub const launch = ProgramSource.launch;
 69         pub const scheduleSnapshot = ProgramSource.scheduleSnapshot;
 70         pub const createPlan = ProgramSource.createPlan;
 71         pub const createCheckedPlan = ProgramSource.createCheckedPlan;
 72         pub const compileFragment = ProgramSource.compileFragment;
 73         pub const createKernelArtifact = ProgramSource.createKernelArtifact;
 74         pub const createKernelCallArtifact = ProgramSource.createKernelCallArtifact;
 75         pub const runCpu = ProgramSource.runCpu;
 76         pub const runCpuWithDiagnostic = ProgramSource.runCpuWithDiagnostic;
 77         pub const verify = ProgramSource.verify;
 78     };
 79 }
 80 
 81 fn Derived(comptime ProgramType: type, comptime transform_spec: anytype) type {
 82     return struct {
 83         const Self = @This();
 84         const Runtime: type = RuntimeSurface(Self);
 85 
 86         pub const name = ProgramType.name;
 87         pub const parameters = ProgramType.parameters;
 88         pub const Layout: type = ProgramType.Layout;
 89         pub const schedule_policy = ProgramType.schedule_policy;
 90         pub const spec = transform_spec;
 91         pub const Limits: type = kernel.Limits;
 92 
 93         pub fn withSchedule(comptime next_schedule: anytype) type {
 94             return Derived(ProgramType.withSchedule(next_schedule), transform_spec);
 95         }
 96 
 97         pub fn transform(comptime impl: anytype) type {
 98             return Derived(ProgramType, kernel.interpret.stack(.{ transform_spec, kernel.interpret.asSpec(impl) }));
 99         }
100 
101         pub fn analyze(comptime analysis_impl: anytype) type {
102             return Analysis(Self, analysis_impl);
103         }
104 
105         pub fn arg(comptime name_value: anytype) usize {
106             return Layout.index(name_value);
107         }
108 
109         pub fn schema() []const kernel.Param {
110             return ProgramType.schema();
111         }
112 
113         pub fn build(allocator: std.mem.Allocator, limits: Limits) !kernel.Graph {
114             return ProgramType.interpret(allocator, limits, transform_spec);
115         }
116 
117         pub fn interpret(allocator: std.mem.Allocator, limits: Limits, initial: anytype) !InterpretResult(ProgramType, StackedSpec(@TypeOf(transform_spec), @TypeOf(initial))) {
118             return ProgramType.interpret(allocator, limits, kernel.interpret.stack(.{ transform_spec, kernel.interpret.asSpec(initial) }));
119         }
120 
121         pub const launch = Runtime.launch;
122         pub const scheduleSnapshot = Runtime.scheduleSnapshot;
123         pub const createPlan = Runtime.createPlan;
124         pub const createCheckedPlan = Runtime.createCheckedPlan;
125         pub const compileFragment = Runtime.compileFragment;
126         pub const createKernelArtifact = Runtime.createKernelArtifact;
127         pub const createKernelCallArtifact = Runtime.createKernelCallArtifact;
128         pub const runCpu = Runtime.runCpu;
129         pub const runCpuWithDiagnostic = Runtime.runCpuWithDiagnostic;
130         pub const verify = Runtime.verify;
131     };
132 }
133 
134 fn Analysis(comptime ProgramType: type, comptime analysis_impl: anytype) type {
135     return struct {
136         pub const Program: type = ProgramType;
137         pub const name = ProgramType.name;
138         pub const parameters = ProgramType.parameters;
139         pub const Layout: type = ProgramType.Layout;
140         pub const Limits: type = kernel.Limits;
141 
142         pub fn arg(comptime name_value: anytype) usize {
143             return Layout.index(name_value);
144         }
145 
146         pub fn schema() []const kernel.Param {
147             return ProgramType.schema();
148         }
149 
150         pub fn interpret(allocator: std.mem.Allocator, limits: Limits) !InterpretResult(ProgramType, @TypeOf(analysis_impl)) {
151             return ProgramType.interpret(allocator, limits, analysis_impl);
152         }
153 
154         pub fn transform(comptime make_transform: anytype) type {
155             return AnalysisTransform(ProgramType, analysis_impl, make_transform);
156         }
157     };
158 }
159 
160 fn AnalysisTransform(comptime ProgramType: type, comptime analysis_impl: anytype, comptime make_transform: anytype) type {
161     return struct {
162         const Self = @This();
163         const Runtime: type = RuntimeSurface(Self);
164         const TransformSpec: type = AnalysisTransformSpec(
165             ProgramType,
166             analysis_impl,
167             make_transform,
168         );
169 
170         pub const name = ProgramType.name;
171         pub const parameters = ProgramType.parameters;
172         pub const Layout: type = ProgramType.Layout;
173         pub const schedule_policy = ProgramType.schedule_policy;
174         pub const Limits: type = kernel.Limits;
175 
176         pub fn withSchedule(comptime next_schedule: anytype) type {
177             return AnalysisTransform(ProgramType.withSchedule(next_schedule), analysis_impl, make_transform);
178         }
179 
180         pub fn transform(comptime impl: anytype) type {
181             return Derived(Self, kernel.interpret.asSpec(impl));
182         }
183 
184         pub fn analyze(comptime next_analysis_impl: anytype) type {
185             return Analysis(Self, next_analysis_impl);
186         }
187 
188         pub fn arg(comptime name_value: anytype) usize {
189             return Layout.index(name_value);
190         }
191 
192         pub fn schema() []const kernel.Param {
193             return ProgramType.schema();
194         }
195 
196         pub fn build(allocator: std.mem.Allocator, limits: Limits) !kernel.Graph {
197             var analysis = try ProgramType.interpret(allocator, limits, analysis_impl);
198             defer kernel.interpret.deinitIfPresent(&analysis);
199             return ProgramType.interpret(allocator, limits, make_transform(&analysis));
200         }
201 
202         pub fn interpret(allocator: std.mem.Allocator, limits: Limits, initial: anytype) !InterpretResult(ProgramType, StackedSpec(TransformSpec, @TypeOf(initial))) {
203             var analysis = try ProgramType.interpret(allocator, limits, analysis_impl);
204             defer kernel.interpret.deinitIfPresent(&analysis);
205             return ProgramType.interpret(allocator, limits, kernel.interpret.stack(.{ kernel.interpret.asSpec(make_transform(&analysis)), kernel.interpret.asSpec(initial) }));
206         }
207 
208         pub const launch = Runtime.launch;
209         pub const scheduleSnapshot = Runtime.scheduleSnapshot;
210         pub const createPlan = Runtime.createPlan;
211         pub const createCheckedPlan = Runtime.createCheckedPlan;
212         pub const compileFragment = Runtime.compileFragment;
213         pub const createKernelArtifact = Runtime.createKernelArtifact;
214         pub const createKernelCallArtifact = Runtime.createKernelCallArtifact;
215         pub const runCpu = Runtime.runCpu;
216         pub const runCpuWithDiagnostic = Runtime.runCpuWithDiagnostic;
217         pub const verify = Runtime.verify;
218     };
219 }
220 
221 fn StackedSpec(comptime Transform: type, comptime Initial: type) type {
222     return @TypeOf(kernel.interpret.stack(.{ @as(Transform, undefined), @as(kernel.interpret.InterpretSpec(Initial), undefined) }));
223 }
224 
225 fn InterpretResult(comptime ProgramType: type, comptime Initial: type) type {
226     const Return = @TypeOf(ProgramType.interpret(
227         @as(std.mem.Allocator, undefined),
228         @as(kernel.Limits, undefined),
229         @as(Initial, undefined),
230     ));
231     return @typeInfo(Return).error_union.payload;
232 }
233 
234 fn AnalysisTransformSpec(comptime ProgramType: type, comptime analysis_impl: anytype, comptime make_transform: anytype) type {
235     return kernel.interpret.InterpretSpec(@TypeOf(make_transform(@as(*const InterpretResult(ProgramType, @TypeOf(analysis_impl)), undefined))));
236 }
237 
238 fn RuntimeSurface(comptime ProgramType: type) type {
239     return struct {
240         pub fn launch(allocator: std.mem.Allocator, limits: kernel.Limits) !kernel.Launch {
241             var graph = try ProgramType.build(allocator, limits);
242             defer graph.deinit();
243             return graph.launch();
244         }
245 
246         pub fn scheduleSnapshot(allocator: std.mem.Allocator, limits: kernel.Limits) !kernel.ScheduleSnapshot {
247             var graph = try ProgramType.build(allocator, limits);
248             defer graph.deinit();
249             return graph.scheduleSnapshot(allocator);
250         }
251 
252         pub fn createPlan(allocator: std.mem.Allocator, limits: kernel.Limits, options: kernel.PlanOptions) !kernel.Plan {
253             var graph = try ProgramType.build(allocator, limits);
254             defer graph.deinit();
255             return graph.createPlan(allocator, options);
256         }
257 
258         pub fn createCheckedPlan(allocator: std.mem.Allocator, limits: kernel.Limits, options: kernel.PlanOptions) !kernel.Plan {
259             var graph = try ProgramType.build(allocator, limits);
260             defer graph.deinit();
261             return graph.createCheckedPlan(allocator, options);
262         }
263 
264         pub fn compileFragment(
265             allocator: std.mem.Allocator,
266             limits: kernel.Limits,
267             handle: kernel.BackendHandle,
268             options: executable.FragmentCompilerOptions,
269         ) !*executable.CompiledFragment {
270             var graph = try ProgramType.build(allocator, limits);
271             defer graph.deinit();
272             return kernel_compile.compileFragment(allocator, handle, &graph, options);
273         }
274 
275         pub fn createKernelArtifact(
276             allocator: std.mem.Allocator,
277             limits: kernel.Limits,
278             handle: kernel.BackendHandle,
279             options: executable.KernelCompilerOptions,
280         ) !kernel.KernelArtifact {
281             var graph = try ProgramType.build(allocator, limits);
282             defer graph.deinit();
283             return kernel_compile.createArtifact(allocator, handle, &graph, options);
284         }
285 
286         pub fn createKernelCallArtifact(
287             allocator: std.mem.Allocator,
288             limits: kernel.Limits,
289             handle: kernel.BackendHandle,
290             options: kernel.KernelCallArtifactOptions,
291         ) !kernel.OwnedKernelCallArtifact {
292             var graph = try ProgramType.build(allocator, limits);
293             defer graph.deinit();
294             return kernel.createKernelCallArtifact(allocator, handle, &graph, options);
295         }
296 
297         pub fn runCpu(allocator: std.mem.Allocator, limits: kernel.Limits, args: []const kernel.Argument) !void {
298             var graph = try ProgramType.build(allocator, limits);
299             defer graph.deinit();
300             try graph.runCpu(allocator, args);
301         }
302 
303         pub fn runCpuWithDiagnostic(
304             allocator: std.mem.Allocator,
305             limits: kernel.Limits,
306             args: []const kernel.Argument,
307             diagnostic: *kernel.ExecutionDiagnostic,
308         ) !void {
309             var graph = try ProgramType.build(allocator, limits);
310             defer graph.deinit();
311             try graph.runCpuWithDiagnostic(allocator, args, diagnostic);
312         }
313 
314         pub fn verify(allocator: std.mem.Allocator, limits: kernel.Limits) !void {
315             var graph = try ProgramType.build(allocator, limits);
316             defer graph.deinit();
317             try graph.verify();
318         }
319     };
320 }
321 
322 fn ScheduledDefinition(comptime definition: anytype, comptime schedule: anytype) type {
323     return struct {
324         pub const name = definition.name;
325         pub const parameters = definition.parameters;
326 
327         pub fn body(raw_builder: anytype, _: anytype) !void {
328             var logical = logical_builder.wrap(raw_builder, schedule);
329             try call(definition, &logical);
330         }
331     };
332 }
333 
334 fn call(comptime definition: anytype, k: anytype) !void {
335     if (comptime parameter.named(definition.parameters)) {
336         try definition.body(k, parameter.Args(definition.parameters, PointerChild(@TypeOf(k))){ .builder = k });
337     } else {
338         const count = comptime parameter.arity(definition.parameters);
339         var args: [count]kernel.Value = undefined;
340         inline for (0..count) |index_value| {
341             args[index_value] = k.argument(index_value);
342         }
343         try definition.body(k, args[0..]);
344     }
345 }
346 
347 fn PointerChild(comptime Pointer: type) type {
348     return switch (@typeInfo(Pointer)) {
349         .pointer => |info| info.child,
350         else => @compileError("logical kernel Program bodies receive builder pointers"),
351     };
352 }