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

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const gpu = @import("gpu");
  3 const choir = @import("choir");
  4 const sys = @import("sys");
  5 const accy_root = @import("../../../root.zig");
  6 const executable = @import("../../root.zig");
  7 
  8 const Allocator = std.mem.Allocator;
  9 const ChoirComposition = choir.composition;
 10 
 11 pub const RuntimeEvidence = struct {
 12     invocation_count: u64,
 13     completed_invocation_count: u64,
 14     failed_invocation_count: u64,
 15     loaded_artifact_count: usize,
 16     live_buffer_count: usize,
 17 };
 18 
 19 pub const RuntimeEvidenceError = error{ForeignLoadedFragment};
 20 
 21 pub const State = struct {
 22     allocator: Allocator,
 23     fragment_id: ChoirComposition.FragmentId,
 24     mutex: sys.thread.Mutex = .{},
 25     cpu: gpu.cpu.State,
 26     fragment: *executable.LoadedFragment,
 27     input_byte_sizes: []usize,
 28     output_byte_sizes: []usize,
 29     invocation_count: u64 = 0,
 30     completed_invocation_count: u64 = 0,
 31     failed_invocation_count: u64 = 0,
 32 
 33     pub fn init(allocator: Allocator, fragment_id: ChoirComposition.FragmentId) State {
 34         return .{
 35             .allocator = allocator,
 36             .fragment_id = fragment_id,
 37             .cpu = gpu.cpu.State.init(allocator),
 38             .fragment = undefined,
 39             .input_byte_sizes = &.{},
 40             .output_byte_sizes = &.{},
 41         };
 42     }
 43 
 44     pub fn handle(self: *State) gpu.BackendHandle {
 45         return self.cpu.handle();
 46     }
 47 
 48     pub fn install(
 49         self: *State,
 50         fragment: *executable.LoadedFragment,
 51         input_byte_sizes: []usize,
 52         output_byte_sizes: []usize,
 53     ) void {
 54         self.fragment = fragment;
 55         self.input_byte_sizes = input_byte_sizes;
 56         self.output_byte_sizes = output_byte_sizes;
 57     }
 58 
 59     pub fn deinitUnloaded(self: *State) void {
 60         std.debug.assert(self.cpu.buffers.count() == 0);
 61         std.debug.assert(self.cpu.loaded.count() == 0);
 62         self.cpu.deinit();
 63         self.* = undefined;
 64     }
 65 };
 66 
 67 pub const vtable = ChoirComposition.FragmentVTable{
 68     .invoke = invoke,
 69     .deinit = deinit,
 70 };
 71 
 72 pub fn runtimeEvidence(fragment: ChoirComposition.LoadedFragment) RuntimeEvidenceError!RuntimeEvidence {
 73     if (fragment.vtable != &vtable) return error.ForeignLoadedFragment;
 74     const state: *State = @ptrCast(@alignCast(fragment.state));
 75     state.mutex.lock();
 76     defer state.mutex.unlock();
 77     return .{
 78         .invocation_count = state.invocation_count,
 79         .completed_invocation_count = state.completed_invocation_count,
 80         .failed_invocation_count = state.failed_invocation_count,
 81         .loaded_artifact_count = state.cpu.loaded.count(),
 82         .live_buffer_count = state.cpu.buffers.count(),
 83     };
 84 }
 85 
 86 fn invoke(
 87     ptr: *anyopaque,
 88     call_site: *const ChoirComposition.CallSite,
 89     context: *ChoirComposition.abi.Context,
 90     frame: *ChoirComposition.abi.Frame,
 91 ) ChoirComposition.abi.Status {
 92     const state: *State = @ptrCast(@alignCast(ptr));
 93     const context_status = context.validate();
 94     if (context_status != .ok) return context_status;
 95     if (call_site.callee.value != state.fragment_id.value) return .missing_fragment;
 96     if (call_site.abi_version != ChoirComposition.abi.version) return .incompatible_frame;
 97     if (call_site.inputs.len != state.input_byte_sizes.len or call_site.outputs.len != state.output_byte_sizes.len) return .invalid_resource;
 98     const expected_resource_count = std.math.add(usize, call_site.inputs.len, call_site.outputs.len) catch return .invalid_resource;
 99 
100     const resolved = frame.resources(context);
101     const resources = switch (resolved) {
102         .status => |status| return status,
103         .value => |values| values,
104     };
105     if (resources.len != expected_resource_count) return .invalid_resource;
106 
107     var input_values: [ChoirComposition.abi.max_resources][]const u8 = undefined;
108     var output_values: [ChoirComposition.abi.max_resources][]u8 = undefined;
109     for (resources[0..call_site.inputs.len], 0..) |resource, index| {
110         const access = resource.accessValue() orelse return .invalid_resource;
111         if (access != .read and access != .read_write) return .invalid_resource;
112         const bytes = switch (resource.bytes()) {
113             .status => |status| return status,
114             .value => |value| value,
115         };
116         if (bytes.len != state.input_byte_sizes[index]) return .invalid_resource;
117         input_values[index] = bytes;
118     }
119     for (resources[call_site.inputs.len..], 0..) |resource, index| {
120         const access = resource.accessValue() orelse return .invalid_resource;
121         if (access != .write and access != .read_write) return .invalid_resource;
122         const bytes = switch (resource.bytes()) {
123             .status => |status| return status,
124             .value => |value| value,
125         };
126         if (bytes.len != state.output_byte_sizes[index]) return .invalid_resource;
127         output_values[index] = bytes;
128     }
129 
130     state.mutex.lock();
131     defer state.mutex.unlock();
132     state.invocation_count +|= 1;
133     const invocation = executable.prepareInvocation(
134         state.fragment,
135         state.allocator,
136         input_values[0..call_site.inputs.len],
137     ) catch {
138         state.failed_invocation_count +|= 1;
139         return .fragment_failure;
140     };
141     defer invocation.deinit();
142     invocation.launch(state.allocator) catch {
143         state.failed_invocation_count +|= 1;
144         return .fragment_failure;
145     };
146     invocation.readOutputs(output_values[0..call_site.outputs.len]) catch {
147         state.failed_invocation_count +|= 1;
148         return .fragment_failure;
149     };
150     state.completed_invocation_count +|= 1;
151     return .ok;
152 }
153 
154 fn deinit(ptr: *anyopaque, _: Allocator) void {
155     const state: *State = @ptrCast(@alignCast(ptr));
156     const allocator = state.allocator;
157     state.mutex.lock();
158     std.debug.assert(state.cpu.buffers.count() == 0);
159     std.debug.assert(state.cpu.loaded.count() == state.fragment.kernelCount());
160     state.fragment.deinit();
161     std.debug.assert(state.cpu.buffers.count() == 0);
162     std.debug.assert(state.cpu.loaded.count() == 0);
163     state.cpu.deinit();
164     allocator.free(state.output_byte_sizes);
165     allocator.free(state.input_byte_sizes);
166     state.mutex.unlock();
167     allocator.destroy(state);
168 }