lib/accy/src/executable/fragment.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const choir = @import("choir");
5 const sys = @import("sys");
6 const accy_root = @import("../root.zig");
7 const accy_choir = @import("../choir/root.zig");
8 const artifact_product = @import("../artifact/root.zig");
9 const kernel_library = @import("../kernel/library/root.zig");
10 const exec_product = @import("plan.zig");
11 const preparation = @import("../preparation/root.zig");
12 const tuning_mod = @import("tuning.zig");
13 const schedule_tuning = @import("schedule.zig");
14 const binding_mod = @import("binding.zig");
15 const loaded_mod = @import("loaded.zig");
16 const fixture = @import("fixture.zig");
17
18 const ir = choir.ir;
19 const passes = choir.passes;
20 const semantic = accy_choir.semantic;
21 const candidate_mod = @import("candidate.zig");
22 const phase_mod = @import("phase.zig");
23
24 pub const LaunchCandidateMeasurement = tuning_mod.LaunchCandidateMeasurement;
25 pub const LaunchTuningSelection = tuning_mod.LaunchTuningSelection;
26 pub const launch_tuning_cache_record_version = tuning_mod.launch_tuning_cache_record_version;
27 pub const launch_tuning_artifact_magic = tuning_mod.launch_tuning_artifact_magic;
28 pub const launch_tuning_artifact_version = tuning_mod.launch_tuning_artifact_version;
29 pub const LaunchTuningCacheRecord = tuning_mod.LaunchTuningCacheRecord;
30 pub const LaunchTuning = tuning_mod.LaunchTuning;
31 pub const encodeLaunchTuningArtifact = tuning_mod.encodeLaunchTuningArtifact;
32 pub const decodeLaunchTuningArtifact = tuning_mod.decodeLaunchTuningArtifact;
33
34 const LaunchTuningCacheKey = tuning_mod.LaunchTuningCacheKey;
35 const LaunchTuningCache = tuning_mod.LaunchTuningCache;
36 const LaunchGraphNode = exec_product.LaunchGraphNode;
37 const LaunchGraphDependency = exec_product.LaunchGraphDependency;
38 const LaunchGraphLoopCarry = exec_product.LaunchGraphLoopCarry;
39 const LaunchGraphLoop = exec_product.LaunchGraphLoop;
40 const LaunchGraphPlan = exec_product.LaunchGraphPlan;
41 const OwnedLaunchGraphPlan = exec_product.OwnedLaunchGraphPlan;
42 const createDataflowLaunchGraphPlan = exec_product.createDataflowLaunchGraphPlan;
43 pub const LaunchCandidateSynchronization = candidate_mod.LaunchCandidateSynchronization;
44 pub const LaunchCandidateBenchmarkOptions = candidate_mod.LaunchCandidateBenchmarkOptions;
45 pub const LaunchCandidateRecord = candidate_mod.LaunchCandidateRecord;
46 pub const FragmentPhase = phase_mod.FragmentPhase;
47 pub const FragmentPhaseObserver = phase_mod.FragmentPhaseObserver;
48 pub const FragmentInstrumentation = phase_mod.FragmentInstrumentation;
49
50 const launchCandidateRecord = candidate_mod.launchCandidateRecord;
51 const plannedKernelForLaunchCandidateRecord = candidate_mod.plannedKernelForLaunchCandidateRecord;
52 const launchResourceClassName = artifact_product.launchResourceClassName;
53 const copyKernelArtifactToAllocator = candidate_mod.copyKernelArtifactToAllocator;
54
55 const SlotBinding = binding_mod.SlotBinding;
56 const ElementCountBufferBinding = binding_mod.ElementCountBufferBinding;
57 const PreparedLaunchBindings = binding_mod.PreparedLaunchBindings;
58 const LoadedKernels = loaded_mod.LoadedKernels;
59 const loadKernels = loaded_mod.loadKernels;
60 const loadKernelsWithInstrumentation = loaded_mod.loadKernelsWithInstrumentation;
61 const RecordingBackendState = gpu.recording.BackendState;
62 const addChoirModule = fixture.addChoirModule;
63 const addU32ChoirModule = fixture.addU32ChoirModule;
64 const addSemanticModule = fixture.addSemanticModule;
65 const fusedAddMulChoirModule = fixture.fusedAddMulChoirModule;
66 const constantAddChoirModule = fixture.constantAddChoirModule;
67 const dotGeneralChoirModule = fixture.dotGeneralChoirModule;
68 const dotGeneralF16ChoirModule = fixture.dotGeneralF16ChoirModule;
69 const reduceChoirModule = fixture.reduceChoirModule;
70 const reduceI32ChoirModule = fixture.reduceI32ChoirModule;
71 const escapedTwoKernelChoirModule = fixture.escapedTwoKernelChoirModule;
72 const kernelCallChoirModule = fixture.kernelCallChoirModule;
73 const aliasedKernelCallChoirModule = fixture.aliasedKernelCallChoirModule;
74 const createTestBackendArtifactPlan = fixture.createTestBackendArtifactPlan;
75 const bufferBinding = fixture.bufferBinding;
76 const slotBindingsForKernel = fixture.slotBindingsForKernel;
77 const slotBindingsForPlan = fixture.slotBindingsForPlan;
78 const elementCountBindingsForPlan = fixture.elementCountBindingsForPlan;
79 const firstElementCountBinding = fixture.firstElementCountBinding;
80
81 pub const FragmentCompilerOptions = struct {
82 /// Limits on the compiler context that tensor lowering builds its semantic
83 /// module in, for sizing the memory used when a tensor program is lowered
84 /// into the compiler's semantic form. Only tensor lowering reads this
85 /// limit, so compiling from a semantic module the caller already built
86 /// ignores it. A limit too small for the program makes the refresh fail
87 /// with `error.OutOfMemory`.
88 semantic_context_limits: semantic.Builder.ContextLimits = .standard,
89 artifact_format: ?gpu.ArtifactFormat = null,
90 math_tier: gpu.BackendMathTier = .exact,
91 kernel_call_registry: ?*const artifact_product.KernelCallRegistry = null,
92 matrix_product_schedule: ?kernel_library.MatrixProductSchedule = null,
93 gather_schedule: ?kernel_library.GatherSchedule = null,
94 scatter_schedule: ?kernel_library.ScatterSchedule = null,
95 scatter_add_schedule: ?kernel_library.ScatterAddSchedule = null,
96 row_sparse_cross_entropy_schedule: ?kernel_library.RowSparseCrossEntropySchedule = null,
97 generated_scan_schedule: ?preparation.target.GeneratedScanSchedule = null,
98 generated_scan_schedule_tuning_artifact: []const u8 = &.{},
99 generated_row_pipeline_schedule: ?preparation.target.GeneratedRowPipelineSchedule = null,
100 generated_row_pipeline_schedule_tuning_artifact: []const u8 = &.{},
101 family_tuning: ?*const kernel_library.tuning.FamilyTuningReader = null,
102 family_tuning_artifact: []const u8 = &.{},
103 matrix_product_schedule_tuning_artifact: []const u8 = &.{},
104 authored_kernel_diagnostic_id: ?[]const u8 = null,
105 launch_tuning_artifact: []const u8 = &.{},
106 preparation_failure: ?*preparation.BackendPreparationFailure = null,
107 instrumentation: FragmentInstrumentation = .{},
108 };
109
110 pub const KernelCompilerOptions = struct {
111 artifact_format: ?gpu.ArtifactFormat = null,
112 authored_kernel_diagnostic_id: ?[]const u8 = null,
113 instrumentation: FragmentInstrumentation = .{},
114 };
115
116 pub const LaunchOptions = exec_product.LaunchOptions;
117
118 const LoadedFragmentState = struct {
119 allocator: std.mem.Allocator,
120 compiled_fragment: *exec_product.CompiledFragment,
121 loaded_kernels: LoadedKernels,
122 tuning_cache: LaunchTuningCache,
123 prepared_plan: ?OwnedLaunchGraphPlan = null,
124
125 fn deinit(self: *LoadedFragmentState) void {
126 if (self.prepared_plan) |*prepared| prepared.deinit();
127 self.tuning_cache.deinit();
128 self.loaded_kernels.deinit();
129 self.compiled_fragment.deinit();
130 self.* = undefined;
131 }
132 };
133
134 pub const LoadedFragment = opaque {
135 fn stateConst(self: *const LoadedFragment) *const LoadedFragmentState {
136 return @ptrCast(@alignCast(self));
137 }
138
139 fn stateMut(self: *LoadedFragment) *LoadedFragmentState {
140 return @ptrCast(@alignCast(self));
141 }
142
143 fn artifactPlan(self: *const LoadedFragment) *const artifact_product.BackendArtifactPlan {
144 const state = self.stateConst();
145 return state.compiled_fragment.artifactPlan();
146 }
147
148 fn loadedKernelCount(self: *const LoadedFragment) usize {
149 const state = self.stateConst();
150 return state.loaded_kernels.kernelCount();
151 }
152
153 pub fn deinit(self: *LoadedFragment) void {
154 const state = self.stateMut();
155 const allocator = state.allocator;
156 state.deinit();
157 allocator.destroy(state);
158 }
159
160 pub fn tuningRecordCount(self: *const LoadedFragment) usize {
161 return self.stateConst().tuning_cache.count();
162 }
163
164 pub fn compiledFingerprint(self: *const LoadedFragment) u64 {
165 return self.stateConst().compiled_fragment.fingerprint();
166 }
167
168 pub fn kernelCount(self: *const LoadedFragment) usize {
169 return self.artifactPlan().kernelCount();
170 }
171
172 pub fn kernelSummary(
173 self: *const LoadedFragment,
174 kernel_index: usize,
175 ) gpu.BackendError!artifact_product.KernelSummary {
176 return try artifact_product.summarizePlannedKernel(try self.plannedKernel(kernel_index));
177 }
178
179 pub fn kernelSummaryForWork(
180 self: *const LoadedFragment,
181 work_item_id: usize,
182 ) gpu.BackendError!artifact_product.KernelSummary {
183 for (self.artifactPlan().kernels.items) |kernel| {
184 if (kernel.work_item_id == work_item_id) return try artifact_product.summarizePlannedKernel(kernel);
185 }
186 return error.InvalidArtifact;
187 }
188
189 pub fn copyKernelSummaries(
190 self: *const LoadedFragment,
191 result_allocator: std.mem.Allocator,
192 ) gpu.BackendError!artifact_product.KernelSummaries {
193 const kernels = self.artifactPlan().kernels.items;
194 const items = result_allocator.alloc(artifact_product.KernelSummary, kernels.len) catch return error.OutOfMemory;
195 var copied: usize = 0;
196 errdefer {
197 for (items[0..copied]) |summary| {
198 result_allocator.free(summary.entry_name);
199 }
200 result_allocator.free(items);
201 }
202
203 for (items, kernels) |*item, kernel| {
204 item.* = try artifact_product.copyKernelSummary(result_allocator, try artifact_product.summarizePlannedKernel(kernel));
205 copied += 1;
206 }
207
208 return .{
209 .allocator = result_allocator,
210 .items = items,
211 };
212 }
213
214 pub fn copyKernelArtifact(
215 self: *const LoadedFragment,
216 allocator: std.mem.Allocator,
217 kernel_index: usize,
218 ) gpu.BackendError!gpu.KernelArtifact {
219 const planned = try self.plannedKernel(kernel_index);
220 return try copyKernelArtifactToAllocator(allocator, planned.artifact);
221 }
222
223 fn plannedKernel(
224 self: *const LoadedFragment,
225 kernel_index: usize,
226 ) gpu.BackendError!artifact_product.PlannedKernel {
227 const artifact_plan = self.artifactPlan();
228 if (kernel_index >= artifact_plan.kernels.items.len) return error.InvalidArtifact;
229 return artifact_plan.kernels.items[kernel_index];
230 }
231
232 fn launchResourceCandidate(
233 self: *const LoadedFragment,
234 kernel_index: usize,
235 candidate_index: usize,
236 ) gpu.BackendError!artifact_product.LaunchResourceCandidate {
237 const kernel = try self.plannedKernel(kernel_index);
238 if (candidate_index >= kernel.launch_resources.candidate_count) return error.LaunchArgumentMismatch;
239 return kernel.launch_resources.candidates[candidate_index];
240 }
241
242 pub fn importLaunchTuningArtifact(
243 self: *LoadedFragment,
244 scratch: std.mem.Allocator,
245 bytes: []const u8,
246 ) gpu.BackendError!void {
247 const state = self.stateMut();
248 const records = try decodeLaunchTuningArtifact(scratch, bytes);
249 defer scratch.free(records);
250 try state.tuning_cache.importRecords(records);
251 try self.refreshPreparedLaunchGraphPlan();
252 }
253
254 pub fn replaceLaunchTuningArtifact(
255 self: *LoadedFragment,
256 scratch: std.mem.Allocator,
257 bytes: []const u8,
258 ) gpu.BackendError!void {
259 const state = self.stateMut();
260 var tuning_cache = LaunchTuningCache.init(state.allocator);
261 errdefer tuning_cache.deinit();
262 if (bytes.len != 0) {
263 const records = try decodeLaunchTuningArtifact(scratch, bytes);
264 defer scratch.free(records);
265 try tuning_cache.importRecords(records);
266 }
267 state.tuning_cache.deinit();
268 state.tuning_cache = tuning_cache;
269 try self.refreshPreparedLaunchGraphPlan();
270 }
271
272 pub fn exportLaunchTuningArtifact(
273 self: *const LoadedFragment,
274 result_allocator: std.mem.Allocator,
275 ) gpu.BackendError![]u8 {
276 const state = self.stateConst();
277 const records = try state.tuning_cache.exportRecords(result_allocator);
278 defer result_allocator.free(records);
279 return try encodeLaunchTuningArtifact(result_allocator, records);
280 }
281
282 pub fn recordLaunchCandidateRecords(
283 self: *LoadedFragment,
284 records: []const LaunchCandidateRecord,
285 ) gpu.BackendError!void {
286 const state = self.stateMut();
287 const artifact_plan = state.compiled_fragment.artifactPlan();
288 const caps = try state.loaded_kernels.handle.queryCapabilities();
289 for (records) |record| _ = try plannedKernelForLaunchCandidateRecord(artifact_plan, record);
290 for (records) |record| {
291 const planned = try plannedKernelForLaunchCandidateRecord(artifact_plan, record);
292 try state.tuning_cache.recordMeasuredSelection(caps, planned, .{
293 .kernel_id = record.kernel.kernel_id,
294 .candidate_index = record.candidate_index,
295 .median_ns = record.median_ns,
296 .sample_count = record.sample_count,
297 });
298 }
299 try self.refreshPreparedLaunchGraphPlan();
300 }
301
302 pub fn createLaunchGraphPlan(
303 self: *const LoadedFragment,
304 allocator: std.mem.Allocator,
305 launch_options: LaunchOptions,
306 ) gpu.BackendError!OwnedLaunchGraphPlan {
307 const state = self.stateConst();
308 var graph = try state.compiled_fragment.createLaunchGraphPlan(allocator, launch_options);
309 errdefer graph.deinit();
310 if (state.tuning_cache.count() != 0 and tuning_mod.launchTuningIsEmpty(launch_options.tuning)) {
311 try graph.applyCachedLaunchTuning(
312 try state.loaded_kernels.handle.queryCapabilities(),
313 state.compiled_fragment.artifactPlan(),
314 &state.tuning_cache,
315 );
316 graph.validated = false;
317 const view = graph.plan();
318 try exec_product.validateLaunchGraph(
319 allocator,
320 state.compiled_fragment.artifactPlan(),
321 view,
322 exec_product.launchGraphNeedsDependencyEvents(view),
323 );
324 graph.validated = true;
325 }
326 return graph;
327 }
328
329 fn refreshPreparedLaunchGraphPlan(self: *LoadedFragment) gpu.BackendError!void {
330 const state = self.stateMut();
331 var next = try self.createLaunchGraphPlan(state.allocator, .{});
332 errdefer next.deinit();
333 if (state.prepared_plan) |*existing| existing.deinit();
334 state.prepared_plan = next;
335 }
336
337 fn invokeInputsForTest(
338 self: *const LoadedFragment,
339 allocator: std.mem.Allocator,
340 scratch: std.mem.Allocator,
341 inputs: []const []const u8,
342 ) gpu.BackendError!void {
343 try self.invokeInputsWithOptionsForTest(allocator, scratch, inputs, .{});
344 }
345
346 fn invokeInputsWithOptionsForTest(
347 self: *const LoadedFragment,
348 allocator: std.mem.Allocator,
349 scratch: std.mem.Allocator,
350 inputs: []const []const u8,
351 launch_options: LaunchOptions,
352 ) gpu.BackendError!void {
353 const bindings = try self.prepareInvocationBindings(allocator, inputs);
354 defer bindings.deinit();
355 try self.submitInvocationWithOptions(scratch, bindings, launch_options);
356 }
357
358 pub fn prepareInvocationBindings(
359 self: *const LoadedFragment,
360 allocator: std.mem.Allocator,
361 inputs: []const []const u8,
362 ) !*PreparedLaunchBindings {
363 const state = self.stateConst();
364 return try binding_mod.prepare(
365 allocator,
366 state.loaded_kernels.handle,
367 state.compiled_fragment.artifactPlan(),
368 inputs,
369 );
370 }
371
372 fn submitPreparedForTest(
373 self: *const LoadedFragment,
374 scratch: std.mem.Allocator,
375 bindings: *const PreparedLaunchBindings,
376 ) gpu.BackendError!void {
377 const state = self.stateConst();
378 if (state.prepared_plan) |*prepared| {
379 try self.launchGraph(
380 scratch,
381 binding_mod.slotBindings(bindings),
382 binding_mod.elementCountBindings(bindings),
383 prepared.plan(),
384 );
385 return;
386 }
387 try self.submitInvocationWithOptions(scratch, bindings, .{});
388 }
389
390 pub fn submitInvocationWithOptions(
391 self: *const LoadedFragment,
392 scratch: std.mem.Allocator,
393 bindings: *const PreparedLaunchBindings,
394 launch_options: LaunchOptions,
395 ) gpu.BackendError!void {
396 const state = self.stateConst();
397 if (launchOptionsUsePreparedPlan(launch_options)) {
398 const prepared = if (state.prepared_plan) |*plan| plan else return error.InvalidArtifact;
399 try self.launchGraph(
400 scratch,
401 binding_mod.slotBindings(bindings),
402 binding_mod.elementCountBindings(bindings),
403 prepared.plan(),
404 );
405 return;
406 }
407 try self.launchAllWithOptions(
408 scratch,
409 binding_mod.slotBindings(bindings),
410 binding_mod.elementCountBindings(bindings),
411 launch_options,
412 );
413 }
414
415 pub fn submitInvocationWithGraph(
416 self: *const LoadedFragment,
417 scratch: std.mem.Allocator,
418 bindings: *const PreparedLaunchBindings,
419 graph: LaunchGraphPlan,
420 ) gpu.BackendError!void {
421 try self.launchGraph(
422 scratch,
423 binding_mod.slotBindings(bindings),
424 binding_mod.elementCountBindings(bindings),
425 graph,
426 );
427 }
428
429 pub fn completeInvocationWithOptions(
430 self: *const LoadedFragment,
431 options: LaunchOptions,
432 ) gpu.BackendError!void {
433 const state = self.stateConst();
434 if (options.signal_event) |event| {
435 try state.loaded_kernels.handle.synchronize(.{ .scope = .event, .event = event });
436 } else if (options.stream) |stream| {
437 try state.loaded_kernels.handle.synchronize(.{ .scope = .stream, .stream = stream });
438 } else {
439 try state.loaded_kernels.handle.synchronize(.{ .scope = .default_stream });
440 }
441 }
442
443 pub fn completeInvocationGraph(
444 self: *const LoadedFragment,
445 graph: LaunchGraphPlan,
446 ) gpu.BackendError!void {
447 const state = self.stateConst();
448 if (graph.nodes.len == 0) return error.InvalidArtifact;
449 var reverse_index = graph.nodes.len;
450 while (reverse_index != 0) {
451 reverse_index -= 1;
452 const node = graph.nodes[reverse_index];
453 var later_same_stream = false;
454 for (graph.nodes[reverse_index + 1 ..]) |later| {
455 if (sameOptionalStream(node.stream, later.stream)) {
456 later_same_stream = true;
457 break;
458 }
459 }
460 if (later_same_stream) continue;
461 if (node.signal_event) |event| {
462 try state.loaded_kernels.handle.synchronize(.{ .scope = .event, .event = event });
463 } else if (node.stream) |stream| {
464 try state.loaded_kernels.handle.synchronize(.{ .scope = .stream, .stream = stream });
465 } else {
466 try state.loaded_kernels.handle.synchronize(.{ .scope = .default_stream });
467 }
468 }
469 }
470
471 pub fn outputCount(self: *const LoadedFragment) usize {
472 return self.artifactPlan().output_slot_ids.len;
473 }
474
475 pub fn outputSlot(
476 self: *const LoadedFragment,
477 index: usize,
478 ) ?*const artifact_product.PlannedSlot {
479 const plan = self.artifactPlan();
480 if (index >= plan.output_slot_ids.len) return null;
481 return plan.slotById(plan.output_slot_ids[index]);
482 }
483
484 pub fn readInvocationOutput(
485 self: *const LoadedFragment,
486 bindings: *const PreparedLaunchBindings,
487 index: usize,
488 host_bytes: []u8,
489 ) gpu.BackendError!void {
490 const state = self.stateConst();
491 const plan = state.compiled_fragment.artifactPlan();
492 if (index >= plan.output_slot_ids.len) return error.InvalidArtifact;
493 const slot_id = plan.output_slot_ids[index];
494 const slot = plan.slotById(slot_id) orelse return error.InvalidArtifact;
495 const byte_size_u64 = slot.byte_size orelse return error.UnsupportedOperation;
496 const byte_size = std.math.cast(usize, byte_size_u64) orelse return error.InvalidArtifact;
497 if (host_bytes.len < byte_size) return error.ReadBufferDestinationTooSmall;
498 if (host_bytes.len > byte_size) return error.InvalidBuffer;
499 const binding = try binding_mod.bindingForSlot(binding_mod.slotBindings(bindings), slot_id, .read_only);
500 try state.loaded_kernels.handle.readBuffer(.{
501 .handle = binding.handle,
502 .bytes = host_bytes,
503 });
504 }
505
506 fn invokeAndReadForTest(
507 self: *const LoadedFragment,
508 allocator: std.mem.Allocator,
509 scratch: std.mem.Allocator,
510 inputs: []const []const u8,
511 outputs: []const []u8,
512 ) gpu.BackendError!void {
513 if (outputs.len != self.outputCount()) return error.InvalidArtifact;
514 const bindings = try self.prepareInvocationBindings(allocator, inputs);
515 defer bindings.deinit();
516 try self.submitPreparedForTest(scratch, bindings);
517 try self.completeInvocationWithOptions(.{});
518 for (outputs, 0..) |host_bytes, index| {
519 try self.readInvocationOutput(bindings, index, host_bytes);
520 }
521 }
522
523 pub fn measureInvocationLaunchCandidates(
524 self: *const LoadedFragment,
525 result_allocator: std.mem.Allocator,
526 scratch: std.mem.Allocator,
527 kernel_index: usize,
528 bindings: *const PreparedLaunchBindings,
529 options: LaunchCandidateBenchmarkOptions,
530 ) gpu.BackendError![]LaunchCandidateRecord {
531 const state = self.stateConst();
532 const kernel = try self.plannedKernel(kernel_index);
533 const artifact_plan = state.compiled_fragment.artifactPlan();
534 const measurements = try state.loaded_kernels.measureLaunchCandidates(
535 scratch,
536 scratch,
537 artifact_plan,
538 kernel_index,
539 binding_mod.slotBindings(bindings),
540 binding_mod.elementCountBindingForKernel(bindings, kernel.kernel_id),
541 options,
542 );
543 defer scratch.free(measurements);
544
545 const records = result_allocator.alloc(LaunchCandidateRecord, measurements.len) catch return error.OutOfMemory;
546 errdefer result_allocator.free(records);
547 for (measurements, 0..) |measurement, index| {
548 const candidate = try self.launchResourceCandidate(kernel_index, measurement.candidate_index);
549 records[index] = try launchCandidateRecord(kernel, candidate, measurement);
550 }
551 return records;
552 }
553
554 pub fn measureAndRecordInvocationLaunchCandidates(
555 self: *LoadedFragment,
556 result_allocator: std.mem.Allocator,
557 scratch: std.mem.Allocator,
558 bindings: *const PreparedLaunchBindings,
559 options: LaunchCandidateBenchmarkOptions,
560 ) gpu.BackendError![]LaunchCandidateRecord {
561 var records = std.ArrayListUnmanaged(LaunchCandidateRecord).empty;
562 errdefer records.deinit(result_allocator);
563
564 var kernel_index: usize = 0;
565 while (kernel_index < self.kernelCount()) : (kernel_index += 1) {
566 const kernel = try self.plannedKernel(kernel_index);
567 if (kernel.launch_resources.candidate_count <= 1) continue;
568
569 const measured = try self.measureInvocationLaunchCandidates(
570 scratch,
571 scratch,
572 kernel_index,
573 bindings,
574 options,
575 );
576 defer scratch.free(measured);
577 records.appendSlice(result_allocator, measured) catch return error.OutOfMemory;
578 }
579
580 const owned_records = records.toOwnedSlice(result_allocator) catch return error.OutOfMemory;
581 errdefer result_allocator.free(owned_records);
582 try self.recordLaunchCandidateRecords(owned_records);
583 return owned_records;
584 }
585
586 fn launchAll(
587 self: *const LoadedFragment,
588 scratch: std.mem.Allocator,
589 slot_bindings: []const SlotBinding,
590 element_count_buffers: []const ElementCountBufferBinding,
591 ) gpu.BackendError!void {
592 try self.launchAllWithOptions(
593 scratch,
594 slot_bindings,
595 element_count_buffers,
596 .{},
597 );
598 }
599
600 fn launchAllWithOptions(
601 self: *const LoadedFragment,
602 scratch: std.mem.Allocator,
603 slot_bindings: []const SlotBinding,
604 element_count_buffers: []const ElementCountBufferBinding,
605 launch_options: LaunchOptions,
606 ) gpu.BackendError!void {
607 var graph = try self.createLaunchGraphPlan(scratch, launch_options);
608 defer graph.deinit();
609 try self.launchGraph(scratch, slot_bindings, element_count_buffers, graph.plan());
610 }
611
612 fn launchGraph(
613 self: *const LoadedFragment,
614 scratch: std.mem.Allocator,
615 slot_bindings: []const SlotBinding,
616 element_count_buffers: []const ElementCountBufferBinding,
617 graph: LaunchGraphPlan,
618 ) gpu.BackendError!void {
619 const state = self.stateConst();
620 const artifact_plan = state.compiled_fragment.artifactPlan();
621 try state.loaded_kernels.launchGraph(
622 scratch,
623 artifact_plan,
624 slot_bindings,
625 element_count_buffers,
626 graph,
627 );
628 }
629
630 fn launchGraphWithDependencyEvents(
631 self: *const LoadedFragment,
632 scratch: std.mem.Allocator,
633 slot_bindings: []const SlotBinding,
634 element_count_buffers: []const ElementCountBufferBinding,
635 graph: LaunchGraphPlan,
636 ) gpu.BackendError!void {
637 const state = self.stateConst();
638 const artifact_plan = state.compiled_fragment.artifactPlan();
639 try state.loaded_kernels.launchGraphWithDependencyEvents(
640 scratch,
641 artifact_plan,
642 slot_bindings,
643 element_count_buffers,
644 graph,
645 );
646 }
647
648 pub fn launchKernelWithArguments(
649 self: *const LoadedFragment,
650 kernel_index: usize,
651 buffers: []const gpu.BufferBinding,
652 scalar_arguments: []const choir_abi.ScalarArgument,
653 launch_options: LaunchOptions,
654 ) gpu.BackendError!void {
655 const state = self.stateConst();
656 const artifact_plan = state.compiled_fragment.artifactPlan();
657 try state.loaded_kernels.launchKernelWithArguments(
658 artifact_plan,
659 kernel_index,
660 buffers,
661 scalar_arguments,
662 launch_options,
663 );
664 }
665 };
666
667 fn launchOptionsUsePreparedPlan(options: LaunchOptions) bool {
668 return options.stream == null and
669 options.wait_events.len == 0 and
670 options.signal_event == null and
671 tuning_mod.launchTuningIsEmpty(options.tuning) and
672 options.runtime_scalar_arguments.len == 0;
673 }
674
675 fn sameOptionalStream(a: ?gpu.StreamHandle, b: ?gpu.StreamHandle) bool {
676 if (a == null or b == null) return a == null and b == null;
677 return a.?.id == b.?.id and a.?.backend == b.?.backend;
678 }
679
680 fn artifactPlanOptions(options: FragmentCompilerOptions) artifact_product.ArtifactPlanOptions {
681 return .{
682 .format = options.artifact_format,
683 .kernel_call_registry = options.kernel_call_registry,
684 };
685 }
686
687 pub fn backendTargetProfile(
688 handle: gpu.BackendHandle,
689 requested_format: ?gpu.ArtifactFormat,
690 math_tier: gpu.BackendMathTier,
691 ) gpu.BackendError!preparation.BackendTargetProfile {
692 const caps = try handle.queryCapabilities();
693 const kind = handle.backendKind() orelse caps.identity.backend;
694 const format = requested_format orelse artifact_product.defaultArtifactFormat(kind) orelse return error.UnsupportedOperation;
695 return try preparation.BackendTargetProfile.initWithMathTier(caps, kind, format, math_tier);
696 }
697
698 pub fn recordBackendPreparationRun(
699 instrumentation: FragmentInstrumentation,
700 run: preparation.BackendPreparationRun,
701 ) !void {
702 try instrumentation.recordElapsed(.run_contract_pipeline, run.contract_ns);
703 try instrumentation.recordElapsed(.run_tensor_pipeline, run.tensor_ns);
704 try instrumentation.recordElapsed(.run_dispatch_pipeline, run.dispatch_ns);
705 try instrumentation.recordElapsed(.run_memory_pipeline, run.memory_ns);
706 try instrumentation.recordElapsed(.run_kernel_pipeline, run.kernel_ns);
707 try instrumentation.recordElapsed(.run_target_pipeline, run.target_ns);
708 }
709
710 pub fn loadFragment(
711 allocator: std.mem.Allocator,
712 handle: gpu.BackendHandle,
713 compiled_fragment: *exec_product.CompiledFragment,
714 options: FragmentCompilerOptions,
715 ) !*LoadedFragment {
716 var compiled_owned = true;
717 defer if (compiled_owned) compiled_fragment.deinit();
718
719 var tuning_cache = LaunchTuningCache.init(allocator);
720 var tuning_owned = true;
721 errdefer if (tuning_owned) tuning_cache.deinit();
722 const tuning_start = nowNs();
723 if (options.launch_tuning_artifact.len != 0) {
724 const records = try decodeLaunchTuningArtifact(allocator, options.launch_tuning_artifact);
725 defer allocator.free(records);
726 try tuning_cache.importRecords(records);
727 try options.instrumentation.record(.import_launch_tuning, tuning_start);
728 }
729
730 var loaded_kernels = try loadKernelsWithInstrumentation(
731 allocator,
732 handle,
733 compiled_fragment.artifactPlan(),
734 options.instrumentation,
735 );
736 var kernels_owned = true;
737 errdefer if (kernels_owned) loaded_kernels.deinit();
738
739 const state = allocator.create(LoadedFragmentState) catch return error.OutOfMemory;
740 var state_owned = true;
741 errdefer if (state_owned) allocator.destroy(state);
742 state.* = .{
743 .allocator = allocator,
744 .compiled_fragment = compiled_fragment,
745 .loaded_kernels = loaded_kernels,
746 .tuning_cache = tuning_cache,
747 };
748 compiled_owned = false;
749 tuning_owned = false;
750 kernels_owned = false;
751 state_owned = false;
752 const fragment: *LoadedFragment = @ptrCast(state);
753 errdefer fragment.deinit();
754 try fragment.refreshPreparedLaunchGraphPlan();
755 return fragment;
756 }
757
758 /// Compiles the stored stage results of `prepared`, rebuilding their plans
759 /// inside the caller's `workspace`, into a compiled fragment so the caller can
760 /// turn a finished preparation into device code it can later load. The returned
761 /// fragment copies the kernel list and its plan, so it stays valid after the
762 /// caller releases `workspace` and `prepared`. The step records its time as
763 /// `compile_fragment` in the instrumentation.
764 pub fn compileFragmentFromPreparedModule(
765 allocator: std.mem.Allocator,
766 handle: gpu.BackendHandle,
767 prepared: *const preparation.pipeline.BackendPreparedModule,
768 options: FragmentCompilerOptions,
769 workspace: []u8,
770 comptime configuration: choir.product.operation.Configuration,
771 ) !*exec_product.CompiledFragment {
772 const job = try createArtifactJobFromPreparedModule(
773 allocator,
774 handle,
775 prepared,
776 options,
777 workspace,
778 configuration,
779 );
780 defer job.deinit();
781 const start = nowNs();
782 const fragment = try exec_product.compileFragmentFromArtifactJob(allocator, job);
783 errdefer fragment.deinit();
784 try options.instrumentation.record(.compile_fragment, start);
785 return fragment;
786 }
787
788 /// Compiles the stored stage results of `prepared`, using `workspace` for the
789 /// rebuilt plans, into a new artifact job that the caller frees with `deinit`,
790 /// so the caller can obtain the compiled kernels by themselves before building
791 /// a runnable fragment. The job is built from the stored records of `prepared`
792 /// alone and holds no live compile job of any stage. The step records its time
793 /// as `plan_create_backend_artifacts`.
794 pub fn createArtifactJobFromPreparedModule(
795 allocator: std.mem.Allocator,
796 handle: gpu.BackendHandle,
797 prepared: *const preparation.pipeline.BackendPreparedModule,
798 options: FragmentCompilerOptions,
799 workspace: []u8,
800 comptime configuration: choir.product.operation.Configuration,
801 ) !*artifact_product.ArtifactJob {
802 const start = nowNs();
803 var plan = try artifact_product.createBackendArtifactPlanFromPreparedModule(
804 allocator,
805 handle,
806 prepared,
807 artifactPlanOptions(options),
808 workspace,
809 configuration,
810 );
811 var plan_owned = true;
812 errdefer if (plan_owned) plan.deinit();
813 const job = try artifact_product.ArtifactJob.init(allocator, plan);
814 plan_owned = false;
815 errdefer job.deinit();
816 try options.instrumentation.record(.plan_create_backend_artifacts, start);
817 return job;
818 }
819
820 pub fn compileFragmentFromPreparedJob(
821 allocator: std.mem.Allocator,
822 handle: gpu.BackendHandle,
823 prepared: *preparation.BackendPreparedJob,
824 options: FragmentCompilerOptions,
825 ) !*exec_product.CompiledFragment {
826 const artifact_module = try createArtifactJobFromPreparedJob(allocator, handle, prepared, options);
827 defer artifact_module.deinit();
828 const exec_start = nowNs();
829 const fragment = try exec_product.compileFragmentFromArtifactJob(allocator, artifact_module);
830 errdefer fragment.deinit();
831 try options.instrumentation.record(.compile_fragment, exec_start);
832 return fragment;
833 }
834
835 pub fn createArtifactJobFromPreparedJob(
836 allocator: std.mem.Allocator,
837 handle: gpu.BackendHandle,
838 prepared: *preparation.BackendPreparedJob,
839 options: FragmentCompilerOptions,
840 ) !*artifact_product.ArtifactJob {
841 const target_module = try prepared.targetModule();
842 const artifact_start = nowNs();
843 const artifact_module = try artifact_product.createArtifactJobFromTargetJob(
844 allocator,
845 handle,
846 target_module,
847 artifactPlanOptions(options),
848 );
849 errdefer artifact_module.deinit();
850 try options.instrumentation.record(.plan_create_backend_artifacts, artifact_start);
851 return artifact_module;
852 }
853
854 pub const FragmentPreparationPlan = struct {
855 allocator: std.mem.Allocator,
856 decoded_family_tuning: ?kernel_library.tuning.OwnedFamilyTuningRecords,
857 decoded_family_tuning_reader: kernel_library.tuning.FamilyTuningReader,
858 matrix_product_tuning_records: []kernel_library.tuning.MatrixProductFamilyScheduleTuningRecord,
859 generated_scan_schedule_decisions: []preparation.target.GeneratedScanScheduleDecision,
860 generated_row_pipeline_schedule_decisions: []preparation.target.GeneratedRowPipelineScheduleDecision,
861 run_options: preparation.BackendPreparationRunOptions,
862
863 pub fn init(
864 self: *FragmentPreparationPlan,
865 allocator: std.mem.Allocator,
866 handle: gpu.BackendHandle,
867 options: FragmentCompilerOptions,
868 ) !void {
869 const target_profile = try backendTargetProfile(handle, options.artifact_format, options.math_tier);
870 self.allocator = allocator;
871 self.decoded_family_tuning = null;
872 self.matrix_product_tuning_records = &.{};
873 self.generated_scan_schedule_decisions = &.{};
874 self.generated_row_pipeline_schedule_decisions = &.{};
875 errdefer self.deinit();
876
877 var family_tuning = options.family_tuning;
878 if (family_tuning == null and options.family_tuning_artifact.len != 0) {
879 self.decoded_family_tuning = try kernel_library.tuning.decodeFamilyTuningArtifact(
880 allocator,
881 options.family_tuning_artifact,
882 );
883 self.decoded_family_tuning_reader = kernel_library.tuning.FamilyTuningReader.init(
884 try handle.queryCapabilities(),
885 self.decoded_family_tuning.?.table(),
886 );
887 family_tuning = &self.decoded_family_tuning_reader;
888 }
889
890 var matrix_product_tuning: ?kernel_library.linalg.MatrixProductScheduleReader = null;
891 if (options.matrix_product_schedule == null and options.matrix_product_schedule_tuning_artifact.len != 0) {
892 const decoded = try schedule_tuning.decodeMatrixProductFamilyScheduleTuningArtifact(
893 allocator,
894 options.matrix_product_schedule_tuning_artifact,
895 );
896 defer allocator.free(decoded);
897 var cache = schedule_tuning.MatrixProductFamilyScheduleTuningCache.init(allocator);
898 defer cache.deinit();
899 try cache.importRecords(decoded);
900 self.matrix_product_tuning_records = try cache.exportRecords(allocator);
901 matrix_product_tuning = .{
902 .device = (try handle.queryCapabilities()).identity,
903 .format = target_profile.artifact_format,
904 .records = self.matrix_product_tuning_records,
905 };
906 }
907
908 if (options.generated_scan_schedule) |schedule| {
909 const decisions = allocator.alloc(preparation.target.GeneratedScanScheduleDecision, 1) catch return error.OutOfMemory;
910 decisions[0] = .{ .schedule = schedule };
911 self.generated_scan_schedule_decisions = decisions;
912 } else if (options.generated_scan_schedule_tuning_artifact.len != 0) {
913 const records = try schedule_tuning.decodeGeneratedScanScheduleTuningArtifact(
914 allocator,
915 options.generated_scan_schedule_tuning_artifact,
916 );
917 defer allocator.free(records);
918 const caps = try handle.queryCapabilities();
919 var decisions = std.ArrayListUnmanaged(preparation.target.GeneratedScanScheduleDecision).empty;
920 errdefer decisions.deinit(allocator);
921 for (records) |record| {
922 if (record.key.format != target_profile.artifact_format) continue;
923 if (record.key.dtype != .f32) continue;
924 var candidate_buffer: [preparation.kernelization.max_scan_schedule_candidates]preparation.target.GeneratedScanSchedule = undefined;
925 const candidates = preparation.kernelization.scanScheduleCandidates(
926 record.key.total,
927 target_profile.artifact_format,
928 &candidate_buffer,
929 );
930 if (candidates.len < 2) continue;
931 const key = try schedule_tuning.GeneratedScanScheduleTuningKey.init(caps, .{
932 .format = target_profile.artifact_format,
933 .total = record.key.total,
934 .dtype = record.key.dtype,
935 .schedule_version = preparation.kernelization.generated_scan_schedule_version,
936 .candidates = candidates,
937 });
938 if (!key.eql(record.key)) continue;
939 decisions.append(allocator, .{
940 .total = record.key.total,
941 .schedule = record.selection.schedule,
942 }) catch return error.OutOfMemory;
943 }
944 self.generated_scan_schedule_decisions = decisions.toOwnedSlice(allocator) catch return error.OutOfMemory;
945 }
946
947 if (options.generated_row_pipeline_schedule) |schedule| {
948 const decisions = allocator.alloc(preparation.target.GeneratedRowPipelineScheduleDecision, 1) catch return error.OutOfMemory;
949 decisions[0] = .{ .schedule = schedule };
950 self.generated_row_pipeline_schedule_decisions = decisions;
951 } else if (options.generated_row_pipeline_schedule_tuning_artifact.len != 0) {
952 const records = try schedule_tuning.decodeGeneratedRowPipelineScheduleTuningArtifact(
953 allocator,
954 options.generated_row_pipeline_schedule_tuning_artifact,
955 );
956 defer allocator.free(records);
957 const caps = try handle.queryCapabilities();
958 var decisions = std.ArrayListUnmanaged(preparation.target.GeneratedRowPipelineScheduleDecision).empty;
959 errdefer decisions.deinit(allocator);
960 for (records) |record| {
961 if (record.key.format != target_profile.artifact_format) continue;
962 if (record.key.dtype != .f32) continue;
963 var candidate_buffer: [preparation.kernelization.max_row_pipeline_schedule_candidates]preparation.target.GeneratedRowPipelineSchedule = undefined;
964 const candidates = preparation.kernelization.rowPipelineScheduleCandidates(
965 record.key.cols,
966 target_profile.artifact_format,
967 &candidate_buffer,
968 );
969 if (candidates.len < 2) continue;
970 const key = try schedule_tuning.GeneratedRowPipelineScheduleTuningKey.init(caps, .{
971 .format = target_profile.artifact_format,
972 .rows = record.key.rows,
973 .cols = record.key.cols,
974 .dtype = record.key.dtype,
975 .schedule_version = preparation.kernelization.generated_row_pipeline_schedule_version,
976 .candidates = candidates,
977 });
978 if (!key.eql(record.key)) continue;
979 decisions.append(allocator, .{
980 .shape = .{ .rows = record.key.rows, .cols = record.key.cols },
981 .schedule = record.selection.schedule,
982 }) catch return error.OutOfMemory;
983 }
984 self.generated_row_pipeline_schedule_decisions = decisions.toOwnedSlice(allocator) catch return error.OutOfMemory;
985 }
986
987 const kernel_library_lowering: preparation.KernelLibraryLowering = if (options.kernel_call_registry != null) .enabled else .disabled;
988 self.run_options = .{
989 .failure = options.preparation_failure,
990 .target_profile = target_profile,
991 .generated_scan_schedules = self.generated_scan_schedule_decisions,
992 .generated_row_pipeline_schedules = self.generated_row_pipeline_schedule_decisions,
993 .tensor = .{
994 .activation = .{ .kernel_library = kernel_library_lowering },
995 .einsum = .{
996 .kernel_library = kernel_library_lowering,
997 .matrix_product_schedule = options.matrix_product_schedule,
998 .matrix_product_tuning = matrix_product_tuning,
999 .family_tuning = family_tuning,
1000 },
1001 .indexing = .{
1002 .kernel_library = kernel_library_lowering,
1003 .gather_schedule = options.gather_schedule,
1004 .scatter_schedule = options.scatter_schedule,
1005 .scatter_add_schedule = options.scatter_add_schedule,
1006 .family_tuning = family_tuning,
1007 },
1008 .loss = .{
1009 .kernel_library = kernel_library_lowering,
1010 .row_sparse_cross_entropy_schedule = options.row_sparse_cross_entropy_schedule,
1011 },
1012 },
1013 };
1014 }
1015
1016 pub fn deinit(self: *FragmentPreparationPlan) void {
1017 if (self.generated_row_pipeline_schedule_decisions.len != 0) self.allocator.free(self.generated_row_pipeline_schedule_decisions);
1018 if (self.generated_scan_schedule_decisions.len != 0) self.allocator.free(self.generated_scan_schedule_decisions);
1019 self.allocator.free(self.matrix_product_tuning_records);
1020 if (self.decoded_family_tuning) |*owned| owned.deinit();
1021 self.* = undefined;
1022 }
1023 };
1024
1025 pub fn prepareFragmentFromSemanticModule(
1026 allocator: std.mem.Allocator,
1027 handle: gpu.BackendHandle,
1028 module: *semantic.SemanticModule,
1029 options: FragmentCompilerOptions,
1030 ) !preparation.BackendPreparedJob {
1031 var module_owned = true;
1032 errdefer if (module_owned) module.deinit();
1033
1034 var plan: FragmentPreparationPlan = undefined;
1035 try plan.init(allocator, handle, options);
1036 defer plan.deinit();
1037
1038 module_owned = false;
1039 return try preparation.prepareBackendJobFromSemanticModule(allocator, module, plan.run_options);
1040 }
1041
1042 pub fn compileFragmentFromSemanticModule(
1043 allocator: std.mem.Allocator,
1044 handle: gpu.BackendHandle,
1045 module: *semantic.SemanticModule,
1046 options: FragmentCompilerOptions,
1047 ) !*exec_product.CompiledFragment {
1048 var prepared = try prepareFragmentFromSemanticModule(allocator, handle, module, options);
1049 defer prepared.deinit();
1050
1051 try recordBackendPreparationRun(options.instrumentation, prepared.run);
1052 return try compileFragmentFromPreparedJob(allocator, handle, &prepared, options);
1053 }
1054
1055 fn nowNs() i128 {
1056 return sys.time.nanoTimestamp();
1057 }
1058
1059 const testing = std.testing;
1060
1061 fn createTestLoadedFragment(
1062 allocator: std.mem.Allocator,
1063 handle: gpu.BackendHandle,
1064 pass_ctx: *passes.PassContext,
1065 choir_module: *ir.Operation,
1066 options: FragmentCompilerOptions,
1067 ) !*LoadedFragment {
1068 const artifact_start = nowNs();
1069 var artifact_plan = try createTestBackendArtifactPlan(
1070 allocator,
1071 handle,
1072 pass_ctx,
1073 choir_module,
1074 artifactPlanOptions(options),
1075 );
1076 var plan_owned = true;
1077 errdefer if (plan_owned) artifact_plan.deinit();
1078
1079 const artifact_module = try artifact_product.ArtifactJob.init(
1080 allocator,
1081 artifact_plan,
1082 );
1083 plan_owned = false;
1084 defer artifact_module.deinit();
1085 try options.instrumentation.record(.plan_create_backend_artifacts, artifact_start);
1086 const exec_start = nowNs();
1087 const compiled = try exec_product.compileFragmentFromArtifactJob(allocator, artifact_module);
1088 try options.instrumentation.record(.compile_fragment, exec_start);
1089 return try loadFragment(allocator, handle, compiled, options);
1090 }
1091
1092 fn compileAndLoadTestSemanticModule(
1093 allocator: std.mem.Allocator,
1094 handle: gpu.BackendHandle,
1095 module: *semantic.SemanticModule,
1096 options: FragmentCompilerOptions,
1097 ) !*LoadedFragment {
1098 const compiled = try compileFragmentFromSemanticModule(allocator, handle, module, options);
1099 return try loadFragment(allocator, handle, compiled, options);
1100 }
1101
1102 fn compileAndLoadTestPreparedJob(
1103 allocator: std.mem.Allocator,
1104 handle: gpu.BackendHandle,
1105 prepared: *preparation.BackendPreparedJob,
1106 options: FragmentCompilerOptions,
1107 ) !*LoadedFragment {
1108 const compiled = try compileFragmentFromPreparedJob(allocator, handle, prepared, options);
1109 return try loadFragment(allocator, handle, compiled, options);
1110 }
1111
1112 test "loaded fragment consumes compiled fragment and partial backend loads on failure" {
1113 const allocator = testing.allocator;
1114
1115 var owned = try escapedTwoKernelChoirModule(allocator, "loaded_fragment_failed_load");
1116 defer owned.deinit();
1117
1118 var cache = passes.AnalysisCache.init(allocator, null);
1119 defer cache.deinit();
1120 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
1121 defer pass_ctx.deinit();
1122
1123 var state = RecordingBackendState{
1124 .allocator = allocator,
1125 .kind = .cuda,
1126 .format = .cuda_ptx,
1127 .fail_load_after_count = 1,
1128 };
1129 const handle = state.handle();
1130
1131 var artifact_plan = try createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{});
1132 var plan_owned = true;
1133 errdefer if (plan_owned) artifact_plan.deinit();
1134 const artifact_module = try artifact_product.ArtifactJob.init(
1135 allocator,
1136 artifact_plan,
1137 );
1138 plan_owned = false;
1139 defer artifact_module.deinit();
1140
1141 const compiled = try exec_product.compileFragmentFromArtifactJob(allocator, artifact_module);
1142 try testing.expectError(
1143 error.RuntimeUnavailable,
1144 loadFragment(allocator, handle, compiled, .{ .artifact_format = .cuda_ptx }),
1145 );
1146 try testing.expectEqual(@as(usize, 1), state.load_count);
1147 try testing.expectEqual(@as(usize, 1), state.destroy_count);
1148 }
1149
1150 test "Choir executable plan loads artifacts and launches CUDA device-count kernels" {
1151 const allocator = testing.allocator;
1152
1153 var owned = try addChoirModule(allocator, "choir_cuda_executable_add");
1154 defer owned.deinit();
1155
1156 var cache = passes.AnalysisCache.init(allocator, null);
1157 defer cache.deinit();
1158 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
1159 defer pass_ctx.deinit();
1160
1161 var state = RecordingBackendState{
1162 .allocator = allocator,
1163 .kind = .cuda,
1164 .format = .cuda_ptx,
1165 };
1166 const handle = state.handle();
1167
1168 var artifact_plan = try createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{});
1169 defer artifact_plan.deinit();
1170
1171 var executable = try loadKernels(allocator, handle, &artifact_plan);
1172 defer executable.deinit();
1173
1174 const kernel = artifact_plan.kernels.items[0];
1175 const slot_bindings = try slotBindingsForKernel(allocator, kernel, .cuda);
1176 defer allocator.free(slot_bindings);
1177 const count_bindings = try elementCountBindingsForPlan(allocator, &artifact_plan, .cuda);
1178 defer allocator.free(count_bindings);
1179
1180 try executable.launchAll(allocator, &artifact_plan, slot_bindings, count_bindings);
1181
1182 try testing.expectEqual(@as(usize, 1), executable.kernelCount());
1183 try testing.expectEqual(@as(usize, 1), state.load_count);
1184 try testing.expectEqual(@as(usize, 1), state.launch_count);
1185 try testing.expectEqual(state.last_loaded_id.?, state.last_launch_loaded_id.?);
1186 try testing.expectEqual(@as(usize, 4), state.last_launch_buffer_count);
1187 try testing.expectEqual(@as(usize, 0), state.last_launch_scalar_count);
1188 try testing.expect(state.last_launch_scalar_u32 == null);
1189 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[0]);
1190 try testing.expectEqual(@as(u32, 32), state.last_launch_threadgroup[0]);
1191 try testing.expectEqual(@as(gpu.BackendObjectId, 100), state.last_buffer_ids[0]);
1192 try testing.expectEqual(@as(gpu.BackendObjectId, 900), state.last_buffer_ids[3]);
1193 try testing.expectEqual(gpu.BufferAccess.write_only, state.last_buffer_access[0]);
1194 try testing.expectEqual(gpu.BufferAccess.read_only, state.last_buffer_access[1]);
1195 try testing.expectEqual(gpu.BufferAccess.read_only, state.last_buffer_access[2]);
1196 try testing.expectEqual(gpu.BufferAccess.read_only, state.last_buffer_access[3]);
1197 }
1198
1199 test "Choir executable plan runs native CPU machine-code prepared kernels" {
1200 const allocator = testing.allocator;
1201
1202 var owned = try addChoirModule(allocator, "choir_cpu_executable_add");
1203 defer owned.deinit();
1204
1205 var cache = passes.AnalysisCache.init(allocator, null);
1206 defer cache.deinit();
1207 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
1208 defer pass_ctx.deinit();
1209
1210 var state = gpu.cpu.State.init(allocator);
1211 defer state.deinit();
1212 const handle = state.handle();
1213
1214 var artifact_plan = createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{}) catch |err| switch (err) {
1215 error.UnsupportedOperation => return error.SkipZigTest,
1216 else => return err,
1217 };
1218 defer artifact_plan.deinit();
1219
1220 var executable = try loadKernels(allocator, handle, &artifact_plan);
1221 defer executable.deinit();
1222
1223 const lhs = [_]f32{ 1.0, 2.0, 3.0, 4.0, -1.5, 8.0, 0.25, 16.0 };
1224 const rhs = [_]f32{ 5.0, -2.0, 0.5, 6.0, 2.5, -3.0, 0.75, 4.0 };
1225 var bindings = try binding_mod.prepare(allocator, handle, &artifact_plan, &.{
1226 std.mem.sliceAsBytes(lhs[0..]),
1227 std.mem.sliceAsBytes(rhs[0..]),
1228 });
1229 defer bindings.deinit();
1230
1231 try executable.launchAll(
1232 allocator,
1233 &artifact_plan,
1234 binding_mod.slotBindings(bindings),
1235 binding_mod.elementCountBindings(bindings),
1236 );
1237
1238 const output_slot_id = artifact_plan.output_slot_ids[0];
1239 const output_binding = try binding_mod.bindingForSlot(binding_mod.slotBindings(bindings), output_slot_id, .read_only);
1240 var output = @as([8]f32, @splat(0.0));
1241 try handle.readBuffer(.{
1242 .handle = output_binding.handle,
1243 .bytes = std.mem.sliceAsBytes(output[0..]),
1244 });
1245
1246 for (output, 0..) |value, index| {
1247 try testing.expectEqual(lhs[index] + rhs[index], value);
1248 }
1249 }
1250
1251 test "Choir executable fragment runs native CPU object prepared kernels" {
1252 const allocator = testing.allocator;
1253
1254 var owned = try addChoirModule(allocator, "choir_cpu_loaded_fragment_object_add");
1255 defer owned.deinit();
1256
1257 var cache = passes.AnalysisCache.init(allocator, null);
1258 defer cache.deinit();
1259 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
1260 defer pass_ctx.deinit();
1261
1262 var state = gpu.cpu.State.init(allocator);
1263 defer state.deinit();
1264 const handle = state.handle();
1265
1266 var fragment = createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{
1267 .artifact_format = .cpu_object,
1268 }) catch |err| switch (err) {
1269 error.UnsupportedOperation => return error.SkipZigTest,
1270 else => return err,
1271 };
1272 defer fragment.deinit();
1273
1274 const summary = try fragment.kernelSummary(0);
1275 try testing.expectEqual(gpu.ArtifactFormat.cpu_object, summary.artifact_format);
1276 try testing.expectEqual(artifact_product.KernelSource.tensor, summary.source);
1277 try testing.expectEqual(@as(usize, 1), fragment.outputCount());
1278
1279 const lhs = [_]f32{ 1.0, 2.0, 3.0, 4.0, -1.5, 8.0, 0.25, 16.0 };
1280 const rhs = [_]f32{ 5.0, -2.0, 0.5, 6.0, 2.5, -3.0, 0.75, 4.0 };
1281 const inputs = [_][]const u8{
1282 std.mem.sliceAsBytes(lhs[0..]),
1283 std.mem.sliceAsBytes(rhs[0..]),
1284 };
1285 const bindings = try fragment.prepareInvocationBindings(allocator, inputs[0..]);
1286 defer bindings.deinit();
1287
1288 try fragment.submitPreparedForTest(allocator, bindings);
1289 try fragment.completeInvocationWithOptions(.{});
1290
1291 var output = @as([8]f32, @splat(0.0));
1292 try fragment.readInvocationOutput(bindings, 0, std.mem.sliceAsBytes(output[0..]));
1293
1294 for (output, 0..) |value, index| {
1295 try testing.expectEqual(lhs[index] + rhs[index], value);
1296 }
1297 }
1298
1299 fn runDotGeneralNativeCpuFragment(format: gpu.ArtifactFormat) !void {
1300 const allocator = testing.allocator;
1301
1302 var state = gpu.cpu.State.init(allocator);
1303 defer state.deinit();
1304 const handle = state.handle();
1305
1306 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1307 errdefer builder.deinit();
1308 const lhs_ty = try builder.tensor(.f32, &.{ 2, 3 });
1309 const rhs_ty = try builder.tensor(.f32, &.{ 3, 2 });
1310 const out_ty = try builder.tensor(.f32, &.{ 2, 2 });
1311 const module_name = switch (format) {
1312 .cpu_machine_code => "choir_cpu_dot_general_machine_code",
1313 .cpu_object => "choir_cpu_dot_general_object",
1314 else => return error.UnsupportedArtifactFormat,
1315 };
1316 var fb = try builder.beginFunction(module_name, &.{ lhs_ty, rhs_ty }, &.{out_ty});
1317 const product = try fb.dotGeneral(
1318 fb.parameter(0),
1319 fb.parameter(1),
1320 out_ty,
1321 &.{1},
1322 &.{0},
1323 &.{},
1324 &.{},
1325 );
1326 try fb.return_(&.{product});
1327 try fb.finish();
1328 const module = try builder.finish();
1329
1330 var fragment = compileAndLoadTestSemanticModule(allocator, handle, module, .{
1331 .artifact_format = format,
1332 }) catch |err| switch (err) {
1333 error.UnsupportedOperation => return error.SkipZigTest,
1334 else => return err,
1335 };
1336 defer fragment.deinit();
1337
1338 const summary = try fragment.kernelSummary(0);
1339 try testing.expectEqual(format, summary.artifact_format);
1340 try testing.expectEqual(artifact_product.KernelSource.tensor, summary.source);
1341 try testing.expectEqual(artifact_product.PlannedKernelCompileLaunch.dot_general, summary.compile_launch);
1342 try testing.expectEqual(@as(usize, 1), fragment.outputCount());
1343
1344 const lhs = [_]f32{ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0 };
1345 const rhs = [_]f32{ 7.0, 8.0, 9.0, 10.0, 11.0, 12.0 };
1346 var output = @as([4]f32, @splat(0.0));
1347 const output_bytes = std.mem.sliceAsBytes(output[0..]);
1348 var outputs = [_][]u8{output_bytes};
1349 try fragment.invokeAndReadForTest(
1350 allocator,
1351 allocator,
1352 &.{ std.mem.sliceAsBytes(lhs[0..]), std.mem.sliceAsBytes(rhs[0..]) },
1353 &outputs,
1354 );
1355
1356 try testing.expectEqualSlices(f32, &.{ 58.0, 64.0, 139.0, 154.0 }, output[0..]);
1357 }
1358
1359 test "Choir executable fragment runs dot_general native CPU machine-code" {
1360 try runDotGeneralNativeCpuFragment(.cpu_machine_code);
1361 }
1362
1363 test "Choir executable fragment runs dot_general native CPU object" {
1364 try runDotGeneralNativeCpuFragment(.cpu_object);
1365 }
1366
1367 test "Choir executable fragment materializes i1 outputs as single bytes on native CPU" {
1368 const allocator = testing.allocator;
1369
1370 var state = gpu.cpu.State.init(allocator);
1371 defer state.deinit();
1372 const handle = state.handle();
1373
1374 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1375 errdefer builder.deinit();
1376 const f32_8 = try builder.tensor(.f32, &.{8});
1377 const i1_8 = try builder.tensor(.i1, &.{8});
1378 var fb = try builder.beginFunction("choir_cpu_bool_byte_flags", &.{ f32_8, f32_8 }, &.{i1_8});
1379 const flags = try fb.compare(fb.parameter(0), fb.parameter(1), i1_8, .gt);
1380 try fb.return_(&.{flags});
1381 try fb.finish();
1382 const module = try builder.finish();
1383
1384 var fragment = compileAndLoadTestSemanticModule(allocator, handle, module, .{
1385 .artifact_format = .cpu_object,
1386 }) catch |err| switch (err) {
1387 error.UnsupportedOperation => return error.SkipZigTest,
1388 else => return err,
1389 };
1390 defer fragment.deinit();
1391
1392 const lhs = [_]f32{ 1.0, -1.0, 2.0, -2.0, 0.5, -0.5, 3.0, 0.0 };
1393 const rhs = [_]f32{ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0 };
1394 var flags_out = @as([8]u8, @splat(0xaa));
1395 var outputs = [_][]u8{flags_out[0..]};
1396 try fragment.invokeAndReadForTest(
1397 allocator,
1398 allocator,
1399 &.{ std.mem.sliceAsBytes(lhs[0..]), std.mem.sliceAsBytes(rhs[0..]) },
1400 &outputs,
1401 );
1402
1403 try testing.expectEqualSlices(u8, &.{ 1, 0, 1, 0, 1, 0, 1, 0 }, flags_out[0..]);
1404 }
1405
1406 test "Choir executable fragment runs kernel library native CPU object" {
1407 const allocator = testing.allocator;
1408
1409 var state = gpu.cpu.State.init(allocator);
1410 defer state.deinit();
1411 const handle = state.handle();
1412
1413 const VectorAdd8F32 = kernel_library.elementwise.VectorAdd8F32;
1414 const options = FragmentCompilerOptions{
1415 .artifact_format = .cpu_object,
1416 .authored_kernel_diagnostic_id = "accy/executable/kernel-library-native-cpu-object",
1417 };
1418 const compiled = VectorAdd8F32.compileFragment(
1419 allocator,
1420 VectorAdd8F32.Limits.testing,
1421 handle,
1422 options,
1423 ) catch |err| switch (err) {
1424 error.UnsupportedOperation => return error.SkipZigTest,
1425 else => return err,
1426 };
1427 var fragment = loadFragment(allocator, handle, compiled, options) catch |err| switch (err) {
1428 error.UnsupportedOperation => return error.SkipZigTest,
1429 else => return err,
1430 };
1431 defer fragment.deinit();
1432
1433 const plan = fragment.artifactPlan();
1434 const planned = plan.kernels.items[0];
1435 const summary = try fragment.kernelSummary(0);
1436 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
1437 try testing.expectEqual(gpu.ArtifactFormat.cpu_object, summary.artifact_format);
1438 try testing.expectEqual(artifact_product.KernelSource.choir_kernel, summary.source);
1439 try testing.expectEqual(@as(u32, 10), summary.compile_argument_count);
1440 try testing.expectEqual(.authored, summary.compile_launch);
1441 try testing.expect(summary.fixed_threadgroup);
1442 try testing.expectEqual(@as(u32, 2), summary.launch_geometry.grid[0]);
1443 try testing.expectEqual(@as(u32, 4), summary.launch_geometry.threadgroup[0]);
1444 try testing.expectEqual(@as(usize, 7), planned.static_arguments.len);
1445
1446 const lhs = [_]f32{ 1.0, 2.0, -3.0, 4.5, 8.0, -1.0, 0.25, 16.0 };
1447 const rhs = [_]f32{ 4.0, -2.0, 6.0, 0.5, -3.0, 5.0, 0.75, -8.0 };
1448 const dst_buffer = try handle.allocateBuffer(.{
1449 .byte_size = 8 * @sizeOf(f32),
1450 .alignment = @alignOf(f32),
1451 .dtype = .f32,
1452 .element_count = 8,
1453 });
1454 defer handle.destroyObject(dst_buffer.id);
1455
1456 const lhs_buffer = try handle.allocateBuffer(.{
1457 .byte_size = 8 * @sizeOf(f32),
1458 .alignment = @alignOf(f32),
1459 .dtype = .f32,
1460 .element_count = 8,
1461 });
1462 defer handle.destroyObject(lhs_buffer.id);
1463
1464 const rhs_buffer = try handle.allocateBuffer(.{
1465 .byte_size = 8 * @sizeOf(f32),
1466 .alignment = @alignOf(f32),
1467 .dtype = .f32,
1468 .element_count = 8,
1469 });
1470 defer handle.destroyObject(rhs_buffer.id);
1471
1472 try handle.writeBuffer(.{
1473 .handle = lhs_buffer,
1474 .bytes = std.mem.sliceAsBytes(lhs[0..]),
1475 });
1476 try handle.writeBuffer(.{
1477 .handle = rhs_buffer,
1478 .bytes = std.mem.sliceAsBytes(rhs[0..]),
1479 });
1480
1481 const bindings = [_]gpu.BufferBinding{
1482 .{
1483 .handle = dst_buffer,
1484 .access = .write_only,
1485 .ownership = .backend,
1486 .byte_size = dst_buffer.byte_size,
1487 },
1488 .{
1489 .handle = lhs_buffer,
1490 .access = .read_only,
1491 .ownership = .backend,
1492 .byte_size = lhs_buffer.byte_size,
1493 },
1494 .{
1495 .handle = rhs_buffer,
1496 .access = .read_only,
1497 .ownership = .backend,
1498 .byte_size = rhs_buffer.byte_size,
1499 },
1500 };
1501
1502 try fragment.launchKernelWithArguments(0, bindings[0..], planned.static_arguments, .{});
1503
1504 var dst = @as([8]f32, @splat(0.0));
1505 try handle.readBuffer(.{
1506 .handle = dst_buffer,
1507 .bytes = std.mem.sliceAsBytes(dst[0..]),
1508 });
1509 try testing.expectEqualSlices(f32, &.{ 5.0, 0.0, 3.0, 5.0, 5.0, 4.0, 1.0, 8.0 }, dst[0..]);
1510 }
1511
1512 fn runKernelCallNativeCpuRuntimeScalarFragment(format: gpu.ArtifactFormat) !void {
1513 const allocator = testing.allocator;
1514
1515 var state = gpu.cpu.State.init(allocator);
1516 defer state.deinit();
1517 const handle = state.handle();
1518
1519 const Axpy8F32 = kernel_library.elementwise.Axpy8F32;
1520 var call_artifact = Axpy8F32.createKernelCallArtifact(allocator, handle, .{
1521 .limits = .standard,
1522 .format = format,
1523 .runtime_scalar_argument_count = 1,
1524 }) catch |err| switch (err) {
1525 error.UnsupportedOperation => return error.SkipZigTest,
1526 else => return err,
1527 };
1528 defer call_artifact.deinit();
1529
1530 const entry = call_artifact.entry();
1531 try testing.expectEqual(format, entry.format);
1532 try testing.expectEqual(@as(u32, 11), entry.argument_count);
1533 try testing.expectEqual(@as(u32, 1), entry.runtime_scalar_argument_count);
1534 try testing.expectEqual(@as(usize, 7), entry.static_arguments.len);
1535
1536 const registry = call_artifact.registry();
1537 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1538 errdefer builder.deinit();
1539 const f32_8 = try builder.tensor(.f32, &.{8});
1540 const module_name = switch (format) {
1541 .cpu_machine_code => "choir_cpu_machine_code_kernel_call_runtime_scalar_axpy",
1542 .cpu_object => "choir_cpu_object_kernel_call_runtime_scalar_axpy",
1543 else => return error.UnsupportedArtifactFormat,
1544 };
1545 var fb = try builder.beginFunction(module_name, &.{ f32_8, f32_8 }, &.{f32_8});
1546 const call = try fb.kernelCall(
1547 &.{ fb.parameter(0), fb.parameter(1) },
1548 &.{f32_8},
1549 .{
1550 .target = Axpy8F32.target,
1551 .operand_effects = &.{ .read, .read },
1552 .result_aliases = &.{null},
1553 },
1554 );
1555 try fb.return_(&.{call.getFirstResult()});
1556 try fb.finish();
1557 const module = try builder.finish();
1558
1559 var fragment = compileAndLoadTestSemanticModule(allocator, handle, module, .{
1560 .artifact_format = format,
1561 .kernel_call_registry = ®istry,
1562 }) catch |err| switch (err) {
1563 error.UnsupportedOperation => return error.SkipZigTest,
1564 else => return err,
1565 };
1566 defer fragment.deinit();
1567
1568 const summary = try fragment.kernelSummary(0);
1569 try testing.expectEqual(format, summary.artifact_format);
1570 try testing.expectEqual(artifact_product.KernelSource.kernel_call, summary.source);
1571 try testing.expectEqual(@as(u32, 11), summary.compile_argument_count);
1572 try testing.expectEqual(@as(u32, 1), summary.runtime_scalar_argument_count);
1573 try testing.expectEqual(@as(usize, 1), fragment.outputCount());
1574
1575 const x = [_]f32{ 1.0, 2.0, -3.0, 4.0, 0.5, -1.5, 8.0, 16.0 };
1576 const y = [_]f32{ 10.0, -4.0, 1.0, 2.0, 3.0, 6.0, -8.0, 0.0 };
1577 const inputs = [_][]const u8{
1578 std.mem.sliceAsBytes(x[0..]),
1579 std.mem.sliceAsBytes(y[0..]),
1580 };
1581 const bindings = try fragment.prepareInvocationBindings(allocator, inputs[0..]);
1582 defer bindings.deinit();
1583
1584 var graph = try fragment.createLaunchGraphPlan(allocator, .{});
1585 defer graph.deinit();
1586 const runtime_args = [_]choir_abi.ScalarArgument{.{ .f32 = 2.0 }};
1587 graph.nodes[0].runtime_scalar_arguments = runtime_args[0..];
1588
1589 try fragment.submitInvocationWithGraph(allocator, bindings, graph.plan());
1590 try fragment.completeInvocationGraph(graph.plan());
1591
1592 var output = @as([8]f32, @splat(0.0));
1593 try fragment.readInvocationOutput(bindings, 0, std.mem.sliceAsBytes(output[0..]));
1594 try testing.expectEqualSlices(f32, &.{ 12.0, 0.0, -5.0, 10.0, 4.0, 3.0, 8.0, 32.0 }, output[0..]);
1595 }
1596
1597 test "Choir executable fragment runs kernel_call native CPU machine-code with runtime scalar" {
1598 try runKernelCallNativeCpuRuntimeScalarFragment(.cpu_machine_code);
1599 }
1600
1601 test "Choir executable fragment runs kernel_call native CPU object with runtime scalar" {
1602 try runKernelCallNativeCpuRuntimeScalarFragment(.cpu_object);
1603 }
1604
1605 test "Choir executable artifact planning accepts CUDA u32 kernels" {
1606 const allocator = testing.allocator;
1607
1608 var owned = try addU32ChoirModule(allocator, "choir_cuda_executable_add_u32");
1609 defer owned.deinit();
1610
1611 var cache = passes.AnalysisCache.init(allocator, null);
1612 defer cache.deinit();
1613 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
1614 defer pass_ctx.deinit();
1615
1616 var state = RecordingBackendState{
1617 .allocator = allocator,
1618 .kind = .cuda,
1619 .format = .cuda_ptx,
1620 };
1621 const handle = state.handle();
1622
1623 var artifact_plan = try createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{});
1624 defer artifact_plan.deinit();
1625
1626 try testing.expectEqual(@as(usize, 1), artifact_plan.kernelCount());
1627 const kernel = artifact_plan.kernels.items[0];
1628 try testing.expect(kernel.compile.required_dtypes.contains(.u32));
1629 try testing.expectEqual(@as(usize, 1), state.create_count);
1630 try testing.expectEqual(kernel.compile.required_dtypes.bits, state.last_create_required_dtype_bits);
1631 }
1632
1633 test "Choir executable fragment launches registered kernel_call with full launch geometry" {
1634 const allocator = testing.allocator;
1635
1636 var owned = try kernelCallChoirModule(allocator, "choir_cuda_executable_kernel_call_geometry");
1637 defer owned.deinit();
1638
1639 var cache = passes.AnalysisCache.init(allocator, null);
1640 defer cache.deinit();
1641 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
1642 defer pass_ctx.deinit();
1643
1644 var state = RecordingBackendState{
1645 .allocator = allocator,
1646 .kind = .cuda,
1647 .format = .cuda_ptx,
1648 };
1649 const handle = state.handle();
1650
1651 const source = ".visible .entry accy_custom_scale() { ret; }";
1652 const registry = artifact_product.KernelCallRegistry{ .entries = &.{.{
1653 .target = "accy.custom.scale",
1654 .version = 1,
1655 .format = .cuda_ptx,
1656 .entry_name = "accy_custom_scale",
1657 .argument_count = 5,
1658 .required_dtypes = gpu.DTypeSet.init(&.{.f32}),
1659 .payload = .{ .text = source },
1660 .launch = .{ .fixed = .{
1661 .grid = .{ 2, 3, 1 },
1662 .threadgroup = .{ 8, 2, 1 },
1663 .dynamic_shared_memory_bytes = 2048,
1664 } },
1665 .element_count_argument = .scalar_u32,
1666 .static_arguments = &.{.{ .u32 = 7 }},
1667 }} };
1668
1669 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{
1670 .kernel_call_registry = ®istry,
1671 });
1672 defer fragment.deinit();
1673
1674 const plan = fragment.artifactPlan();
1675 const kernel = plan.kernels.items[0];
1676 const slot_bindings = try slotBindingsForPlan(allocator, plan, .cuda);
1677 defer allocator.free(slot_bindings);
1678 const count_bindings = try elementCountBindingsForPlan(allocator, plan, .cuda);
1679 defer allocator.free(count_bindings);
1680
1681 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
1682 try testing.expectEqual(artifact_product.PlannedKernelSource.kernel_call, kernel.compile.source);
1683 try testing.expectEqual(@as(u32, 2048), kernel.launch_resources.geometry.dynamic_shared_memory_bytes);
1684
1685 try fragment.launchAll(allocator, slot_bindings, count_bindings);
1686
1687 try testing.expectEqual(@as(usize, 1), state.launch_count);
1688 try testing.expectEqual(@as(u32, 2), state.last_launch_grid[0]);
1689 try testing.expectEqual(@as(u32, 3), state.last_launch_grid[1]);
1690 try testing.expectEqual(@as(u32, 8), state.last_launch_threadgroup[0]);
1691 try testing.expectEqual(@as(u32, 2), state.last_launch_threadgroup[1]);
1692 try testing.expectEqual(@as(u32, 2048), state.last_launch_dynamic_shared_memory_bytes);
1693 try testing.expectEqual(@as(usize, 2), state.last_launch_scalar_count);
1694 try testing.expectEqual(@as(u32, 8), state.last_launch_scalar_u32_values[0]);
1695 try testing.expectEqual(@as(u32, 7), state.last_launch_scalar_u32_values[1]);
1696 }
1697
1698 test "Choir executable fragment launches registered kernel_call with runtime derived geometry" {
1699 const allocator = testing.allocator;
1700
1701 var owned = try kernelCallChoirModule(allocator, "choir_cuda_executable_kernel_call_derived_geometry");
1702 defer owned.deinit();
1703
1704 var cache = passes.AnalysisCache.init(allocator, null);
1705 defer cache.deinit();
1706 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
1707 defer pass_ctx.deinit();
1708
1709 var state = RecordingBackendState{
1710 .allocator = allocator,
1711 .kind = .cuda,
1712 .format = .cuda_ptx,
1713 };
1714 const handle = state.handle();
1715
1716 const source = ".visible .entry accy_custom_scale() { ret; }";
1717 const registry = artifact_product.KernelCallRegistry{ .entries = &.{.{
1718 .target = "accy.custom.scale",
1719 .version = 1,
1720 .format = .cuda_ptx,
1721 .entry_name = "accy_custom_scale",
1722 .argument_count = 6,
1723 .required_dtypes = gpu.DTypeSet.init(&.{.f32}),
1724 .payload = .{ .text = source },
1725 .launch = .{ .derived = .{
1726 .grid = .{
1727 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = 8 } },
1728 .{ .runtime_u32_ceil_div = .{ .argument_index = 1, .divisor = 4 } },
1729 .{ .fixed = 1 },
1730 },
1731 .threadgroup = .{ 8, 4, 1 },
1732 .dynamic_shared_memory_bytes = 1024,
1733 } },
1734 .runtime_scalar_argument_count = 2,
1735 .static_arguments = &.{.{ .u32 = 7 }},
1736 }} };
1737
1738 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{
1739 .kernel_call_registry = ®istry,
1740 });
1741 defer fragment.deinit();
1742
1743 const plan = fragment.artifactPlan();
1744 const kernel = plan.kernels.items[0];
1745 const slot_bindings = try slotBindingsForPlan(allocator, plan, .cuda);
1746 defer allocator.free(slot_bindings);
1747
1748 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
1749 try testing.expectEqual(artifact_product.PlannedKernelSource.kernel_call, kernel.compile.source);
1750 try testing.expectEqual(@as(u32, 2), kernel.runtime_scalar_argument_count);
1751
1752 try fragment.launchAllWithOptions(allocator, slot_bindings, &.{}, .{
1753 .runtime_scalar_arguments = &.{ .{ .u32 = 17 }, .{ .u32 = 9 } },
1754 });
1755
1756 try testing.expectEqual(@as(usize, 1), state.launch_count);
1757 try testing.expectEqual(@as(u32, 3), state.last_launch_grid[0]);
1758 try testing.expectEqual(@as(u32, 3), state.last_launch_grid[1]);
1759 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[2]);
1760 try testing.expectEqual(@as(u32, 8), state.last_launch_threadgroup[0]);
1761 try testing.expectEqual(@as(u32, 4), state.last_launch_threadgroup[1]);
1762 try testing.expectEqual(@as(u32, 1024), state.last_launch_dynamic_shared_memory_bytes);
1763 try testing.expectEqual(@as(usize, 3), state.last_launch_scalar_count);
1764 try testing.expectEqual(@as(u32, 17), state.last_launch_scalar_u32_values[0]);
1765 try testing.expectEqual(@as(u32, 9), state.last_launch_scalar_u32_values[1]);
1766 try testing.expectEqual(@as(u32, 7), state.last_launch_scalar_u32_values[2]);
1767 }
1768
1769 test "Choir executable fragment lets launch scalars override the kernel_call count default" {
1770 const allocator = testing.allocator;
1771
1772 var owned = try kernelCallChoirModule(allocator, "choir_cuda_executable_kernel_call_count_override");
1773 defer owned.deinit();
1774
1775 var cache = passes.AnalysisCache.init(allocator, null);
1776 defer cache.deinit();
1777 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
1778 defer pass_ctx.deinit();
1779
1780 var state = RecordingBackendState{
1781 .allocator = allocator,
1782 .kind = .cuda,
1783 .format = .cuda_ptx,
1784 };
1785 const handle = state.handle();
1786
1787 const source = ".visible .entry accy_custom_scale() { ret; }";
1788 const registry = artifact_product.KernelCallRegistry{ .entries = &.{.{
1789 .target = "accy.custom.scale",
1790 .version = 1,
1791 .format = .cuda_ptx,
1792 .entry_name = "accy_custom_scale",
1793 .argument_count = 4,
1794 .required_dtypes = gpu.DTypeSet.init(&.{.f32}),
1795 .payload = .{ .text = source },
1796 .launch = .{ .derived = .{
1797 .grid = .{
1798 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = 4 } },
1799 .{ .fixed = 1 },
1800 .{ .fixed = 1 },
1801 },
1802 .threadgroup = .{ 4, 1, 1 },
1803 .dynamic_shared_memory_bytes = 0,
1804 } },
1805 .runtime_scalar_argument_count = 1,
1806 .shape_family_fingerprint = 0xfeed,
1807 .shape_profile = .{
1808 .name = "rows",
1809 .fingerprint = 0xfeed,
1810 .dimensions = &.{.{ .name = "n", .runtime_scalar_argument_index = 0, .bounds = .{ .max = 4096 } }},
1811 },
1812 }} };
1813
1814 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{
1815 .kernel_call_registry = ®istry,
1816 });
1817 defer fragment.deinit();
1818
1819 const plan = fragment.artifactPlan();
1820 const kernel = plan.kernels.items[0];
1821 const slot_bindings = try slotBindingsForPlan(allocator, plan, .cuda);
1822 defer allocator.free(slot_bindings);
1823
1824 try testing.expectEqual(@as(u32, 1), kernel.runtime_scalar_argument_count);
1825 try testing.expectEqual(@as(usize, 1), kernel.runtime_scalar_defaults.len);
1826 try testing.expectEqual(@as(u32, 8), kernel.runtime_scalar_defaults[0].u32);
1827 try testing.expectEqual(artifact_product.ElementCountArgument.none, kernel.element_count_argument);
1828
1829 try fragment.launchAllWithOptions(allocator, slot_bindings, &.{}, .{
1830 .runtime_scalar_arguments = &.{.{ .u32 = 13 }},
1831 });
1832 try testing.expectEqual(@as(usize, 1), state.last_launch_scalar_count);
1833 try testing.expectEqual(@as(u32, 13), state.last_launch_scalar_u32_values[0]);
1834 try testing.expectEqual(@as(u32, 4), state.last_launch_grid[0]);
1835
1836 try fragment.launchAllWithOptions(allocator, slot_bindings, &.{}, .{});
1837 try testing.expectEqual(@as(u32, 8), state.last_launch_scalar_u32_values[0]);
1838 try testing.expectEqual(@as(u32, 2), state.last_launch_grid[0]);
1839 }
1840
1841 test "Choir executable fragment launches aliased kernel_call outputs read-write" {
1842 const allocator = testing.allocator;
1843
1844 var owned = try aliasedKernelCallChoirModule(allocator, "choir_cuda_executable_kernel_call_alias_access");
1845 defer owned.deinit();
1846
1847 var cache = passes.AnalysisCache.init(allocator, null);
1848 defer cache.deinit();
1849 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
1850 defer pass_ctx.deinit();
1851
1852 var state = RecordingBackendState{
1853 .allocator = allocator,
1854 .kind = .cuda,
1855 .format = .cuda_ptx,
1856 };
1857 const handle = state.handle();
1858
1859 const source = ".visible .entry accy_custom_update() { ret; }";
1860 const registry = artifact_product.KernelCallRegistry{ .entries = &.{.{
1861 .target = "accy.custom.update",
1862 .version = 1,
1863 .format = .cuda_ptx,
1864 .entry_name = "accy_custom_update",
1865 .argument_count = 2,
1866 .required_dtypes = gpu.DTypeSet.init(&.{.f32}),
1867 .payload = .{ .text = source },
1868 .launch = .{ .fixed = .{
1869 .grid = .{ 1, 1, 1 },
1870 .threadgroup = .{ 8, 1, 1 },
1871 } },
1872 }} };
1873
1874 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{
1875 .kernel_call_registry = ®istry,
1876 });
1877 defer fragment.deinit();
1878
1879 const plan = fragment.artifactPlan();
1880 const kernel = plan.kernels.items[0];
1881 const slot_bindings = try slotBindingsForPlan(allocator, plan, .cuda);
1882 defer allocator.free(slot_bindings);
1883 const count_bindings = try elementCountBindingsForPlan(allocator, plan, .cuda);
1884 defer allocator.free(count_bindings);
1885
1886 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
1887 try testing.expectEqual(artifact_product.PlannedKernelSource.kernel_call, kernel.compile.source);
1888 try testing.expectEqual(@as(usize, 1), kernel.input_slot_ids.len);
1889 try testing.expectEqual(kernel.output_slot_id, kernel.input_slot_ids[0]);
1890
1891 try fragment.launchAll(allocator, slot_bindings, count_bindings);
1892
1893 try testing.expectEqual(@as(usize, 1), state.launch_count);
1894 try testing.expectEqual(@as(usize, 2), state.last_launch_buffer_count);
1895 try testing.expectEqual(state.last_buffer_ids[0], state.last_buffer_ids[1]);
1896 try testing.expectEqual(gpu.BufferAccess.read_write, state.last_buffer_access[0]);
1897 try testing.expectEqual(gpu.BufferAccess.read_only, state.last_buffer_access[1]);
1898 }
1899
1900 test "Choir executable fragment launches prepared graph with runtime scalars" {
1901 const allocator = testing.allocator;
1902
1903 var owned = try kernelCallChoirModule(allocator, "choir_cuda_executable_kernel_call_prepared_graph_runtime");
1904 defer owned.deinit();
1905
1906 var cache = passes.AnalysisCache.init(allocator, null);
1907 defer cache.deinit();
1908 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
1909 defer pass_ctx.deinit();
1910
1911 var state = RecordingBackendState{
1912 .allocator = allocator,
1913 .kind = .cuda,
1914 .format = .cuda_ptx,
1915 };
1916 const handle = state.handle();
1917
1918 const source = ".visible .entry accy_custom_scale() { ret; }";
1919 const registry = artifact_product.KernelCallRegistry{ .entries = &.{.{
1920 .target = "accy.custom.scale",
1921 .version = 1,
1922 .format = .cuda_ptx,
1923 .entry_name = "accy_custom_scale",
1924 .argument_count = 6,
1925 .required_dtypes = gpu.DTypeSet.init(&.{.f32}),
1926 .payload = .{ .text = source },
1927 .launch = .{ .derived = .{
1928 .grid = .{
1929 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = 8 } },
1930 .{ .runtime_u32_ceil_div = .{ .argument_index = 1, .divisor = 4 } },
1931 .{ .fixed = 1 },
1932 },
1933 .threadgroup = .{ 8, 4, 1 },
1934 .dynamic_shared_memory_bytes = 1024,
1935 } },
1936 .runtime_scalar_argument_count = 2,
1937 .static_arguments = &.{.{ .u32 = 7 }},
1938 }} };
1939
1940 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{
1941 .kernel_call_registry = ®istry,
1942 });
1943 defer fragment.deinit();
1944
1945 var lhs = @as([8]f32, @splat(1.0));
1946 var rhs = @as([8]f32, @splat(2.0));
1947 const input_bytes = [_][]const u8{
1948 std.mem.sliceAsBytes(lhs[0..]),
1949 std.mem.sliceAsBytes(rhs[0..]),
1950 };
1951 const bindings = try fragment.prepareInvocationBindings(allocator, input_bytes[0..]);
1952 defer bindings.deinit();
1953
1954 var graph = try fragment.createLaunchGraphPlan(allocator, .{});
1955 defer graph.deinit();
1956 const runtime_args = [_]choir_abi.ScalarArgument{ .{ .u32 = 17 }, .{ .u32 = 9 } };
1957 graph.nodes[0].runtime_scalar_arguments = runtime_args[0..];
1958
1959 try fragment.submitInvocationWithGraph(allocator, bindings, graph.plan());
1960
1961 try testing.expectEqual(@as(usize, 1), state.launch_count);
1962 try testing.expectEqual(@as(u32, 3), state.last_launch_grid[0]);
1963 try testing.expectEqual(@as(u32, 3), state.last_launch_grid[1]);
1964 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[2]);
1965 try testing.expectEqual(@as(u32, 8), state.last_launch_threadgroup[0]);
1966 try testing.expectEqual(@as(u32, 4), state.last_launch_threadgroup[1]);
1967 try testing.expectEqual(@as(u32, 1024), state.last_launch_dynamic_shared_memory_bytes);
1968 try testing.expectEqual(@as(usize, 3), state.last_launch_scalar_count);
1969 try testing.expectEqual(@as(u32, 17), state.last_launch_scalar_u32_values[0]);
1970 try testing.expectEqual(@as(u32, 9), state.last_launch_scalar_u32_values[1]);
1971 try testing.expectEqual(@as(u32, 7), state.last_launch_scalar_u32_values[2]);
1972 }
1973
1974 test "Choir executable plan launches CUDA choir dot_general without static shape scalars" {
1975 const allocator = testing.allocator;
1976
1977 var owned = try dotGeneralChoirModule(allocator, "choir_cuda_executable_dot_general");
1978 defer owned.deinit();
1979
1980 var cache = passes.AnalysisCache.init(allocator, null);
1981 defer cache.deinit();
1982 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
1983 defer pass_ctx.deinit();
1984
1985 var state = RecordingBackendState{
1986 .allocator = allocator,
1987 .kind = .cuda,
1988 .format = .cuda_ptx,
1989 };
1990 const handle = state.handle();
1991
1992 var artifact_plan = try createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{});
1993 defer artifact_plan.deinit();
1994
1995 var executable = try loadKernels(allocator, handle, &artifact_plan);
1996 defer executable.deinit();
1997
1998 const kernel = artifact_plan.kernels.items[0];
1999 const caps = try handle.queryCapabilities();
2000 const key = LaunchTuningCacheKey.init(caps, kernel);
2001 try testing.expectEqual(kernel.output_layout_fingerprint, key.output_layout_fingerprint);
2002 try testing.expectEqual(kernel.input_layout_fingerprint, key.input_layout_fingerprint);
2003 try testing.expect(key.output_layout_fingerprint != 0);
2004 try testing.expect(key.input_layout_fingerprint != 0);
2005 try testing.expectEqual(artifact_product.LaunchTileKind.dot_general, key.tile_kind);
2006 try testing.expectEqual(@as(u32, 16), key.tile_m);
2007 try testing.expectEqual(@as(u32, 16), key.tile_n);
2008 try testing.expectEqual(@as(u32, 16), key.tile_k);
2009 try testing.expectEqual(@as(u32, 1), key.tile_batch);
2010 try testing.expect(key.tile_has_input_dtype);
2011 try testing.expectEqual(choir_abi.DType.f32, key.tile_input_dtype);
2012 try testing.expect(key.tile_has_output_dtype);
2013 try testing.expectEqual(choir_abi.DType.f32, key.tile_output_dtype);
2014 try testing.expectEqual(@as(u32, 2048), key.tile_input_tile_bytes);
2015 try testing.expectEqual(@as(u32, 1024), key.tile_output_tile_bytes);
2016 try testing.expectEqual(@as(u32, 0), key.tile_scratch_memory_bytes);
2017 try testing.expectEqual(@as(u32, 4), key.candidate_count);
2018 try testing.expectEqual(artifact_product.launchResourcePlanFingerprint(kernel.launch_resources), key.launch_resource_fingerprint);
2019 try testing.expect(key.launch_resource_fingerprint != 0);
2020 const slot_bindings = try slotBindingsForPlan(allocator, &artifact_plan, .cuda);
2021 defer allocator.free(slot_bindings);
2022 const count_bindings = try elementCountBindingsForPlan(allocator, &artifact_plan, .cuda);
2023 defer allocator.free(count_bindings);
2024
2025 try executable.launchAll(allocator, &artifact_plan, slot_bindings, count_bindings);
2026
2027 try testing.expectEqual(@as(usize, 1), executable.kernelCount());
2028 try testing.expectEqual(@as(usize, 1), state.load_count);
2029 try testing.expectEqual(@as(usize, 1), state.launch_count);
2030 try testing.expectEqual(@as(usize, 3), state.last_launch_buffer_count);
2031 try testing.expectEqual(@as(usize, 0), state.last_launch_scalar_count);
2032 try testing.expect(state.last_launch_scalar_u32 == null);
2033 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[0]);
2034 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[1]);
2035 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[2]);
2036 try testing.expectEqual(@as(u32, 16), state.last_launch_threadgroup[0]);
2037 try testing.expectEqual(@as(u32, 16), state.last_launch_threadgroup[1]);
2038 try testing.expectEqual(@as(u32, 1), state.last_launch_threadgroup[2]);
2039 try testing.expectEqual(@as(usize, 2), kernel.input_slot_ids.len);
2040 }
2041
2042 test "Choir executable plan rejects CUDA f16 dot_general until choir lowering supports it" {
2043 const allocator = testing.allocator;
2044
2045 var owned = try dotGeneralF16ChoirModule(allocator, "choir_cuda_executable_dot_general_f16");
2046 defer owned.deinit();
2047
2048 var cache = passes.AnalysisCache.init(allocator, null);
2049 defer cache.deinit();
2050 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2051 defer pass_ctx.deinit();
2052
2053 var state = RecordingBackendState{
2054 .allocator = allocator,
2055 .kind = .cuda,
2056 .format = .cuda_ptx,
2057 };
2058 const handle = state.handle();
2059
2060 try testing.expectError(
2061 error.CapabilityMismatch,
2062 createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{}),
2063 );
2064 }
2065
2066 test "Choir executable plan launches Vulkan dot_general without element-count bindings" {
2067 const allocator = testing.allocator;
2068
2069 var owned = try dotGeneralChoirModule(allocator, "choir_vulkan_executable_dot_general");
2070 defer owned.deinit();
2071
2072 var cache = passes.AnalysisCache.init(allocator, null);
2073 defer cache.deinit();
2074 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2075 defer pass_ctx.deinit();
2076
2077 var state = RecordingBackendState{
2078 .allocator = allocator,
2079 .kind = .vulkan,
2080 .format = .vulkan_spirv,
2081 };
2082 const handle = state.handle();
2083
2084 var artifact_plan = try createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{});
2085 defer artifact_plan.deinit();
2086
2087 var executable = try loadKernels(allocator, handle, &artifact_plan);
2088 defer executable.deinit();
2089
2090 try testing.expect(state.last_create_had_payload);
2091 try testing.expectEqual(@as(u32, 3), state.last_create_argument_count);
2092
2093 const slot_bindings = try slotBindingsForPlan(allocator, &artifact_plan, .vulkan);
2094 defer allocator.free(slot_bindings);
2095
2096 try executable.launchAll(allocator, &artifact_plan, slot_bindings, &.{});
2097
2098 try testing.expectEqual(@as(usize, 1), state.launch_count);
2099 try testing.expectEqual(@as(usize, 3), state.last_launch_buffer_count);
2100 try testing.expectEqual(@as(usize, 0), state.last_launch_scalar_count);
2101 try testing.expectEqual(@as(u32, 2), state.last_launch_grid[0]);
2102 try testing.expectEqual(@as(u32, 2), state.last_launch_grid[1]);
2103 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[2]);
2104 try testing.expectEqual(@as(u32, 8), state.last_launch_threadgroup[0]);
2105 try testing.expectEqual(@as(u32, 8), state.last_launch_threadgroup[1]);
2106 try testing.expectEqual(@as(u32, 1), state.last_launch_threadgroup[2]);
2107 }
2108
2109 test "Choir executable plan launches CUDA reduction with static input count" {
2110 const allocator = testing.allocator;
2111
2112 var owned = try reduceChoirModule(allocator, "choir_cuda_executable_reduce");
2113 defer owned.deinit();
2114
2115 var cache = passes.AnalysisCache.init(allocator, null);
2116 defer cache.deinit();
2117 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2118 defer pass_ctx.deinit();
2119
2120 var state = RecordingBackendState{
2121 .allocator = allocator,
2122 .kind = .cuda,
2123 .format = .cuda_ptx,
2124 };
2125 const handle = state.handle();
2126
2127 var artifact_plan = try createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{});
2128 defer artifact_plan.deinit();
2129
2130 var executable = try loadKernels(allocator, handle, &artifact_plan);
2131 defer executable.deinit();
2132
2133 const kernel = artifact_plan.kernels.items[0];
2134 const caps = try handle.queryCapabilities();
2135 const key = LaunchTuningCacheKey.init(caps, kernel);
2136 try testing.expectEqual(artifact_product.LaunchTileKind.reduction, key.tile_kind);
2137 try testing.expectEqual(@as(u32, 32), key.tile_m);
2138 try testing.expectEqual(@as(u32, 256), key.tile_n);
2139 try testing.expectEqual(@as(u32, 0), key.tile_k);
2140 try testing.expectEqual(@as(u32, 1), key.tile_batch);
2141 try testing.expect(key.tile_has_input_dtype);
2142 try testing.expectEqual(choir_abi.DType.f32, key.tile_input_dtype);
2143 try testing.expect(key.tile_has_output_dtype);
2144 try testing.expectEqual(choir_abi.DType.f32, key.tile_output_dtype);
2145 try testing.expectEqual(@as(u32, 32768), key.tile_input_tile_bytes);
2146 try testing.expectEqual(@as(u32, 128), key.tile_output_tile_bytes);
2147 try testing.expectEqual(@as(u32, 0), key.tile_scratch_memory_bytes);
2148 try testing.expectEqual(artifact_product.LaunchReductionKind.sum, key.tile_reduction_kind);
2149 try testing.expectEqual(@as(u32, 1), key.tile_reduction_rank);
2150 try testing.expectEqual(@as(u32, 0), key.tile_reduction_axis);
2151 try testing.expectEqual(@as(u32, 256), key.tile_reduction_extent);
2152 try testing.expect(key.launch_resource_fingerprint != 0);
2153 const slot_bindings = try slotBindingsForPlan(allocator, &artifact_plan, .cuda);
2154 defer allocator.free(slot_bindings);
2155 const count_bindings = try elementCountBindingsForPlan(allocator, &artifact_plan, .cuda);
2156 defer allocator.free(count_bindings);
2157
2158 try executable.launchAll(allocator, &artifact_plan, slot_bindings, count_bindings);
2159
2160 try testing.expectEqual(@as(usize, 1), executable.kernelCount());
2161 try testing.expectEqual(@as(usize, 1), state.load_count);
2162 try testing.expectEqual(@as(usize, 1), state.launch_count);
2163 try testing.expectEqual(@as(usize, 3), state.last_launch_buffer_count);
2164 try testing.expectEqual(@as(usize, 0), state.last_launch_scalar_count);
2165 try testing.expect(state.last_launch_scalar_u32 == null);
2166 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[0]);
2167 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[1]);
2168 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[2]);
2169 try testing.expectEqual(@as(u32, 32), state.last_launch_threadgroup[0]);
2170 try testing.expectEqual(@as(u32, 1), state.last_launch_threadgroup[1]);
2171 try testing.expectEqual(@as(u32, 1), state.last_launch_threadgroup[2]);
2172 try testing.expectEqual(@as(usize, 1), kernel.input_slot_ids.len);
2173 try testing.expectEqual(@as(gpu.BackendObjectId, 900), state.last_buffer_ids[2]);
2174 try testing.expectEqual(gpu.BufferAccess.read_only, state.last_buffer_access[2]);
2175 }
2176
2177 test "Choir executable plan creates CUDA plans from i32 reductions" {
2178 const allocator = testing.allocator;
2179
2180 var owned = try reduceI32ChoirModule(allocator, "choir_cuda_executable_reduce_i32");
2181 defer owned.deinit();
2182
2183 var cache = passes.AnalysisCache.init(allocator, null);
2184 defer cache.deinit();
2185 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2186 defer pass_ctx.deinit();
2187
2188 var state = RecordingBackendState{
2189 .allocator = allocator,
2190 .kind = .cuda,
2191 .format = .cuda_ptx,
2192 };
2193 const handle = state.handle();
2194
2195 var plan = try createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{});
2196 defer plan.deinit();
2197
2198 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
2199 try testing.expect(plan.kernels.items[0].compile.required_dtypes.contains(.i32));
2200 }
2201
2202 test "Choir executable plan launches Vulkan reduction with device count binding" {
2203 const allocator = testing.allocator;
2204
2205 var owned = try reduceChoirModule(allocator, "choir_vulkan_executable_reduce");
2206 defer owned.deinit();
2207
2208 var cache = passes.AnalysisCache.init(allocator, null);
2209 defer cache.deinit();
2210 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2211 defer pass_ctx.deinit();
2212
2213 var state = RecordingBackendState{
2214 .allocator = allocator,
2215 .kind = .vulkan,
2216 .format = .vulkan_spirv,
2217 };
2218 const handle = state.handle();
2219
2220 var artifact_plan = try createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{});
2221 defer artifact_plan.deinit();
2222
2223 var executable = try loadKernels(allocator, handle, &artifact_plan);
2224 defer executable.deinit();
2225
2226 const kernel = artifact_plan.kernels.items[0];
2227 const caps = try handle.queryCapabilities();
2228 const key = LaunchTuningCacheKey.init(caps, kernel);
2229 try testing.expectEqual(artifact_product.LaunchTileKind.reduction, key.tile_kind);
2230 try testing.expectEqual(@as(u32, 64), key.tile_m);
2231 try testing.expectEqual(@as(u32, 256), key.tile_n);
2232 try testing.expectEqual(@as(u32, 0), key.tile_k);
2233 try testing.expectEqual(@as(u32, 1), key.tile_batch);
2234 try testing.expect(key.tile_has_input_dtype);
2235 try testing.expectEqual(choir_abi.DType.f32, key.tile_input_dtype);
2236 try testing.expect(key.tile_has_output_dtype);
2237 try testing.expectEqual(choir_abi.DType.f32, key.tile_output_dtype);
2238 try testing.expectEqual(@as(u32, 65536), key.tile_input_tile_bytes);
2239 try testing.expectEqual(@as(u32, 256), key.tile_output_tile_bytes);
2240 try testing.expectEqual(@as(u32, 0), key.tile_scratch_memory_bytes);
2241 try testing.expectEqual(artifact_product.LaunchReductionKind.sum, key.tile_reduction_kind);
2242 try testing.expectEqual(@as(u32, 1), key.tile_reduction_rank);
2243 try testing.expectEqual(@as(u32, 0), key.tile_reduction_axis);
2244 try testing.expectEqual(@as(u32, 256), key.tile_reduction_extent);
2245 try testing.expect(key.launch_resource_fingerprint != 0);
2246 try testing.expect(state.last_create_had_payload);
2247 try testing.expectEqual(@as(u32, 3), state.last_create_argument_count);
2248
2249 const slot_bindings = try slotBindingsForPlan(allocator, &artifact_plan, .vulkan);
2250 defer allocator.free(slot_bindings);
2251 const count_binding = ElementCountBufferBinding{
2252 .kernel_id = kernel.kernel_id,
2253 .binding = bufferBinding(200, .vulkan, 4),
2254 };
2255
2256 try executable.launchAll(allocator, &artifact_plan, slot_bindings, &.{count_binding});
2257
2258 try testing.expectEqual(@as(usize, 1), executable.kernelCount());
2259 try testing.expectEqual(@as(usize, 1), state.load_count);
2260 try testing.expectEqual(@as(usize, 1), state.launch_count);
2261 try testing.expectEqual(@as(usize, 3), state.last_launch_buffer_count);
2262 try testing.expectEqual(@as(usize, 0), state.last_launch_scalar_count);
2263 try testing.expect(state.last_launch_scalar_u32 == null);
2264 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[0]);
2265 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[1]);
2266 try testing.expectEqual(@as(u32, 1), state.last_launch_grid[2]);
2267 try testing.expectEqual(@as(u32, 64), state.last_launch_threadgroup[0]);
2268 try testing.expectEqual(@as(u32, 1), state.last_launch_threadgroup[1]);
2269 try testing.expectEqual(@as(u32, 1), state.last_launch_threadgroup[2]);
2270 try testing.expectEqual(@as(usize, 1), kernel.input_slot_ids.len);
2271 try testing.expectEqual(@as(gpu.BackendObjectId, 200), state.last_buffer_ids[2]);
2272 try testing.expectEqual(gpu.BufferAccess.read_only, state.last_buffer_access[2]);
2273 }
2274
2275 test "Choir executable launch candidate measurement synchronizes stream" {
2276 const allocator = testing.allocator;
2277
2278 var owned = try addChoirModule(allocator, "choir_cuda_executable_measure_stream_sync");
2279 defer owned.deinit();
2280
2281 var cache = passes.AnalysisCache.init(allocator, null);
2282 defer cache.deinit();
2283 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2284 defer pass_ctx.deinit();
2285
2286 var state = RecordingBackendState{
2287 .allocator = allocator,
2288 .kind = .cuda,
2289 .format = .cuda_ptx,
2290 };
2291 const handle = state.handle();
2292
2293 var artifact_plan = try createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{});
2294 defer artifact_plan.deinit();
2295
2296 var executable = try loadKernels(allocator, handle, &artifact_plan);
2297 defer executable.deinit();
2298
2299 const kernel = artifact_plan.kernels.items[0];
2300 const slot_bindings = try slotBindingsForKernel(allocator, kernel, .cuda);
2301 defer allocator.free(slot_bindings);
2302 const count_bindings = try elementCountBindingsForPlan(allocator, &artifact_plan, .cuda);
2303 defer allocator.free(count_bindings);
2304
2305 const measurements = try executable.measureLaunchCandidates(
2306 allocator,
2307 allocator,
2308 &artifact_plan,
2309 0,
2310 slot_bindings,
2311 firstElementCountBinding(count_bindings),
2312 .{
2313 .warmup = 1,
2314 .samples = 2,
2315 .base_options = .{ .stream = .{ .id = 77, .backend = .cuda } },
2316 .synchronize = .stream,
2317 },
2318 );
2319 defer allocator.free(measurements);
2320
2321 try testing.expectEqual(kernel.launch_resources.candidate_count * 3, state.launch_count);
2322 try testing.expectEqual(kernel.launch_resources.candidate_count * 3, state.sync_count);
2323 try testing.expectEqual(gpu.SyncScope.stream, state.last_sync_scope.?);
2324 try testing.expectEqual(@as(gpu.BackendObjectId, 77), state.last_sync_stream.?);
2325 }
2326
2327 test "Choir executable fragment creates artifacts loads handles and launches" {
2328 const allocator = testing.allocator;
2329 const PhaseRecorder = struct {
2330 phases: [3]FragmentPhase = undefined,
2331 count: usize = 0,
2332
2333 fn observe(context: ?*anyopaque, phase_name: []const u8, _: u64) anyerror!void {
2334 const self: *@This() = @ptrCast(@alignCast(context.?));
2335 if (self.count >= self.phases.len) return error.TooManyPhases;
2336 self.phases[self.count] = if (std.mem.eql(u8, phase_name, FragmentPhase.run_contract_pipeline.name()))
2337 .run_contract_pipeline
2338 else if (std.mem.eql(u8, phase_name, FragmentPhase.run_target_pipeline.name()))
2339 .run_target_pipeline
2340 else if (std.mem.eql(u8, phase_name, FragmentPhase.plan_create_backend_artifacts.name()))
2341 .plan_create_backend_artifacts
2342 else if (std.mem.eql(u8, phase_name, FragmentPhase.compile_fragment.name()))
2343 .compile_fragment
2344 else if (std.mem.eql(u8, phase_name, FragmentPhase.load_backend_artifacts.name()))
2345 .load_backend_artifacts
2346 else if (std.mem.eql(u8, phase_name, FragmentPhase.import_launch_tuning.name()))
2347 .import_launch_tuning
2348 else
2349 return error.InvalidPhase;
2350 self.count += 1;
2351 }
2352
2353 fn has(self: *const @This(), phase: FragmentPhase) bool {
2354 for (self.phases[0..self.count]) |recorded| {
2355 if (recorded == phase) return true;
2356 }
2357 return false;
2358 }
2359 };
2360
2361 var owned = try addChoirModule(allocator, "choir_loaded_fragment_add");
2362 defer owned.deinit();
2363
2364 var cache = passes.AnalysisCache.init(allocator, null);
2365 defer cache.deinit();
2366 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2367 defer pass_ctx.deinit();
2368
2369 var state = RecordingBackendState{
2370 .allocator = allocator,
2371 .kind = .cuda,
2372 .format = .cuda_ptx,
2373 };
2374 const handle = state.handle();
2375
2376 var phase_recorder = PhaseRecorder{};
2377 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{
2378 .instrumentation = .{
2379 .context = &phase_recorder,
2380 .observe = PhaseRecorder.observe,
2381 },
2382 });
2383 defer fragment.deinit();
2384
2385 const kernel = fragment.artifactPlan().kernels.items[0];
2386 const slot_bindings = try slotBindingsForKernel(allocator, kernel, .cuda);
2387 defer allocator.free(slot_bindings);
2388 const count_bindings = try elementCountBindingsForPlan(allocator, fragment.artifactPlan(), .cuda);
2389 defer allocator.free(count_bindings);
2390
2391 try fragment.launchAll(allocator, slot_bindings, count_bindings);
2392
2393 try testing.expectEqual(@as(usize, 1), fragment.artifactPlan().kernelCount());
2394 try testing.expectEqual(@as(usize, 1), fragment.loadedKernelCount());
2395 try testing.expectEqual(@as(usize, 1), state.load_count);
2396 try testing.expectEqual(@as(usize, 1), state.launch_count);
2397 try testing.expect(state.last_launch_scalar_u32 == null);
2398 try testing.expectEqual(@as(usize, 3), phase_recorder.count);
2399 try testing.expectEqual(FragmentPhase.plan_create_backend_artifacts, phase_recorder.phases[0]);
2400 try testing.expectEqual(FragmentPhase.compile_fragment, phase_recorder.phases[1]);
2401 try testing.expectEqual(FragmentPhase.load_backend_artifacts, phase_recorder.phases[2]);
2402 try testing.expect(phase_recorder.has(.plan_create_backend_artifacts));
2403 try testing.expect(phase_recorder.has(.compile_fragment));
2404 try testing.expect(phase_recorder.has(.load_backend_artifacts));
2405 }
2406
2407 test "Choir executable fragment from semantic module uses target product" {
2408 const allocator = testing.allocator;
2409 const PhaseRecorder = struct {
2410 phases: [9]FragmentPhase = undefined,
2411 count: usize = 0,
2412
2413 fn observe(context: ?*anyopaque, phase_name: []const u8, _: u64) anyerror!void {
2414 const self: *@This() = @ptrCast(@alignCast(context.?));
2415 if (self.count >= self.phases.len) return error.TooManyPhases;
2416 self.phases[self.count] = if (std.mem.eql(u8, phase_name, FragmentPhase.run_contract_pipeline.name()))
2417 .run_contract_pipeline
2418 else if (std.mem.eql(u8, phase_name, FragmentPhase.run_tensor_pipeline.name()))
2419 .run_tensor_pipeline
2420 else if (std.mem.eql(u8, phase_name, FragmentPhase.run_dispatch_pipeline.name()))
2421 .run_dispatch_pipeline
2422 else if (std.mem.eql(u8, phase_name, FragmentPhase.run_memory_pipeline.name()))
2423 .run_memory_pipeline
2424 else if (std.mem.eql(u8, phase_name, FragmentPhase.run_kernel_pipeline.name()))
2425 .run_kernel_pipeline
2426 else if (std.mem.eql(u8, phase_name, FragmentPhase.run_target_pipeline.name()))
2427 .run_target_pipeline
2428 else if (std.mem.eql(u8, phase_name, FragmentPhase.plan_create_backend_artifacts.name()))
2429 .plan_create_backend_artifacts
2430 else if (std.mem.eql(u8, phase_name, FragmentPhase.compile_fragment.name()))
2431 .compile_fragment
2432 else if (std.mem.eql(u8, phase_name, FragmentPhase.load_backend_artifacts.name()))
2433 .load_backend_artifacts
2434 else if (std.mem.eql(u8, phase_name, FragmentPhase.import_launch_tuning.name()))
2435 .import_launch_tuning
2436 else
2437 return error.InvalidPhase;
2438 self.count += 1;
2439 }
2440
2441 fn has(self: *const @This(), phase: FragmentPhase) bool {
2442 for (self.phases[0..self.count]) |recorded| {
2443 if (recorded == phase) return true;
2444 }
2445 return false;
2446 }
2447 };
2448
2449 const module = try addSemanticModule(allocator, "choir_executable_from_semantic_add");
2450
2451 var state = RecordingBackendState{
2452 .allocator = allocator,
2453 .kind = .cuda,
2454 .format = .cuda_ptx,
2455 };
2456 const handle = state.handle();
2457
2458 var phase_recorder = PhaseRecorder{};
2459 var fragment = try compileAndLoadTestSemanticModule(allocator, handle, module, .{
2460 .instrumentation = .{
2461 .context = &phase_recorder,
2462 .observe = PhaseRecorder.observe,
2463 },
2464 });
2465 defer fragment.deinit();
2466
2467 try testing.expectEqual(@as(usize, 1), fragment.artifactPlan().kernelCount());
2468 try testing.expectEqual(@as(usize, 1), fragment.loadedKernelCount());
2469 try testing.expectEqual(@as(usize, 1), state.load_count);
2470 try testing.expectEqual(@as(usize, 9), phase_recorder.count);
2471 try testing.expectEqual(FragmentPhase.run_contract_pipeline, phase_recorder.phases[0]);
2472 try testing.expectEqual(FragmentPhase.run_tensor_pipeline, phase_recorder.phases[1]);
2473 try testing.expectEqual(FragmentPhase.run_dispatch_pipeline, phase_recorder.phases[2]);
2474 try testing.expectEqual(FragmentPhase.run_memory_pipeline, phase_recorder.phases[3]);
2475 try testing.expectEqual(FragmentPhase.run_kernel_pipeline, phase_recorder.phases[4]);
2476 try testing.expectEqual(FragmentPhase.run_target_pipeline, phase_recorder.phases[5]);
2477 try testing.expectEqual(FragmentPhase.plan_create_backend_artifacts, phase_recorder.phases[6]);
2478 try testing.expectEqual(FragmentPhase.compile_fragment, phase_recorder.phases[7]);
2479 try testing.expectEqual(FragmentPhase.load_backend_artifacts, phase_recorder.phases[8]);
2480 try testing.expect(phase_recorder.has(.run_contract_pipeline));
2481 try testing.expect(phase_recorder.has(.run_tensor_pipeline));
2482 try testing.expect(phase_recorder.has(.run_dispatch_pipeline));
2483 try testing.expect(phase_recorder.has(.run_memory_pipeline));
2484 try testing.expect(phase_recorder.has(.run_kernel_pipeline));
2485 try testing.expect(phase_recorder.has(.run_target_pipeline));
2486 try testing.expect(phase_recorder.has(.plan_create_backend_artifacts));
2487 try testing.expect(phase_recorder.has(.compile_fragment));
2488 try testing.expect(phase_recorder.has(.load_backend_artifacts));
2489 }
2490
2491 test "Choir executable fragment compiles a prepared job without rerunning preparation" {
2492 const allocator = testing.allocator;
2493 const PhaseRecorder = struct {
2494 phases: [3]FragmentPhase = undefined,
2495 count: usize = 0,
2496
2497 fn observe(context: ?*anyopaque, phase_name: []const u8, _: u64) anyerror!void {
2498 const self: *@This() = @ptrCast(@alignCast(context.?));
2499 if (self.count >= self.phases.len) return error.TooManyPhases;
2500 self.phases[self.count] = if (std.mem.eql(u8, phase_name, FragmentPhase.plan_create_backend_artifacts.name()))
2501 .plan_create_backend_artifacts
2502 else if (std.mem.eql(u8, phase_name, FragmentPhase.compile_fragment.name()))
2503 .compile_fragment
2504 else if (std.mem.eql(u8, phase_name, FragmentPhase.load_backend_artifacts.name()))
2505 .load_backend_artifacts
2506 else if (std.mem.eql(u8, phase_name, FragmentPhase.import_launch_tuning.name()))
2507 .import_launch_tuning
2508 else
2509 return error.InvalidPhase;
2510 self.count += 1;
2511 }
2512
2513 fn has(self: *const @This(), phase: FragmentPhase) bool {
2514 for (self.phases[0..self.count]) |recorded| {
2515 if (recorded == phase) return true;
2516 }
2517 return false;
2518 }
2519 };
2520
2521 var state = RecordingBackendState{
2522 .allocator = allocator,
2523 .kind = .cuda,
2524 .format = .cuda_ptx,
2525 };
2526 const handle = state.handle();
2527
2528 const target_profile = try backendTargetProfile(handle, .cuda_ptx, .exact);
2529 const module = try addSemanticModule(allocator, "choir_executable_from_prepared_add");
2530 var prepared = try preparation.prepareBackendJobFromSemanticModule(allocator, module, .{
2531 .target_profile = target_profile,
2532 });
2533 defer prepared.deinit();
2534
2535 var phase_recorder = PhaseRecorder{};
2536 var fragment = try compileAndLoadTestPreparedJob(allocator, handle, &prepared, .{
2537 .artifact_format = .cuda_ptx,
2538 .instrumentation = .{
2539 .context = &phase_recorder,
2540 .observe = PhaseRecorder.observe,
2541 },
2542 });
2543 defer fragment.deinit();
2544
2545 try testing.expectEqual(@as(usize, 1), fragment.artifactPlan().kernelCount());
2546 try testing.expectEqual(@as(usize, 1), fragment.loadedKernelCount());
2547 try testing.expectEqual(@as(usize, 1), state.load_count);
2548 try testing.expectEqual(@as(u64, preparation.target_pass_count), prepared.run.target_stats.pass_runs);
2549 try testing.expectEqual(@as(usize, 1), try prepared.generatedKernelCount());
2550 const generated_summary = try prepared.generatedKernelSummary(0);
2551 try testing.expectEqualStrings("accy_choir_elementwise_add_0_v4", generated_summary.entry_name);
2552 try testing.expectEqualStrings("accy.artifact", artifact_product.product_name);
2553 try testing.expectEqual(@as(usize, 3), phase_recorder.count);
2554 try testing.expectEqual(FragmentPhase.plan_create_backend_artifacts, phase_recorder.phases[0]);
2555 try testing.expectEqual(FragmentPhase.compile_fragment, phase_recorder.phases[1]);
2556 try testing.expectEqual(FragmentPhase.load_backend_artifacts, phase_recorder.phases[2]);
2557 try testing.expect(phase_recorder.has(.plan_create_backend_artifacts));
2558 try testing.expect(phase_recorder.has(.compile_fragment));
2559 try testing.expect(phase_recorder.has(.load_backend_artifacts));
2560 }
2561
2562 test "Choir executable fragment prepareInvocationBindings allocates buffers and count bindings" {
2563 const allocator = testing.allocator;
2564
2565 var owned = try addChoirModule(allocator, "choir_loaded_fragment_prepared_bindings");
2566 defer owned.deinit();
2567
2568 var cache = passes.AnalysisCache.init(allocator, null);
2569 defer cache.deinit();
2570 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2571 defer pass_ctx.deinit();
2572
2573 var state = RecordingBackendState{
2574 .allocator = allocator,
2575 .kind = .cuda,
2576 .format = .cuda_ptx,
2577 };
2578 const handle = state.handle();
2579
2580 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
2581 defer fragment.deinit();
2582
2583 const lhs = @as([8]f32, @splat(1.0));
2584 const rhs = @as([8]f32, @splat(2.0));
2585 const bindings = try fragment.prepareInvocationBindings(allocator, &.{
2586 std.mem.asBytes(&lhs),
2587 std.mem.asBytes(&rhs),
2588 });
2589 var bindings_live = true;
2590 defer if (bindings_live) bindings.deinit();
2591
2592 const kernel = fragment.artifactPlan().kernels.items[0];
2593 const live_buffer_count = binding_mod.liveBufferCount(bindings);
2594 try testing.expectEqual(@as(usize, 4), live_buffer_count);
2595 try testing.expectEqual(@as(usize, 3), binding_mod.slotBindings(bindings).len);
2596 try testing.expectEqual(@as(usize, 1), binding_mod.elementCountBindings(bindings).len);
2597 try testing.expect(binding_mod.elementCountBindingForKernel(bindings, kernel.kernel_id) != null);
2598 const count_binding = binding_mod.elementCountBindingForKernel(bindings, kernel.kernel_id).?;
2599 try testing.expectEqual(@as(usize, @sizeOf(u32)), count_binding.byte_size);
2600 try testing.expectEqual(@as(usize, 4), state.buffer_allocate_count);
2601 try testing.expectEqual(@as(usize, 3), state.write_count);
2602 try testing.expectEqual(@as(usize, @sizeOf(u32)), state.last_write_byte_count);
2603 try testing.expectEqual(@as(u32, 8), state.last_write_u32.?);
2604
2605 try fragment.submitPreparedForTest(allocator, bindings);
2606 try testing.expectEqual(@as(usize, 1), state.launch_count);
2607 try testing.expectEqual(@as(usize, 4), state.last_launch_buffer_count);
2608 try testing.expectEqual(@as(usize, 0), state.destroy_count);
2609
2610 bindings.deinit();
2611 bindings_live = false;
2612 try testing.expectEqual(live_buffer_count, state.destroy_count);
2613 }
2614
2615 test "Choir executable fragment prepareInvocationBindings skips reduction init constants" {
2616 const allocator = testing.allocator;
2617
2618 var owned = try reduceChoirModule(allocator, "choir_loaded_fragment_prepared_reduction_bindings");
2619 defer owned.deinit();
2620
2621 var cache = passes.AnalysisCache.init(allocator, null);
2622 defer cache.deinit();
2623 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2624 defer pass_ctx.deinit();
2625
2626 var state = RecordingBackendState{
2627 .allocator = allocator,
2628 .kind = .cuda,
2629 .format = .cuda_ptx,
2630 };
2631 const handle = state.handle();
2632
2633 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
2634 defer fragment.deinit();
2635
2636 const input = @as([256]f32, @splat(1.0));
2637 const bindings = try fragment.prepareInvocationBindings(allocator, &.{std.mem.asBytes(&input)});
2638 defer bindings.deinit();
2639
2640 const plan = fragment.artifactPlan();
2641 var constant_slot_count: usize = 0;
2642 for (plan.slots) |slot| {
2643 if (slot.role.constant) constant_slot_count += 1;
2644 }
2645
2646 try testing.expectEqual(@as(usize, 1), constant_slot_count);
2647 try testing.expectEqual(@as(usize, 3), binding_mod.liveBufferCount(bindings));
2648 try testing.expectEqual(@as(usize, 2), binding_mod.slotBindings(bindings).len);
2649 try testing.expectEqual(@as(usize, 1), binding_mod.elementCountBindings(bindings).len);
2650 for (binding_mod.slotBindings(bindings)) |slot_binding| {
2651 const slot = plan.slotById(slot_binding.slot_id) orelse return error.InvalidArtifact;
2652 try testing.expect(!slot.role.constant);
2653 }
2654
2655 const kernel = plan.kernels.items[0];
2656 try testing.expect(binding_mod.elementCountBindingForKernel(bindings, kernel.kernel_id) != null);
2657 try testing.expectEqual(@as(usize, 3), state.buffer_allocate_count);
2658 try testing.expectEqual(@as(usize, 2), state.write_count);
2659 try testing.expectEqual(@as(usize, @sizeOf(u32)), state.last_write_byte_count);
2660 try testing.expectEqual(@as(u32, 1), state.last_write_u32.?);
2661
2662 try fragment.submitPreparedForTest(allocator, bindings);
2663 try testing.expectEqual(@as(usize, 1), state.launch_count);
2664 try testing.expectEqual(@as(usize, 3), state.last_launch_buffer_count);
2665 }
2666
2667 test "Choir executable fragment prepareInvocationBindings writes kernel constant buffers" {
2668 const allocator = testing.allocator;
2669
2670 var owned = try constantAddChoirModule(allocator, "choir_loaded_fragment_constant_bindings");
2671 defer owned.deinit();
2672
2673 var cache = passes.AnalysisCache.init(allocator, null);
2674 defer cache.deinit();
2675 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2676 defer pass_ctx.deinit();
2677
2678 var state = RecordingBackendState{
2679 .allocator = allocator,
2680 .kind = .cuda,
2681 .format = .cuda_ptx,
2682 };
2683 const handle = state.handle();
2684
2685 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
2686 defer fragment.deinit();
2687
2688 const input = @as([8]f32, @splat(1.0));
2689 const bindings = try fragment.prepareInvocationBindings(allocator, &.{std.mem.asBytes(&input)});
2690 defer bindings.deinit();
2691
2692 const plan = fragment.artifactPlan();
2693 var constant_slot: ?artifact_product.PlannedSlot = null;
2694 for (plan.slots) |slot| {
2695 if (slot.role.constant) constant_slot = slot;
2696 }
2697 const constant = constant_slot orelse return error.TestExpectedConstant;
2698
2699 try testing.expectEqual(@as(usize, 4), binding_mod.liveBufferCount(bindings));
2700 try testing.expectEqual(@as(usize, 3), binding_mod.slotBindings(bindings).len);
2701 try testing.expectEqual(@as(usize, 1), binding_mod.elementCountBindings(bindings).len);
2702 try testing.expectEqual(@as(usize, 3), state.write_count);
2703
2704 var found_constant_binding = false;
2705 for (binding_mod.slotBindings(bindings)) |slot_binding| {
2706 if (slot_binding.slot_id != constant.slot_id) continue;
2707 found_constant_binding = true;
2708 try testing.expectEqual(gpu.BufferAccess.read_only, slot_binding.binding.access);
2709 try testing.expectEqual(@as(usize, 8 * @sizeOf(f32)), slot_binding.binding.byte_size);
2710 }
2711 try testing.expect(found_constant_binding);
2712
2713 try fragment.submitPreparedForTest(allocator, bindings);
2714 try testing.expectEqual(@as(usize, 1), state.launch_count);
2715 try testing.expectEqual(@as(usize, 4), state.last_launch_buffer_count);
2716 }
2717
2718 test "Choir executable fragment reads program outputs back to host" {
2719 const allocator = testing.allocator;
2720
2721 var owned = try addChoirModule(allocator, "choir_loaded_fragment_read_outputs");
2722 defer owned.deinit();
2723
2724 var cache = passes.AnalysisCache.init(allocator, null);
2725 defer cache.deinit();
2726 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2727 defer pass_ctx.deinit();
2728
2729 var state = RecordingBackendState{
2730 .allocator = allocator,
2731 .kind = .cuda,
2732 .format = .cuda_ptx,
2733 };
2734 const handle = state.handle();
2735
2736 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
2737 defer fragment.deinit();
2738
2739 try testing.expectEqual(@as(usize, 1), fragment.outputCount());
2740 const slot = fragment.outputSlot(0).?;
2741 try testing.expectEqual(@as(?u64, 8 * @sizeOf(f32)), slot.byte_size);
2742 try testing.expect(fragment.outputSlot(1) == null);
2743
2744 const lhs = @as([8]f32, @splat(1.0));
2745 const rhs = @as([8]f32, @splat(2.0));
2746 const bindings = try fragment.prepareInvocationBindings(allocator, &.{
2747 std.mem.asBytes(&lhs),
2748 std.mem.asBytes(&rhs),
2749 });
2750 defer bindings.deinit();
2751
2752 try fragment.submitPreparedForTest(allocator, bindings);
2753 try fragment.completeInvocationWithOptions(.{});
2754
2755 const output_slot_id = fragment.artifactPlan().output_slot_ids[0];
2756 const expected = try binding_mod.bindingForSlot(binding_mod.slotBindings(bindings), output_slot_id, .read_only);
2757
2758 var host_out = @as([8]f32, @splat(7.0));
2759 try fragment.readInvocationOutput(bindings, 0, std.mem.sliceAsBytes(host_out[0..]));
2760
2761 try testing.expectEqual(@as(usize, 1), state.read_count);
2762 try testing.expectEqual(@as(usize, 8 * @sizeOf(f32)), state.last_read_byte_count);
2763 try testing.expectEqual(expected.handle.id, state.last_read_buffer_id.?);
2764 try testing.expect(state.sync_count >= 1);
2765 try testing.expectEqual(gpu.SyncScope.default_stream, state.last_sync_scope.?);
2766 for (host_out) |value| try testing.expectEqual(@as(f32, 0.0), value);
2767
2768 try testing.expectError(
2769 error.InvalidArtifact,
2770 fragment.readInvocationOutput(bindings, 1, std.mem.sliceAsBytes(host_out[0..])),
2771 );
2772 var small: [4]u8 = undefined;
2773 try testing.expectError(error.ReadBufferDestinationTooSmall, fragment.readInvocationOutput(bindings, 0, small[0..]));
2774 }
2775
2776 test "Choir executable fragment launches inputs and reads outputs" {
2777 const allocator = testing.allocator;
2778
2779 var owned = try addChoirModule(allocator, "choir_loaded_fragment_launch_read");
2780 defer owned.deinit();
2781
2782 var cache = passes.AnalysisCache.init(allocator, null);
2783 defer cache.deinit();
2784 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2785 defer pass_ctx.deinit();
2786
2787 var state = RecordingBackendState{
2788 .allocator = allocator,
2789 .kind = .cuda,
2790 .format = .cuda_ptx,
2791 };
2792 const handle = state.handle();
2793
2794 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
2795 defer fragment.deinit();
2796
2797 const lhs = @as([8]f32, @splat(1.0));
2798 const rhs = @as([8]f32, @splat(2.0));
2799 var host_out = @as([8]f32, @splat(9.0));
2800 var outputs = [_][]u8{std.mem.sliceAsBytes(host_out[0..])};
2801
2802 try fragment.invokeAndReadForTest(allocator, allocator, &.{
2803 std.mem.asBytes(&lhs),
2804 std.mem.asBytes(&rhs),
2805 }, &outputs);
2806
2807 try testing.expectEqual(@as(usize, 1), state.launch_count);
2808 try testing.expectEqual(@as(usize, 1), state.read_count);
2809 try testing.expectEqual(@as(usize, 8 * @sizeOf(f32)), state.last_read_byte_count);
2810 try testing.expect(state.sync_count >= 1);
2811 for (host_out) |value| try testing.expectEqual(@as(f32, 0.0), value);
2812
2813 try testing.expectError(error.InvalidArtifact, fragment.invokeAndReadForTest(allocator, allocator, &.{
2814 std.mem.asBytes(&lhs),
2815 std.mem.asBytes(&rhs),
2816 }, &.{}));
2817 }
2818
2819 test "Choir executable fragment exposes kernel summaries by index and work item" {
2820 const allocator = testing.allocator;
2821
2822 var owned = try addChoirModule(allocator, "choir_loaded_fragment_kernel_summary_lookup");
2823 defer owned.deinit();
2824
2825 var cache = passes.AnalysisCache.init(allocator, null);
2826 defer cache.deinit();
2827 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2828 defer pass_ctx.deinit();
2829
2830 var state = RecordingBackendState{
2831 .allocator = allocator,
2832 .kind = .cuda,
2833 .format = .cuda_ptx,
2834 };
2835 const handle = state.handle();
2836
2837 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
2838 defer fragment.deinit();
2839
2840 try testing.expectEqual(@as(usize, 1), fragment.kernelCount());
2841 const indexed = try fragment.kernelSummary(0);
2842 const by_work = try fragment.kernelSummaryForWork(indexed.work_item_id);
2843 try testing.expect(artifact_product.kernelSummariesEqual(indexed, by_work));
2844 try testing.expectError(error.InvalidArtifact, fragment.kernelSummary(1));
2845 try testing.expectError(error.InvalidArtifact, fragment.kernelSummaryForWork(std.math.maxInt(usize)));
2846
2847 var summaries = try fragment.copyKernelSummaries(allocator);
2848 defer summaries.deinit();
2849 try testing.expectEqual(fragment.kernelCount(), summaries.len());
2850 const copied_indexed = try summaries.summary(0);
2851 const copied_by_work = try summaries.summaryForWork(indexed.work_item_id);
2852 try testing.expect(artifact_product.kernelSummariesEqual(indexed, copied_indexed));
2853 try testing.expect(artifact_product.kernelSummariesEqual(indexed, copied_by_work));
2854 try testing.expectEqualStrings(indexed.entry_name, copied_indexed.entry_name);
2855 try testing.expectError(error.InvalidArtifact, summaries.summary(fragment.kernelCount()));
2856 try testing.expectError(error.InvalidArtifact, summaries.summaryForWork(std.math.maxInt(usize)));
2857 }
2858
2859 test "Choir executable fragment measures launch candidates through prepared bindings" {
2860 const allocator = testing.allocator;
2861
2862 var owned = try addChoirModule(allocator, "choir_loaded_fragment_measure_candidates");
2863 defer owned.deinit();
2864
2865 var cache = passes.AnalysisCache.init(allocator, null);
2866 defer cache.deinit();
2867 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2868 defer pass_ctx.deinit();
2869
2870 var state = RecordingBackendState{
2871 .allocator = allocator,
2872 .kind = .cuda,
2873 .format = .cuda_ptx,
2874 };
2875 const handle = state.handle();
2876
2877 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
2878 defer fragment.deinit();
2879
2880 const lhs = @as([8]f32, @splat(1.0));
2881 const rhs = @as([8]f32, @splat(2.0));
2882 const bindings = try fragment.prepareInvocationBindings(allocator, &.{
2883 std.mem.asBytes(&lhs),
2884 std.mem.asBytes(&rhs),
2885 });
2886 defer bindings.deinit();
2887
2888 try testing.expectEqual(@as(usize, 1), fragment.kernelCount());
2889 const kernel = try fragment.plannedKernel(0);
2890 try testing.expect(kernel.launch_resources.candidate_count > 1);
2891 const summary = try fragment.kernelSummary(0);
2892 try testing.expectEqual(artifact_product.KernelSource.tensor, summary.source);
2893 try testing.expectEqual(kernel.kernel_id, summary.kernel_id);
2894 try testing.expectEqual(kernel.work_item_id, summary.work_item_id);
2895 try testing.expectEqual(kernel.element_count, summary.element_count);
2896 try testing.expectEqual(kernel.op_count, summary.op_count);
2897 try testing.expectEqual(kernel.output_layout_fingerprint, summary.output_layout_fingerprint);
2898 try testing.expectEqual(kernel.input_layout_fingerprint, summary.input_layout_fingerprint);
2899 try testing.expectEqual(kernel.launch_resources.candidate_count, summary.launch_candidate_count);
2900 try testing.expectEqualStrings(launchResourceClassName(kernel.launch_resources.resource_class), summary.launch_resource_class);
2901 var copied_artifact = try fragment.copyKernelArtifact(allocator, 0);
2902 defer copied_artifact.deinit();
2903 try testing.expectEqual(kernel.artifact.format, copied_artifact.format);
2904 try testing.expectEqualStrings(kernel.artifact.entry_name, copied_artifact.entry_name);
2905
2906 const records = try fragment.measureInvocationLaunchCandidates(
2907 allocator,
2908 allocator,
2909 0,
2910 bindings,
2911 .{ .warmup = 1, .samples = 2 },
2912 );
2913 defer allocator.free(records);
2914 try testing.expectEqual(kernel.launch_resources.candidate_count, records.len);
2915 try testing.expectEqual(kernel.launch_resources.candidate_count * 3, state.launch_count);
2916 try testing.expectEqual(kernel.kernel_id, records[0].kernel.kernel_id);
2917 try testing.expectEqual(kernel.output_layout_fingerprint, records[0].kernel.output_layout_fingerprint);
2918 try testing.expectEqual(kernel.input_layout_fingerprint, records[0].kernel.input_layout_fingerprint);
2919 try testing.expectEqual(@as(usize, 0), records[0].candidate_index);
2920 try testing.expectEqual(kernel.launch_resources.candidates[0].geometry.threadgroup[0], records[0].geometry.threadgroup[0]);
2921 try testing.expectEqual(@as(u32, 2), records[0].sample_count);
2922
2923 const candidate = try fragment.launchResourceCandidate(0, records[0].candidate_index);
2924 try testing.expectEqual(kernel.launch_resources.candidates[0].geometry.threadgroup[0], candidate.geometry.threadgroup[0]);
2925 try testing.expectError(error.InvalidArtifact, fragment.plannedKernel(fragment.kernelCount()));
2926 try testing.expectError(
2927 error.LaunchArgumentMismatch,
2928 fragment.launchResourceCandidate(0, kernel.launch_resources.candidate_count),
2929 );
2930 }
2931
2932 test "Choir executable fragment measures and records launch candidates" {
2933 const allocator = testing.allocator;
2934
2935 var owned = try addChoirModule(allocator, "choir_loaded_fragment_measure_record_candidates");
2936 defer owned.deinit();
2937
2938 var cache = passes.AnalysisCache.init(allocator, null);
2939 defer cache.deinit();
2940 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
2941 defer pass_ctx.deinit();
2942
2943 var state = RecordingBackendState{
2944 .allocator = allocator,
2945 .kind = .cuda,
2946 .format = .cuda_ptx,
2947 };
2948 const handle = state.handle();
2949
2950 var measured = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
2951 defer measured.deinit();
2952
2953 const lhs = @as([8]f32, @splat(1.0));
2954 const rhs = @as([8]f32, @splat(2.0));
2955 const bindings = try measured.prepareInvocationBindings(allocator, &.{
2956 std.mem.asBytes(&lhs),
2957 std.mem.asBytes(&rhs),
2958 });
2959 defer bindings.deinit();
2960
2961 const measured_kernel = measured.artifactPlan().kernels.items[0];
2962 try testing.expect(measured_kernel.launch_resources.candidate_count > 1);
2963
2964 const records = try measured.measureAndRecordInvocationLaunchCandidates(
2965 allocator,
2966 allocator,
2967 bindings,
2968 .{ .warmup = 0, .samples = 1 },
2969 );
2970 defer allocator.free(records);
2971
2972 try testing.expectEqual(measured_kernel.launch_resources.candidate_count, records.len);
2973 try testing.expectEqual(measured_kernel.launch_resources.candidate_count, state.launch_count);
2974 try testing.expectEqual(@as(usize, 1), measured.tuningRecordCount());
2975 for (records, 0..) |record, index| {
2976 try testing.expectEqual(measured_kernel.kernel_id, record.kernel.kernel_id);
2977 try testing.expectEqual(index, record.candidate_index);
2978 try testing.expectEqual(@as(u32, 1), record.sample_count);
2979 }
2980
2981 const artifact_bytes = try measured.exportLaunchTuningArtifact(allocator);
2982 defer allocator.free(artifact_bytes);
2983 const decoded = try decodeLaunchTuningArtifact(allocator, artifact_bytes);
2984 defer allocator.free(decoded);
2985 try testing.expectEqual(@as(usize, 1), decoded.len);
2986 try testing.expectEqual(measured_kernel.kernel_id, decoded[0].selection.kernel_id);
2987 try testing.expect(decoded[0].selection.candidate_index < measured_kernel.launch_resources.candidate_count);
2988
2989 var replay = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{
2990 .launch_tuning_artifact = artifact_bytes,
2991 });
2992 defer replay.deinit();
2993
2994 var graph = try replay.createLaunchGraphPlan(allocator, .{});
2995 defer graph.deinit();
2996 try testing.expectEqual(@as(usize, 1), graph.tuning_selections.len);
2997 try testing.expectEqual(decoded[0].selection.candidate_index, graph.tuning_selections[0].candidate_index);
2998
2999 const replay_kernel = replay.artifactPlan().kernels.items[0];
3000 const slot_bindings = try slotBindingsForKernel(allocator, replay_kernel, .cuda);
3001 defer allocator.free(slot_bindings);
3002 const count_bindings = try elementCountBindingsForPlan(allocator, replay.artifactPlan(), .cuda);
3003 defer allocator.free(count_bindings);
3004
3005 try replay.launchAll(allocator, slot_bindings, count_bindings);
3006
3007 const selected = replay_kernel.launch_resources.candidates[graph.tuning_selections[0].candidate_index];
3008 try testing.expectEqual(selected.geometry.grid[0], state.last_launch_grid[0]);
3009 try testing.expectEqual(selected.geometry.threadgroup[0], state.last_launch_threadgroup[0]);
3010 }
3011
3012 test "Choir executable fragment measured tuning skips fixed launch candidates" {
3013 const allocator = testing.allocator;
3014
3015 var owned = try kernelCallChoirModule(allocator, "choir_loaded_fragment_measure_record_fixed_candidates");
3016 defer owned.deinit();
3017
3018 var cache = passes.AnalysisCache.init(allocator, null);
3019 defer cache.deinit();
3020 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
3021 defer pass_ctx.deinit();
3022
3023 var state = RecordingBackendState{
3024 .allocator = allocator,
3025 .kind = .cuda,
3026 .format = .cuda_ptx,
3027 };
3028 const handle = state.handle();
3029
3030 const source = ".visible .entry accy_custom_scale() { ret; }";
3031 const registry = artifact_product.KernelCallRegistry{ .entries = &.{.{
3032 .target = "accy.custom.scale",
3033 .version = 1,
3034 .format = .cuda_ptx,
3035 .entry_name = "accy_custom_scale",
3036 .argument_count = 5,
3037 .required_dtypes = gpu.DTypeSet.init(&.{.f32}),
3038 .payload = .{ .text = source },
3039 .launch = .{ .fixed = .{
3040 .grid = .{ 2, 3, 1 },
3041 .threadgroup = .{ 8, 2, 1 },
3042 } },
3043 .element_count_argument = .scalar_u32,
3044 .static_arguments = &.{.{ .u32 = 7 }},
3045 }} };
3046
3047 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{
3048 .kernel_call_registry = ®istry,
3049 });
3050 defer fragment.deinit();
3051
3052 const kernel = fragment.artifactPlan().kernels.items[0];
3053 try testing.expectEqual(@as(usize, 1), kernel.launch_resources.candidate_count);
3054
3055 const lhs = @as([8]f32, @splat(1.0));
3056 const rhs = @as([8]f32, @splat(2.0));
3057 const bindings = try fragment.prepareInvocationBindings(allocator, &.{
3058 std.mem.asBytes(&lhs),
3059 std.mem.asBytes(&rhs),
3060 });
3061 defer bindings.deinit();
3062
3063 const records = try fragment.measureAndRecordInvocationLaunchCandidates(
3064 allocator,
3065 allocator,
3066 bindings,
3067 .{ .warmup = 0, .samples = 1 },
3068 );
3069 defer allocator.free(records);
3070
3071 try testing.expectEqual(@as(usize, 0), records.len);
3072 try testing.expectEqual(@as(usize, 0), fragment.tuningRecordCount());
3073 try testing.expectEqual(@as(usize, 0), state.launch_count);
3074 }
3075
3076 test "Choir executable fragment launches fused CUDA kernels" {
3077 const allocator = testing.allocator;
3078
3079 var owned = try fusedAddMulChoirModule(allocator, "choir_loaded_fragment_fused_add_mul");
3080 defer owned.deinit();
3081
3082 var cache = passes.AnalysisCache.init(allocator, null);
3083 defer cache.deinit();
3084 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
3085 defer pass_ctx.deinit();
3086
3087 var state = RecordingBackendState{
3088 .allocator = allocator,
3089 .kind = .cuda,
3090 .format = .cuda_ptx,
3091 };
3092 const handle = state.handle();
3093
3094 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
3095 defer fragment.deinit();
3096
3097 const kernel = fragment.artifactPlan().kernels.items[0];
3098 const slot_bindings = try slotBindingsForKernel(allocator, kernel, .cuda);
3099 defer allocator.free(slot_bindings);
3100 const count_bindings = try elementCountBindingsForPlan(allocator, fragment.artifactPlan(), .cuda);
3101 defer allocator.free(count_bindings);
3102
3103 try fragment.launchAll(allocator, slot_bindings, count_bindings);
3104
3105 try testing.expectEqual(@as(usize, 1), fragment.artifactPlan().kernelCount());
3106 try testing.expectEqual(@as(usize, 2), kernel.op_count);
3107 try testing.expectEqual(@as(usize, 5), state.last_launch_buffer_count);
3108 try testing.expectEqual(@as(usize, 0), state.last_launch_scalar_count);
3109 try testing.expect(state.last_launch_scalar_u32 == null);
3110 try testing.expectEqual(gpu.BufferAccess.write_only, state.last_buffer_access[0]);
3111 try testing.expectEqual(gpu.BufferAccess.read_only, state.last_buffer_access[1]);
3112 try testing.expectEqual(gpu.BufferAccess.read_only, state.last_buffer_access[2]);
3113 try testing.expectEqual(gpu.BufferAccess.read_only, state.last_buffer_access[3]);
3114 try testing.expectEqual(gpu.BufferAccess.read_only, state.last_buffer_access[4]);
3115 }
3116
3117 test "Choir executable fragment records launch candidate records and applies artifact at creation" {
3118 const allocator = testing.allocator;
3119
3120 var owned = try addChoirModule(allocator, "choir_loaded_fragment_tuning_records");
3121 defer owned.deinit();
3122
3123 var cache = passes.AnalysisCache.init(allocator, null);
3124 defer cache.deinit();
3125 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
3126 defer pass_ctx.deinit();
3127
3128 var state = RecordingBackendState{
3129 .allocator = allocator,
3130 .kind = .cuda,
3131 .format = .cuda_ptx,
3132 };
3133 const handle = state.handle();
3134
3135 var measured = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
3136 defer measured.deinit();
3137
3138 const measured_kernel = measured.artifactPlan().kernels.items[0];
3139 try testing.expect(measured_kernel.launch_resources.candidate_count > 1);
3140 const records = [_]LaunchCandidateRecord{
3141 try launchCandidateRecord(
3142 measured_kernel,
3143 measured_kernel.launch_resources.candidates[0],
3144 .{
3145 .kernel_id = measured_kernel.kernel_id,
3146 .candidate_index = 0,
3147 .median_ns = 500,
3148 .sample_count = 8,
3149 },
3150 ),
3151 try launchCandidateRecord(
3152 measured_kernel,
3153 measured_kernel.launch_resources.candidates[1],
3154 .{
3155 .kernel_id = measured_kernel.kernel_id,
3156 .candidate_index = 1,
3157 .median_ns = 200,
3158 .sample_count = 4,
3159 },
3160 ),
3161 };
3162 try measured.recordLaunchCandidateRecords(&records);
3163 try testing.expectEqual(@as(usize, 1), measured.tuningRecordCount());
3164
3165 const artifact_bytes = try measured.exportLaunchTuningArtifact(allocator);
3166 defer allocator.free(artifact_bytes);
3167
3168 var replay = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{
3169 .launch_tuning_artifact = artifact_bytes,
3170 });
3171 defer replay.deinit();
3172 try testing.expectEqual(@as(usize, 1), replay.tuningRecordCount());
3173
3174 const replay_kernel = replay.artifactPlan().kernels.items[0];
3175 const slot_bindings = try slotBindingsForKernel(allocator, replay_kernel, .cuda);
3176 defer allocator.free(slot_bindings);
3177 const count_bindings = try elementCountBindingsForPlan(allocator, replay.artifactPlan(), .cuda);
3178 defer allocator.free(count_bindings);
3179
3180 var graph = try replay.createLaunchGraphPlan(allocator, .{});
3181 defer graph.deinit();
3182 try testing.expectEqual(@as(usize, 1), graph.tuning_selections.len);
3183 try testing.expectEqual(replay_kernel.kernel_id, graph.tuning_selections[0].kernel_id);
3184 try testing.expectEqual(@as(usize, 1), graph.tuning_selections[0].candidate_index);
3185
3186 try replay.launchAll(allocator, slot_bindings, count_bindings);
3187
3188 const selected = replay_kernel.launch_resources.candidates[1];
3189 try testing.expectEqual(@as(usize, 1), state.launch_count);
3190 try testing.expectEqual(selected.geometry.grid[0], state.last_launch_grid[0]);
3191 try testing.expectEqual(selected.geometry.threadgroup[0], state.last_launch_threadgroup[0]);
3192 }
3193
3194 test "Choir executable fragment rejects mismatched launch candidate records" {
3195 const allocator = testing.allocator;
3196
3197 var owned = try addChoirModule(allocator, "choir_loaded_fragment_bad_candidate_records");
3198 defer owned.deinit();
3199
3200 var cache = passes.AnalysisCache.init(allocator, null);
3201 defer cache.deinit();
3202 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
3203 defer pass_ctx.deinit();
3204
3205 var state = RecordingBackendState{
3206 .allocator = allocator,
3207 .kind = .cuda,
3208 .format = .cuda_ptx,
3209 };
3210 const handle = state.handle();
3211
3212 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
3213 defer fragment.deinit();
3214
3215 const kernel = fragment.artifactPlan().kernels.items[0];
3216 var record = try launchCandidateRecord(
3217 kernel,
3218 kernel.launch_resources.candidates[0],
3219 .{
3220 .kernel_id = kernel.kernel_id,
3221 .candidate_index = 0,
3222 .median_ns = 500,
3223 .sample_count = 8,
3224 },
3225 );
3226 record.geometry.threadgroup[0] += 1;
3227 try testing.expectError(error.LaunchArgumentMismatch, fragment.recordLaunchCandidateRecords(&.{record}));
3228 try testing.expectEqual(@as(usize, 0), fragment.tuningRecordCount());
3229
3230 record = try launchCandidateRecord(
3231 kernel,
3232 kernel.launch_resources.candidates[0],
3233 .{
3234 .kernel_id = kernel.kernel_id,
3235 .candidate_index = 0,
3236 .median_ns = 500,
3237 .sample_count = 8,
3238 },
3239 );
3240 record.geometry.dynamic_shared_memory_bytes += 1;
3241 try testing.expectError(error.LaunchArgumentMismatch, fragment.recordLaunchCandidateRecords(&.{record}));
3242 try testing.expectEqual(@as(usize, 0), fragment.tuningRecordCount());
3243
3244 var layout_record = try launchCandidateRecord(
3245 kernel,
3246 kernel.launch_resources.candidates[0],
3247 .{
3248 .kernel_id = kernel.kernel_id,
3249 .candidate_index = 0,
3250 .median_ns = 500,
3251 .sample_count = 8,
3252 },
3253 );
3254 layout_record.kernel.output_layout_fingerprint ^= 1;
3255 try testing.expectError(error.LaunchArgumentMismatch, fragment.recordLaunchCandidateRecords(&.{layout_record}));
3256 try testing.expectEqual(@as(usize, 0), fragment.tuningRecordCount());
3257
3258 var compile_record = try launchCandidateRecord(
3259 kernel,
3260 kernel.launch_resources.candidates[0],
3261 .{
3262 .kernel_id = kernel.kernel_id,
3263 .candidate_index = 0,
3264 .median_ns = 500,
3265 .sample_count = 8,
3266 },
3267 );
3268 compile_record.kernel.compile_required_features = .{ .tensor_cores = true };
3269 try testing.expectError(error.LaunchArgumentMismatch, fragment.recordLaunchCandidateRecords(&.{compile_record}));
3270 try testing.expectEqual(@as(usize, 0), fragment.tuningRecordCount());
3271
3272 compile_record = try launchCandidateRecord(
3273 kernel,
3274 kernel.launch_resources.candidates[0],
3275 .{
3276 .kernel_id = kernel.kernel_id,
3277 .candidate_index = 0,
3278 .median_ns = 500,
3279 .sample_count = 8,
3280 },
3281 );
3282 compile_record.kernel.compile_required_subgroup = .{ .supported = true, .shuffle = true };
3283 try testing.expectError(error.LaunchArgumentMismatch, fragment.recordLaunchCandidateRecords(&.{compile_record}));
3284 try testing.expectEqual(@as(usize, 0), fragment.tuningRecordCount());
3285 }
3286
3287 test "Choir executable fragment imports launch tuning artifact at creation" {
3288 const allocator = testing.allocator;
3289
3290 var owned = try addChoirModule(allocator, "choir_loaded_fragment_tuning_artifact");
3291 defer owned.deinit();
3292
3293 var cache = passes.AnalysisCache.init(allocator, null);
3294 defer cache.deinit();
3295 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
3296 defer pass_ctx.deinit();
3297
3298 var state = RecordingBackendState{
3299 .allocator = allocator,
3300 .kind = .cuda,
3301 .format = .cuda_ptx,
3302 };
3303 const handle = state.handle();
3304
3305 var measured = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
3306 defer measured.deinit();
3307
3308 const measured_kernel = measured.artifactPlan().kernels.items[0];
3309 try testing.expect(measured_kernel.launch_resources.candidate_count > 1);
3310 const records = [_]LaunchCandidateRecord{
3311 try launchCandidateRecord(
3312 measured_kernel,
3313 measured_kernel.launch_resources.candidates[0],
3314 .{
3315 .kernel_id = measured_kernel.kernel_id,
3316 .candidate_index = 0,
3317 .median_ns = 500,
3318 .sample_count = 8,
3319 },
3320 ),
3321 try launchCandidateRecord(
3322 measured_kernel,
3323 measured_kernel.launch_resources.candidates[1],
3324 .{
3325 .kernel_id = measured_kernel.kernel_id,
3326 .candidate_index = 1,
3327 .median_ns = 200,
3328 .sample_count = 4,
3329 },
3330 ),
3331 };
3332 try measured.recordLaunchCandidateRecords(&records);
3333
3334 const artifact_bytes = try measured.exportLaunchTuningArtifact(allocator);
3335 defer allocator.free(artifact_bytes);
3336
3337 var replay = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{
3338 .launch_tuning_artifact = artifact_bytes,
3339 });
3340 defer replay.deinit();
3341 try testing.expectEqual(@as(usize, 1), replay.tuningRecordCount());
3342
3343 const replay_kernel = replay.artifactPlan().kernels.items[0];
3344 const slot_bindings = try slotBindingsForKernel(allocator, replay_kernel, .cuda);
3345 defer allocator.free(slot_bindings);
3346 const count_bindings = try elementCountBindingsForPlan(allocator, replay.artifactPlan(), .cuda);
3347 defer allocator.free(count_bindings);
3348
3349 var graph = try replay.createLaunchGraphPlan(allocator, .{});
3350 defer graph.deinit();
3351 try testing.expectEqual(@as(usize, 1), graph.tuning_selections.len);
3352 try testing.expectEqual(replay_kernel.kernel_id, graph.tuning_selections[0].kernel_id);
3353 try testing.expectEqual(@as(usize, 1), graph.tuning_selections[0].candidate_index);
3354
3355 try replay.launchAll(allocator, slot_bindings, count_bindings);
3356
3357 const selected = replay_kernel.launch_resources.candidates[1];
3358 try testing.expectEqual(@as(usize, 1), state.launch_count);
3359 try testing.expectEqual(selected.geometry.grid[0], state.last_launch_grid[0]);
3360 try testing.expectEqual(selected.geometry.threadgroup[0], state.last_launch_threadgroup[0]);
3361 }
3362
3363 test "Choir executable fragment rejects invalid launch tuning artifacts without mutation" {
3364 const allocator = testing.allocator;
3365
3366 var owned = try addChoirModule(allocator, "choir_loaded_fragment_bad_tuning_artifact");
3367 defer owned.deinit();
3368
3369 var cache = passes.AnalysisCache.init(allocator, null);
3370 defer cache.deinit();
3371 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
3372 defer pass_ctx.deinit();
3373
3374 var state = RecordingBackendState{
3375 .allocator = allocator,
3376 .kind = .cuda,
3377 .format = .cuda_ptx,
3378 };
3379 const handle = state.handle();
3380 const caps = try handle.queryCapabilities();
3381
3382 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
3383 defer fragment.deinit();
3384
3385 const kernel = fragment.artifactPlan().kernels.items[0];
3386 const record = LaunchTuningCacheRecord{
3387 .key = LaunchTuningCacheKey.init(caps, kernel),
3388 .selection = .{
3389 .kernel_id = kernel.kernel_id,
3390 .candidate_index = 0,
3391 .median_ns = 500,
3392 .sample_count = 8,
3393 },
3394 };
3395 const artifact_bytes = try encodeLaunchTuningArtifact(allocator, &.{record});
3396 defer allocator.free(artifact_bytes);
3397 try fragment.importLaunchTuningArtifact(allocator, artifact_bytes);
3398 try testing.expectEqual(@as(usize, 1), fragment.tuningRecordCount());
3399
3400 var bad_magic = try allocator.dupe(u8, artifact_bytes);
3401 defer allocator.free(bad_magic);
3402 bad_magic[0] ^= 0xff;
3403 try testing.expectError(error.InvalidArtifact, fragment.importLaunchTuningArtifact(allocator, bad_magic));
3404 try testing.expectEqual(@as(usize, 1), fragment.tuningRecordCount());
3405 }
3406
3407 test "Choir executable fragment replaces launch tuning artifacts" {
3408 const allocator = testing.allocator;
3409
3410 var owned = try addChoirModule(allocator, "choir_loaded_fragment_replace_tuning_artifact");
3411 defer owned.deinit();
3412
3413 var cache = passes.AnalysisCache.init(allocator, null);
3414 defer cache.deinit();
3415 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
3416 defer pass_ctx.deinit();
3417
3418 var state = RecordingBackendState{
3419 .allocator = allocator,
3420 .kind = .cuda,
3421 .format = .cuda_ptx,
3422 };
3423 const handle = state.handle();
3424
3425 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
3426 defer fragment.deinit();
3427
3428 const kernel = fragment.artifactPlan().kernels.items[0];
3429 try testing.expect(kernel.launch_resources.candidate_count > 1);
3430 const records = [_]LaunchCandidateRecord{try launchCandidateRecord(
3431 kernel,
3432 kernel.launch_resources.candidates[1],
3433 .{
3434 .kernel_id = kernel.kernel_id,
3435 .candidate_index = 1,
3436 .median_ns = 200,
3437 .sample_count = 4,
3438 },
3439 )};
3440 try fragment.recordLaunchCandidateRecords(&records);
3441 try testing.expectEqual(@as(usize, 1), fragment.tuningRecordCount());
3442
3443 const artifact_bytes = try fragment.exportLaunchTuningArtifact(allocator);
3444 defer allocator.free(artifact_bytes);
3445
3446 try fragment.replaceLaunchTuningArtifact(allocator, &.{});
3447 try testing.expectEqual(@as(usize, 0), fragment.tuningRecordCount());
3448
3449 try fragment.replaceLaunchTuningArtifact(allocator, artifact_bytes);
3450 try testing.expectEqual(@as(usize, 1), fragment.tuningRecordCount());
3451 }
3452
3453 test "Choir executable fragment explicit launch tuning overrides cached records" {
3454 const allocator = testing.allocator;
3455
3456 var owned = try addChoirModule(allocator, "choir_loaded_fragment_tuning_override");
3457 defer owned.deinit();
3458
3459 var cache = passes.AnalysisCache.init(allocator, null);
3460 defer cache.deinit();
3461 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
3462 defer pass_ctx.deinit();
3463
3464 var state = RecordingBackendState{
3465 .allocator = allocator,
3466 .kind = .cuda,
3467 .format = .cuda_ptx,
3468 };
3469 const handle = state.handle();
3470 var fragment = try createTestLoadedFragment(allocator, handle, &pass_ctx, owned.choir_module, .{});
3471 defer fragment.deinit();
3472
3473 const kernel = fragment.artifactPlan().kernels.items[0];
3474 try testing.expect(kernel.launch_resources.candidate_count > 1);
3475 const records = [_]LaunchCandidateRecord{try launchCandidateRecord(
3476 kernel,
3477 kernel.launch_resources.candidates[1],
3478 .{
3479 .kernel_id = kernel.kernel_id,
3480 .candidate_index = 1,
3481 .median_ns = 200,
3482 .sample_count = 4,
3483 },
3484 )};
3485 try fragment.recordLaunchCandidateRecords(&records);
3486 try testing.expectEqual(@as(usize, 1), fragment.tuningRecordCount());
3487
3488 const explicit = [_]LaunchCandidateMeasurement{.{
3489 .kernel_id = kernel.kernel_id,
3490 .candidate_index = 0,
3491 .median_ns = 1,
3492 .sample_count = 1,
3493 }};
3494 const slot_bindings = try slotBindingsForKernel(allocator, kernel, .cuda);
3495 defer allocator.free(slot_bindings);
3496 const count_bindings = try elementCountBindingsForPlan(allocator, fragment.artifactPlan(), .cuda);
3497 defer allocator.free(count_bindings);
3498
3499 var graph = try fragment.createLaunchGraphPlan(allocator, .{
3500 .tuning = .{ .measurements = &explicit },
3501 });
3502 defer graph.deinit();
3503 try testing.expectEqual(@as(usize, 0), graph.tuning_selections.len);
3504 try testing.expectEqual(@as(usize, 1), graph.nodes[0].tuning.measurements.len);
3505
3506 try fragment.launchAllWithOptions(allocator, slot_bindings, count_bindings, .{
3507 .tuning = .{ .measurements = &explicit },
3508 });
3509
3510 const selected = kernel.launch_resources.candidates[0];
3511 try testing.expectEqual(@as(usize, 1), state.launch_count);
3512 try testing.expectEqual(selected.geometry.grid[0], state.last_launch_grid[0]);
3513 try testing.expectEqual(selected.geometry.threadgroup[0], state.last_launch_threadgroup[0]);
3514 }
3515
3516 test "Choir executable plan launches Vulkan device-count kernels with count buffer" {
3517 const allocator = testing.allocator;
3518
3519 var owned = try addChoirModule(allocator, "choir_vulkan_executable_add");
3520 defer owned.deinit();
3521
3522 var cache = passes.AnalysisCache.init(allocator, null);
3523 defer cache.deinit();
3524 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
3525 defer pass_ctx.deinit();
3526
3527 var state = RecordingBackendState{
3528 .allocator = allocator,
3529 .kind = .vulkan,
3530 .format = .vulkan_spirv,
3531 };
3532 const handle = state.handle();
3533
3534 var artifact_plan = try createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{});
3535 defer artifact_plan.deinit();
3536
3537 var executable = try loadKernels(allocator, handle, &artifact_plan);
3538 defer executable.deinit();
3539
3540 const kernel = artifact_plan.kernels.items[0];
3541 const slot_bindings = try slotBindingsForKernel(allocator, kernel, .vulkan);
3542 defer allocator.free(slot_bindings);
3543 const count_binding = ElementCountBufferBinding{
3544 .kernel_id = kernel.kernel_id,
3545 .binding = bufferBinding(200, .vulkan, 4),
3546 };
3547
3548 try executable.launchAll(allocator, &artifact_plan, slot_bindings, &.{count_binding});
3549
3550 try testing.expectEqual(@as(usize, 1), state.load_count);
3551 try testing.expectEqual(@as(usize, 1), state.launch_count);
3552 try testing.expectEqual(@as(usize, 4), state.last_launch_buffer_count);
3553 try testing.expectEqual(@as(usize, 0), state.last_launch_scalar_count);
3554 try testing.expect(state.last_launch_scalar_u32 == null);
3555 try testing.expectEqual(@as(u32, 8), state.last_launch_threadgroup[0]);
3556 try testing.expectEqual(@as(gpu.BackendObjectId, 200), state.last_buffer_ids[3]);
3557 try testing.expectEqual(gpu.BufferAccess.read_only, state.last_buffer_access[3]);
3558 }
3559
3560 test "Choir executable plan requires a device-count buffer for Vulkan launches" {
3561 const allocator = testing.allocator;
3562
3563 var owned = try addChoirModule(allocator, "choir_vulkan_executable_missing_count");
3564 defer owned.deinit();
3565
3566 var cache = passes.AnalysisCache.init(allocator, null);
3567 defer cache.deinit();
3568 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
3569 defer pass_ctx.deinit();
3570
3571 var state = RecordingBackendState{
3572 .allocator = allocator,
3573 .kind = .vulkan,
3574 .format = .vulkan_spirv,
3575 };
3576 const handle = state.handle();
3577
3578 var artifact_plan = try createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{});
3579 defer artifact_plan.deinit();
3580
3581 var executable = try loadKernels(allocator, handle, &artifact_plan);
3582 defer executable.deinit();
3583
3584 const kernel = artifact_plan.kernels.items[0];
3585 const slot_bindings = try slotBindingsForKernel(allocator, kernel, .vulkan);
3586 defer allocator.free(slot_bindings);
3587
3588 try testing.expectError(
3589 error.LaunchArgumentMismatch,
3590 executable.launchAll(allocator, &artifact_plan, slot_bindings, &.{}),
3591 );
3592 try testing.expectEqual(@as(usize, 0), state.launch_count);
3593 }
3594
3595 test "Choir executable plan destroys loaded artifacts on deinit" {
3596 const allocator = testing.allocator;
3597
3598 var owned = try addChoirModule(allocator, "choir_executable_destroy_loaded");
3599 defer owned.deinit();
3600
3601 var cache = passes.AnalysisCache.init(allocator, null);
3602 defer cache.deinit();
3603 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.ctx, allocator, &cache);
3604 defer pass_ctx.deinit();
3605
3606 var state = RecordingBackendState{
3607 .allocator = allocator,
3608 .kind = .cuda,
3609 .format = .cuda_ptx,
3610 };
3611 const handle = state.handle();
3612
3613 var artifact_plan = try createTestBackendArtifactPlan(allocator, handle, &pass_ctx, owned.choir_module, .{});
3614 defer artifact_plan.deinit();
3615
3616 var executable = try loadKernels(allocator, handle, &artifact_plan);
3617 const loaded_id = state.last_loaded_id.?;
3618 executable.deinit();
3619
3620 try testing.expectEqual(@as(usize, 1), state.destroy_count);
3621 try testing.expectEqual(loaded_id, state.last_destroyed_id.?);
3622 }
3623
3624 fn familyTuningBlobTestModule(allocator: std.mem.Allocator, name: []const u8) !*semantic.SemanticModule {
3625 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
3626 errdefer builder.deinit();
3627 const lhs_ty = try builder.tensor(.f32, &.{ 5, 3 });
3628 const rhs_ty = try builder.tensor(.f32, &.{ 3, 7 });
3629 const out_ty = try builder.tensor(.f32, &.{ 5, 7 });
3630 var fb = try builder.beginFunction(name, &.{ lhs_ty, rhs_ty }, &.{out_ty});
3631 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "mk,kn->mn");
3632 try fb.return_(&.{out});
3633 try fb.finish();
3634 return try builder.finish();
3635 }
3636
3637 fn findKernelCallTarget(op: *ir.Operation) ?[]const u8 {
3638 if (std.mem.eql(u8, op.name.name, accy_choir.dialect.AccyDialect.KernelCallOp.operation_name)) {
3639 const target_attr = op.getAttr("target") orelse return null;
3640 const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return null;
3641 return target.payload;
3642 }
3643 for (op.regions.items) |*region| {
3644 var block_iter = region.getBlocks();
3645 while (block_iter.next()) |block| {
3646 var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
3647 while (current) |current_op| {
3648 if (findKernelCallTarget(current_op)) |found| return found;
3649 current = current_op.next_op;
3650 }
3651 }
3652 }
3653 return null;
3654 }
3655
3656 test "Choir executable fragment consults an embedded family tuning artifact" {
3657 const allocator = std.testing.allocator;
3658 var state = gpu.recording.BackendState{
3659 .allocator = allocator,
3660 .kind = .cuda,
3661 .format = .cuda_ptx,
3662 };
3663 const handle = state.handle();
3664 const library = kernel_library;
3665 const tuning = library.tuning;
3666
3667 var registry = try matrixProductTuningTestRegistry(allocator, handle);
3668 defer registry.deinit();
3669 const registry_value = registry.registry();
3670
3671 const probe = library.linalg.MatrixProduct{ .m = 5, .n = 7, .k = 3 };
3672 const caps = try handle.queryCapabilities();
3673 const device = tuning.deviceFingerprint(caps);
3674 const thread_candidates = library.linalg.matrixProductThreadCandidatesForExtents(probe.m, probe.n);
3675 var winner = probe;
3676 winner.threads = thread_candidates.slice()[0];
3677 const winner_target = try library.linalg.matrixProductFamilyTarget(allocator, winner);
3678 defer allocator.free(winner_target);
3679
3680 const records = [_]tuning.FamilyTuningRecord{.{
3681 .key = try library.linalg.matrixProductFamilyTuningKey(allocator, device, probe),
3682 .target = winner_target,
3683 .winner_median_ns = 800,
3684 .runner_up_median_ns = 1200,
3685 .sample_count = 30,
3686 }};
3687 const encoded = try tuning.encodeFamilyTuningArtifact(allocator, records[0..]);
3688 defer allocator.free(encoded);
3689
3690 const module = try familyTuningBlobTestModule(allocator, "family_tuning_blob_prepare");
3691 var prepared = try prepareFragmentFromSemanticModule(
3692 allocator,
3693 handle,
3694 module,
3695 .{
3696 .kernel_call_registry = ®istry_value,
3697 .family_tuning_artifact = encoded,
3698 },
3699 );
3700 defer prepared.deinit();
3701
3702 const selected_target = findKernelCallTarget(prepared.choir_module) orelse
3703 return error.TestExpectedKernelCall;
3704 try std.testing.expectEqualStrings(winner_target, selected_target);
3705 }
3706
3707 fn matrixProductScheduleTuningArtifactForThreads(
3708 allocator: std.mem.Allocator,
3709 handle: gpu.BackendHandle,
3710 probe: kernel_library.linalg.MatrixProduct,
3711 selected_threads: kernel_library.Threads2D,
3712 ) ![]u8 {
3713 const thread_candidates = kernel_library.linalg.matrixProductThreadCandidatesForExtents(probe.m, probe.n);
3714 const candidate_slice = thread_candidates.slice();
3715 try std.testing.expect(candidate_slice.len >= 2);
3716 const capacity = kernel_library.tuning.matrix_product_family_schedule_tuning_max_candidates;
3717 var schedule_candidates: [capacity]kernel_library.tuning.MatrixProductFamilyScheduleThreads =
3718 undefined;
3719 try std.testing.expect(candidate_slice.len <= schedule_candidates.len);
3720 for (candidate_slice, 0..) |candidate, index| {
3721 schedule_candidates[index] = .{ .x = candidate.x, .y = candidate.y };
3722 }
3723 const problem = kernel_library.tuning.MatrixProductFamilyScheduleTuningProblem{
3724 .format = .cuda_ptx,
3725 .m = probe.m,
3726 .n = probe.n,
3727 .k = probe.k,
3728 .dtype = probe.dtype,
3729 .accumulation_dtype = probe.accumulation_dtype,
3730 .family_version = kernel_library.linalg.matrix_product_family_version,
3731 .candidates = schedule_candidates[0..candidate_slice.len],
3732 };
3733 const record = kernel_library.tuning.MatrixProductFamilyScheduleTuningRecord{
3734 .key = try kernel_library.tuning.MatrixProductFamilyScheduleTuningKey.init(
3735 (try handle.queryCapabilities()).identity,
3736 problem,
3737 ),
3738 .selection = .{
3739 .threads = .{ .x = selected_threads.x, .y = selected_threads.y },
3740 .winner_median_ns = 800,
3741 .runner_up_median_ns = 1200,
3742 .sample_count = 30,
3743 },
3744 };
3745 return try schedule_tuning.encodeMatrixProductFamilyScheduleTuningArtifact(allocator, &.{record});
3746 }
3747
3748 fn nonHeuristicMatrixProductThreads(probe: kernel_library.linalg.MatrixProduct) !kernel_library.Threads2D {
3749 const heuristic = kernel_library.linalg.matrixProductThreadsForExtents(probe.m, probe.n);
3750 const thread_candidates = kernel_library.linalg.matrixProductThreadCandidatesForExtents(probe.m, probe.n);
3751 for (thread_candidates.slice()) |candidate| {
3752 if (candidate.x != heuristic.x or candidate.y != heuristic.y) return candidate;
3753 }
3754 return error.TestExpectedMatrixProductThreadCandidate;
3755 }
3756
3757 fn matrixProductTuningTestDescriptor(
3758 allocator: std.mem.Allocator,
3759 threads: kernel_library.Threads2D,
3760 ) !kernel_library.OwnedCatalogDescriptor {
3761 const lhs_dims = [_]i64{ 5, 3 };
3762 const rhs_dims = [_]i64{ 3, 7 };
3763 const out_dims = [_]i64{ 5, 7 };
3764 return (try kernel_library.selectOwned(allocator, .{ .matrix_product = .{
3765 .dtype = .f32,
3766 .lhs_indices = "mk",
3767 .rhs_indices = "kn",
3768 .output_indices = "mn",
3769 .lhs_dims = lhs_dims[0..],
3770 .rhs_dims = rhs_dims[0..],
3771 .output_dims = out_dims[0..],
3772 .schedule = .{ .thread_blocks = threads },
3773 } })) orelse return error.TestExpectedMatrixProductThreadCandidate;
3774 }
3775
3776 fn matrixProductTuningTestRegistry(
3777 allocator: std.mem.Allocator,
3778 handle: gpu.BackendHandle,
3779 ) !kernel_library.OwnedKernelCallArtifactRegistry {
3780 const probe = kernel_library.linalg.MatrixProduct{ .m = 5, .n = 7, .k = 3 };
3781 const heuristic_threads = kernel_library.linalg.matrixProductThreadsForExtents(probe.m, probe.n);
3782 const alternate_threads = try nonHeuristicMatrixProductThreads(probe);
3783 var descriptors: [2]kernel_library.OwnedCatalogDescriptor = undefined;
3784 var descriptor_count: usize = 0;
3785 defer for (descriptors[0..descriptor_count]) |*descriptor| descriptor.deinit();
3786 descriptors[0] = try matrixProductTuningTestDescriptor(allocator, heuristic_threads);
3787 descriptor_count = 1;
3788 descriptors[1] = try matrixProductTuningTestDescriptor(allocator, alternate_threads);
3789 descriptor_count = 2;
3790 return try kernel_library.createOwnedKernelCallArtifactRegistry(
3791 allocator,
3792 handle,
3793 descriptors[0..],
3794 .{ .limits = .testing },
3795 );
3796 }
3797
3798 test "Choir executable fragment consumes matrix product schedule tuning artifacts" {
3799 const allocator = std.testing.allocator;
3800 var state = gpu.recording.BackendState{
3801 .allocator = allocator,
3802 .kind = .cuda,
3803 .format = .cuda_ptx,
3804 };
3805 const handle = state.handle();
3806 const probe = kernel_library.linalg.MatrixProduct{ .m = 5, .n = 7, .k = 3 };
3807 const selected_threads = try nonHeuristicMatrixProductThreads(probe);
3808 const encoded = try matrixProductScheduleTuningArtifactForThreads(allocator, handle, probe, selected_threads);
3809 defer allocator.free(encoded);
3810
3811 var registry = try matrixProductTuningTestRegistry(allocator, handle);
3812 defer registry.deinit();
3813 const registry_value = registry.registry();
3814
3815 var winner = probe;
3816 winner.threads = selected_threads;
3817 const winner_target = try kernel_library.linalg.matrixProductFamilyTarget(allocator, winner);
3818 defer allocator.free(winner_target);
3819
3820 const module = try familyTuningBlobTestModule(allocator, "matrix_product_schedule_tuning_blob_prepare");
3821 var prepared = try prepareFragmentFromSemanticModule(
3822 allocator,
3823 handle,
3824 module,
3825 .{
3826 .kernel_call_registry = ®istry_value,
3827 .matrix_product_schedule_tuning_artifact = encoded,
3828 },
3829 );
3830 defer prepared.deinit();
3831
3832 const selected_target = findKernelCallTarget(prepared.choir_module) orelse
3833 return error.TestExpectedKernelCall;
3834 try std.testing.expectEqualStrings(winner_target, selected_target);
3835 }
3836
3837 test "Choir executable fragment prefers explicit matrix product schedule over tuning artifact" {
3838 const allocator = std.testing.allocator;
3839 var state = gpu.recording.BackendState{
3840 .allocator = allocator,
3841 .kind = .cuda,
3842 .format = .cuda_ptx,
3843 };
3844 const handle = state.handle();
3845 const probe = kernel_library.linalg.MatrixProduct{ .m = 5, .n = 7, .k = 3 };
3846 const selected_threads = try nonHeuristicMatrixProductThreads(probe);
3847 const explicit_threads = kernel_library.linalg.matrixProductThreadsForExtents(probe.m, probe.n);
3848 try std.testing.expect(selected_threads.x != explicit_threads.x or selected_threads.y != explicit_threads.y);
3849 const encoded = try matrixProductScheduleTuningArtifactForThreads(allocator, handle, probe, selected_threads);
3850 defer allocator.free(encoded);
3851
3852 var registry = try matrixProductTuningTestRegistry(allocator, handle);
3853 defer registry.deinit();
3854 const registry_value = registry.registry();
3855
3856 var explicit = probe;
3857 explicit.threads = explicit_threads;
3858 const explicit_target = try kernel_library.linalg.matrixProductFamilyTarget(allocator, explicit);
3859 defer allocator.free(explicit_target);
3860
3861 const module = try familyTuningBlobTestModule(allocator, "matrix_product_schedule_tuning_explicit_prepare");
3862 var prepared = try prepareFragmentFromSemanticModule(
3863 allocator,
3864 handle,
3865 module,
3866 .{
3867 .kernel_call_registry = ®istry_value,
3868 .matrix_product_schedule = .{ .thread_blocks = explicit_threads },
3869 .matrix_product_schedule_tuning_artifact = encoded,
3870 },
3871 );
3872 defer prepared.deinit();
3873
3874 const selected_target = findKernelCallTarget(prepared.choir_module) orelse
3875 return error.TestExpectedKernelCall;
3876 try std.testing.expectEqualStrings(explicit_target, selected_target);
3877 }
3878
3879 test "Choir executable fragment rejects malformed matrix product schedule tuning artifacts" {
3880 const allocator = std.testing.allocator;
3881 var state = gpu.recording.BackendState{
3882 .allocator = allocator,
3883 .kind = .cuda,
3884 .format = .cuda_ptx,
3885 };
3886 const module = try familyTuningBlobTestModule(allocator, "matrix_product_schedule_tuning_blob_invalid");
3887 const bad_bytes = [_]u8{ 1, 2, 3 };
3888 try std.testing.expectError(error.InvalidArtifact, prepareFragmentFromSemanticModule(
3889 allocator,
3890 state.handle(),
3891 module,
3892 .{ .matrix_product_schedule_tuning_artifact = bad_bytes[0..] },
3893 ));
3894
3895 const probe = kernel_library.linalg.MatrixProduct{ .m = 5, .n = 7, .k = 3 };
3896 const bad_selection = try matrixProductScheduleTuningArtifactForThreads(
3897 allocator,
3898 state.handle(),
3899 probe,
3900 .{ .x = 99, .y = 99 },
3901 );
3902 defer allocator.free(bad_selection);
3903 var registry = try matrixProductTuningTestRegistry(allocator, state.handle());
3904 defer registry.deinit();
3905 const registry_value = registry.registry();
3906 const selected_module = try familyTuningBlobTestModule(allocator, "matrix_product_schedule_tuning_blob_bad_selection");
3907 var failure: preparation.BackendPreparationFailure = .{};
3908 defer failure.deinit(allocator);
3909 try std.testing.expectError(error.PassFailed, prepareFragmentFromSemanticModule(
3910 allocator,
3911 state.handle(),
3912 selected_module,
3913 .{
3914 .kernel_call_registry = ®istry_value,
3915 .matrix_product_schedule_tuning_artifact = bad_selection,
3916 .preparation_failure = &failure,
3917 },
3918 ));
3919 try std.testing.expectEqualStrings(preparation.tensor_pipeline_name, failure.pipeline_name.?);
3920 try std.testing.expectEqual(passes.PassFailureKind.pass, failure.failure_kind.?);
3921 try std.testing.expectEqualStrings(preparation.einsum_lowering_pass_name, failure.pass_name.?);
3922 try std.testing.expect(failure.target_op_name != null);
3923 try std.testing.expect(failure.worker_count > 0);
3924 }
3925
3926 test "Choir executable fragment rejects malformed family tuning artifacts" {
3927 const allocator = std.testing.allocator;
3928 var state = gpu.recording.BackendState{
3929 .allocator = allocator,
3930 .kind = .cuda,
3931 .format = .cuda_ptx,
3932 };
3933 const module = try familyTuningBlobTestModule(allocator, "family_tuning_blob_invalid");
3934 const bad_bytes = [_]u8{ 1, 2, 3 };
3935 try std.testing.expectError(error.InvalidArtifact, prepareFragmentFromSemanticModule(
3936 allocator,
3937 state.handle(),
3938 module,
3939 .{ .family_tuning_artifact = bad_bytes[0..] },
3940 ));
3941 }
3942
3943 fn matrixProductScheduleRecipe(
3944 allocator: std.mem.Allocator,
3945 handle: gpu.BackendHandle,
3946 artifact: []const u8,
3947 root: *ir.Operation,
3948 ) ![]u8 {
3949 var plan: FragmentPreparationPlan = undefined;
3950 try plan.init(allocator, handle, .{ .matrix_product_schedule_tuning_artifact = artifact });
3951 defer plan.deinit();
3952 return preparation.recipe.encode(allocator, .tensor, root, plan.run_options);
3953 }
3954
3955 test "matrix product schedule recipe restores the artifact reader after plan destruction" {
3956 const allocator = std.testing.allocator;
3957 var state = gpu.recording.BackendState{
3958 .allocator = allocator,
3959 .kind = .cuda,
3960 .format = .cuda_ptx,
3961 };
3962 const probe = kernel_library.linalg.MatrixProduct{ .m = 5, .n = 7, .k = 3 };
3963 const winner = try nonHeuristicMatrixProductThreads(probe);
3964 const bytes = block: {
3965 const artifact = try matrixProductScheduleTuningArtifactForThreads(
3966 allocator,
3967 state.handle(),
3968 probe,
3969 winner,
3970 );
3971 defer allocator.free(artifact);
3972 const module = try familyTuningBlobTestModule(allocator, "matrix_recipe_input");
3973 defer module.deinit();
3974 break :block try matrixProductScheduleRecipe(
3975 allocator,
3976 state.handle(),
3977 artifact,
3978 module.choir_module,
3979 );
3980 };
3981 defer allocator.free(bytes);
3982 var decoded = try accy_choir.record.codec.decode(
3983 allocator,
3984 preparation.recipe.Record(.tensor),
3985 .tensor,
3986 bytes,
3987 );
3988 defer decoded.deinit();
3989 const reader = decoded.value.options.einsum.matrix_product_tuning.?;
3990 try std.testing.expectEqual(winner, (try reader.resolve(probe)).?);
3991 var missing = probe;
3992 missing.k += 1;
3993 try std.testing.expectEqual(null, try reader.resolve(missing));
3994 }
3995
3996 test "matrix product schedule recipe normalizes artifact duplicates through the cache" {
3997 const allocator = std.testing.allocator;
3998 var state = gpu.recording.BackendState{
3999 .allocator = allocator,
4000 .kind = .cuda,
4001 .format = .cuda_ptx,
4002 };
4003 const probe = kernel_library.linalg.MatrixProduct{ .m = 5, .n = 7, .k = 3 };
4004 const winner = try nonHeuristicMatrixProductThreads(probe);
4005 const artifact = try matrixProductScheduleTuningArtifactForThreads(
4006 allocator,
4007 state.handle(),
4008 probe,
4009 winner,
4010 );
4011 defer allocator.free(artifact);
4012 const decoded = try schedule_tuning.decodeMatrixProductFamilyScheduleTuningArtifact(
4013 allocator,
4014 artifact,
4015 );
4016 defer allocator.free(decoded);
4017 var entries = [_]kernel_library.tuning.MatrixProductFamilyScheduleTuningRecord{
4018 decoded[0], decoded[0],
4019 };
4020 const slow = kernel_library.linalg.matrixProductThreadsForExtents(probe.m, probe.n);
4021 entries[0].selection.threads = .{ .x = slow.x, .y = slow.y };
4022 entries[0].selection.winner_median_ns += 1;
4023 const module = try familyTuningBlobTestModule(allocator, "matrix_recipe_duplicates");
4024 defer module.deinit();
4025 var recipes: [2][]u8 = undefined;
4026 var count: usize = 0;
4027 defer for (recipes[0..count]) |bytes| allocator.free(bytes);
4028 for (&recipes) |*bytes| {
4029 const encoded = try schedule_tuning.encodeMatrixProductFamilyScheduleTuningArtifact(
4030 allocator,
4031 &entries,
4032 );
4033 defer allocator.free(encoded);
4034 bytes.* = try matrixProductScheduleRecipe(
4035 allocator,
4036 state.handle(),
4037 encoded,
4038 module.choir_module,
4039 );
4040 count += 1;
4041 std.mem.swap(
4042 kernel_library.tuning.MatrixProductFamilyScheduleTuningRecord,
4043 &entries[0],
4044 &entries[1],
4045 );
4046 }
4047 try std.testing.expectEqualStrings(recipes[0], recipes[1]);
4048 var restored = try accy_choir.record.codec.decode(
4049 allocator,
4050 preparation.recipe.Record(.tensor),
4051 .tensor,
4052 recipes[0],
4053 );
4054 defer restored.deinit();
4055 const reader = restored.value.options.einsum.matrix_product_tuning.?;
4056 try std.testing.expectEqual(@as(usize, 1), reader.records.len);
4057 try std.testing.expectEqual(winner, (try reader.resolve(probe)).?);
4058 }
4059
4060 fn matrixProductSchedulePlanFailure(
4061 allocator: std.mem.Allocator,
4062 handle: gpu.BackendHandle,
4063 artifact: []const u8,
4064 ) !void {
4065 var plan: FragmentPreparationPlan = undefined;
4066 try plan.init(allocator, handle, .{ .matrix_product_schedule_tuning_artifact = artifact });
4067 defer plan.deinit();
4068 const reader = plan.run_options.tensor.einsum.matrix_product_tuning.?;
4069 try std.testing.expectEqual(@as(usize, 1), reader.records.len);
4070 }
4071
4072 test "matrix product schedule snapshot releases every failed acquisition" {
4073 const allocator = std.testing.allocator;
4074 var state = gpu.recording.BackendState{
4075 .allocator = allocator,
4076 .kind = .cuda,
4077 .format = .cuda_ptx,
4078 };
4079 const probe = kernel_library.linalg.MatrixProduct{ .m = 5, .n = 7, .k = 3 };
4080 const winner = try nonHeuristicMatrixProductThreads(probe);
4081 const artifact = try matrixProductScheduleTuningArtifactForThreads(
4082 allocator,
4083 state.handle(),
4084 probe,
4085 winner,
4086 );
4087 defer allocator.free(artifact);
4088 try std.testing.checkAllAllocationFailures(allocator, matrixProductSchedulePlanFailure, .{
4089 state.handle(), artifact,
4090 });
4091 }