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

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const choir = @import("choir");
  3 const executable = @import("../../root.zig");
  4 const composition = @import("../root.zig");
  5 const cpu = @import("root.zig");
  6 
  7 const Allocator = std.mem.Allocator;
  8 const ChoirComposition = choir.composition;
  9 
 10 pub fn materializeCpuObject(
 11     compilation: *cpu.CpuObjectCompilation,
 12     allocator: Allocator,
 13     module: *const ChoirComposition.CompositionModule,
 14     fragment: *const ChoirComposition.Fragment,
 15 ) !ChoirComposition.MaterializedFragment {
 16     try verifyDurableFragment(compilation, module, fragment);
 17     const compiled_fragment = compilation.compiled_fragment orelse return error.AlreadyMaterialized;
 18     const state = try allocator.create(cpu.runtime.State);
 19     state.* = cpu.runtime.State.init(allocator, fragment.id);
 20     var state_owned = true;
 21     errdefer if (state_owned) {
 22         state.deinitUnloaded();
 23         allocator.destroy(state);
 24     };
 25 
 26     const runtime_input_sizes = try retainByteSizes(allocator, compilation.metadata.input_boundaries);
 27     var input_sizes_owned = true;
 28     errdefer if (input_sizes_owned) allocator.free(runtime_input_sizes);
 29     const runtime_output_sizes = try retainByteSizes(allocator, compilation.metadata.output_boundaries);
 30     var output_sizes_owned = true;
 31     errdefer if (output_sizes_owned) allocator.free(runtime_output_sizes);
 32 
 33     compilation.compiled_fragment = null;
 34     const load_options = executable.FragmentCompilerOptions{
 35         .artifact_format = .cpu_object,
 36         .launch_tuning_artifact = compilation.launch_tuning_artifact,
 37     };
 38     const loaded_fragment = try executable.loadFragment(
 39         allocator,
 40         state.handle(),
 41         compiled_fragment,
 42         load_options,
 43     );
 44     var loaded_owned = true;
 45     errdefer if (loaded_owned) loaded_fragment.deinit();
 46 
 47     state.install(loaded_fragment, runtime_input_sizes, runtime_output_sizes);
 48     state_owned = false;
 49     loaded_owned = false;
 50     input_sizes_owned = false;
 51     output_sizes_owned = false;
 52 
 53     return .{
 54         .state = state,
 55         .vtable = &cpu.runtime.vtable,
 56     };
 57 }
 58 
 59 fn verifyDurableFragment(
 60     compilation: *const cpu.CpuObjectCompilation,
 61     module: *const ChoirComposition.CompositionModule,
 62     fragment: *const ChoirComposition.Fragment,
 63 ) !void {
 64     if (fragment.id.value != compilation.fragment_id.value or fragment.pipeline != .accy) return error.FragmentMismatch;
 65     if (!fragment.pipeline_input.eql(compilation.input_product)) return error.ProductMismatch;
 66     if (!artifactsEqual(fragment.artifacts, compilation.metadata.artifacts)) return error.ArtifactMismatch;
 67     const partition = module.source_module.partition(fragment.partition) orelse return error.FragmentMismatch;
 68     if (partition.pipeline != .accy) return error.FragmentMismatch;
 69     if (partition.inputs.len != compilation.metadata.input_boundaries.len or partition.outputs.len != compilation.metadata.output_boundaries.len) {
 70         return error.BoundaryMismatch;
 71     }
 72     for (partition.inputs, compilation.metadata.input_boundaries) |boundary_id, metadata| {
 73         const boundary = module.source_module.boundary(boundary_id) orelse return error.BoundaryMismatch;
 74         if (!boundaryMatches(boundary, metadata)) return error.BoundaryMismatch;
 75     }
 76     for (partition.outputs, compilation.metadata.output_boundaries) |boundary_id, metadata| {
 77         const boundary = module.source_module.boundary(boundary_id) orelse return error.BoundaryMismatch;
 78         if (!boundaryMatches(boundary, metadata)) return error.BoundaryMismatch;
 79     }
 80     const entry_symbol = compilation.metadata.entrySymbol() orelse return error.ArtifactMismatch;
 81     for (fragment.exports) |export_value| {
 82         if (std.mem.eql(u8, export_value.symbol, entry_symbol) and export_value.abi_version == ChoirComposition.abi.version) return;
 83     }
 84     return error.ArtifactMismatch;
 85 }
 86 
 87 fn boundaryMatches(boundary: *const ChoirComposition.Boundary, metadata: composition.BoundaryMetadata) bool {
 88     const byte_size = std.math.cast(u64, metadata.byte_size) orelse return false;
 89     return boundary.byte_size == byte_size and
 90         boundary.element_type == metadata.element_type and
 91         std.mem.eql(u64, boundary.dimensions, metadata.dimensions);
 92 }
 93 
 94 fn retainByteSizes(allocator: Allocator, boundaries: []const composition.BoundaryMetadata) Allocator.Error![]usize {
 95     const byte_sizes = try allocator.alloc(usize, boundaries.len);
 96     for (boundaries, byte_sizes) |boundary, *byte_size| byte_size.* = boundary.byte_size;
 97     return byte_sizes;
 98 }
 99 
100 fn artifactsEqual(actual: []const choir.backends.artifact.Artifact, expected: []const choir.backends.artifact.Artifact) bool {
101     if (actual.len != expected.len) return false;
102     for (actual, expected) |actual_artifact, expected_artifact| {
103         if (!actual_artifact.eql(expected_artifact)) return false;
104     }
105     return true;
106 }