lib/accy/src/executable/pipeline.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4
5 const artifact_product = @import("../artifact/model/root.zig");
6
7 pub const product_name = "accy.exec.pipeline";
8
9 pub const PipelineLaunch = struct {
10 pipeline: artifact_product.KernelCallPipeline,
11 registry: artifact_product.KernelCallRegistry,
12 format: gpu.ArtifactFormat,
13 operands: []const gpu.BufferBinding = &.{},
14 results: []const gpu.BufferBinding = &.{},
15 runtime_scalar_arguments: []const choir_abi.ScalarArgument = &.{},
16 intermediates: ?[]const gpu.BufferBinding = null,
17 artifacts: ?*const PipelineArtifactPool = null,
18 stream: ?gpu.StreamHandle = null,
19 diagnostic_id: ?[]const u8 = null,
20 };
21
22 /// The compiled and loaded kernels of one kernel pipeline, one per stage,
23 /// labeled with the chain's target, version and artifact format, and owned by
24 /// the caller. The caller keeps a pool to launch the same chain many times
25 /// without loading its kernels again. `deinit` destroys every loaded kernel and
26 /// frees every compiled one. `deinit` waits for nothing, so the caller first
27 /// makes sure no queued launch still uses the kernels.
28 pub const PipelineArtifactPool = struct {
29 allocator: std.mem.Allocator,
30 handle: gpu.BackendHandle,
31 target: []const u8,
32 version: u32,
33 format: gpu.ArtifactFormat,
34 entries: []PoolEntry,
35
36 pub const PoolEntry = struct {
37 artifact: gpu.KernelArtifact,
38 loaded: gpu.LoadedArtifact,
39 };
40
41 pub fn deinit(self: *PipelineArtifactPool) void {
42 deinitPipelinePoolEntries(self.handle, self.entries);
43 if (self.entries.len != 0) self.allocator.free(self.entries);
44 self.allocator.free(self.target);
45 self.* = undefined;
46 }
47 };
48
49 fn deinitPipelinePoolEntries(handle: gpu.BackendHandle, entries: []PipelineArtifactPool.PoolEntry) void {
50 for (entries) |*pool_entry| {
51 handle.destroyObject(pool_entry.loaded.id);
52 pool_entry.artifact.deinit();
53 }
54 }
55
56 /// Prepares a reusable set of loaded kernels once for a kernel pipeline. The
57 /// call checks the whole chain against the registry first, then compiles and
58 /// loads each stage's registry entry into a new pool. A stage whose entry is
59 /// missing for `format` gives `error.UnsupportedOperation`, and a failure part
60 /// way through unloads and frees the stages already loaded.
61 pub fn loadPipelineArtifacts(
62 allocator: std.mem.Allocator,
63 handle: gpu.BackendHandle,
64 pipeline: artifact_product.KernelCallPipeline,
65 registry: artifact_product.KernelCallRegistry,
66 format: gpu.ArtifactFormat,
67 ) gpu.BackendError!PipelineArtifactPool {
68 try pipeline.validate(registry, format);
69 const target = allocator.dupe(u8, pipeline.target) catch return error.OutOfMemory;
70 errdefer allocator.free(target);
71 const entries = allocator.alloc(PipelineArtifactPool.PoolEntry, pipeline.stages.len) catch return error.OutOfMemory;
72 var loaded_count: usize = 0;
73 errdefer {
74 deinitPipelinePoolEntries(handle, entries[0..loaded_count]);
75 allocator.free(entries);
76 }
77 for (pipeline.stages, entries) |stage, *pool_entry| {
78 const entry = registry.find(stage.target, stage.version, format) orelse return error.UnsupportedOperation;
79 var stage_artifact = try handle.createArtifact(.{
80 .kernel_name = entry.entry_name,
81 .requested_format = format,
82 .argument_count = entry.argument_count,
83 .scalar_argument_count = try entry.scalarArgumentCount(),
84 .required_dtypes = entry.required_dtypes,
85 .required_features = entry.required_features,
86 .required_subgroup = entry.required_subgroup,
87 .push_constants = entry.push_constants,
88 .payload = entry.payload,
89 });
90 errdefer stage_artifact.deinit();
91 const loaded = try handle.loadArtifact(&stage_artifact);
92 pool_entry.* = .{ .artifact = stage_artifact, .loaded = loaded };
93 loaded_count += 1;
94 }
95 return .{
96 .allocator = allocator,
97 .handle = handle,
98 .target = target,
99 .version = pipeline.version,
100 .format = format,
101 .entries = entries,
102 };
103 }
104
105 pub fn intermediateByteSize(
106 spec: artifact_product.PipelineIntermediate,
107 runtime_scalar_arguments: []const choir_abi.ScalarArgument,
108 ) gpu.BackendError!usize {
109 const extent = try spec.extent.resolveExtent(runtime_scalar_arguments);
110 if (extent == 0) return error.InvalidArtifact;
111 return std.math.mul(usize, extent, spec.dtype.sizeOf()) catch return error.InvalidArtifact;
112 }
113
114 /// Allocates each intermediate, a scratch buffer between stages, once so the
115 /// caller can pass the buffers to many launches. The call checks each runtime
116 /// scalar argument, then allocates one readable and writable device buffer for
117 /// each intermediate at the size those arguments give, aligned to 256 bytes.
118 /// The caller releases the buffers with `deinitPipelineIntermediates`, and only
119 /// after the work queued on them has finished.
120 pub fn allocatePipelineIntermediates(
121 allocator: std.mem.Allocator,
122 handle: gpu.BackendHandle,
123 pipeline: artifact_product.KernelCallPipeline,
124 runtime_scalar_arguments: []const choir_abi.ScalarArgument,
125 ) gpu.BackendError![]gpu.BufferBinding {
126 try pipeline.validateRuntimeScalarArguments(runtime_scalar_arguments);
127 const bindings = allocator.alloc(gpu.BufferBinding, pipeline.intermediates.len) catch return error.OutOfMemory;
128 var initialized: usize = 0;
129 errdefer {
130 destroyPipelineIntermediateObjects(handle, bindings[0..initialized]);
131 allocator.free(bindings);
132 }
133 for (pipeline.intermediates, bindings) |spec, *binding| {
134 const extent = try spec.extent.resolveExtent(runtime_scalar_arguments);
135 const byte_size = try intermediateByteSize(spec, runtime_scalar_arguments);
136 const buffer = try handle.allocateBuffer(.{
137 .byte_size = byte_size,
138 .alignment = 256,
139 .dtype = spec.dtype,
140 .element_count = extent,
141 });
142 binding.* = .{
143 .handle = buffer,
144 .access = .read_write,
145 .ownership = buffer.ownership,
146 .byte_size = buffer.byte_size,
147 };
148 initialized += 1;
149 }
150 return bindings;
151 }
152
153 pub fn deinitPipelineIntermediates(
154 allocator: std.mem.Allocator,
155 handle: gpu.BackendHandle,
156 bindings: []gpu.BufferBinding,
157 ) void {
158 destroyPipelineIntermediateObjects(handle, bindings);
159 allocator.free(bindings);
160 }
161
162 fn destroyPipelineIntermediateObjects(handle: gpu.BackendHandle, bindings: []const gpu.BufferBinding) void {
163 for (bindings) |binding| handle.destroyObject(binding.handle.id);
164 }
165
166 /// Runs a kernel pipeline once for the caller. The call checks the chain, the
167 /// operand and result counts, each runtime scalar argument, and any pool and
168 /// scratch buffers the caller passed, before launching any stage. A pool must
169 /// match the chain's stage count, version, format and target, and each scratch
170 /// buffer must be large enough and disjoint from every other scratch, operand
171 /// or result buffer, else `error.LaunchArgumentMismatch`. The stages launch in
172 /// their listed order on `request.stream`. When the call allocated scratch or
173 /// loaded kernels itself, the call waits for the stream on success before
174 /// destroying them. When the caller passed both a pool and scratch buffers, the
175 /// call returns without waiting, and the caller synchronizes before reusing or
176 /// freeing them. An error after launching starts a best-effort wait, and the
177 /// objects the call owns are destroyed even when that wait fails.
178 pub fn launchPipeline(
179 scratch: std.mem.Allocator,
180 handle: gpu.BackendHandle,
181 request: PipelineLaunch,
182 ) gpu.BackendError!void {
183 const pipeline = request.pipeline;
184 try pipeline.validate(request.registry, request.format);
185 if (request.operands.len != pipeline.operand_count) return error.LaunchArgumentMismatch;
186 if (request.results.len != pipeline.result_count) return error.LaunchArgumentMismatch;
187 try pipeline.validateRuntimeScalarArguments(request.runtime_scalar_arguments);
188
189 if (request.artifacts) |pool| {
190 if (pool.entries.len != pipeline.stages.len) return error.LaunchArgumentMismatch;
191 if (pool.version != pipeline.version) return error.LaunchArgumentMismatch;
192 if (pool.format != request.format) return error.LaunchArgumentMismatch;
193 if (!std.mem.eql(u8, pool.target, pipeline.target)) return error.LaunchArgumentMismatch;
194 }
195
196 var owned_intermediates: []gpu.BufferBinding = &.{};
197 const intermediates = if (request.intermediates) |provided| intermediates: {
198 try validatePipelineIntermediates(
199 pipeline,
200 request.runtime_scalar_arguments,
201 request.operands,
202 request.results,
203 provided,
204 );
205 break :intermediates provided;
206 } else intermediates: {
207 owned_intermediates = try allocatePipelineIntermediates(
208 scratch,
209 handle,
210 pipeline,
211 request.runtime_scalar_arguments,
212 );
213 break :intermediates owned_intermediates;
214 };
215 defer if (request.intermediates == null) deinitPipelineIntermediates(scratch, handle, owned_intermediates);
216
217 var no_loaded_artifacts: [0]gpu.LoadedArtifact = .{};
218 const loaded_artifacts = if (request.artifacts == null)
219 scratch.alloc(gpu.LoadedArtifact, pipeline.stages.len) catch return error.OutOfMemory
220 else
221 no_loaded_artifacts[0..];
222 defer if (request.artifacts == null) scratch.free(loaded_artifacts);
223 var loaded_count: usize = 0;
224 defer if (request.artifacts == null) destroyLoadedArtifacts(handle, loaded_artifacts[0..loaded_count]);
225
226 const needs_sync = request.intermediates == null or request.artifacts == null;
227 errdefer if (needs_sync) synchronizePipelineLaunch(handle, request.stream) catch {};
228 try launchStages(scratch, handle, request, intermediates, loaded_artifacts, &loaded_count);
229 if (needs_sync) try synchronizePipelineLaunch(handle, request.stream);
230 }
231
232 fn validatePipelineIntermediates(
233 pipeline: artifact_product.KernelCallPipeline,
234 runtime_scalar_arguments: []const choir_abi.ScalarArgument,
235 operands: []const gpu.BufferBinding,
236 results: []const gpu.BufferBinding,
237 provided: []const gpu.BufferBinding,
238 ) gpu.BackendError!void {
239 if (provided.len != pipeline.intermediates.len) return error.LaunchArgumentMismatch;
240 for (pipeline.intermediates, provided, 0..) |spec, binding, index| {
241 const byte_size = try intermediateByteSize(spec, runtime_scalar_arguments);
242 if (binding.byte_size < byte_size) return error.LaunchArgumentMismatch;
243 for (provided[0..index]) |previous| {
244 if (bufferBindingsAlias(binding, previous)) return error.LaunchArgumentMismatch;
245 }
246 for (operands) |operand| {
247 if (bufferBindingsAlias(binding, operand)) return error.LaunchArgumentMismatch;
248 }
249 for (results) |result| {
250 if (bufferBindingsAlias(binding, result)) return error.LaunchArgumentMismatch;
251 }
252 }
253 }
254
255 fn bufferBindingsAlias(lhs: gpu.BufferBinding, rhs: gpu.BufferBinding) bool {
256 return lhs.handle.backend == rhs.handle.backend and lhs.handle.id == rhs.handle.id;
257 }
258
259 fn destroyLoadedArtifacts(handle: gpu.BackendHandle, loaded_artifacts: []const gpu.LoadedArtifact) void {
260 for (loaded_artifacts) |loaded| handle.destroyObject(loaded.id);
261 }
262
263 fn synchronizePipelineLaunch(handle: gpu.BackendHandle, stream: ?gpu.StreamHandle) gpu.BackendError!void {
264 if (stream) |stream_handle| {
265 try handle.synchronize(.{ .scope = .stream, .stream = stream_handle });
266 } else {
267 try handle.synchronize(.{ .scope = .device });
268 }
269 }
270
271 fn launchStages(
272 scratch: std.mem.Allocator,
273 handle: gpu.BackendHandle,
274 request: PipelineLaunch,
275 intermediates: []const gpu.BufferBinding,
276 loaded_artifacts: []gpu.LoadedArtifact,
277 loaded_count: *usize,
278 ) gpu.BackendError!void {
279 for (request.pipeline.stages, 0..) |stage, stage_index| {
280 const entry = request.registry.find(stage.target, stage.version, request.format) orelse {
281 return error.UnsupportedOperation;
282 };
283 try launchStage(scratch, handle, request, entry, stage, stage_index, intermediates, loaded_artifacts, loaded_count);
284 }
285 }
286
287 fn launchStage(
288 scratch: std.mem.Allocator,
289 handle: gpu.BackendHandle,
290 request: PipelineLaunch,
291 entry: artifact_product.KernelCallArtifact,
292 stage: artifact_product.PipelineStage,
293 stage_index: usize,
294 intermediates: []const gpu.BufferBinding,
295 loaded_artifacts: []gpu.LoadedArtifact,
296 loaded_count: *usize,
297 ) gpu.BackendError!void {
298 const scalars = scratch.alloc(choir_abi.ScalarArgument, stage.scalars.len) catch return error.OutOfMemory;
299 defer scratch.free(scalars);
300 for (stage.scalars, scalars) |derivation, *value| {
301 value.* = try derivation.resolveScalar(request.runtime_scalar_arguments);
302 }
303
304 const buffers = scratch.alloc(gpu.BufferBinding, stage.buffers.len) catch return error.OutOfMemory;
305 defer scratch.free(buffers);
306 for (stage.buffers, buffers) |ref, *binding| {
307 binding.* = switch (ref) {
308 .operand => |index| request.operands[index],
309 .result => |index| request.results[index],
310 .intermediate => |index| intermediates[index],
311 };
312 }
313
314 const geometry = switch (entry.launch) {
315 .derived => |derived| try derived.geometry(scalars),
316 .fixed => |fixed| fixed,
317 };
318
319 if (request.artifacts) |pool| {
320 const pool_entry = &pool.entries[stage_index];
321 try handle.launch(.{
322 .artifact = &pool_entry.artifact,
323 .loaded_artifact = pool_entry.loaded,
324 .buffers = buffers,
325 .scalar_arguments = scalars,
326 .geometry = geometry,
327 .stream = request.stream,
328 .diagnostic_id = request.diagnostic_id,
329 });
330 return;
331 }
332
333 var stage_artifact = try handle.createArtifact(.{
334 .kernel_name = entry.entry_name,
335 .requested_format = request.format,
336 .argument_count = entry.argument_count,
337 .scalar_argument_count = try entry.scalarArgumentCount(),
338 .required_dtypes = entry.required_dtypes,
339 .required_features = entry.required_features,
340 .required_subgroup = entry.required_subgroup,
341 .push_constants = entry.push_constants,
342 .diagnostic_id = request.diagnostic_id,
343 .payload = entry.payload,
344 });
345 defer stage_artifact.deinit();
346 const loaded = try handle.loadArtifact(&stage_artifact);
347 errdefer handle.destroyObject(loaded.id);
348
349 try handle.launch(.{
350 .artifact = &stage_artifact,
351 .loaded_artifact = loaded,
352 .buffers = buffers,
353 .scalar_arguments = scalars,
354 .geometry = geometry,
355 .stream = request.stream,
356 .diagnostic_id = request.diagnostic_id,
357 });
358 if (loaded_count.* >= loaded_artifacts.len) return error.InvalidArtifact;
359 loaded_artifacts[loaded_count.*] = loaded;
360 loaded_count.* += 1;
361 }
362
363 const testing = std.testing;
364
365 fn pipelineTestEntry(
366 comptime target: []const u8,
367 argument_count: u32,
368 runtime_scalars: u32,
369 launch: artifact_product.KernelCallLaunch,
370 ) artifact_product.KernelCallArtifact {
371 return .{
372 .target = target,
373 .version = 1,
374 .format = .cuda_ptx,
375 .entry_name = target,
376 .argument_count = argument_count,
377 .payload = .{ .text = "// " ++ target },
378 .launch = launch,
379 .runtime_scalar_argument_count = runtime_scalars,
380 };
381 }
382
383 const test_entries = [_]artifact_product.KernelCallArtifact{
384 pipelineTestEntry("accy.kernel.test.block_scan", 4, 1, .{ .derived = .{
385 .grid = .{
386 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = 64 } },
387 .{ .fixed = 1 },
388 .{ .fixed = 1 },
389 },
390 .threadgroup = .{ 64, 1, 1 },
391 } }),
392 pipelineTestEntry("accy.kernel.test.sums_scan", 3, 1, .{ .derived = .{
393 .grid = .{
394 .{ .fixed = 1 },
395 .{ .fixed = 1 },
396 .{ .fixed = 1 },
397 },
398 .threadgroup = .{ 96, 1, 1 },
399 } }),
400 pipelineTestEntry("accy.kernel.test.add_base", 3, 1, .{ .derived = .{
401 .grid = .{
402 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = 64 } },
403 .{ .fixed = 1 },
404 .{ .fixed = 1 },
405 },
406 .threadgroup = .{ 64, 1, 1 },
407 } }),
408 };
409
410 const test_pipeline = artifact_product.KernelCallPipeline{
411 .target = "accy.kernel.test.device_scan",
412 .version = 1,
413 .operand_count = 1,
414 .result_count = 1,
415 .runtime_scalar_argument_count = 1,
416 .intermediates = &.{
417 .{ .dtype = .f32, .extent = .{ .ceil_div = .{ .argument_index = 0, .divisor = 64 } } },
418 .{ .dtype = .f32, .extent = .{ .ceil_div = .{ .argument_index = 0, .divisor = 64 } } },
419 },
420 .stages = &.{
421 .{
422 .target = "accy.kernel.test.block_scan",
423 .version = 1,
424 .buffers = &.{ .{ .result = 0 }, .{ .operand = 0 }, .{ .intermediate = 0 } },
425 .scalars = &.{.{ .forward = 0 }},
426 },
427 .{
428 .target = "accy.kernel.test.sums_scan",
429 .version = 1,
430 .buffers = &.{ .{ .intermediate = 1 }, .{ .intermediate = 0 } },
431 .scalars = &.{.{ .ceil_div = .{ .argument_index = 0, .divisor = 64 } }},
432 },
433 .{
434 .target = "accy.kernel.test.add_base",
435 .version = 1,
436 .buffers = &.{ .{ .result = 0 }, .{ .intermediate = 1 } },
437 .scalars = &.{.{ .forward = 0 }},
438 },
439 },
440 };
441
442 fn testBinding(id: gpu.BackendObjectId) gpu.BufferBinding {
443 return .{
444 .handle = .{ .id = id, .backend = .cuda, .byte_size = 20000, .ownership = .backend },
445 .access = .read_write,
446 .ownership = .backend,
447 .byte_size = 20000,
448 };
449 }
450
451 fn recordingDestroyedId(state: *const gpu.recording.BackendState, id: gpu.BackendObjectId) bool {
452 const count = @min(state.destroy_count, state.destroyed_ids.len);
453 for (state.destroyed_ids[0..count]) |destroyed_id| {
454 if (destroyed_id == id) return true;
455 }
456 return false;
457 }
458
459 test "pipeline executor derives intermediates scalars and geometry per stage" {
460 var state = gpu.recording.BackendState{
461 .allocator = testing.allocator,
462 .kind = .cuda,
463 .format = .cuda_ptx,
464 };
465 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
466
467 const operands = [_]gpu.BufferBinding{testBinding(1001)};
468 const results = [_]gpu.BufferBinding{testBinding(1002)};
469 const args = [_]choir_abi.ScalarArgument{.{ .u32 = 5000 }};
470
471 try launchPipeline(testing.allocator, state.handle(), .{
472 .pipeline = test_pipeline,
473 .registry = registry,
474 .format = .cuda_ptx,
475 .operands = operands[0..],
476 .results = results[0..],
477 .runtime_scalar_arguments = args[0..],
478 });
479
480 try testing.expectEqual(@as(usize, 2), state.buffer_allocate_count);
481 try testing.expectEqual(@as(usize, 79 * 4), state.allocated_buffer_byte_sizes[0]);
482 try testing.expectEqual(@as(usize, 79 * 4), state.allocated_buffer_byte_sizes[1]);
483
484 try testing.expectEqual(@as(usize, 3), state.launch_count);
485 try testing.expectEqual([3]u32{ 79, 1, 1 }, state.launch_grids[0]);
486 try testing.expectEqual([3]u32{ 1, 1, 1 }, state.launch_grids[1]);
487 try testing.expectEqual([3]u32{ 79, 1, 1 }, state.launch_grids[2]);
488 try testing.expectEqual(@as(?u32, 5000), state.launch_scalar_u32s[0]);
489 try testing.expectEqual(@as(?u32, 79), state.launch_scalar_u32s[1]);
490 try testing.expectEqual(@as(?u32, 5000), state.launch_scalar_u32s[2]);
491
492 try testing.expectEqual(@as(usize, 2), state.last_launch_buffer_count);
493 try testing.expectEqual(@as(gpu.BackendObjectId, 1002), state.last_buffer_ids[0]);
494 try testing.expectEqual(@as(usize, 3), state.load_count);
495 try testing.expectEqual(@as(usize, 5), state.destroy_count);
496 try testing.expect(recordingDestroyedId(&state, state.allocated_buffer_ids[0]));
497 try testing.expect(recordingDestroyedId(&state, state.allocated_buffer_ids[1]));
498 }
499
500 test "pipeline executor rejects binding count mismatches" {
501 var state = gpu.recording.BackendState{
502 .allocator = testing.allocator,
503 .kind = .cuda,
504 .format = .cuda_ptx,
505 };
506 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
507 const args = [_]choir_abi.ScalarArgument{.{ .u32 = 5000 }};
508 const operands = [_]gpu.BufferBinding{testBinding(1001)};
509 const results = [_]gpu.BufferBinding{testBinding(1002)};
510
511 try testing.expectError(error.LaunchArgumentMismatch, launchPipeline(testing.allocator, state.handle(), .{
512 .pipeline = test_pipeline,
513 .registry = registry,
514 .format = .cuda_ptx,
515 .operands = &.{},
516 .results = results[0..],
517 .runtime_scalar_arguments = args[0..],
518 }));
519 try testing.expectError(error.LaunchArgumentMismatch, launchPipeline(testing.allocator, state.handle(), .{
520 .pipeline = test_pipeline,
521 .registry = registry,
522 .format = .cuda_ptx,
523 .operands = operands[0..],
524 .results = results[0..],
525 .runtime_scalar_arguments = &.{},
526 }));
527 try testing.expectEqual(@as(usize, 0), state.launch_count);
528 }
529
530 test "pipeline executor rejects runtime scalar bounds before allocation" {
531 var state = gpu.recording.BackendState{
532 .allocator = testing.allocator,
533 .kind = .cuda,
534 .format = .cuda_ptx,
535 };
536 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
537 var bounded = test_pipeline;
538 bounded.runtime_scalar_bounds = &.{.{ .argument_index = 0, .max_u32 = 5000 }};
539 const args = [_]choir_abi.ScalarArgument{.{ .u32 = 5001 }};
540 const operands = [_]gpu.BufferBinding{testBinding(1001)};
541 const results = [_]gpu.BufferBinding{testBinding(1002)};
542
543 try testing.expectError(error.LaunchArgumentMismatch, launchPipeline(testing.allocator, state.handle(), .{
544 .pipeline = bounded,
545 .registry = registry,
546 .format = .cuda_ptx,
547 .operands = operands[0..],
548 .results = results[0..],
549 .runtime_scalar_arguments = args[0..],
550 }));
551 try testing.expectEqual(@as(usize, 0), state.buffer_allocate_count);
552 try testing.expectEqual(@as(usize, 0), state.launch_count);
553
554 try testing.expectError(
555 error.LaunchArgumentMismatch,
556 allocatePipelineIntermediates(testing.allocator, state.handle(), bounded, args[0..]),
557 );
558 try testing.expectEqual(@as(usize, 0), state.buffer_allocate_count);
559 }
560
561 test "pipeline executor refuses unresolvable stages before any launch" {
562 var state = gpu.recording.BackendState{
563 .allocator = testing.allocator,
564 .kind = .cuda,
565 .format = .cuda_ptx,
566 };
567 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..1] };
568 const args = [_]choir_abi.ScalarArgument{.{ .u32 = 5000 }};
569 const operands = [_]gpu.BufferBinding{testBinding(1001)};
570 const results = [_]gpu.BufferBinding{testBinding(1002)};
571
572 try testing.expectError(error.InvalidArtifact, launchPipeline(testing.allocator, state.handle(), .{
573 .pipeline = test_pipeline,
574 .registry = registry,
575 .format = .cuda_ptx,
576 .operands = operands[0..],
577 .results = results[0..],
578 .runtime_scalar_arguments = args[0..],
579 }));
580 try testing.expectEqual(@as(usize, 0), state.launch_count);
581 try testing.expectEqual(@as(usize, 0), state.destroy_count);
582 }
583
584 test "pipeline executor releases intermediates after partial allocation failure" {
585 var state = gpu.recording.BackendState{
586 .allocator = testing.allocator,
587 .kind = .cuda,
588 .format = .cuda_ptx,
589 .fail_buffer_allocate_after_count = 1,
590 };
591 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
592 const args = [_]choir_abi.ScalarArgument{.{ .u32 = 5000 }};
593 const operands = [_]gpu.BufferBinding{testBinding(1001)};
594 const results = [_]gpu.BufferBinding{testBinding(1002)};
595
596 try testing.expectError(error.OutOfMemory, launchPipeline(testing.allocator, state.handle(), .{
597 .pipeline = test_pipeline,
598 .registry = registry,
599 .format = .cuda_ptx,
600 .operands = operands[0..],
601 .results = results[0..],
602 .runtime_scalar_arguments = args[0..],
603 }));
604 try testing.expectEqual(@as(usize, 1), state.buffer_allocate_count);
605 try testing.expectEqual(@as(usize, 1), state.destroy_count);
606 try testing.expect(recordingDestroyedId(&state, state.allocated_buffer_ids[0]));
607 }
608
609 test "pipeline executor releases loaded artifacts after staged load failure" {
610 var state = gpu.recording.BackendState{
611 .allocator = testing.allocator,
612 .kind = .cuda,
613 .format = .cuda_ptx,
614 .fail_load_after_count = 1,
615 };
616 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
617 const args = [_]choir_abi.ScalarArgument{.{ .u32 = 5000 }};
618 const operands = [_]gpu.BufferBinding{testBinding(1001)};
619 const results = [_]gpu.BufferBinding{testBinding(1002)};
620
621 try testing.expectError(error.RuntimeUnavailable, launchPipeline(testing.allocator, state.handle(), .{
622 .pipeline = test_pipeline,
623 .registry = registry,
624 .format = .cuda_ptx,
625 .operands = operands[0..],
626 .results = results[0..],
627 .runtime_scalar_arguments = args[0..],
628 }));
629 try testing.expectEqual(@as(usize, 2), state.buffer_allocate_count);
630 try testing.expectEqual(@as(usize, 1), state.load_count);
631 try testing.expectEqual(@as(usize, 1), state.launch_count);
632 try testing.expectEqual(@as(usize, 3), state.destroy_count);
633 try testing.expect(recordingDestroyedId(&state, state.allocated_buffer_ids[0]));
634 try testing.expect(recordingDestroyedId(&state, state.allocated_buffer_ids[1]));
635 }
636
637 test "pipeline executor reuses provided intermediates without allocating" {
638 var state = gpu.recording.BackendState{
639 .allocator = testing.allocator,
640 .kind = .cuda,
641 .format = .cuda_ptx,
642 };
643 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
644 const args = [_]choir_abi.ScalarArgument{.{ .u32 = 5000 }};
645 const operands = [_]gpu.BufferBinding{testBinding(1001)};
646 const results = [_]gpu.BufferBinding{testBinding(1002)};
647
648 const provided = try allocatePipelineIntermediates(testing.allocator, state.handle(), test_pipeline, args[0..]);
649 defer deinitPipelineIntermediates(testing.allocator, state.handle(), provided);
650 try testing.expectEqual(@as(usize, 2), state.buffer_allocate_count);
651
652 try launchPipeline(testing.allocator, state.handle(), .{
653 .pipeline = test_pipeline,
654 .registry = registry,
655 .format = .cuda_ptx,
656 .operands = operands[0..],
657 .results = results[0..],
658 .runtime_scalar_arguments = args[0..],
659 .intermediates = provided,
660 });
661 try testing.expectEqual(@as(usize, 2), state.buffer_allocate_count);
662 try testing.expectEqual(@as(usize, 3), state.launch_count);
663
664 try launchPipeline(testing.allocator, state.handle(), .{
665 .pipeline = test_pipeline,
666 .registry = registry,
667 .format = .cuda_ptx,
668 .operands = operands[0..],
669 .results = results[0..],
670 .runtime_scalar_arguments = args[0..],
671 .intermediates = provided,
672 });
673 try testing.expectEqual(@as(usize, 2), state.buffer_allocate_count);
674 try testing.expectEqual(@as(usize, 6), state.launch_count);
675 }
676
677 test "pipeline executor rejects undersized or miscounted provided intermediates" {
678 var state = gpu.recording.BackendState{
679 .allocator = testing.allocator,
680 .kind = .cuda,
681 .format = .cuda_ptx,
682 };
683 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
684 const args = [_]choir_abi.ScalarArgument{.{ .u32 = 5000 }};
685 const operands = [_]gpu.BufferBinding{testBinding(1001)};
686 const results = [_]gpu.BufferBinding{testBinding(1002)};
687
688 const miscounted = [_]gpu.BufferBinding{testBinding(2001)};
689 try testing.expectError(error.LaunchArgumentMismatch, launchPipeline(testing.allocator, state.handle(), .{
690 .pipeline = test_pipeline,
691 .registry = registry,
692 .format = .cuda_ptx,
693 .operands = operands[0..],
694 .results = results[0..],
695 .runtime_scalar_arguments = args[0..],
696 .intermediates = miscounted[0..],
697 }));
698
699 var undersized = [_]gpu.BufferBinding{ testBinding(2001), testBinding(2002) };
700 undersized[1].byte_size = 4;
701 try testing.expectError(error.LaunchArgumentMismatch, launchPipeline(testing.allocator, state.handle(), .{
702 .pipeline = test_pipeline,
703 .registry = registry,
704 .format = .cuda_ptx,
705 .operands = operands[0..],
706 .results = results[0..],
707 .runtime_scalar_arguments = args[0..],
708 .intermediates = undersized[0..],
709 }));
710
711 const duplicate = [_]gpu.BufferBinding{ testBinding(2001), testBinding(2001) };
712 try testing.expectError(error.LaunchArgumentMismatch, launchPipeline(testing.allocator, state.handle(), .{
713 .pipeline = test_pipeline,
714 .registry = registry,
715 .format = .cuda_ptx,
716 .operands = operands[0..],
717 .results = results[0..],
718 .runtime_scalar_arguments = args[0..],
719 .intermediates = duplicate[0..],
720 }));
721
722 const operand_alias = [_]gpu.BufferBinding{ operands[0], testBinding(2002) };
723 try testing.expectError(error.LaunchArgumentMismatch, launchPipeline(testing.allocator, state.handle(), .{
724 .pipeline = test_pipeline,
725 .registry = registry,
726 .format = .cuda_ptx,
727 .operands = operands[0..],
728 .results = results[0..],
729 .runtime_scalar_arguments = args[0..],
730 .intermediates = operand_alias[0..],
731 }));
732
733 const result_alias = [_]gpu.BufferBinding{ testBinding(2001), results[0] };
734 try testing.expectError(error.LaunchArgumentMismatch, launchPipeline(testing.allocator, state.handle(), .{
735 .pipeline = test_pipeline,
736 .registry = registry,
737 .format = .cuda_ptx,
738 .operands = operands[0..],
739 .results = results[0..],
740 .runtime_scalar_arguments = args[0..],
741 .intermediates = result_alias[0..],
742 }));
743 try testing.expectEqual(@as(usize, 0), state.launch_count);
744 }
745
746 test "pipeline executor launches from a loaded artifact pool without driver round trips" {
747 var state = gpu.recording.BackendState{
748 .allocator = testing.allocator,
749 .kind = .cuda,
750 .format = .cuda_ptx,
751 };
752 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
753 const args = [_]choir_abi.ScalarArgument{.{ .u32 = 5000 }};
754 const operands = [_]gpu.BufferBinding{testBinding(1001)};
755 const results = [_]gpu.BufferBinding{testBinding(1002)};
756
757 var pool = try loadPipelineArtifacts(testing.allocator, state.handle(), test_pipeline, registry, .cuda_ptx);
758 defer pool.deinit();
759 try testing.expectEqual(@as(usize, 3), state.create_count);
760 try testing.expectEqual(@as(usize, 3), state.load_count);
761
762 const provided = try allocatePipelineIntermediates(testing.allocator, state.handle(), test_pipeline, args[0..]);
763 defer deinitPipelineIntermediates(testing.allocator, state.handle(), provided);
764
765 var round: usize = 0;
766 while (round < 3) : (round += 1) {
767 try launchPipeline(testing.allocator, state.handle(), .{
768 .pipeline = test_pipeline,
769 .registry = registry,
770 .format = .cuda_ptx,
771 .operands = operands[0..],
772 .results = results[0..],
773 .runtime_scalar_arguments = args[0..],
774 .intermediates = provided,
775 .artifacts = &pool,
776 });
777 }
778 try testing.expectEqual(@as(usize, 3), state.create_count);
779 try testing.expectEqual(@as(usize, 3), state.load_count);
780 try testing.expectEqual(@as(usize, 9), state.launch_count);
781 }
782
783 test "pipeline artifact pool deinit destroys loaded artifacts" {
784 var state = gpu.recording.BackendState{
785 .allocator = testing.allocator,
786 .kind = .cuda,
787 .format = .cuda_ptx,
788 };
789 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
790
791 var pool = try loadPipelineArtifacts(testing.allocator, state.handle(), test_pipeline, registry, .cuda_ptx);
792 try testing.expectEqual(@as(usize, 3), state.load_count);
793 try testing.expectEqual(@as(usize, 0), state.destroy_count);
794
795 pool.deinit();
796 try testing.expectEqual(@as(usize, 3), state.destroy_count);
797 try testing.expectEqual(@as(?gpu.BackendObjectId, 3), state.last_destroyed_id);
798 }
799
800 test "pipeline stage artifacts carry their entry's scalar count" {
801 var state = gpu.recording.BackendState{
802 .allocator = testing.allocator,
803 .kind = .cuda,
804 .format = .cuda_ptx,
805 };
806 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
807
808 var pool = try loadPipelineArtifacts(testing.allocator, state.handle(), test_pipeline, registry, .cuda_ptx);
809 defer pool.deinit();
810 for (pool.entries, test_pipeline.stages) |pool_entry, stage| {
811 const entry = registry.find(stage.target, stage.version, .cuda_ptx).?;
812 try testing.expectEqual(entry.runtime_scalar_argument_count, pool_entry.artifact.scalar_argument_count);
813 try testing.expectEqual(
814 entry.argument_count - entry.runtime_scalar_argument_count,
815 try pool_entry.artifact.bufferArgumentCount(),
816 );
817 }
818 }
819
820 test "pipeline artifact pool cleans loaded artifacts after staged load failure" {
821 var state = gpu.recording.BackendState{
822 .allocator = testing.allocator,
823 .kind = .cuda,
824 .format = .cuda_ptx,
825 .fail_load_after_count = 1,
826 };
827 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
828
829 try testing.expectError(
830 error.RuntimeUnavailable,
831 loadPipelineArtifacts(testing.allocator, state.handle(), test_pipeline, registry, .cuda_ptx),
832 );
833 try testing.expectEqual(@as(usize, 1), state.load_count);
834 try testing.expectEqual(@as(usize, 1), state.destroy_count);
835 try testing.expectEqual(@as(?gpu.BackendObjectId, 1), state.last_destroyed_id);
836 }
837
838 test "pipeline executor rejects mismatched artifact pools before any launch" {
839 var state = gpu.recording.BackendState{
840 .allocator = testing.allocator,
841 .kind = .cuda,
842 .format = .cuda_ptx,
843 };
844 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
845 const args = [_]choir_abi.ScalarArgument{.{ .u32 = 5000 }};
846 const operands = [_]gpu.BufferBinding{testBinding(1001)};
847 const results = [_]gpu.BufferBinding{testBinding(1002)};
848
849 var pool = try loadPipelineArtifacts(testing.allocator, state.handle(), test_pipeline, registry, .cuda_ptx);
850 defer pool.deinit();
851
852 var renamed = test_pipeline;
853 renamed.target = "accy.kernel.test.other_pipeline";
854 try testing.expectError(error.LaunchArgumentMismatch, launchPipeline(testing.allocator, state.handle(), .{
855 .pipeline = renamed,
856 .registry = registry,
857 .format = .cuda_ptx,
858 .operands = operands[0..],
859 .results = results[0..],
860 .runtime_scalar_arguments = args[0..],
861 .artifacts = &pool,
862 }));
863
864 var reversioned = test_pipeline;
865 reversioned.version = 9;
866 try testing.expectError(error.LaunchArgumentMismatch, launchPipeline(testing.allocator, state.handle(), .{
867 .pipeline = reversioned,
868 .registry = registry,
869 .format = .cuda_ptx,
870 .operands = operands[0..],
871 .results = results[0..],
872 .runtime_scalar_arguments = args[0..],
873 .artifacts = &pool,
874 }));
875 try testing.expectEqual(@as(usize, 0), state.launch_count);
876 }
877
878 test "pipeline executor rejects artifact pool format mismatches before any launch" {
879 var state = gpu.recording.BackendState{
880 .allocator = testing.allocator,
881 .kind = .cuda,
882 .format = .cuda_ptx,
883 };
884 const registry = artifact_product.KernelCallRegistry{ .entries = test_entries[0..] };
885 var metal_entries = test_entries;
886 for (&metal_entries) |*entry| entry.format = .metal_msl;
887 const metal_registry = artifact_product.KernelCallRegistry{ .entries = metal_entries[0..] };
888 const args = [_]choir_abi.ScalarArgument{.{ .u32 = 5000 }};
889 const operands = [_]gpu.BufferBinding{testBinding(1001)};
890 const results = [_]gpu.BufferBinding{testBinding(1002)};
891
892 var pool = try loadPipelineArtifacts(testing.allocator, state.handle(), test_pipeline, registry, .cuda_ptx);
893 defer pool.deinit();
894
895 try testing.expectError(error.LaunchArgumentMismatch, launchPipeline(testing.allocator, state.handle(), .{
896 .pipeline = test_pipeline,
897 .registry = metal_registry,
898 .format = .metal_msl,
899 .operands = operands[0..],
900 .results = results[0..],
901 .runtime_scalar_arguments = args[0..],
902 .artifacts = &pool,
903 }));
904 try testing.expectEqual(@as(usize, 0), state.buffer_allocate_count);
905 try testing.expectEqual(@as(usize, 0), state.launch_count);
906 }