lib/accy/src/executable/compiler.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const gpu = @import("gpu");
  3 const choir = @import("choir");
  4 const fragment = @import("fragment.zig");
  5 const preparation = @import("../preparation/root.zig");
  6 const semantic = @import("../choir/root.zig").semantic;
  7 const Prepared = preparation.pipeline.BackendPreparedModule;
  8 
  9 pub const FragmentCompilationRequest = struct {
 10     source: []const u8,
 11     variant: []const u8 = "",
 12     work: choir.product.revision.receipt.Limits,
 13     record_bytes: u32,
 14     artifact_workspace: []u8,
 15 };
 16 
 17 pub const FragmentCompilerCacheUpdate = struct {
 18     preparation: preparation.BackendPreparationCacheUpdate,
 19     /// The cache owns the loaded fragment that this refresh produced. A caller
 20     /// launches the current program through this pointer after a refresh. The
 21     /// pointer stays valid until the next refresh that succeeds or until the
 22     /// cache is destroyed, and both free it.
 23     fragment: *fragment.LoadedFragment,
 24 
 25     pub fn deinit(self: *FragmentCompilerCacheUpdate, allocator: std.mem.Allocator) void {
 26         self.preparation.deinit(allocator);
 27         self.* = undefined;
 28     }
 29 };
 30 
 31 /// The cache holds one preparation cache and the loaded fragment built from it.
 32 /// A caller keeps one of these per program to recompile and reload it cheaply
 33 /// as the program changes. A refresh reuses a compile stage only when the
 34 /// preparation cache's own admission check accepts it. Device code is compiled
 35 /// from the prepared result and loaded again on every refresh. A refresh works
 36 /// on a forked copy of the preparation cache and swaps the copy and the new
 37 /// loaded fragment in only after both succeed, so a failure leaves the earlier
 38 /// preparation and fragment in place.
 39 pub const FragmentCompilerCache = struct {
 40     allocator: std.mem.Allocator,
 41     handle: gpu.BackendHandle,
 42     preparation_cache: preparation.BackendPreparationCache,
 43     loaded: ?*fragment.LoadedFragment = null,
 44 
 45     pub fn init(allocator: std.mem.Allocator, handle: gpu.BackendHandle) FragmentCompilerCache {
 46         return .{
 47             .allocator = allocator,
 48             .handle = handle,
 49             .preparation_cache = preparation.BackendPreparationCache.init(allocator),
 50         };
 51     }
 52 
 53     pub fn deinit(self: *FragmentCompilerCache) void {
 54         if (self.loaded) |loaded| loaded.deinit();
 55         self.preparation_cache.deinit();
 56         self.* = undefined;
 57     }
 58 
 59     pub fn currentPrepared(self: *const FragmentCompilerCache) ?*const Prepared {
 60         return self.preparation_cache.currentPrepared();
 61     }
 62 
 63     pub fn currentFragment(self: *const FragmentCompilerCache) ?*fragment.LoadedFragment {
 64         return self.loaded;
 65     }
 66 
 67     /// The call frees the draft, the caller's semantic module, whether the
 68     /// refresh succeeds or fails. A caller hands a new draft of the program to
 69     /// this call to recompile it. The caller owns `report`, which holds the
 70     /// receipts of the stages the refresh produced.
 71     pub fn refreshFromSemanticModule(
 72         self: *FragmentCompilerCache,
 73         module: *semantic.SemanticModule,
 74         options: fragment.FragmentCompilerOptions,
 75         request: FragmentCompilationRequest,
 76         report: *preparation.publication.PreparationReport,
 77         comptime configuration: choir.product.operation.Configuration,
 78     ) !FragmentCompilerCacheUpdate {
 79         return self.refresh(.{ .draft = module }, options, request, report, configuration);
 80     }
 81 
 82     pub fn refreshFromSemanticRevision(
 83         self: *FragmentCompilerCache,
 84         source: *const choir.product.revision.Revision,
 85         options: fragment.FragmentCompilerOptions,
 86         request: FragmentCompilationRequest,
 87         report: *preparation.publication.PreparationReport,
 88         comptime configuration: choir.product.operation.Configuration,
 89     ) !FragmentCompilerCacheUpdate {
 90         return self.refresh(.{ .retained = source }, options, request, report, configuration);
 91     }
 92 
 93     const Input = union(enum) {
 94         draft: *semantic.SemanticModule,
 95         retained: *const choir.product.revision.Revision,
 96     };
 97 
 98     fn refresh(
 99         self: *FragmentCompilerCache,
100         input: Input,
101         options: fragment.FragmentCompilerOptions,
102         request: FragmentCompilationRequest,
103         report: *preparation.publication.PreparationReport,
104         comptime configuration: choir.product.operation.Configuration,
105     ) !FragmentCompilerCacheUpdate {
106         var draft_owned = input == .draft;
107         defer if (draft_owned) input.draft.deinit();
108         var plan: fragment.FragmentPreparationPlan = undefined;
109         try plan.init(self.allocator, self.handle, options);
110         defer plan.deinit();
111         var staged = try self.preparation_cache.fork();
112         var staged_owned = true;
113         defer if (staged_owned) staged.deinit();
114         const current = preparation.publication.PreparationRequest{
115             .source = request.source,
116             .variant = request.variant,
117             .work = request.work,
118             .record_bytes = request.record_bytes,
119             .options = plan.run_options,
120         };
121         draft_owned = false;
122         var update = switch (input) {
123             .draft => |module| try staged.refreshFromSemanticModule(module, current, report, configuration),
124             .retained => |source| try staged.refreshFromSemanticRevision(source, current, report, configuration),
125         };
126         errdefer update.deinit(self.allocator);
127         try recordPreparation(options.instrumentation, report);
128         const compiled = try fragment.compileFragmentFromPreparedModule(
129             self.allocator,
130             self.handle,
131             update.prepared,
132             options,
133             request.artifact_workspace,
134             configuration,
135         );
136         const loaded = try fragment.loadFragment(self.allocator, self.handle, compiled, options);
137         self.preparation_cache.deinit();
138         self.preparation_cache = staged;
139         staged_owned = false;
140         if (self.loaded) |previous| previous.deinit();
141         self.loaded = loaded;
142         return .{ .preparation = update, .fragment = loaded };
143     }
144 
145     fn recordPreparation(
146         instrumentation: fragment.FragmentInstrumentation,
147         report: *const preparation.publication.PreparationReport,
148     ) !void {
149         const phases = [_]fragment.FragmentPhase{
150             .run_contract_pipeline, .run_tensor_pipeline, .run_dispatch_pipeline,
151             .run_memory_pipeline,   .run_kernel_pipeline, .run_target_pipeline,
152         };
153         for (phases, report.elapsed_ns[1..]) |phase, elapsed| {
154             try instrumentation.recordElapsed(phase, elapsed);
155         }
156     }
157 };