lib/accy/src/executable/composition/cpu/compilation.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const gpu = @import("gpu");
  3 const choir = @import("choir");
  4 const accy_root = @import("../../../root.zig");
  5 const executable = @import("../../root.zig");
  6 const composition = @import("../root.zig");
  7 const cpu = @import("root.zig");
  8 
  9 const Allocator = std.mem.Allocator;
 10 const accy = accy_root.choir;
 11 const ChoirComposition = choir.composition;
 12 const preparation = accy_root.preparation;
 13 
 14 pub const CpuObjectCompileOptions = struct {
 15     request: executable.FragmentCompilationRequest,
 16     compiler: executable.FragmentCompilerOptions = .{},
 17 };
 18 
 19 pub const CpuObjectCompilation = struct {
 20     allocator: Allocator,
 21     fragment_id: ChoirComposition.FragmentId,
 22     launch_tuning_artifact: []const u8,
 23     input_product: choir.product.ProductKey,
 24     compiled_fragment: ?*executable.CompiledFragment,
 25     metadata: composition.CpuObjectMetadata,
 26 
 27     pub fn deinit(self: *CpuObjectCompilation) void {
 28         if (self.compiled_fragment) |fragment| fragment.deinit();
 29         self.metadata.deinit();
 30         self.input_product.record.release();
 31         if (self.launch_tuning_artifact.len != 0) self.allocator.free(@constCast(self.launch_tuning_artifact));
 32         self.* = undefined;
 33     }
 34 
 35     pub fn materializer(self: *CpuObjectCompilation) ChoirComposition.FragmentMaterializer {
 36         return .{
 37             .id = self.fragment_id,
 38             .context = self,
 39             .materialize = materialize,
 40         };
 41     }
 42 };
 43 
 44 pub fn compileCpuObject(
 45     allocator: Allocator,
 46     fragment_id: ChoirComposition.FragmentId,
 47     module: *accy.SemanticModule,
 48     options: CpuObjectCompileOptions,
 49     report: *preparation.publication.PreparationReport,
 50     comptime configuration: choir.product.operation.Configuration,
 51 ) !CpuObjectCompilation {
 52     defer module.deinit();
 53     const launch_tuning_artifact: []const u8 = if (options.compiler.launch_tuning_artifact.len == 0)
 54         &.{}
 55     else
 56         try allocator.dupe(u8, options.compiler.launch_tuning_artifact);
 57     errdefer if (launch_tuning_artifact.len != 0) allocator.free(@constCast(launch_tuning_artifact));
 58     var compiler_options = options.compiler;
 59     compiler_options.artifact_format = .cpu_object;
 60     var backend_state = gpu.cpu.State.init(allocator);
 61     defer backend_state.deinit();
 62     var plan: executable.fragment.FragmentPreparationPlan = undefined;
 63     try plan.init(allocator, backend_state.handle(), compiler_options);
 64     defer plan.deinit();
 65     const prepared = try preparation.publication.prepare(allocator, .{ .draft = module }, .{
 66         .source = options.request.source,
 67         .variant = options.request.variant,
 68         .work = options.request.work,
 69         .record_bytes = options.request.record_bytes,
 70         .options = plan.run_options,
 71     }, report, configuration);
 72     defer prepared.deinit();
 73     try recordPreparation(compiler_options.instrumentation, report);
 74     const input_record = try prepared.stage(.semantic).metadata().retain();
 75     errdefer input_record.release();
 76     const compiled_fragment = try executable.compileFragmentFromPreparedModule(
 77         allocator,
 78         backend_state.handle(),
 79         prepared,
 80         compiler_options,
 81         options.request.artifact_workspace,
 82         configuration,
 83     );
 84     errdefer compiled_fragment.deinit();
 85     const metadata = try composition.CpuObjectMetadata.init(allocator, compiled_fragment);
 86     return .{
 87         .allocator = allocator,
 88         .fragment_id = fragment_id,
 89         .launch_tuning_artifact = launch_tuning_artifact,
 90         .input_product = choir.product.productKey(input_record),
 91         .compiled_fragment = compiled_fragment,
 92         .metadata = metadata,
 93     };
 94 }
 95 
 96 fn recordPreparation(
 97     instrumentation: executable.FragmentInstrumentation,
 98     report: *const preparation.publication.PreparationReport,
 99 ) !void {
100     const phases = [_]executable.FragmentPhase{
101         .run_contract_pipeline, .run_tensor_pipeline, .run_dispatch_pipeline,
102         .run_memory_pipeline,   .run_kernel_pipeline, .run_target_pipeline,
103     };
104     for (phases, report.elapsed_ns[1..]) |phase, elapsed| {
105         try instrumentation.recordElapsed(phase, elapsed);
106     }
107 }
108 
109 fn materialize(
110     ptr: *anyopaque,
111     allocator: Allocator,
112     module: *const ChoirComposition.CompositionModule,
113     fragment: *const ChoirComposition.Fragment,
114 ) anyerror!ChoirComposition.MaterializedFragment {
115     const compilation: *CpuObjectCompilation = @ptrCast(@alignCast(ptr));
116     return cpu.materializeCpuObject(compilation, allocator, module, fragment);
117 }