lib/accy/src/executable/invocation.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const accy_root = @import("../root.zig");
4 const fragment_mod = @import("fragment.zig");
5 const binding_mod = @import("binding.zig");
6 const plan_mod = @import("plan.zig");
7
8 pub const InvocationState = enum {
9 prepared,
10 running,
11 completed,
12 cancelled,
13 failed,
14 };
15
16 pub const InvocationError = gpu.BackendError || error{InvalidInvocationState};
17
18 const InvocationStorage = struct {
19 allocator: std.mem.Allocator,
20 fragment: *fragment_mod.LoadedFragment,
21 bindings: *binding_mod.PreparedLaunchBindings,
22 status: InvocationState = .prepared,
23 failure: ?gpu.BackendError = null,
24
25 fn deinit(self: *InvocationStorage) void {
26 self.bindings.deinit();
27 self.* = undefined;
28 }
29 };
30
31 pub const Invocation = opaque {
32 fn storageConst(self: *const Invocation) *const InvocationStorage {
33 return @ptrCast(@alignCast(self));
34 }
35
36 fn storageMut(self: *Invocation) *InvocationStorage {
37 return @ptrCast(@alignCast(self));
38 }
39
40 pub fn deinit(self: *Invocation) void {
41 const storage = self.storageMut();
42 const allocator = storage.allocator;
43 storage.deinit();
44 allocator.destroy(storage);
45 }
46
47 pub fn state(self: *const Invocation) InvocationState {
48 return self.storageConst().status;
49 }
50
51 pub fn failure(self: *const Invocation) ?gpu.BackendError {
52 return self.storageConst().failure;
53 }
54
55 pub fn cancel(self: *Invocation) InvocationError!void {
56 const storage = self.storageMut();
57 if (storage.status != .prepared) return error.InvalidInvocationState;
58 storage.status = .cancelled;
59 }
60
61 pub fn launch(
62 self: *Invocation,
63 scratch: std.mem.Allocator,
64 ) InvocationError!void {
65 try self.launchWithOptions(scratch, .{});
66 }
67
68 pub fn launchWithOptions(
69 self: *Invocation,
70 scratch: std.mem.Allocator,
71 options: fragment_mod.LaunchOptions,
72 ) InvocationError!void {
73 const storage = self.storageMut();
74 if (storage.status != .prepared) return error.InvalidInvocationState;
75 storage.status = .running;
76 storage.fragment.submitInvocationWithOptions(scratch, storage.bindings, options) catch |err| {
77 storage.failure = err;
78 storage.status = .failed;
79 return err;
80 };
81 storage.fragment.completeInvocationWithOptions(options) catch |err| {
82 storage.failure = err;
83 storage.status = .failed;
84 return err;
85 };
86 storage.status = .completed;
87 }
88
89 pub fn launchWithGraph(
90 self: *Invocation,
91 scratch: std.mem.Allocator,
92 graph: plan_mod.LaunchGraphPlan,
93 ) InvocationError!void {
94 const storage = self.storageMut();
95 if (storage.status != .prepared) return error.InvalidInvocationState;
96 storage.status = .running;
97 storage.fragment.submitInvocationWithGraph(scratch, storage.bindings, graph) catch |err| {
98 storage.failure = err;
99 storage.status = .failed;
100 return err;
101 };
102 storage.fragment.completeInvocationGraph(graph) catch |err| {
103 storage.failure = err;
104 storage.status = .failed;
105 return err;
106 };
107 storage.status = .completed;
108 }
109
110 pub fn outputCount(self: *const Invocation) usize {
111 return self.storageConst().fragment.outputCount();
112 }
113
114 pub fn readOutput(
115 self: *const Invocation,
116 index: usize,
117 host_bytes: []u8,
118 ) InvocationError!void {
119 const storage = self.storageConst();
120 if (storage.status != .completed) return error.InvalidInvocationState;
121 try storage.fragment.readInvocationOutput(storage.bindings, index, host_bytes);
122 }
123
124 pub fn readOutputs(
125 self: *const Invocation,
126 outputs: []const []u8,
127 ) InvocationError!void {
128 if (outputs.len != self.outputCount()) return error.InvalidArtifact;
129 for (outputs, 0..) |host_bytes, index| {
130 try self.readOutput(index, host_bytes);
131 }
132 }
133
134 pub fn measureLaunchCandidateRecords(
135 self: *const Invocation,
136 result_allocator: std.mem.Allocator,
137 scratch: std.mem.Allocator,
138 kernel_index: usize,
139 options: fragment_mod.LaunchCandidateBenchmarkOptions,
140 ) gpu.BackendError![]fragment_mod.LaunchCandidateRecord {
141 const storage = self.storageConst();
142 if (storage.status != .prepared) return error.LaunchArgumentMismatch;
143 return try storage.fragment.measureInvocationLaunchCandidates(
144 result_allocator,
145 scratch,
146 kernel_index,
147 storage.bindings,
148 options,
149 );
150 }
151
152 pub fn measureAndRecordLaunchCandidateRecords(
153 self: *Invocation,
154 result_allocator: std.mem.Allocator,
155 scratch: std.mem.Allocator,
156 options: fragment_mod.LaunchCandidateBenchmarkOptions,
157 ) gpu.BackendError![]fragment_mod.LaunchCandidateRecord {
158 const storage = self.storageMut();
159 if (storage.status != .prepared) return error.LaunchArgumentMismatch;
160 return try storage.fragment.measureAndRecordInvocationLaunchCandidates(
161 result_allocator,
162 scratch,
163 storage.bindings,
164 options,
165 );
166 }
167 };
168
169 pub fn prepare(
170 fragment: *fragment_mod.LoadedFragment,
171 allocator: std.mem.Allocator,
172 inputs: []const []const u8,
173 ) !*Invocation {
174 const bindings = try fragment.prepareInvocationBindings(allocator, inputs);
175 errdefer bindings.deinit();
176
177 const storage = allocator.create(InvocationStorage) catch return error.OutOfMemory;
178 storage.* = .{
179 .allocator = allocator,
180 .fragment = fragment,
181 .bindings = bindings,
182 };
183 return @ptrCast(storage);
184 }
185
186 pub fn run(
187 fragment: *fragment_mod.LoadedFragment,
188 allocator: std.mem.Allocator,
189 scratch: std.mem.Allocator,
190 inputs: []const []const u8,
191 outputs: []const []u8,
192 ) InvocationError!void {
193 var invocation = try prepare(fragment, allocator, inputs);
194 defer invocation.deinit();
195 try invocation.launch(scratch);
196 try invocation.readOutputs(outputs);
197 }
198
199 test "invocation cancellation is terminal before launch" {
200 const fixture = @import("fixture.zig");
201
202 const allocator = std.testing.allocator;
203 var backend_state = gpu.recording.BackendState{
204 .allocator = allocator,
205 .kind = .cuda,
206 .format = .cuda_ptx,
207 };
208 const handle = backend_state.handle();
209 const module = try fixture.addSemanticModule(allocator, "invocation_cancel");
210 const compiled = try fragment_mod.compileFragmentFromSemanticModule(
211 allocator,
212 handle,
213 module,
214 .{ .artifact_format = .cuda_ptx },
215 );
216 var fragment = try fragment_mod.loadFragment(
217 allocator,
218 handle,
219 compiled,
220 .{ .artifact_format = .cuda_ptx },
221 );
222 defer fragment.deinit();
223
224 const input = [_]f32{ 1, 2, 3, 4, 5, 6, 7, 8 };
225 var invocation = try prepare(fragment, allocator, &.{ std.mem.sliceAsBytes(&input), std.mem.sliceAsBytes(&input) });
226 defer invocation.deinit();
227
228 try std.testing.expectEqual(InvocationState.prepared, invocation.state());
229 try invocation.cancel();
230 try std.testing.expectEqual(InvocationState.cancelled, invocation.state());
231 try std.testing.expectError(error.InvalidInvocationState, invocation.launch(allocator));
232 var output: [8]f32 = undefined;
233 try std.testing.expectError(error.InvalidInvocationState, invocation.readOutput(0, std.mem.sliceAsBytes(&output)));
234 }
235
236 test "synchronous invocation reaches completion before output access" {
237 const fixture = @import("fixture.zig");
238
239 const allocator = std.testing.allocator;
240 var backend_state = gpu.recording.BackendState{
241 .allocator = allocator,
242 .kind = .cuda,
243 .format = .cuda_ptx,
244 };
245 const handle = backend_state.handle();
246 const module = try fixture.addSemanticModule(allocator, "invocation_complete");
247 const compiled = try fragment_mod.compileFragmentFromSemanticModule(
248 allocator,
249 handle,
250 module,
251 .{ .artifact_format = .cuda_ptx },
252 );
253 var fragment = try fragment_mod.loadFragment(
254 allocator,
255 handle,
256 compiled,
257 .{ .artifact_format = .cuda_ptx },
258 );
259 defer fragment.deinit();
260
261 const input = [_]f32{ 1, 2, 3, 4, 5, 6, 7, 8 };
262 var invocation = try prepare(fragment, allocator, &.{ std.mem.sliceAsBytes(&input), std.mem.sliceAsBytes(&input) });
263 defer invocation.deinit();
264
265 var output: [8]f32 = undefined;
266 try std.testing.expectError(error.InvalidInvocationState, invocation.readOutput(0, std.mem.sliceAsBytes(&output)));
267 try invocation.launch(allocator);
268 try std.testing.expectEqual(InvocationState.completed, invocation.state());
269 try std.testing.expect(invocation.failure() == null);
270 try invocation.readOutput(0, std.mem.sliceAsBytes(&output));
271 }
272
273 test "default invocation scale reuses retained plans and synchronizes only the default stream" {
274 const fixture = @import("fixture.zig");
275
276 const allocator = std.testing.allocator;
277 var backend_state = gpu.recording.BackendState{
278 .allocator = allocator,
279 .kind = .cuda,
280 .format = .cuda_ptx,
281 };
282 const handle = backend_state.handle();
283 const module = try fixture.addSemanticModule(allocator, "invocation_retained_plan_scale");
284 const compiled = try fragment_mod.compileFragmentFromSemanticModule(
285 allocator,
286 handle,
287 module,
288 .{ .artifact_format = .cuda_ptx },
289 );
290 var fragment = try fragment_mod.loadFragment(
291 allocator,
292 handle,
293 compiled,
294 .{ .artifact_format = .cuda_ptx },
295 );
296 defer fragment.deinit();
297
298 const input = [_]f32{ 1, 2, 3, 4, 5, 6, 7, 8 };
299 var failing = std.testing.FailingAllocator.init(allocator, .{});
300 failing.fail_index = failing.alloc_index;
301 failing.resize_fail_index = failing.resize_index;
302 const invocation_count = 12;
303 for (0..invocation_count) |_| {
304 var invocation = try prepare(fragment, allocator, &.{ std.mem.sliceAsBytes(&input), std.mem.sliceAsBytes(&input) });
305 defer invocation.deinit();
306 try invocation.launch(failing.allocator());
307 try std.testing.expectEqual(InvocationState.completed, invocation.state());
308 }
309
310 try std.testing.expect(!failing.has_induced_failure);
311 try std.testing.expectEqual(@as(usize, invocation_count), backend_state.launch_count);
312 try std.testing.expectEqual(@as(usize, invocation_count), backend_state.sync_count);
313 try std.testing.expectEqual(gpu.SyncScope.default_stream, backend_state.last_sync_scope.?);
314 try std.testing.expect(backend_state.last_sync_stream == null);
315 try std.testing.expect(backend_state.last_sync_event == null);
316 }
317
318 test "invocation launch failure is terminal" {
319 const fixture = @import("fixture.zig");
320
321 const allocator = std.testing.allocator;
322 var backend_state = gpu.recording.BackendState{
323 .allocator = allocator,
324 .kind = .cuda,
325 .format = .cuda_ptx,
326 .fail_launch_after_count = 0,
327 };
328 const handle = backend_state.handle();
329 const module = try fixture.addSemanticModule(allocator, "invocation_failure");
330 const compiled = try fragment_mod.compileFragmentFromSemanticModule(
331 allocator,
332 handle,
333 module,
334 .{ .artifact_format = .cuda_ptx },
335 );
336 var fragment = try fragment_mod.loadFragment(
337 allocator,
338 handle,
339 compiled,
340 .{ .artifact_format = .cuda_ptx },
341 );
342 defer fragment.deinit();
343
344 const input = [_]f32{ 1, 2, 3, 4, 5, 6, 7, 8 };
345 var invocation = try prepare(fragment, allocator, &.{ std.mem.sliceAsBytes(&input), std.mem.sliceAsBytes(&input) });
346 defer invocation.deinit();
347
348 try std.testing.expectError(error.RuntimeUnavailable, invocation.launch(allocator));
349 try std.testing.expectEqual(InvocationState.failed, invocation.state());
350 try std.testing.expectEqual(error.RuntimeUnavailable, invocation.failure().?);
351 try std.testing.expectError(error.InvalidInvocationState, invocation.launch(allocator));
352 }