lib/accy/src/artifact/plan.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 gpu_codegen = choir.backends.gpu;
6 const GpuDialect = choir.dialects.gpu.GpuDialect;
7 const kernel_program = @import("../kernel/model/root.zig").program;
8 const accy_choir = @import("../choir/root.zig");
9 const artifact_model = @import("model/root.zig");
10 const preparation = @import("../preparation/root.zig");
11 const kernelization_model = @import("../preparation/kernelization/model/root.zig");
12 const backend_legalization = preparation.backend;
13 const bufferization = preparation.bufferization;
14 const dialect_mod = accy_choir.dialect;
15 const kernelization = preparation.kernelization;
16 const kernel_outlining = preparation.outlining;
17 const layout_planning = preparation.layout;
18 const memory_space = preparation.memory;
19 const schedule_planning = preparation.schedule;
20 const semantic = accy_choir.semantic;
21 const target_profile = preparation.target;
22 const target_product = @import("../target/root.zig");
23
24 const ir = choir.ir;
25 const passes = choir.passes;
26
27 const ArtifactPlanOptions = artifact_model.ArtifactPlanOptions;
28 const ElementCountArgument = artifact_model.ElementCountArgument;
29 const KernelCallArtifact = artifact_model.KernelCallArtifact;
30 const KernelCallDerivedLaunch = artifact_model.KernelCallDerivedLaunch;
31 const KernelCallDerivedLaunchAxis = artifact_model.KernelCallDerivedLaunchAxis;
32 const KernelCallLaunch = artifact_model.KernelCallLaunch;
33 const KernelCallRegistry = artifact_model.KernelCallRegistry;
34 const KernelCallShapeProfile = artifact_model.KernelCallShapeProfile;
35 const StandaloneKernelOptions = artifact_model.StandaloneKernelOptions;
36 const defaultArtifactFormat = artifact_model.defaultArtifactFormat;
37 const validateThreadgroup = artifact_model.validateThreadgroup;
38
39 pub const BackendArtifactPlanSource = struct {
40 pass_ctx: *passes.PassContext,
41 choir_module: *ir.Operation,
42 lowered_kernels: *const kernelization.KernelizationAnalysis,
43 };
44
45 const CompilePlan = struct {
46 entry_name: []const u8,
47 entry_name_owned: ?[]u8 = null,
48 argument_count: u32 = 0,
49 payload: gpu.CompilePayload = .none,
50 launch: CompileLaunchPlan = .generic,
51 lowered_kernel: ?*const kernelization.LoweredKernel = null,
52 required_dtypes: gpu.DTypeSet = .{},
53 required_features: choir_abi.Features = .{},
54 required_subgroup: choir_abi.SubgroupRequirements = .{},
55 push_constants: choir_abi.PushConstants = .{},
56 shape_family_fingerprint: ?u64 = null,
57 element_count_argument: ?ElementCountArgument = null,
58 runtime_scalar_argument_count: u32 = 0,
59 runtime_scalar_defaults: []choir_abi.ScalarArgument = &.{},
60 static_arguments: []choir_abi.ScalarArgument = &.{},
61
62 fn deinit(self: *CompilePlan, allocator: std.mem.Allocator) void {
63 if (self.entry_name_owned) |entry_name| allocator.free(entry_name);
64 if (self.runtime_scalar_defaults.len != 0) allocator.free(self.runtime_scalar_defaults);
65 if (self.static_arguments.len != 0) allocator.free(self.static_arguments);
66 deinitCompilePayload(allocator, self.payload);
67 self.* = undefined;
68 }
69 };
70
71 const CompileLaunchPlan = union(enum) {
72 generic,
73 dot_general: DotGeneralLaunchPlan,
74 reduction: ReductionLaunchPlan,
75 kernel_call: KernelCallLaunch,
76 };
77
78 pub const PlannedKernelCompilePayload = enum {
79 none,
80 bytes,
81 words_u32,
82 text,
83 };
84
85 pub const PlannedKernelCompileLaunch = enum {
86 authored,
87 generic,
88 dot_general,
89 reduction,
90 kernel_call,
91 };
92
93 pub const PlannedKernelCompileContract = struct {
94 source: PlannedKernelSource,
95 launch: PlannedKernelCompileLaunch,
96 format: gpu.ArtifactFormat,
97 entry_name: []u8,
98 argument_count: u32,
99 required_dtypes: gpu.DTypeSet,
100 required_features: choir_abi.Features = .{},
101 required_subgroup: choir_abi.SubgroupRequirements = .{},
102 shape_family_fingerprint: ?u64 = null,
103 payload: PlannedKernelCompilePayload,
104 payload_byte_count: usize,
105
106 pub fn init(
107 allocator: std.mem.Allocator,
108 source: PlannedKernelSource,
109 launch: PlannedKernelCompileLaunch,
110 format: gpu.ArtifactFormat,
111 entry_name: []const u8,
112 argument_count: u32,
113 required_dtypes: gpu.DTypeSet,
114 required_features: choir_abi.Features,
115 required_subgroup: choir_abi.SubgroupRequirements,
116 shape_family_fingerprint: ?u64,
117 payload: gpu.CompilePayload,
118 ) !PlannedKernelCompileContract {
119 return .{
120 .source = source,
121 .launch = launch,
122 .format = format,
123 .entry_name = try allocator.dupe(u8, entry_name),
124 .argument_count = argument_count,
125 .required_dtypes = required_dtypes,
126 .required_features = required_features,
127 .required_subgroup = required_subgroup,
128 .shape_family_fingerprint = shape_family_fingerprint,
129 .payload = compilePayloadKind(payload),
130 .payload_byte_count = try compilePayloadByteCount(payload),
131 };
132 }
133
134 pub fn deinit(self: *PlannedKernelCompileContract, allocator: std.mem.Allocator) void {
135 allocator.free(self.entry_name);
136 self.* = undefined;
137 }
138
139 fn copy(self: PlannedKernelCompileContract, allocator: std.mem.Allocator) gpu.BackendError!PlannedKernelCompileContract {
140 return .{
141 .source = self.source,
142 .launch = self.launch,
143 .format = self.format,
144 .entry_name = allocator.dupe(u8, self.entry_name) catch return error.OutOfMemory,
145 .argument_count = self.argument_count,
146 .required_dtypes = self.required_dtypes,
147 .required_features = self.required_features,
148 .required_subgroup = self.required_subgroup,
149 .shape_family_fingerprint = self.shape_family_fingerprint,
150 .payload = self.payload,
151 .payload_byte_count = self.payload_byte_count,
152 };
153 }
154 };
155
156 const DotGeneralLaunchPlan = struct {
157 input_dtype: choir_abi.DType,
158 output_dtype: choir_abi.DType,
159 m: u32,
160 n: u32,
161 k: u32,
162 batch: u32 = 1,
163 };
164
165 const ReductionLaunchPlan = struct {
166 kind: LaunchReductionKind,
167 input_dtype: choir_abi.DType,
168 output_dtype: choir_abi.DType,
169 input_rank: u8,
170 axis: u8,
171 input_element_count: u32,
172 output_element_count: u32,
173 reduction_extent: u32,
174 };
175
176 pub const PlannedKernel = struct {
177 compile: PlannedKernelCompileContract,
178 kernel_id: usize,
179 work_item_id: usize,
180 output_slot_id: usize,
181 input_slot_ids: []usize,
182 output_layout_fingerprint: u64,
183 input_layout_fingerprint: u64,
184 element_count: u64,
185 op_count: usize,
186 resources: schedule_planning.ScheduleResourceEstimate,
187 artifact: gpu.KernelArtifact,
188 launch_resources: LaunchResourcePlan,
189 kernel_call_launch: ?KernelCallLaunch = null,
190 element_count_argument: ElementCountArgument,
191 element_count_argument_value: u64,
192 runtime_scalar_argument_count: u32,
193 runtime_scalar_defaults: []choir_abi.ScalarArgument,
194 static_arguments: []choir_abi.ScalarArgument,
195 output_fill_pattern: ?u32 = null,
196 scratch_fill_pattern: ?u32 = null,
197
198 fn deinit(self: *PlannedKernel, allocator: std.mem.Allocator) void {
199 self.artifact.deinit();
200 self.compile.deinit(allocator);
201 allocator.free(self.input_slot_ids);
202 if (self.runtime_scalar_defaults.len != 0) allocator.free(self.runtime_scalar_defaults);
203 if (self.static_arguments.len != 0) allocator.free(self.static_arguments);
204 self.* = undefined;
205 }
206
207 fn copy(self: PlannedKernel, allocator: std.mem.Allocator) gpu.BackendError!PlannedKernel {
208 var compile = try self.compile.copy(allocator);
209 errdefer compile.deinit(allocator);
210
211 const input_slot_ids = allocator.dupe(usize, self.input_slot_ids) catch return error.OutOfMemory;
212 errdefer allocator.free(input_slot_ids);
213
214 var artifact = try copyKernelArtifact(allocator, self.artifact);
215 errdefer artifact.deinit();
216
217 const runtime_scalar_defaults = allocator.dupe(choir_abi.ScalarArgument, self.runtime_scalar_defaults) catch return error.OutOfMemory;
218 errdefer if (runtime_scalar_defaults.len != 0) allocator.free(runtime_scalar_defaults);
219
220 const static_arguments = allocator.dupe(choir_abi.ScalarArgument, self.static_arguments) catch return error.OutOfMemory;
221 errdefer if (static_arguments.len != 0) allocator.free(static_arguments);
222
223 return .{
224 .compile = compile,
225 .kernel_id = self.kernel_id,
226 .work_item_id = self.work_item_id,
227 .output_slot_id = self.output_slot_id,
228 .input_slot_ids = input_slot_ids,
229 .output_layout_fingerprint = self.output_layout_fingerprint,
230 .input_layout_fingerprint = self.input_layout_fingerprint,
231 .element_count = self.element_count,
232 .op_count = self.op_count,
233 .resources = self.resources,
234 .artifact = artifact,
235 .launch_resources = self.launch_resources,
236 .kernel_call_launch = self.kernel_call_launch,
237 .element_count_argument = self.element_count_argument,
238 .element_count_argument_value = self.element_count_argument_value,
239 .runtime_scalar_argument_count = self.runtime_scalar_argument_count,
240 .runtime_scalar_defaults = runtime_scalar_defaults,
241 .static_arguments = static_arguments,
242 .output_fill_pattern = self.output_fill_pattern,
243 .scratch_fill_pattern = self.scratch_fill_pattern,
244 };
245 }
246 };
247
248 pub const PlannedKernelSource = enum {
249 tensor,
250 kernel_call,
251 choir_kernel,
252 };
253
254 pub const PlannedSlot = struct {
255 slot_id: usize,
256 role: bufferization.BufferRole,
257 dtype: choir_abi.DType,
258 memory_space: memory_space.MemorySpace,
259 memory_access: memory_space.MemoryAccess,
260 boundary_transfer: memory_space.BoundaryTransfer,
261 layout_kind: layout_planning.LayoutKind,
262 dims: []i64,
263 element_strides: ?[]u64,
264 minor_to_major: []usize,
265 element_count: ?u64,
266 byte_size: ?u64,
267 alignment: u64,
268 contiguous: bool,
269 static_layout: bool,
270 layout_fingerprint: u64,
271 constant_payload: []u8,
272
273 fn init(
274 allocator: std.mem.Allocator,
275 slot: bufferization.BufferSlot,
276 memory_assignment: memory_space.MemorySpaceAssignment,
277 layout_assignment: layout_planning.LayoutAssignment,
278 ) !PlannedSlot {
279 if (slot.id != memory_assignment.slot_id or slot.id != layout_assignment.slot_id) {
280 return error.InvalidArtifact;
281 }
282 const dims = try allocator.dupe(i64, slot.dims);
283 errdefer allocator.free(dims);
284
285 var element_strides: ?[]u64 = null;
286 if (layout_assignment.element_strides) |strides| {
287 element_strides = try allocator.dupe(u64, strides);
288 errdefer if (element_strides) |owned| allocator.free(owned);
289 }
290
291 const minor_to_major = try allocator.dupe(usize, layout_assignment.minor_to_major);
292 errdefer allocator.free(minor_to_major);
293
294 const constant_payload = try constantPayloadForSlot(allocator, slot);
295 errdefer if (constant_payload.len != 0) allocator.free(constant_payload);
296
297 const layout_fingerprint = layoutFingerprint(
298 memory_assignment.space,
299 layout_assignment.kind,
300 dims,
301 element_strides,
302 minor_to_major,
303 layout_assignment.alignment,
304 layout_assignment.contiguous,
305 layout_assignment.static_layout,
306 );
307
308 return .{
309 .slot_id = slot.id,
310 .role = slot.role,
311 .dtype = slot.dtype,
312 .memory_space = memory_assignment.space,
313 .memory_access = memory_assignment.access,
314 .boundary_transfer = memory_assignment.transfer,
315 .layout_kind = layout_assignment.kind,
316 .dims = dims,
317 .element_strides = element_strides,
318 .minor_to_major = minor_to_major,
319 .element_count = slot.element_count,
320 .byte_size = slot.byte_size,
321 .alignment = layout_assignment.alignment,
322 .contiguous = layout_assignment.contiguous,
323 .static_layout = layout_assignment.static_layout,
324 .layout_fingerprint = layout_fingerprint,
325 .constant_payload = constant_payload,
326 };
327 }
328
329 fn deinit(self: *PlannedSlot, allocator: std.mem.Allocator) void {
330 allocator.free(self.dims);
331 if (self.element_strides) |strides| allocator.free(strides);
332 allocator.free(self.minor_to_major);
333 if (self.constant_payload.len != 0) allocator.free(self.constant_payload);
334 self.* = undefined;
335 }
336
337 fn copy(self: PlannedSlot, allocator: std.mem.Allocator) gpu.BackendError!PlannedSlot {
338 const dims = allocator.dupe(i64, self.dims) catch return error.OutOfMemory;
339 errdefer allocator.free(dims);
340
341 var element_strides: ?[]u64 = null;
342 if (self.element_strides) |strides| {
343 element_strides = allocator.dupe(u64, strides) catch return error.OutOfMemory;
344 errdefer if (element_strides) |owned| allocator.free(owned);
345 }
346
347 const minor_to_major = allocator.dupe(usize, self.minor_to_major) catch return error.OutOfMemory;
348 errdefer allocator.free(minor_to_major);
349
350 const constant_payload = allocator.dupe(u8, self.constant_payload) catch return error.OutOfMemory;
351 errdefer if (constant_payload.len != 0) allocator.free(constant_payload);
352
353 return .{
354 .slot_id = self.slot_id,
355 .role = self.role,
356 .dtype = self.dtype,
357 .memory_space = self.memory_space,
358 .memory_access = self.memory_access,
359 .boundary_transfer = self.boundary_transfer,
360 .layout_kind = self.layout_kind,
361 .dims = dims,
362 .element_strides = element_strides,
363 .minor_to_major = minor_to_major,
364 .element_count = self.element_count,
365 .byte_size = self.byte_size,
366 .alignment = self.alignment,
367 .contiguous = self.contiguous,
368 .static_layout = self.static_layout,
369 .layout_fingerprint = self.layout_fingerprint,
370 .constant_payload = constant_payload,
371 };
372 }
373
374 pub fn hasStaticSize(self: PlannedSlot) bool {
375 return self.byte_size != null;
376 }
377
378 pub fn constantBytes(self: PlannedSlot) ?[]const u8 {
379 if (!self.role.constant) return null;
380 return self.constant_payload;
381 }
382 };
383
384 pub const BackendArtifactPlan = struct {
385 allocator: std.mem.Allocator,
386 backend_kind: gpu.BackendKind,
387 format: gpu.ArtifactFormat,
388 target_profile: target_profile.BackendTargetProfile,
389 slots: []PlannedSlot,
390 input_slot_ids: []usize,
391 output_slot_ids: []usize,
392 kernels: std.ArrayListUnmanaged(PlannedKernel),
393 total_kernel_ops: usize = 0,
394 total_static_bytes: u64 = 0,
395
396 pub fn init(
397 allocator: std.mem.Allocator,
398 profile: target_profile.BackendTargetProfile,
399 ) BackendArtifactPlan {
400 return .{
401 .allocator = allocator,
402 .backend_kind = profile.backend_kind,
403 .format = profile.artifact_format,
404 .target_profile = profile,
405 .slots = &.{},
406 .input_slot_ids = &.{},
407 .output_slot_ids = &.{},
408 .kernels = .empty,
409 };
410 }
411
412 pub fn deinit(self: *BackendArtifactPlan) void {
413 for (self.kernels.items) |*kernel| {
414 kernel.deinit(self.allocator);
415 }
416 self.kernels.deinit(self.allocator);
417 for (self.slots) |*slot| {
418 slot.deinit(self.allocator);
419 }
420 if (self.slots.len != 0) self.allocator.free(self.slots);
421 if (self.input_slot_ids.len != 0) self.allocator.free(self.input_slot_ids);
422 if (self.output_slot_ids.len != 0) self.allocator.free(self.output_slot_ids);
423 self.* = undefined;
424 }
425
426 pub fn copy(self: *const BackendArtifactPlan, allocator: std.mem.Allocator) gpu.BackendError!BackendArtifactPlan {
427 var copied = BackendArtifactPlan.init(allocator, self.target_profile);
428 errdefer copied.deinit();
429
430 const slots = allocator.alloc(PlannedSlot, self.slots.len) catch return error.OutOfMemory;
431 var slots_owned = true;
432 var initialized_slots: usize = 0;
433 errdefer {
434 if (slots_owned) {
435 for (slots[0..initialized_slots]) |*slot| slot.deinit(allocator);
436 if (slots.len != 0) allocator.free(slots);
437 }
438 }
439 for (self.slots, 0..) |slot, index| {
440 slots[index] = try slot.copy(allocator);
441 initialized_slots += 1;
442 }
443 copied.slots = slots;
444 slots_owned = false;
445
446 copied.input_slot_ids = allocator.dupe(usize, self.input_slot_ids) catch return error.OutOfMemory;
447
448 copied.output_slot_ids = allocator.dupe(usize, self.output_slot_ids) catch return error.OutOfMemory;
449
450 try copied.kernels.ensureTotalCapacity(allocator, self.kernels.items.len);
451 for (self.kernels.items) |kernel| {
452 var kernel_copy = try kernel.copy(allocator);
453 var kernel_owned = true;
454 errdefer if (kernel_owned) kernel_copy.deinit(allocator);
455 copied.kernels.appendAssumeCapacity(kernel_copy);
456 kernel_owned = false;
457 }
458
459 copied.total_kernel_ops = self.total_kernel_ops;
460 copied.total_static_bytes = self.total_static_bytes;
461 return copied;
462 }
463
464 pub fn slotCount(self: BackendArtifactPlan) usize {
465 return self.slots.len;
466 }
467
468 pub fn kernelCount(self: BackendArtifactPlan) usize {
469 return self.kernels.items.len;
470 }
471
472 pub fn slotById(self: *const BackendArtifactPlan, slot_id: usize) ?*const PlannedSlot {
473 if (slot_id < self.slots.len and self.slots[slot_id].slot_id == slot_id) {
474 return &self.slots[slot_id];
475 }
476 for (self.slots) |*slot| {
477 if (slot.slot_id == slot_id) return slot;
478 }
479 return null;
480 }
481
482 fn setSlots(
483 self: *BackendArtifactPlan,
484 buffer_plan: *const bufferization.BufferPlanAnalysis,
485 memory_plan: *const memory_space.MemorySpacePlanAnalysis,
486 layout_plan: *const layout_planning.LayoutPlanAnalysis,
487 choir_module: *ir.Operation,
488 ) !void {
489 std.debug.assert(self.slots.len == 0);
490 std.debug.assert(self.input_slot_ids.len == 0);
491 std.debug.assert(self.output_slot_ids.len == 0);
492
493 const slots = try self.allocator.alloc(PlannedSlot, buffer_plan.slots.items.len);
494 var initialized: usize = 0;
495 errdefer {
496 for (slots[0..initialized]) |*slot| {
497 slot.deinit(self.allocator);
498 }
499 if (slots.len != 0) self.allocator.free(slots);
500 }
501 for (buffer_plan.slots.items, 0..) |slot, index| {
502 const memory_assignment = memory_plan.getAssignmentForSlot(slot.id) orelse {
503 return error.MissingMemorySpaceAssignment;
504 };
505 const layout_assignment = layout_plan.getAssignmentForSlot(slot.id) orelse {
506 return error.MissingLayoutAssignment;
507 };
508 slots[index] = try PlannedSlot.init(
509 self.allocator,
510 slot,
511 memory_assignment.*,
512 layout_assignment.*,
513 );
514 initialized += 1;
515 }
516
517 const input_slot_ids = try collectInputSlotIds(self.allocator, buffer_plan);
518 errdefer if (input_slot_ids.len != 0) self.allocator.free(input_slot_ids);
519
520 const output_slot_ids = try collectOutputSlotIds(self.allocator, choir_module, buffer_plan);
521 errdefer if (output_slot_ids.len != 0) self.allocator.free(output_slot_ids);
522
523 self.slots = slots;
524 self.input_slot_ids = input_slot_ids;
525 self.output_slot_ids = output_slot_ids;
526 }
527
528 fn addKernel(
529 self: *BackendArtifactPlan,
530 outline: kernelization_model.KernelOutline,
531 resources: schedule_planning.ScheduleResourceEstimate,
532 legal: backend_legalization.BackendKernelLegalization,
533 artifact: gpu.KernelArtifact,
534 launch_resources: LaunchResourcePlan,
535 kernel_call_launch: ?KernelCallLaunch,
536 element_count_argument: ElementCountArgument,
537 element_count_argument_value: u64,
538 runtime_scalar_argument_count: u32,
539 runtime_scalar_defaults: []const choir_abi.ScalarArgument,
540 static_arguments: []const choir_abi.ScalarArgument,
541 compile: PlannedKernelCompileContract,
542 output_fill_pattern: ?u32,
543 scratch_fill_pattern: ?u32,
544 ) !void {
545 var owned_artifact = artifact;
546 errdefer owned_artifact.deinit();
547 var owned_compile = compile;
548 errdefer owned_compile.deinit(self.allocator);
549
550 const input_slot_ids = try self.allocator.dupe(usize, outline.input_slot_ids);
551 errdefer self.allocator.free(input_slot_ids);
552 const owned_runtime_scalar_defaults = try self.allocator.dupe(choir_abi.ScalarArgument, runtime_scalar_defaults);
553 errdefer if (owned_runtime_scalar_defaults.len != 0) self.allocator.free(owned_runtime_scalar_defaults);
554 const owned_static_arguments = try self.allocator.dupe(choir_abi.ScalarArgument, static_arguments);
555 errdefer if (owned_static_arguments.len != 0) self.allocator.free(owned_static_arguments);
556 const output_layout_fingerprint = try self.layoutFingerprintForSlot(outline.output_slot_id);
557 const input_layout_fingerprint = try self.layoutFingerprintForInputs(input_slot_ids);
558
559 try self.kernels.append(self.allocator, .{
560 .compile = owned_compile,
561 .kernel_id = outline.id,
562 .work_item_id = outline.work_item_id,
563 .output_slot_id = outline.output_slot_id,
564 .input_slot_ids = input_slot_ids,
565 .output_layout_fingerprint = output_layout_fingerprint,
566 .input_layout_fingerprint = input_layout_fingerprint,
567 .element_count = outline.element_count,
568 .op_count = outline.op_count,
569 .resources = resources,
570 .artifact = owned_artifact,
571 .launch_resources = launch_resources,
572 .kernel_call_launch = kernel_call_launch,
573 .element_count_argument = element_count_argument,
574 .element_count_argument_value = element_count_argument_value,
575 .runtime_scalar_argument_count = runtime_scalar_argument_count,
576 .runtime_scalar_defaults = owned_runtime_scalar_defaults,
577 .static_arguments = owned_static_arguments,
578 .output_fill_pattern = output_fill_pattern,
579 .scratch_fill_pattern = scratch_fill_pattern,
580 });
581 self.total_kernel_ops += outline.op_count;
582 self.total_static_bytes += legal.static_bytes;
583 }
584
585 pub fn addStandaloneKernel(
586 self: *BackendArtifactPlan,
587 artifact: gpu.KernelArtifact,
588 launch_resources: LaunchResourcePlan,
589 compile: PlannedKernelCompileContract,
590 options: StandaloneKernelOptions,
591 ) !void {
592 var owned_artifact = artifact;
593 errdefer owned_artifact.deinit();
594 var owned_compile = compile;
595 errdefer owned_compile.deinit(self.allocator);
596
597 const input_slot_ids = try self.allocator.alloc(usize, 0);
598 errdefer self.allocator.free(input_slot_ids);
599 const owned_static_arguments = try self.allocator.dupe(choir_abi.ScalarArgument, options.static_arguments);
600 errdefer if (owned_static_arguments.len != 0) self.allocator.free(owned_static_arguments);
601
602 try self.kernels.append(self.allocator, .{
603 .compile = owned_compile,
604 .kernel_id = 0,
605 .work_item_id = 0,
606 .output_slot_id = 0,
607 .input_slot_ids = input_slot_ids,
608 .output_layout_fingerprint = 0,
609 .input_layout_fingerprint = 0,
610 .element_count = launch_resources.element_count,
611 .op_count = 0,
612 .resources = .{
613 .element_count = launch_resources.element_count,
614 .element_size = 1,
615 .op_count = 0,
616 },
617 .artifact = owned_artifact,
618 .launch_resources = launch_resources,
619 .kernel_call_launch = null,
620 .element_count_argument = .none,
621 .element_count_argument_value = 0,
622 .runtime_scalar_argument_count = options.runtime_scalar_argument_count,
623 .runtime_scalar_defaults = &.{},
624 .static_arguments = owned_static_arguments,
625 });
626 }
627
628 fn layoutFingerprintForSlot(
629 self: *const BackendArtifactPlan,
630 slot_id: usize,
631 ) !u64 {
632 const slot = self.slotById(slot_id) orelse return error.InvalidArtifact;
633 return slot.layout_fingerprint;
634 }
635
636 fn layoutFingerprintForInputs(
637 self: *const BackendArtifactPlan,
638 input_slot_ids: []const usize,
639 ) !u64 {
640 var builder = choir.product.incremental.FingerprintBuilder{};
641 builder.updateBytes("accy.artifact.input_layouts");
642 builder.updateUsize(input_slot_ids.len);
643 for (input_slot_ids) |slot_id| {
644 const slot = self.slotById(slot_id) orelse return error.InvalidArtifact;
645 builder.updateU64(slot.layout_fingerprint);
646 }
647 return builder.finish();
648 }
649 };
650
651 test "artifact plan copy owns standalone kernel data" {
652 const allocator = std.testing.allocator;
653
654 var artifact_plan = BackendArtifactPlan.init(allocator, .{
655 .backend_kind = .cuda,
656 .artifact_format = .cuda_ptx,
657 .dtype_bits = gpu.DTypeSet.init(&.{.f32}).bits,
658 });
659 var plan_owned = true;
660 defer if (plan_owned) artifact_plan.deinit();
661
662 var artifact = try gpu.KernelArtifact.init(allocator, .{
663 .backend = .cuda,
664 .format = .cuda_ptx,
665 .entry_name = "kernel0",
666 .argument_count = 0,
667 });
668 var artifact_owned = true;
669 errdefer if (artifact_owned) artifact.deinit();
670 try artifact.setOwnedText("payload");
671
672 var compile = try PlannedKernelCompileContract.init(
673 allocator,
674 .choir_kernel,
675 .authored,
676 .cuda_ptx,
677 "kernel0",
678 0,
679 gpu.DTypeSet.init(&.{.f32}),
680 .{},
681 .{},
682 null,
683 .{ .text = "payload" },
684 );
685 var compile_owned = true;
686 errdefer if (compile_owned) compile.deinit(allocator);
687
688 try artifact_plan.addStandaloneKernel(
689 artifact,
690 .{
691 .format = .cuda_ptx,
692 .element_count = 1,
693 .geometry = .{
694 .grid = .{ 1, 1, 1 },
695 .threadgroup = .{ 1, 1, 1 },
696 },
697 },
698 compile,
699 .{},
700 );
701 artifact_owned = false;
702 compile_owned = false;
703
704 var copied = try artifact_plan.copy(allocator);
705 defer copied.deinit();
706
707 artifact_plan.deinit();
708 plan_owned = false;
709
710 try std.testing.expectEqual(@as(usize, 1), copied.kernelCount());
711 const kernel = copied.kernels.items[0];
712 try std.testing.expectEqualStrings("kernel0", kernel.artifact.entry_name);
713 try std.testing.expectEqual(gpu.PayloadOwnership.owned, kernel.artifact.payload_ownership);
714 switch (kernel.artifact.payload) {
715 .text => |text| try std.testing.expectEqualStrings("payload", text),
716 else => return error.InvalidArtifact,
717 }
718 try std.testing.expectEqualStrings("kernel0", kernel.compile.entry_name);
719 try std.testing.expectEqual(PlannedKernelCompilePayload.text, kernel.compile.payload);
720 }
721
722 pub const BackendSlotLifetime = struct {
723 slot_id: usize,
724 role: bufferization.BufferRole,
725 dtype: choir_abi.DType,
726 memory_space: memory_space.MemorySpace,
727 layout_kind: layout_planning.LayoutKind,
728 element_count: ?u64,
729 byte_size: ?u64,
730 alignment: u64,
731 layout_fingerprint: u64,
732 first_kernel_index: usize = std.math.maxInt(usize),
733 last_kernel_index: usize = 0,
734
735 pub fn hasKernelUse(self: BackendSlotLifetime) bool {
736 return self.first_kernel_index != std.math.maxInt(usize);
737 }
738
739 pub fn isBackendAllocatable(self: BackendSlotLifetime) bool {
740 return self.hasKernelUse() and !self.role.input and !self.role.constant and self.byte_size != null;
741 }
742 };
743
744 pub const BackendBufferAllocationPlan = struct {
745 allocation_id: usize,
746 byte_size: u64,
747 dtype: choir_abi.DType,
748 memory_space: memory_space.MemorySpace,
749 layout_kind: layout_planning.LayoutKind,
750 element_count: ?u64,
751 alignment: u64,
752 layout_fingerprint: u64,
753 last_kernel_index: usize,
754 };
755
756 pub const BackendBufferReuseAssignment = struct {
757 slot_id: usize,
758 allocation_id: usize,
759 };
760
761 pub const BackendMemoryPlan = struct {
762 allocator: std.mem.Allocator,
763 lifetimes: []BackendSlotLifetime,
764 allocations: []BackendBufferAllocationPlan,
765 assignments: []BackendBufferReuseAssignment,
766 total_static_slot_bytes: u64 = 0,
767 allocated_static_bytes: u64 = 0,
768 peak_static_live_bytes: u64 = 0,
769
770 pub fn deinit(self: *BackendMemoryPlan) void {
771 if (self.lifetimes.len != 0) self.allocator.free(self.lifetimes);
772 if (self.allocations.len != 0) self.allocator.free(self.allocations);
773 if (self.assignments.len != 0) self.allocator.free(self.assignments);
774 self.* = undefined;
775 }
776
777 pub fn assignmentForSlot(
778 self: BackendMemoryPlan,
779 slot_id: usize,
780 ) ?BackendBufferReuseAssignment {
781 for (self.assignments) |assignment| {
782 if (assignment.slot_id == slot_id) return assignment;
783 }
784 return null;
785 }
786
787 pub fn allocationForId(
788 self: BackendMemoryPlan,
789 allocation_id: usize,
790 ) ?BackendBufferAllocationPlan {
791 for (self.allocations) |allocation| {
792 if (allocation.allocation_id == allocation_id) return allocation;
793 }
794 return null;
795 }
796 };
797
798 pub const LaunchResourceClass = enum {
799 unknown,
800 memory_bound,
801 balanced,
802 compute_weighted,
803 };
804
805 pub const LaunchTileKind = enum {
806 none,
807 dot_general,
808 reduction,
809 elementwise_rank2,
810 };
811
812 pub const LaunchReductionKind = enum {
813 none,
814 sum,
815 max,
816 min,
817 };
818
819 pub const LaunchTilePlan = struct {
820 kind: LaunchTileKind = .none,
821 m: u32 = 0,
822 n: u32 = 0,
823 k: u32 = 0,
824 batch: u32 = 1,
825 input_dtype: ?choir_abi.DType = null,
826 output_dtype: ?choir_abi.DType = null,
827 input_tile_bytes: u32 = 0,
828 output_tile_bytes: u32 = 0,
829 scratch_memory_bytes: u32 = 0,
830 reduction_kind: LaunchReductionKind = .none,
831 reduction_rank: u32 = 0,
832 reduction_axis: u32 = 0,
833 reduction_extent: u32 = 0,
834
835 pub fn active(self: LaunchTilePlan) bool {
836 return self.kind != .none;
837 }
838 };
839
840 pub const max_launch_resource_candidates = 8;
841
842 pub const LaunchResourceCandidate = struct {
843 geometry: choir_abi.LaunchGeometry = .{},
844 score: u32 = std.math.maxInt(u32),
845 estimated_static_bytes_per_threadgroup: u64 = 0,
846 estimated_element_ops_per_threadgroup: u64 = 0,
847 tile: LaunchTilePlan = .{},
848 };
849
850 pub const LaunchResourcePlan = struct {
851 format: gpu.ArtifactFormat,
852 element_count: u64,
853 geometry: choir_abi.LaunchGeometry,
854 subgroup_size: ?u32 = null,
855 subgroup_aligned: bool = false,
856 fixed_threadgroup: bool = false,
857 resource_class: LaunchResourceClass = .unknown,
858 element_ops_per_kib: u64 = 0,
859 estimated_static_bytes_per_threadgroup: u64 = 0,
860 estimated_element_ops_per_threadgroup: u64 = 0,
861 static_bytes_complete: bool = false,
862 tile: LaunchTilePlan = .{},
863 candidate_count: usize = 0,
864 candidates: [max_launch_resource_candidates]LaunchResourceCandidate = @as([max_launch_resource_candidates]LaunchResourceCandidate, @splat(.{})),
865
866 pub fn selectedCandidate(self: LaunchResourcePlan) ?LaunchResourceCandidate {
867 if (self.candidate_count == 0) return null;
868 return self.candidates[0];
869 }
870 };
871
872 pub fn createLaunchResourcePlan(
873 caps: gpu.BackendCapabilities,
874 format: gpu.ArtifactFormat,
875 resources: schedule_planning.ScheduleResourceEstimate,
876 ) gpu.BackendError!LaunchResourcePlan {
877 const subgroup_size = subgroupSizeFor(caps);
878 const fixed_threadgroup = fixedThreadgroupSizeFor(format);
879 const resource_class = classifyLaunchResources(resources);
880 const candidates = if (fixed_threadgroup) |fixed|
881 try fixedLaunchResourceCandidates(caps, format, resources, resource_class, fixed)
882 else
883 try dynamicLaunchResourceCandidates(caps, format, resources, resource_class, subgroup_size);
884 if (candidates.count == 0) return error.CapabilityMismatch;
885 const selected = candidates.items[0];
886 const plan = LaunchResourcePlan{
887 .format = format,
888 .element_count = resources.element_count,
889 .geometry = selected.geometry,
890 .subgroup_size = subgroup_size,
891 .subgroup_aligned = if (subgroup_size) |size| blk: {
892 const thread_count = try threadgroupThreadCount(selected.geometry);
893 break :blk thread_count >= size and thread_count % size == 0;
894 } else false,
895 .fixed_threadgroup = fixed_threadgroup != null,
896 .resource_class = resource_class,
897 .element_ops_per_kib = resources.elementOpsPerKiB(),
898 .estimated_static_bytes_per_threadgroup = selected.estimated_static_bytes_per_threadgroup,
899 .estimated_element_ops_per_threadgroup = selected.estimated_element_ops_per_threadgroup,
900 .static_bytes_complete = resources.static_bytes_complete,
901 .candidate_count = candidates.count,
902 .candidates = candidates.items,
903 };
904 try caps.validateLaunchGeometry(plan.geometry);
905 return plan;
906 }
907
908 fn createLaunchResourcePlanForWork(
909 caps: gpu.BackendCapabilities,
910 format: gpu.ArtifactFormat,
911 work: schedule_planning.ScheduleWorkItem,
912 compile_plan: CompilePlan,
913 ) gpu.BackendError!LaunchResourcePlan {
914 const plan = switch (compile_plan.launch) {
915 .dot_general => |dot| try createDotGeneralLaunchResourcePlan(caps, format, work.resources, dot, loweredKernelBody(compile_plan)),
916 .reduction => |reduction| blk: {
917 var resources = work.resources;
918 resources.element_count = work.element_count;
919 break :blk try createReductionLaunchResourcePlan(caps, format, resources, reduction, loweredKernelBody(compile_plan));
920 },
921 .kernel_call => |kernel_call| return createKernelCallLaunchResourcePlan(caps, format, work.resources, kernel_call),
922 .generic => blk: {
923 switch (loweredKernelBody(compile_plan)) {
924 .row_pipeline => |pipeline| break :blk try createFixedLaunchResourcePlan(caps, format, work.resources, .{
925 .grid = .{ pipeline.rows, 1, 1 },
926 .threadgroup = .{ pipeline.threads, 1, 1 },
927 }),
928 .flash_attention => |flash| break :blk try createFixedLaunchResourcePlan(caps, format, work.resources, .{
929 .grid = .{ flash.seq / flash.br, 1, 1 },
930 .threadgroup = .{ flash.threads_x, flash.threads_y, 1 },
931 }),
932 .scan => |scan_plan| break :blk try createFixedLaunchResourcePlan(caps, format, work.resources, .{
933 .grid = .{ scan_plan.blocks, 1, 1 },
934 .threadgroup = .{ scan_plan.threads, 1, 1 },
935 }),
936 .elementwise_rank2 => |rank2| break :blk try createElementwiseRank2LaunchResourcePlan(caps, format, work.resources, rank2),
937 .elementwise_vector => |vector_plan| {
938 var resources = work.resources;
939 resources.element_count = vector_plan.quads;
940 break :blk try createLaunchResourcePlan(caps, format, resources);
941 },
942 else => {},
943 }
944 if (fixedThreadgroupSizeFor(format) != null) {
945 if (compile_plan.lowered_kernel) |lowered| {
946 if (lowered.launchGeometry()) |geometry| {
947 break :blk try createFixedLaunchResourcePlan(caps, format, work.resources, geometry);
948 }
949 }
950 }
951 break :blk try createLaunchResourcePlan(caps, format, work.resources);
952 },
953 };
954 return attachGeneratedDynamicSharedMemory(caps, plan, compile_plan);
955 }
956
957 fn createKernelCallLaunchResourcePlan(
958 caps: gpu.BackendCapabilities,
959 format: gpu.ArtifactFormat,
960 resources: schedule_planning.ScheduleResourceEstimate,
961 launch: KernelCallLaunch,
962 ) gpu.BackendError!LaunchResourcePlan {
963 return switch (launch) {
964 .derived => |derived| blk: {
965 try validateKernelCallDerivedLaunch(derived);
966 break :blk createFixedLaunchResourcePlan(caps, format, resources, .{
967 .threadgroup = derived.threadgroup,
968 .dynamic_shared_memory_bytes = derived.dynamic_shared_memory_bytes,
969 });
970 },
971 .fixed => |geometry| createFixedLaunchResourcePlan(caps, format, resources, geometry),
972 };
973 }
974
975 fn createFixedLaunchResourcePlan(
976 caps: gpu.BackendCapabilities,
977 format: gpu.ArtifactFormat,
978 resources: schedule_planning.ScheduleResourceEstimate,
979 geometry: choir_abi.LaunchGeometry,
980 ) gpu.BackendError!LaunchResourcePlan {
981 try caps.validateLaunchGeometry(geometry);
982 const thread_count = try threadgroupThreadCount(geometry);
983 const subgroup_size = subgroupSizeFor(caps);
984 const resource_class = classifyLaunchResources(resources);
985 const candidate = LaunchResourceCandidate{
986 .geometry = geometry,
987 .score = 0,
988 .estimated_static_bytes_per_threadgroup = scaleResourceToThreadgroup(
989 resources.static_total_bytes,
990 resources.element_count,
991 thread_count,
992 ),
993 .estimated_element_ops_per_threadgroup = scaleResourceToThreadgroup(
994 resources.estimated_element_ops,
995 resources.element_count,
996 thread_count,
997 ),
998 };
999 var candidates = @as([max_launch_resource_candidates]LaunchResourceCandidate, @splat(.{}));
1000 candidates[0] = candidate;
1001 return .{
1002 .format = format,
1003 .element_count = resources.element_count,
1004 .geometry = geometry,
1005 .subgroup_size = subgroup_size,
1006 .subgroup_aligned = if (subgroup_size) |size| thread_count >= size and thread_count % size == 0 else false,
1007 .fixed_threadgroup = true,
1008 .resource_class = resource_class,
1009 .element_ops_per_kib = resources.elementOpsPerKiB(),
1010 .estimated_static_bytes_per_threadgroup = candidate.estimated_static_bytes_per_threadgroup,
1011 .estimated_element_ops_per_threadgroup = candidate.estimated_element_ops_per_threadgroup,
1012 .static_bytes_complete = resources.static_bytes_complete,
1013 .candidate_count = 1,
1014 .candidates = candidates,
1015 };
1016 }
1017
1018 fn loweredKernelBody(compile_plan: CompilePlan) kernelization.LoweredKernelBody {
1019 const lowered = compile_plan.lowered_kernel orelse return .generic;
1020 return lowered.body;
1021 }
1022
1023 fn attachGeneratedDynamicSharedMemory(
1024 caps: gpu.BackendCapabilities,
1025 plan_value: LaunchResourcePlan,
1026 compile_plan: CompilePlan,
1027 ) gpu.BackendError!LaunchResourcePlan {
1028 const bytes = generatedDynamicSharedMemoryBytes(compile_plan);
1029 if (bytes == 0) return plan_value;
1030
1031 var plan = plan_value;
1032 plan.geometry.dynamic_shared_memory_bytes = bytes;
1033 try caps.validateLaunchGeometry(plan.geometry);
1034 var candidate_index: usize = 0;
1035 while (candidate_index < plan.candidate_count) : (candidate_index += 1) {
1036 plan.candidates[candidate_index].geometry.dynamic_shared_memory_bytes = bytes;
1037 try caps.validateLaunchGeometry(plan.candidates[candidate_index].geometry);
1038 }
1039 return plan;
1040 }
1041
1042 fn generatedDynamicSharedMemoryBytes(compile_plan: CompilePlan) u32 {
1043 const lowered = compile_plan.lowered_kernel orelse return 0;
1044 return lowered.dynamic_shared_memory_bytes;
1045 }
1046
1047 fn createElementwiseRank2LaunchResourcePlan(
1048 caps: gpu.BackendCapabilities,
1049 format: gpu.ArtifactFormat,
1050 resources: schedule_planning.ScheduleResourceEstimate,
1051 rank2: kernelization.product.ElementwiseRank2Plan,
1052 ) gpu.BackendError!LaunchResourcePlan {
1053 const default_geometry = choir_abi.LaunchGeometry{
1054 .grid = .{
1055 try gridSizeFor(rank2.cols, rank2.threads_x),
1056 try gridSizeFor(rank2.rows, rank2.threads_y),
1057 1,
1058 },
1059 .threadgroup = .{ rank2.threads_x, rank2.threads_y, 1 },
1060 };
1061 if (fixedThreadgroupSizeFor(format) != null) {
1062 var plan = try createFixedLaunchResourcePlan(caps, format, resources, default_geometry);
1063 plan.tile = elementwiseRank2TilePlan(rank2);
1064 plan.candidates[0].tile = plan.tile;
1065 return plan;
1066 }
1067
1068 const resource_class = classifyLaunchResources(resources);
1069 const subgroup_size = subgroupSizeFor(caps);
1070 const target_thread_count = try threadgroupThreadCount(default_geometry);
1071 var candidates = LaunchCandidateBuffer{};
1072 try appendElementwiseRank2LaunchCandidate(&candidates, caps, resources, rank2, rank2.threads_x, rank2.threads_y, target_thread_count);
1073 for (elementwise_rank2_launch_shapes) |shape| {
1074 try appendElementwiseRank2LaunchCandidate(&candidates, caps, resources, rank2, shape.x, shape.y, target_thread_count);
1075 }
1076 if (candidates.count == 0) return error.CapabilityMismatch;
1077
1078 const selected = candidates.items[0];
1079 return .{
1080 .format = format,
1081 .element_count = resources.element_count,
1082 .geometry = selected.geometry,
1083 .subgroup_size = subgroup_size,
1084 .subgroup_aligned = if (subgroup_size) |size| blk: {
1085 const thread_count = try threadgroupThreadCount(selected.geometry);
1086 break :blk thread_count >= size and thread_count % size == 0;
1087 } else false,
1088 .fixed_threadgroup = false,
1089 .resource_class = resource_class,
1090 .element_ops_per_kib = resources.elementOpsPerKiB(),
1091 .estimated_static_bytes_per_threadgroup = selected.estimated_static_bytes_per_threadgroup,
1092 .estimated_element_ops_per_threadgroup = selected.estimated_element_ops_per_threadgroup,
1093 .static_bytes_complete = resources.static_bytes_complete,
1094 .tile = selected.tile,
1095 .candidate_count = candidates.count,
1096 .candidates = candidates.items,
1097 };
1098 }
1099
1100 const ElementwiseRank2LaunchShape = struct {
1101 x: u32,
1102 y: u32,
1103 };
1104
1105 const elementwise_rank2_launch_shapes = [_]ElementwiseRank2LaunchShape{
1106 .{ .x = 16, .y = 16 },
1107 .{ .x = 8, .y = 32 },
1108 .{ .x = 64, .y = 4 },
1109 .{ .x = 128, .y = 2 },
1110 .{ .x = 256, .y = 1 },
1111 .{ .x = 32, .y = 4 },
1112 .{ .x = 16, .y = 8 },
1113 };
1114
1115 fn appendElementwiseRank2LaunchCandidate(
1116 candidates: *LaunchCandidateBuffer,
1117 caps: gpu.BackendCapabilities,
1118 resources: schedule_planning.ScheduleResourceEstimate,
1119 rank2: kernelization.product.ElementwiseRank2Plan,
1120 threads_x: u32,
1121 threads_y: u32,
1122 target_thread_count: u32,
1123 ) gpu.BackendError!void {
1124 if (candidates.count >= max_launch_resource_candidates) return;
1125 if (threads_x == 0 or threads_y == 0) return;
1126 const geometry = choir_abi.LaunchGeometry{
1127 .grid = .{
1128 try gridSizeFor(rank2.cols, threads_x),
1129 try gridSizeFor(rank2.rows, threads_y),
1130 1,
1131 },
1132 .threadgroup = .{ threads_x, threads_y, 1 },
1133 };
1134 caps.validateLaunchGeometry(geometry) catch |err| switch (err) {
1135 error.CapabilityMismatch => return,
1136 else => return err,
1137 };
1138 for (candidates.items[0..candidates.count]) |existing| {
1139 if (sameLaunchGeometry(existing.geometry, geometry)) return;
1140 }
1141 const thread_count = try threadgroupThreadCount(geometry);
1142 candidates.items[candidates.count] = .{
1143 .geometry = geometry,
1144 .score = candidateScore(thread_count, target_thread_count),
1145 .estimated_static_bytes_per_threadgroup = scaleResourceToThreadgroup(
1146 resources.static_total_bytes,
1147 resources.element_count,
1148 thread_count,
1149 ),
1150 .estimated_element_ops_per_threadgroup = scaleResourceToThreadgroup(
1151 resources.estimated_element_ops,
1152 resources.element_count,
1153 thread_count,
1154 ),
1155 .tile = elementwiseRank2TilePlan(rank2),
1156 };
1157 candidates.count += 1;
1158 }
1159
1160 fn elementwiseRank2TilePlan(rank2: kernelization.product.ElementwiseRank2Plan) LaunchTilePlan {
1161 return .{
1162 .kind = .elementwise_rank2,
1163 .m = rank2.rows,
1164 .n = rank2.cols,
1165 };
1166 }
1167
1168 fn createDotGeneralLaunchResourcePlan(
1169 caps: gpu.BackendCapabilities,
1170 format: gpu.ArtifactFormat,
1171 resources: schedule_planning.ScheduleResourceEstimate,
1172 dot: DotGeneralLaunchPlan,
1173 body: kernelization.LoweredKernelBody,
1174 ) gpu.BackendError!LaunchResourcePlan {
1175 if (gpu.artifactFormatUsesHostLoopLaunch(format)) {
1176 return createDotGeneralHostLoopLaunchResourcePlan(caps, format, resources, dot);
1177 }
1178 switch (body) {
1179 .dot_block_tile => |block_tile| return createDotGeneralBlockTileLaunchResourcePlan(caps, format, resources, dot, block_tile),
1180 .dot_mma_tile => |mma_tile| return createDotGeneralMmaTileLaunchResourcePlan(caps, format, resources, dot, mma_tile),
1181 else => {},
1182 }
1183 var candidates = LaunchCandidateBuffer{};
1184 for (try dotGeneralTileShapes(format), 0..) |shape, index| {
1185 try appendDotGeneralLaunchCandidate(&candidates, caps, resources, dot, shape, @intCast(index));
1186 }
1187 if (candidates.count == 0) return error.CapabilityMismatch;
1188 const candidate = candidates.items[0];
1189 return .{
1190 .format = format,
1191 .element_count = resources.element_count,
1192 .geometry = candidate.geometry,
1193 .fixed_threadgroup = true,
1194 .resource_class = classifyLaunchResources(resources),
1195 .element_ops_per_kib = resources.elementOpsPerKiB(),
1196 .estimated_static_bytes_per_threadgroup = candidate.estimated_static_bytes_per_threadgroup,
1197 .estimated_element_ops_per_threadgroup = candidate.estimated_element_ops_per_threadgroup,
1198 .static_bytes_complete = resources.static_bytes_complete,
1199 .tile = candidate.tile,
1200 .candidate_count = candidates.count,
1201 .candidates = candidates.items,
1202 };
1203 }
1204
1205 fn createDotGeneralHostLoopLaunchResourcePlan(
1206 caps: gpu.BackendCapabilities,
1207 format: gpu.ArtifactFormat,
1208 resources: schedule_planning.ScheduleResourceEstimate,
1209 dot: DotGeneralLaunchPlan,
1210 ) gpu.BackendError!LaunchResourcePlan {
1211 const geometry = choir_abi.LaunchGeometry{
1212 .grid = .{ try gridSizeFor(resources.element_count, 1), 1, 1 },
1213 .threadgroup = .{ 1, 1, 1 },
1214 };
1215 try caps.validateLaunchGeometry(geometry);
1216 var candidates = LaunchCandidateBuffer{};
1217 candidates.items[0] = .{
1218 .geometry = geometry,
1219 .score = 0,
1220 .estimated_static_bytes_per_threadgroup = scaleResourceToThreadgroup(
1221 resources.static_total_bytes,
1222 resources.element_count,
1223 1,
1224 ),
1225 .estimated_element_ops_per_threadgroup = scaleResourceToThreadgroup(
1226 resources.estimated_element_ops,
1227 resources.element_count,
1228 1,
1229 ),
1230 .tile = try dotGeneralTilePlan(dot, .{ .m = 1, .n = 1 }),
1231 };
1232 candidates.count = 1;
1233 const candidate = candidates.items[0];
1234 return .{
1235 .format = format,
1236 .element_count = resources.element_count,
1237 .geometry = candidate.geometry,
1238 .fixed_threadgroup = true,
1239 .resource_class = classifyLaunchResources(resources),
1240 .element_ops_per_kib = resources.elementOpsPerKiB(),
1241 .estimated_static_bytes_per_threadgroup = candidate.estimated_static_bytes_per_threadgroup,
1242 .estimated_element_ops_per_threadgroup = candidate.estimated_element_ops_per_threadgroup,
1243 .static_bytes_complete = resources.static_bytes_complete,
1244 .tile = candidate.tile,
1245 .candidate_count = candidates.count,
1246 .candidates = candidates.items,
1247 };
1248 }
1249
1250 fn createDotGeneralBlockTileLaunchResourcePlan(
1251 caps: gpu.BackendCapabilities,
1252 format: gpu.ArtifactFormat,
1253 resources: schedule_planning.ScheduleResourceEstimate,
1254 dot: DotGeneralLaunchPlan,
1255 block_tile: kernelization.DotGeneralBlockTile,
1256 ) gpu.BackendError!LaunchResourcePlan {
1257 const geometry = choir_abi.LaunchGeometry{
1258 .grid = .{
1259 try gridSizeFor(dot.n, block_tile.bn),
1260 try gridSizeFor(dot.m, block_tile.bm),
1261 dot.batch * block_tile.splits,
1262 },
1263 .threadgroup = .{ block_tile.threadsX(), block_tile.threadsY(), 1 },
1264 };
1265 try caps.validateLaunchGeometry(geometry);
1266 const thread_count = try threadgroupThreadCount(geometry);
1267 var candidates = LaunchCandidateBuffer{};
1268 candidates.items[0] = .{
1269 .geometry = geometry,
1270 .score = 0,
1271 .estimated_static_bytes_per_threadgroup = scaleResourceToThreadgroup(
1272 resources.static_total_bytes,
1273 resources.element_count,
1274 thread_count,
1275 ),
1276 .estimated_element_ops_per_threadgroup = scaleResourceToThreadgroup(
1277 resources.estimated_element_ops,
1278 resources.element_count,
1279 thread_count,
1280 ),
1281 .tile = try dotGeneralTilePlan(dot, .{ .m = block_tile.bm, .n = block_tile.bn }),
1282 };
1283 candidates.count = 1;
1284 const candidate = candidates.items[0];
1285 return .{
1286 .format = format,
1287 .element_count = resources.element_count,
1288 .geometry = candidate.geometry,
1289 .fixed_threadgroup = true,
1290 .resource_class = classifyLaunchResources(resources),
1291 .element_ops_per_kib = resources.elementOpsPerKiB(),
1292 .estimated_static_bytes_per_threadgroup = candidate.estimated_static_bytes_per_threadgroup,
1293 .estimated_element_ops_per_threadgroup = candidate.estimated_element_ops_per_threadgroup,
1294 .static_bytes_complete = resources.static_bytes_complete,
1295 .tile = candidate.tile,
1296 .candidate_count = candidates.count,
1297 .candidates = candidates.items,
1298 };
1299 }
1300
1301 fn createDotGeneralMmaTileLaunchResourcePlan(
1302 caps: gpu.BackendCapabilities,
1303 format: gpu.ArtifactFormat,
1304 resources: schedule_planning.ScheduleResourceEstimate,
1305 dot: DotGeneralLaunchPlan,
1306 mma_tile: kernelization.DotGeneralMmaTile,
1307 ) gpu.BackendError!LaunchResourcePlan {
1308 const geometry = choir_abi.LaunchGeometry{
1309 .grid = .{
1310 try gridSizeFor(dot.n, mma_tile.bn),
1311 try gridSizeFor(dot.m, mma_tile.bm),
1312 dot.batch * mma_tile.splits,
1313 },
1314 .threadgroup = .{ 32, mma_tile.warps(), 1 },
1315 };
1316 try caps.validateLaunchGeometry(geometry);
1317 const thread_count = try threadgroupThreadCount(geometry);
1318 var candidates = LaunchCandidateBuffer{};
1319 candidates.items[0] = .{
1320 .geometry = geometry,
1321 .score = 0,
1322 .estimated_static_bytes_per_threadgroup = scaleResourceToThreadgroup(
1323 resources.static_total_bytes,
1324 resources.element_count,
1325 thread_count,
1326 ),
1327 .estimated_element_ops_per_threadgroup = scaleResourceToThreadgroup(
1328 resources.estimated_element_ops,
1329 resources.element_count,
1330 thread_count,
1331 ),
1332 .tile = try dotGeneralTilePlan(dot, .{ .m = mma_tile.bm, .n = mma_tile.bn }),
1333 };
1334 candidates.count = 1;
1335 const candidate = candidates.items[0];
1336 return .{
1337 .format = format,
1338 .element_count = resources.element_count,
1339 .geometry = candidate.geometry,
1340 .fixed_threadgroup = true,
1341 .resource_class = classifyLaunchResources(resources),
1342 .element_ops_per_kib = resources.elementOpsPerKiB(),
1343 .estimated_static_bytes_per_threadgroup = candidate.estimated_static_bytes_per_threadgroup,
1344 .estimated_element_ops_per_threadgroup = candidate.estimated_element_ops_per_threadgroup,
1345 .static_bytes_complete = resources.static_bytes_complete,
1346 .tile = candidate.tile,
1347 .candidate_count = candidates.count,
1348 .candidates = candidates.items,
1349 };
1350 }
1351
1352 fn dotGeneralTilePlan(
1353 dot: DotGeneralLaunchPlan,
1354 shape: DotGeneralTileShape,
1355 ) gpu.BackendError!LaunchTilePlan {
1356 return .{
1357 .kind = .dot_general,
1358 .m = shape.m,
1359 .n = shape.n,
1360 .k = dot.k,
1361 .batch = dot.batch,
1362 .input_dtype = dot.input_dtype,
1363 .output_dtype = dot.output_dtype,
1364 .input_tile_bytes = try dotGeneralInputTileBytes(shape.m, shape.n, dot.k, dot.input_dtype),
1365 .output_tile_bytes = try dotGeneralOutputTileBytes(shape.m, shape.n, dot.output_dtype),
1366 };
1367 }
1368
1369 const DotGeneralTileShape = struct {
1370 m: u32,
1371 n: u32,
1372 };
1373
1374 const cuda_dot_general_tile_shapes = [_]DotGeneralTileShape{
1375 .{ .m = 16, .n = 16 },
1376 .{ .m = 8, .n = 32 },
1377 .{ .m = 32, .n = 8 },
1378 .{ .m = 8, .n = 8 },
1379 };
1380
1381 const vulkan_dot_general_tile_shapes = [_]DotGeneralTileShape{
1382 .{ .m = 8, .n = 8 },
1383 .{ .m = 8, .n = 16 },
1384 .{ .m = 16, .n = 8 },
1385 .{ .m = 16, .n = 16 },
1386 };
1387
1388 const metal_dot_general_tile_shapes = [_]DotGeneralTileShape{
1389 .{ .m = 8, .n = 8 },
1390 .{ .m = 8, .n = 16 },
1391 .{ .m = 16, .n = 8 },
1392 .{ .m = 16, .n = 16 },
1393 };
1394
1395 fn dotGeneralTileShapes(format: gpu.ArtifactFormat) gpu.BackendError![]const DotGeneralTileShape {
1396 return switch (format) {
1397 .cuda_ptx => &cuda_dot_general_tile_shapes,
1398 .vulkan_spirv => &vulkan_dot_general_tile_shapes,
1399 .metal_msl => &metal_dot_general_tile_shapes,
1400 else => error.UnsupportedArtifactFormat,
1401 };
1402 }
1403
1404 fn appendDotGeneralLaunchCandidate(
1405 candidates: *LaunchCandidateBuffer,
1406 caps: gpu.BackendCapabilities,
1407 resources: schedule_planning.ScheduleResourceEstimate,
1408 dot: DotGeneralLaunchPlan,
1409 shape: DotGeneralTileShape,
1410 score: u32,
1411 ) gpu.BackendError!void {
1412 if (candidates.count >= max_launch_resource_candidates) return;
1413 const geometry = choir_abi.LaunchGeometry{
1414 .grid = .{ try gridSizeFor(dot.n, shape.n), try gridSizeFor(dot.m, shape.m), dot.batch },
1415 .threadgroup = .{ shape.n, shape.m, 1 },
1416 };
1417 caps.validateLaunchGeometry(geometry) catch |err| switch (err) {
1418 error.CapabilityMismatch => return,
1419 else => return err,
1420 };
1421 for (candidates.items[0..candidates.count]) |existing| {
1422 if (sameLaunchGeometry(existing.geometry, geometry)) return;
1423 }
1424 const thread_count = try threadgroupThreadCount(geometry);
1425 candidates.items[candidates.count] = .{
1426 .geometry = geometry,
1427 .score = score,
1428 .estimated_static_bytes_per_threadgroup = scaleResourceToThreadgroup(
1429 resources.static_total_bytes,
1430 resources.element_count,
1431 thread_count,
1432 ),
1433 .estimated_element_ops_per_threadgroup = scaleResourceToThreadgroup(
1434 resources.estimated_element_ops,
1435 resources.element_count,
1436 thread_count,
1437 ),
1438 .tile = try dotGeneralTilePlan(dot, shape),
1439 };
1440 candidates.count += 1;
1441 }
1442
1443 fn dotGeneralInputTileBytes(
1444 m: u32,
1445 n: u32,
1446 k: u32,
1447 dtype: choir_abi.DType,
1448 ) gpu.BackendError!u32 {
1449 const lhs_elements = std.math.mul(u32, m, k) catch return error.CapabilityMismatch;
1450 const rhs_elements = std.math.mul(u32, k, n) catch return error.CapabilityMismatch;
1451 const elements = std.math.add(u32, lhs_elements, rhs_elements) catch return error.CapabilityMismatch;
1452 return std.math.mul(u32, elements, @as(u32, dtype.sizeOf())) catch return error.CapabilityMismatch;
1453 }
1454
1455 fn dotGeneralOutputTileBytes(
1456 m: u32,
1457 n: u32,
1458 dtype: choir_abi.DType,
1459 ) gpu.BackendError!u32 {
1460 const elements = std.math.mul(u32, m, n) catch return error.CapabilityMismatch;
1461 return std.math.mul(u32, elements, @as(u32, dtype.sizeOf())) catch return error.CapabilityMismatch;
1462 }
1463
1464 fn createReductionLaunchResourcePlan(
1465 caps: gpu.BackendCapabilities,
1466 format: gpu.ArtifactFormat,
1467 resources: schedule_planning.ScheduleResourceEstimate,
1468 reduction: ReductionLaunchPlan,
1469 body: kernelization.LoweredKernelBody,
1470 ) gpu.BackendError!LaunchResourcePlan {
1471 try validateReductionLaunchPlan(resources, reduction);
1472 switch (body) {
1473 .reduction_atomic => |atomic_plan| return createFixedGeometryReductionLaunchResourcePlan(caps, format, resources, reduction, .{
1474 .grid = .{ atomic_plan.blocks, 1, 1 },
1475 .threadgroup = .{ atomic_plan.threads, 1, 1 },
1476 }),
1477 .reduction_single_block => |threads| return createFixedGeometryReductionLaunchResourcePlan(caps, format, resources, reduction, .{
1478 .grid = .{ 1, 1, 1 },
1479 .threadgroup = .{ threads, 1, 1 },
1480 }),
1481 .reduction_warp_rows => |warp_plan| {
1482 const lanes = std.math.mul(u32, warp_plan.rows, 32) catch return error.CapabilityMismatch;
1483 return createFixedGeometryReductionLaunchResourcePlan(caps, format, resources, reduction, .{
1484 .grid = .{ (lanes + warp_plan.threads - 1) / warp_plan.threads, 1, 1 },
1485 .threadgroup = .{ warp_plan.threads, 1, 1 },
1486 });
1487 },
1488 else => {},
1489 }
1490 var plan = try createLaunchResourcePlan(caps, format, resources);
1491 var kept: usize = 0;
1492 var candidate_index: usize = 0;
1493 while (candidate_index < plan.candidate_count) : (candidate_index += 1) {
1494 const tile = reductionTilePlan(
1495 reduction,
1496 plan.candidates[candidate_index].geometry.threadgroup[0],
1497 ) catch |err| switch (err) {
1498 error.CapabilityMismatch => continue,
1499 else => return err,
1500 };
1501 plan.candidates[kept] = plan.candidates[candidate_index];
1502 plan.candidates[kept].tile = tile;
1503 kept += 1;
1504 }
1505 if (kept == 0) return error.CapabilityMismatch;
1506 plan.candidate_count = kept;
1507 plan.geometry = plan.candidates[0].geometry;
1508 plan.tile = plan.candidates[0].tile;
1509 return plan;
1510 }
1511
1512 fn createFixedGeometryReductionLaunchResourcePlan(
1513 caps: gpu.BackendCapabilities,
1514 format: gpu.ArtifactFormat,
1515 resources: schedule_planning.ScheduleResourceEstimate,
1516 reduction: ReductionLaunchPlan,
1517 geometry: choir_abi.LaunchGeometry,
1518 ) gpu.BackendError!LaunchResourcePlan {
1519 try caps.validateLaunchGeometry(geometry);
1520 var candidates = LaunchCandidateBuffer{};
1521 candidates.items[0] = .{
1522 .geometry = geometry,
1523 .score = 0,
1524 .estimated_static_bytes_per_threadgroup = resources.static_total_bytes,
1525 .estimated_element_ops_per_threadgroup = resources.estimated_element_ops,
1526 .tile = try reductionTilePlan(reduction, 1),
1527 };
1528 candidates.count = 1;
1529 const candidate = candidates.items[0];
1530 return .{
1531 .format = format,
1532 .element_count = resources.element_count,
1533 .geometry = candidate.geometry,
1534 .fixed_threadgroup = true,
1535 .resource_class = classifyLaunchResources(resources),
1536 .element_ops_per_kib = resources.elementOpsPerKiB(),
1537 .estimated_static_bytes_per_threadgroup = candidate.estimated_static_bytes_per_threadgroup,
1538 .estimated_element_ops_per_threadgroup = candidate.estimated_element_ops_per_threadgroup,
1539 .static_bytes_complete = resources.static_bytes_complete,
1540 .tile = candidate.tile,
1541 .candidate_count = candidates.count,
1542 .candidates = candidates.items,
1543 };
1544 }
1545
1546 fn validateReductionLaunchPlan(
1547 resources: schedule_planning.ScheduleResourceEstimate,
1548 reduction: ReductionLaunchPlan,
1549 ) gpu.BackendError!void {
1550 if (resources.element_count != @as(u64, reduction.output_element_count)) return error.InvalidArtifact;
1551 const covered_elements = std.math.mul(
1552 u32,
1553 reduction.output_element_count,
1554 reduction.reduction_extent,
1555 ) catch return error.CapabilityMismatch;
1556 if (covered_elements != reduction.input_element_count) return error.InvalidArtifact;
1557 }
1558
1559 fn reductionTilePlan(
1560 reduction: ReductionLaunchPlan,
1561 output_tile_elements: u32,
1562 ) gpu.BackendError!LaunchTilePlan {
1563 return .{
1564 .kind = .reduction,
1565 .m = output_tile_elements,
1566 .n = reduction.reduction_extent,
1567 .k = @intCast(reduction.axis),
1568 .batch = 1,
1569 .input_dtype = reduction.input_dtype,
1570 .output_dtype = reduction.output_dtype,
1571 .input_tile_bytes = try reductionInputTileBytes(
1572 output_tile_elements,
1573 reduction.reduction_extent,
1574 reduction.input_dtype,
1575 ),
1576 .output_tile_bytes = try reductionOutputTileBytes(
1577 output_tile_elements,
1578 reduction.output_dtype,
1579 ),
1580 .reduction_kind = reduction.kind,
1581 .reduction_rank = @intCast(reduction.input_rank),
1582 .reduction_axis = @intCast(reduction.axis),
1583 .reduction_extent = reduction.reduction_extent,
1584 };
1585 }
1586
1587 fn reductionInputTileBytes(
1588 output_tile_elements: u32,
1589 reduction_extent: u32,
1590 dtype: choir_abi.DType,
1591 ) gpu.BackendError!u32 {
1592 const elements = std.math.mul(u32, output_tile_elements, reduction_extent) catch return error.CapabilityMismatch;
1593 return std.math.mul(u32, elements, @as(u32, dtype.sizeOf())) catch return error.CapabilityMismatch;
1594 }
1595
1596 fn reductionOutputTileBytes(
1597 output_tile_elements: u32,
1598 dtype: choir_abi.DType,
1599 ) gpu.BackendError!u32 {
1600 return std.math.mul(u32, output_tile_elements, @as(u32, dtype.sizeOf())) catch return error.CapabilityMismatch;
1601 }
1602
1603 fn threadgroupThreadCount(geometry: choir_abi.LaunchGeometry) gpu.BackendError!u32 {
1604 const xy = std.math.mul(u32, geometry.threadgroup[0], geometry.threadgroup[1]) catch return error.CapabilityMismatch;
1605 return std.math.mul(u32, xy, geometry.threadgroup[2]) catch return error.CapabilityMismatch;
1606 }
1607
1608 pub fn createBackendMemoryPlan(
1609 allocator: std.mem.Allocator,
1610 artifact_plan: *const BackendArtifactPlan,
1611 ) !BackendMemoryPlan {
1612 const lifetimes = try allocator.alloc(BackendSlotLifetime, artifact_plan.slots.len);
1613 errdefer allocator.free(lifetimes);
1614 for (artifact_plan.slots, 0..) |slot, index| {
1615 lifetimes[index] = .{
1616 .slot_id = slot.slot_id,
1617 .role = slot.role,
1618 .dtype = slot.dtype,
1619 .memory_space = slot.memory_space,
1620 .layout_kind = slot.layout_kind,
1621 .element_count = slot.element_count,
1622 .byte_size = slot.byte_size,
1623 .alignment = slot.alignment,
1624 .layout_fingerprint = slot.layout_fingerprint,
1625 };
1626 }
1627
1628 for (artifact_plan.kernels.items, 0..) |kernel, kernel_index| {
1629 try recordSlotLifetimeUse(lifetimes, kernel.output_slot_id, kernel_index);
1630 for (kernel.input_slot_ids) |slot_id| {
1631 try recordSlotLifetimeUse(lifetimes, slot_id, kernel_index);
1632 }
1633 }
1634
1635 const end_index = artifact_plan.kernels.items.len;
1636 for (artifact_plan.output_slot_ids) |slot_id| {
1637 const lifetime = lifetimeForSlot(lifetimes, slot_id) orelse return error.InvalidArtifact;
1638 if (lifetime.hasKernelUse()) {
1639 lifetime.last_kernel_index = @max(lifetime.last_kernel_index, end_index);
1640 }
1641 }
1642
1643 var allocation_list: std.ArrayListUnmanaged(BackendBufferAllocationPlan) = .empty;
1644 errdefer allocation_list.deinit(allocator);
1645 var assignment_list: std.ArrayListUnmanaged(BackendBufferReuseAssignment) = .empty;
1646 errdefer assignment_list.deinit(allocator);
1647 var candidate_indices: std.ArrayListUnmanaged(usize) = .empty;
1648 defer candidate_indices.deinit(allocator);
1649
1650 var total_static_slot_bytes: u64 = 0;
1651 for (lifetimes, 0..) |lifetime, index| {
1652 if (!lifetime.isBackendAllocatable()) continue;
1653 total_static_slot_bytes += lifetime.byte_size.?;
1654 try candidate_indices.append(allocator, index);
1655 }
1656 std.mem.sort(usize, candidate_indices.items, lifetimes, lessLifetimeIndex);
1657
1658 var allocated_static_bytes: u64 = 0;
1659 for (candidate_indices.items) |lifetime_index| {
1660 const lifetime = lifetimes[lifetime_index];
1661 const allocation_index = reusableAllocationIndex(allocation_list.items, lifetime) orelse blk: {
1662 const allocation_id = allocation_list.items.len;
1663 const allocation = BackendBufferAllocationPlan{
1664 .allocation_id = allocation_id,
1665 .byte_size = lifetime.byte_size.?,
1666 .dtype = lifetime.dtype,
1667 .memory_space = lifetime.memory_space,
1668 .layout_kind = lifetime.layout_kind,
1669 .element_count = lifetime.element_count,
1670 .alignment = lifetime.alignment,
1671 .layout_fingerprint = lifetime.layout_fingerprint,
1672 .last_kernel_index = lifetime.last_kernel_index,
1673 };
1674 try allocation_list.append(allocator, allocation);
1675 allocated_static_bytes += allocation.byte_size;
1676 break :blk allocation_id;
1677 };
1678 allocation_list.items[allocation_index].last_kernel_index = lifetime.last_kernel_index;
1679 try assignment_list.append(allocator, .{
1680 .slot_id = lifetime.slot_id,
1681 .allocation_id = allocation_list.items[allocation_index].allocation_id,
1682 });
1683 }
1684
1685 const allocations = try allocation_list.toOwnedSlice(allocator);
1686 errdefer if (allocations.len != 0) allocator.free(allocations);
1687 const assignments = try assignment_list.toOwnedSlice(allocator);
1688 errdefer if (assignments.len != 0) allocator.free(assignments);
1689
1690 return .{
1691 .allocator = allocator,
1692 .lifetimes = lifetimes,
1693 .allocations = allocations,
1694 .assignments = assignments,
1695 .total_static_slot_bytes = total_static_slot_bytes,
1696 .allocated_static_bytes = allocated_static_bytes,
1697 .peak_static_live_bytes = peakStaticLiveBytes(lifetimes, end_index),
1698 };
1699 }
1700
1701 fn recordSlotLifetimeUse(
1702 lifetimes: []BackendSlotLifetime,
1703 slot_id: usize,
1704 kernel_index: usize,
1705 ) !void {
1706 const lifetime = lifetimeForSlot(lifetimes, slot_id) orelse return error.InvalidArtifact;
1707 if (!lifetime.hasKernelUse()) {
1708 lifetime.first_kernel_index = kernel_index;
1709 lifetime.last_kernel_index = kernel_index;
1710 return;
1711 }
1712 lifetime.first_kernel_index = @min(lifetime.first_kernel_index, kernel_index);
1713 lifetime.last_kernel_index = @max(lifetime.last_kernel_index, kernel_index);
1714 }
1715
1716 fn lifetimeForSlot(
1717 lifetimes: []BackendSlotLifetime,
1718 slot_id: usize,
1719 ) ?*BackendSlotLifetime {
1720 if (slot_id < lifetimes.len and lifetimes[slot_id].slot_id == slot_id) {
1721 return &lifetimes[slot_id];
1722 }
1723 for (lifetimes) |*lifetime| {
1724 if (lifetime.slot_id == slot_id) return lifetime;
1725 }
1726 return null;
1727 }
1728
1729 fn lessLifetimeIndex(
1730 lifetimes: []BackendSlotLifetime,
1731 lhs_index: usize,
1732 rhs_index: usize,
1733 ) bool {
1734 const lhs = lifetimes[lhs_index];
1735 const rhs = lifetimes[rhs_index];
1736 if (lhs.first_kernel_index != rhs.first_kernel_index) {
1737 return lhs.first_kernel_index < rhs.first_kernel_index;
1738 }
1739 return lhs.slot_id < rhs.slot_id;
1740 }
1741
1742 fn reusableAllocationIndex(
1743 allocations: []const BackendBufferAllocationPlan,
1744 lifetime: BackendSlotLifetime,
1745 ) ?usize {
1746 for (allocations, 0..) |allocation, index| {
1747 if (allocation.last_kernel_index >= lifetime.first_kernel_index) continue;
1748 if (allocation.byte_size != lifetime.byte_size.?) continue;
1749 if (allocation.dtype != lifetime.dtype) continue;
1750 if (allocation.memory_space != lifetime.memory_space) continue;
1751 if (allocation.layout_kind != lifetime.layout_kind) continue;
1752 if (allocation.element_count != lifetime.element_count) continue;
1753 if (allocation.alignment != lifetime.alignment) continue;
1754 if (allocation.layout_fingerprint != lifetime.layout_fingerprint) continue;
1755 return index;
1756 }
1757 return null;
1758 }
1759
1760 fn peakStaticLiveBytes(lifetimes: []const BackendSlotLifetime, end_index: usize) u64 {
1761 var peak: u64 = 0;
1762 for (0..end_index + 1) |kernel_index| {
1763 var live: u64 = 0;
1764 for (lifetimes) |lifetime| {
1765 if (!lifetime.isBackendAllocatable()) continue;
1766 if (lifetime.first_kernel_index > kernel_index) continue;
1767 if (lifetime.last_kernel_index < kernel_index) continue;
1768 live += lifetime.byte_size.?;
1769 }
1770 peak = @max(peak, live);
1771 }
1772 return peak;
1773 }
1774
1775 pub fn createBackendArtifactPlan(
1776 allocator: std.mem.Allocator,
1777 handle: gpu.BackendHandle,
1778 source: BackendArtifactPlanSource,
1779 options: ArtifactPlanOptions,
1780 ) !BackendArtifactPlan {
1781 const pass_ctx = source.pass_ctx;
1782 const choir_module = source.choir_module;
1783 const schedule_plan = try schedule_planning.getSchedulePlanAnalysis(pass_ctx, choir_module);
1784 const buffer_plan = try bufferization.getBufferPlanAnalysis(pass_ctx, choir_module);
1785 const memory_plan = try memory_space.getMemorySpacePlanAnalysis(pass_ctx, choir_module);
1786 const layout_plan = try layout_planning.getLayoutPlanAnalysis(pass_ctx, choir_module);
1787 const outline_plan = try kernel_outlining.getKernelOutlinePlanAnalysis(pass_ctx, choir_module);
1788 const legal_plan = try backend_legalization.getBackendLegalizationAnalysis(pass_ctx, choir_module);
1789
1790 return createFromPlans(allocator, handle, .{
1791 .root = choir_module,
1792 .profile = target_profile.readBackendTargetProfile(choir_module),
1793 .schedule = schedule_plan,
1794 .buffers = buffer_plan,
1795 .spaces = memory_plan,
1796 .layouts = layout_plan,
1797 .outlines = outline_plan,
1798 .generated = source.lowered_kernels,
1799 .legal = legal_plan,
1800 }, options);
1801 }
1802
1803 pub const PlanInputs = struct {
1804 root: *ir.Operation,
1805 profile: ?preparation.BackendTargetProfile,
1806 schedule: *const schedule_planning.SchedulePlanAnalysis,
1807 buffers: *const bufferization.BufferPlanAnalysis,
1808 spaces: *const memory_space.MemorySpacePlanAnalysis,
1809 layouts: *const layout_planning.LayoutPlanAnalysis,
1810 outlines: *const kernelization_model.KernelOutlinePlanAnalysis,
1811 generated: *const kernelization.KernelizationAnalysis,
1812 legal: *const backend_legalization.BackendLegalizationAnalysis,
1813 target: ?*const accy_choir.record.target.Record = null,
1814 };
1815
1816 fn createFromPlans(
1817 allocator: std.mem.Allocator,
1818 handle: gpu.BackendHandle,
1819 source: PlanInputs,
1820 options: ArtifactPlanOptions,
1821 ) !BackendArtifactPlan {
1822 const caps = try handle.queryCapabilities();
1823 const existing_profile = source.profile;
1824 const backend_kind = handle.backendKind() orelse caps.identity.backend;
1825 const format = options.format orelse if (existing_profile) |profile| profile.artifact_format else defaultArtifactFormat(backend_kind) orelse {
1826 return error.UnsupportedOperation;
1827 };
1828 if (!caps.supportsArtifactFormat(format)) return error.UnsupportedArtifactFormat;
1829 const profile = try selectedBackendTargetProfile(caps, backend_kind, format, existing_profile);
1830
1831 const choir_module = source.root;
1832 const schedule_plan = source.schedule;
1833 const buffer_plan = source.buffers;
1834 const memory_plan = source.spaces;
1835 const layout_plan = source.layouts;
1836 const outline_plan = source.outlines;
1837 const legal_plan = source.legal;
1838 if (source.target != null and (existing_profile == null or
1839 existing_profile.?.artifact_format != format)) return error.TargetRecipeMismatch;
1840
1841 var plan = BackendArtifactPlan.init(allocator, profile);
1842 errdefer plan.deinit();
1843 try plan.setSlots(buffer_plan, memory_plan, layout_plan, choir_module);
1844
1845 for (outline_plan.kernels.items) |outline| {
1846 const legal = legalizationForKernel(legal_plan, outline.id) orelse {
1847 return error.MissingKernelLegalization;
1848 };
1849 if (!legal.isLegal()) return backend_legalization.backendKernelStatusError(legal.status);
1850
1851 const work = workItemById(schedule_plan, outline.work_item_id) orelse {
1852 return error.MissingScheduleWorkItem;
1853 };
1854 var compile_plan = try compilePlanForWork(
1855 allocator,
1856 format,
1857 outline,
1858 work.*,
1859 buffer_plan,
1860 source.generated,
1861 options.kernel_call_registry,
1862 source.target,
1863 );
1864 defer compile_plan.deinit(allocator);
1865
1866 const element_count_argument = try elementCountArgumentForWork(format, work.*, compile_plan);
1867 const element_count_argument_value = elementCountArgumentValueForWork(work.*, compile_plan);
1868 const launch_resources = try createLaunchResourcePlanForWork(caps, format, work.*, compile_plan);
1869
1870 const required_dtypes = requiredDTypesForCompilePlan(work.dtype, &compile_plan);
1871 var compile = try PlannedKernelCompileContract.init(
1872 allocator,
1873 plannedCompileSource(compile_plan),
1874 plannedCompileLaunch(compile_plan),
1875 format,
1876 compile_plan.entry_name,
1877 compile_plan.argument_count,
1878 required_dtypes,
1879 compile_plan.required_features,
1880 compile_plan.required_subgroup,
1881 compile_plan.shape_family_fingerprint,
1882 compile_plan.payload,
1883 );
1884 var compile_owned = true;
1885 errdefer if (compile_owned) compile.deinit(allocator);
1886
1887 const artifact = try createKernelArtifactFromCompilePlan(
1888 handle,
1889 format,
1890 &compile_plan,
1891 outline.name,
1892 required_dtypes,
1893 compile_plan.required_features,
1894 compile_plan.required_subgroup,
1895 );
1896
1897 compile_owned = false;
1898 try plan.addKernel(
1899 outline,
1900 work.resources,
1901 legal,
1902 artifact,
1903 launch_resources,
1904 kernelCallLaunchForCompilePlan(compile_plan),
1905 element_count_argument,
1906 element_count_argument_value,
1907 compile_plan.runtime_scalar_argument_count,
1908 compile_plan.runtime_scalar_defaults,
1909 compile_plan.static_arguments,
1910 compile,
1911 outputFillPatternForCompilePlan(compile_plan),
1912 scratchFillPatternForCompilePlan(compile_plan),
1913 );
1914 }
1915
1916 return plan;
1917 }
1918
1919 fn scratchFillPatternForCompilePlan(compile_plan: CompilePlan) ?u32 {
1920 const lowered = compile_plan.lowered_kernel orelse return null;
1921 return lowered.scratch_fill_pattern;
1922 }
1923
1924 fn outputFillPatternForCompilePlan(compile_plan: CompilePlan) ?u32 {
1925 const lowered = compile_plan.lowered_kernel orelse return null;
1926 return lowered.output_fill_pattern;
1927 }
1928
1929 pub fn createBackendArtifactPlanFromTargetJob(
1930 allocator: std.mem.Allocator,
1931 handle: gpu.BackendHandle,
1932 target_module: *target_product.TargetJob,
1933 options: ArtifactPlanOptions,
1934 ) !BackendArtifactPlan {
1935 var pass_ctx = target_module.passContext();
1936 defer pass_ctx.deinit();
1937 return try createBackendArtifactPlan(
1938 allocator,
1939 handle,
1940 .{
1941 .pass_ctx = &pass_ctx,
1942 .choir_module = target_module.choir_module,
1943 .lowered_kernels = target_module.kernelizationProduct(),
1944 },
1945 options,
1946 );
1947 }
1948
1949 fn requiredDTypesForCompilePlan(work_dtype: choir_abi.DType, compile_plan: *const CompilePlan) gpu.DTypeSet {
1950 var dtypes = gpu.DTypeSet.init(&.{work_dtype});
1951 dtypes.bits |= compile_plan.required_dtypes.bits;
1952 if (compile_plan.lowered_kernel) |lowered| {
1953 const lowered_dtypes = lowered.requiredDTypes();
1954 dtypes.bits |= lowered_dtypes.bits;
1955 }
1956 return dtypes;
1957 }
1958
1959 fn requiredDTypesForKernelCallArtifact(
1960 artifact: KernelCallArtifact,
1961 work: schedule_planning.ScheduleWorkItem,
1962 outline: kernelization_model.KernelOutline,
1963 buffer_plan: *const bufferization.BufferPlanAnalysis,
1964 ) gpu.BackendError!gpu.DTypeSet {
1965 var dtypes = artifact.required_dtypes;
1966 dtypes.insert(work.dtype);
1967 for (outline.input_slot_ids) |slot_id| {
1968 const slot = bufferSlotById(buffer_plan, slot_id) orelse return error.InvalidArtifact;
1969 dtypes.insert(slot.dtype);
1970 }
1971 return dtypes;
1972 }
1973
1974 fn bufferSlotById(
1975 buffer_plan: *const bufferization.BufferPlanAnalysis,
1976 slot_id: usize,
1977 ) ?*const bufferization.BufferSlot {
1978 if (slot_id < buffer_plan.slots.items.len and buffer_plan.slots.items[slot_id].id == slot_id) {
1979 return &buffer_plan.slots.items[slot_id];
1980 }
1981 for (buffer_plan.slots.items) |*slot| {
1982 if (slot.id == slot_id) return slot;
1983 }
1984 return null;
1985 }
1986
1987 fn plannedCompileLaunch(compile_plan: CompilePlan) PlannedKernelCompileLaunch {
1988 return switch (compile_plan.launch) {
1989 .generic => .generic,
1990 .dot_general => .dot_general,
1991 .reduction => .reduction,
1992 .kernel_call => .kernel_call,
1993 };
1994 }
1995
1996 fn plannedCompileSource(compile_plan: CompilePlan) PlannedKernelSource {
1997 return switch (compile_plan.launch) {
1998 .kernel_call => .kernel_call,
1999 else => .tensor,
2000 };
2001 }
2002
2003 fn kernelCallLaunchForCompilePlan(compile_plan: CompilePlan) ?KernelCallLaunch {
2004 return switch (compile_plan.launch) {
2005 .kernel_call => |launch| launch,
2006 else => null,
2007 };
2008 }
2009
2010 fn compilePayloadKind(payload: gpu.CompilePayload) PlannedKernelCompilePayload {
2011 return switch (payload) {
2012 .none => .none,
2013 .bytes => .bytes,
2014 .words_u32 => .words_u32,
2015 .text => .text,
2016 };
2017 }
2018
2019 fn compilePayloadByteCount(payload: gpu.CompilePayload) gpu.BackendError!usize {
2020 return switch (payload) {
2021 .none => 0,
2022 .bytes => |bytes| bytes.len,
2023 .words_u32 => |words| std.math.mul(usize, words.len, @sizeOf(u32)) catch return error.CapabilityMismatch,
2024 .text => |text| text.len,
2025 };
2026 }
2027
2028 fn collectInputSlotIds(
2029 allocator: std.mem.Allocator,
2030 buffer_plan: *const bufferization.BufferPlanAnalysis,
2031 ) ![]usize {
2032 var slot_ids: std.ArrayListUnmanaged(usize) = .empty;
2033 errdefer slot_ids.deinit(allocator);
2034 for (buffer_plan.slots.items) |slot| {
2035 if (!slot.role.input) continue;
2036 try slot_ids.append(allocator, slot.id);
2037 }
2038 return try slot_ids.toOwnedSlice(allocator);
2039 }
2040
2041 fn collectOutputSlotIds(
2042 allocator: std.mem.Allocator,
2043 choir_module: *ir.Operation,
2044 buffer_plan: *const bufferization.BufferPlanAnalysis,
2045 ) ![]usize {
2046 var slot_ids: std.ArrayListUnmanaged(usize) = .empty;
2047 errdefer slot_ids.deinit(allocator);
2048 try appendReturnSlotIds(allocator, choir_module, buffer_plan, &slot_ids);
2049 return try slot_ids.toOwnedSlice(allocator);
2050 }
2051
2052 fn appendReturnSlotIds(
2053 allocator: std.mem.Allocator,
2054 op: *ir.Operation,
2055 buffer_plan: *const bufferization.BufferPlanAnalysis,
2056 slot_ids: *std.ArrayListUnmanaged(usize),
2057 ) !void {
2058 if (isReturnOp(op)) {
2059 for (op.getOperandValues()) |value| {
2060 const slot = buffer_plan.getSlot(value) orelse return error.UnsupportedOperation;
2061 try slot_ids.append(allocator, slot.id);
2062 }
2063 }
2064
2065 for (op.regions.items) |*region| {
2066 var block_iter = region.getBlocks();
2067 while (block_iter.next()) |block| {
2068 var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
2069 while (current) |current_op| {
2070 try appendReturnSlotIds(allocator, current_op, buffer_plan, slot_ids);
2071 current = current_op.next_op;
2072 }
2073 }
2074 }
2075 }
2076
2077 fn isReturnOp(op: *ir.Operation) bool {
2078 return isName(op.name.name, "func.return") or
2079 isName(op.name.name, dialect_mod.AccyDialect.ReturnOp.operation_name);
2080 }
2081
2082 fn isKernelCallOp(op: *ir.Operation) bool {
2083 return isName(op.name.name, dialect_mod.AccyDialect.KernelCallOp.operation_name);
2084 }
2085
2086 fn isName(actual: []const u8, expected: []const u8) bool {
2087 return std.mem.eql(u8, actual, expected);
2088 }
2089
2090 fn constantPayloadForSlot(
2091 allocator: std.mem.Allocator,
2092 slot: bufferization.BufferSlot,
2093 ) ![]u8 {
2094 if (!slot.role.constant) return &.{};
2095 const producer = slot.producer orelse return error.InvalidArtifact;
2096 if (!isName(producer.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name)) {
2097 return error.InvalidArtifact;
2098 }
2099 const attr = producer.getAttr("payload") orelse return error.InvalidArtifact;
2100 if (!std.mem.eql(u8, attr.abstract.name, dialect_mod.AccyDialect.ConstantOp.payload_attr_name)) {
2101 return error.InvalidArtifact;
2102 }
2103 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return error.InvalidArtifact;
2104 if (slot.byte_size) |byte_size| {
2105 const expected = std.math.cast(usize, byte_size) orelse return error.InvalidArtifact;
2106 if (dialect_attr.payload.len != expected) return error.InvalidArtifact;
2107 }
2108 return try allocator.dupe(u8, dialect_attr.payload);
2109 }
2110
2111 fn layoutFingerprint(
2112 space: memory_space.MemorySpace,
2113 kind: layout_planning.LayoutKind,
2114 dims: []const i64,
2115 element_strides: ?[]const u64,
2116 minor_to_major: []const usize,
2117 alignment: u64,
2118 contiguous: bool,
2119 static_layout: bool,
2120 ) u64 {
2121 var builder = choir.product.incremental.FingerprintBuilder{};
2122 builder.updateBytes("accy.artifact.layout");
2123 builder.updateEnumTag(space);
2124 builder.updateEnumTag(kind);
2125 builder.updateI64Slice(dims);
2126 builder.updateOptionalU64Slice(element_strides);
2127 builder.updateUsizeSlice(minor_to_major);
2128 builder.updateU64(alignment);
2129 builder.updateBool(contiguous);
2130 builder.updateBool(static_layout);
2131 return builder.finish();
2132 }
2133
2134 fn selectedBackendTargetProfile(
2135 caps: gpu.BackendCapabilities,
2136 backend_kind: gpu.BackendKind,
2137 format: gpu.ArtifactFormat,
2138 existing: ?target_profile.BackendTargetProfile,
2139 ) gpu.BackendError!target_profile.BackendTargetProfile {
2140 const profile = existing orelse try target_profile.BackendTargetProfile.init(caps, backend_kind, format);
2141 if (profile.backend_kind != backend_kind) return error.CapabilityMismatch;
2142 if (profile.artifact_format != format) return error.UnsupportedArtifactFormat;
2143 if (!caps.supportsArtifactFormat(profile.artifact_format)) return error.UnsupportedArtifactFormat;
2144 if (!profile.isSupportedBy(caps)) return error.CapabilityMismatch;
2145 return profile;
2146 }
2147
2148 fn legalizationForKernel(
2149 legal_plan: *const backend_legalization.BackendLegalizationAnalysis,
2150 kernel_id: usize,
2151 ) ?backend_legalization.BackendKernelLegalization {
2152 for (legal_plan.kernels.items) |legal| {
2153 if (legal.kernel_id == kernel_id) return legal;
2154 }
2155 return null;
2156 }
2157
2158 fn workItemById(
2159 schedule_plan: *const schedule_planning.SchedulePlanAnalysis,
2160 work_item_id: usize,
2161 ) ?*const schedule_planning.ScheduleWorkItem {
2162 if (work_item_id < schedule_plan.work_items.items.len) {
2163 const work = &schedule_plan.work_items.items[work_item_id];
2164 if (work.id == work_item_id) return work;
2165 }
2166 for (schedule_plan.work_items.items) |*work| {
2167 if (work.id == work_item_id) return work;
2168 }
2169 return null;
2170 }
2171
2172 fn compilePlanForWork(
2173 allocator: std.mem.Allocator,
2174 format: gpu.ArtifactFormat,
2175 outline: kernelization_model.KernelOutline,
2176 work: schedule_planning.ScheduleWorkItem,
2177 buffer_plan: *const bufferization.BufferPlanAnalysis,
2178 kernelization_plan: *const kernelization.KernelizationAnalysis,
2179 kernel_call_registry: ?*const KernelCallRegistry,
2180 captured: ?*const accy_choir.record.target.Record,
2181 ) gpu.BackendError!CompilePlan {
2182 if (outline.kind == .kernel_call) {
2183 return try compilePlanFromKernelCall(
2184 allocator,
2185 format,
2186 outline,
2187 work,
2188 buffer_plan,
2189 kernel_call_registry,
2190 );
2191 }
2192
2193 if (format == .cuda_ptx or format == .vulkan_spirv or format == .metal_msl or format == .webgpu_wgsl or gpu.artifactFormatUsesHostLoopLaunch(format)) {
2194 if (kernelization_plan.getForWork(work.id)) |lowered| {
2195 var compile_plan = if (captured) |record|
2196 try compilePlanFromCapturedKernel(allocator, format, lowered, work, record)
2197 else
2198 try compilePlanFromLoweredKernel(allocator, format, lowered, work);
2199 errdefer compile_plan.deinit(allocator);
2200 compile_plan.launch = try compileLaunchPlanForWork(outline, work, buffer_plan);
2201 return compile_plan;
2202 }
2203 return kernelizationMissError(outline, work, buffer_plan);
2204 }
2205
2206 return error.UnsupportedOperation;
2207 }
2208
2209 fn compilePlanFromKernelCall(
2210 allocator: std.mem.Allocator,
2211 format: gpu.ArtifactFormat,
2212 outline: kernelization_model.KernelOutline,
2213 work: schedule_planning.ScheduleWorkItem,
2214 buffer_plan: *const bufferization.BufferPlanAnalysis,
2215 kernel_call_registry: ?*const KernelCallRegistry,
2216 ) gpu.BackendError!CompilePlan {
2217 const registry = kernel_call_registry orelse return error.UnsupportedOperation;
2218 if (work.ops.len != 1 or work.ops[0] != work.root) return error.InvalidArtifact;
2219 if (!isKernelCallOp(work.root)) return error.InvalidArtifact;
2220
2221 const target = try kernelCallTarget(work.root);
2222 const version = try kernelCallVersion(work.root);
2223 const artifact = registry.find(target, version, format) orelse return error.UnsupportedOperation;
2224 try validateKernelCallArtifact(format, artifact);
2225
2226 const payload = try duplicateCompilePayload(allocator, artifact.payload);
2227 errdefer deinitCompilePayload(allocator, payload);
2228
2229 const explicit_scalars = dialect_mod.AccyDialect.kernelCallRuntimeScalars(work.root) catch {
2230 return error.InvalidArtifact;
2231 };
2232 var runtime_scalar_defaults = try kernelCallRuntimeScalarDefaults(allocator, artifact, explicit_scalars);
2233 errdefer if (runtime_scalar_defaults.len != 0) allocator.free(runtime_scalar_defaults);
2234
2235 const static_arguments = allocator.dupe(choir_abi.ScalarArgument, artifact.static_arguments) catch return error.OutOfMemory;
2236 errdefer if (static_arguments.len != 0) allocator.free(static_arguments);
2237
2238 const expected_argument_count = try expectedKernelCallArgumentCount(
2239 outline,
2240 artifact.element_count_argument,
2241 artifact.runtime_scalar_argument_count,
2242 artifact.static_arguments.len,
2243 );
2244 if (artifact.argument_count != expected_argument_count) return error.InvalidArtifact;
2245
2246 var element_count_argument = artifact.element_count_argument;
2247 var runtime_scalar_argument_count = artifact.runtime_scalar_argument_count;
2248 var launch = artifact.launch;
2249 if (explicit_scalars == null and kernelCallCountScalarIsElementCount(artifact)) {
2250 if (gpu.artifactFormatUsesHostLoopLaunch(format)) {
2251 element_count_argument = .scalar_u32;
2252 runtime_scalar_argument_count = 0;
2253 launch = try resolveKernelCallLaunchForElementCount(artifact.launch, work.element_count);
2254 const geometry = switch (launch) {
2255 .fixed => |fixed| fixed,
2256 .derived => return error.InvalidArtifact,
2257 };
2258 const shape_count = choir_abi.launch_shape_argument_count;
2259 if (static_arguments.len < shape_count) return error.InvalidArtifact;
2260 const shape_arguments = try choir_abi.launchShapeArguments(allocator, work.element_count, geometry);
2261 defer allocator.free(shape_arguments);
2262 @memcpy(static_arguments[static_arguments.len - shape_count ..], shape_arguments);
2263 } else {
2264 if (work.element_count > std.math.maxInt(u32)) return error.LaunchArgumentMismatch;
2265 const count_default = allocator.alloc(choir_abi.ScalarArgument, 1) catch return error.OutOfMemory;
2266 count_default[0] = .{ .u32 = @intCast(work.element_count) };
2267 runtime_scalar_defaults = count_default;
2268 }
2269 }
2270
2271 return .{
2272 .entry_name = artifact.entry_name,
2273 .argument_count = artifact.argument_count,
2274 .payload = payload,
2275 .launch = .{ .kernel_call = launch },
2276 .required_dtypes = try requiredDTypesForKernelCallArtifact(artifact, work, outline, buffer_plan),
2277 .required_features = artifact.required_features,
2278 .required_subgroup = artifact.required_subgroup,
2279 .push_constants = artifact.push_constants,
2280 .shape_family_fingerprint = artifact.shape_family_fingerprint,
2281 .element_count_argument = element_count_argument,
2282 .runtime_scalar_argument_count = runtime_scalar_argument_count,
2283 .runtime_scalar_defaults = runtime_scalar_defaults,
2284 .static_arguments = static_arguments,
2285 };
2286 }
2287
2288 fn kernelCallRuntimeScalarDefaults(
2289 allocator: std.mem.Allocator,
2290 artifact: KernelCallArtifact,
2291 explicit_scalars: ?dialect_mod.AccyDialect.KernelCallRuntimeScalars,
2292 ) gpu.BackendError![]choir_abi.ScalarArgument {
2293 const scalars = explicit_scalars orelse return &.{};
2294 if (scalars.count != artifact.runtime_scalar_argument_count) return error.InvalidArtifact;
2295 const defaults = allocator.alloc(choir_abi.ScalarArgument, scalars.count) catch return error.OutOfMemory;
2296 for (scalars.slice(), 0..) |scalar, index| {
2297 defaults[index] = backendScalarFromKernelCallScalar(scalar);
2298 }
2299 return defaults;
2300 }
2301
2302 fn backendScalarFromKernelCallScalar(
2303 scalar: dialect_mod.AccyDialect.KernelCallScalar,
2304 ) choir_abi.ScalarArgument {
2305 return switch (scalar.kind) {
2306 .i32 => .{ .i32 = @bitCast(@as(u32, @truncate(scalar.bits))) },
2307 .u32 => .{ .u32 = @truncate(scalar.bits) },
2308 .i64 => .{ .i64 = @bitCast(scalar.bits) },
2309 .u64 => .{ .u64 = scalar.bits },
2310 .f32 => .{ .f32 = @bitCast(@as(u32, @truncate(scalar.bits))) },
2311 .f64 => .{ .f64 = @bitCast(scalar.bits) },
2312 };
2313 }
2314
2315 fn kernelCallCountScalarIsElementCount(artifact: KernelCallArtifact) bool {
2316 if (artifact.runtime_scalar_argument_count != 1) return false;
2317 if (artifact.element_count_argument != .none) return false;
2318 const shape_profile = artifact.shape_profile orelse return false;
2319 if (shape_profile.dimensions.len != 1) return false;
2320 return shape_profile.dimensions[0].runtime_scalar_argument_index == 0;
2321 }
2322
2323 fn resolveKernelCallLaunchForElementCount(
2324 launch: KernelCallLaunch,
2325 element_count: u64,
2326 ) gpu.BackendError!KernelCallLaunch {
2327 const derived = switch (launch) {
2328 .fixed => return launch,
2329 .derived => |derived| derived,
2330 };
2331 const count = std.math.cast(u32, element_count) orelse return error.LaunchArgumentMismatch;
2332 const count_argument = [_]choir_abi.ScalarArgument{.{ .u32 = count }};
2333 return .{ .fixed = try derived.geometry(count_argument[0..]) };
2334 }
2335
2336 fn validateKernelCallArtifact(
2337 format: gpu.ArtifactFormat,
2338 artifact: KernelCallArtifact,
2339 ) gpu.BackendError!void {
2340 if (artifact.target.len == 0) return error.InvalidArtifact;
2341 if (artifact.version == 0) return error.InvalidArtifact;
2342 if (artifact.format != format) return error.UnsupportedArtifactFormat;
2343 if (artifact.entry_name.len == 0) return error.InvalidArtifact;
2344 switch (artifact.payload) {
2345 .none => return error.InvalidArtifact,
2346 .bytes => |bytes| if (bytes.len == 0) return error.InvalidArtifact,
2347 .words_u32 => |words| if (words.len == 0) return error.InvalidArtifact,
2348 .text => |text| if (text.len == 0) return error.InvalidArtifact,
2349 }
2350 const static_argument_count: u32 = std.math.cast(u32, artifact.static_arguments.len) orelse return error.InvalidArtifact;
2351 const scalar_min = std.math.add(u32, artifact.runtime_scalar_argument_count, static_argument_count) catch return error.InvalidArtifact;
2352 if (scalar_min > artifact.argument_count) return error.InvalidArtifact;
2353 if (artifact.shape_profile) |profile| {
2354 const fingerprint = artifact.shape_family_fingerprint orelse return error.InvalidArtifact;
2355 if (profile.fingerprint != fingerprint) return error.InvalidArtifact;
2356 try profile.validate(artifact.runtime_scalar_argument_count);
2357 }
2358 switch (artifact.launch) {
2359 .derived => |derived| try validateKernelCallDerivedLaunch(derived),
2360 .fixed => {},
2361 }
2362 }
2363
2364 fn validateKernelCallDerivedLaunch(launch: KernelCallDerivedLaunch) gpu.BackendError!void {
2365 _ = try validateThreadgroup(launch.threadgroup);
2366 for (launch.grid) |axis| {
2367 switch (axis) {
2368 .fixed => |value| if (value == 0) return error.InvalidArtifact,
2369 .runtime_u32_ceil_div => |runtime| if (runtime.divisor == 0) return error.InvalidArtifact,
2370 }
2371 }
2372 }
2373
2374 fn kernelCallTarget(op: *ir.Operation) gpu.BackendError![]const u8 {
2375 const attr = op.getAttr("target") orelse return error.InvalidArtifact;
2376 if (!std.mem.eql(u8, attr.abstract.name, dialect_mod.AccyDialect.KernelCallOp.target_attr_name)) {
2377 return error.InvalidArtifact;
2378 }
2379 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return error.InvalidArtifact;
2380 if (dialect_attr.payload.len == 0) return error.InvalidArtifact;
2381 return dialect_attr.payload;
2382 }
2383
2384 fn kernelCallVersion(op: *ir.Operation) gpu.BackendError!u32 {
2385 const attr = op.getAttrAs(ir.Attribute.IntegerAttr, "version") orelse return error.InvalidArtifact;
2386 const value = attr.getValue();
2387 if (value < 1 or value > @as(i64, std.math.maxInt(u32))) return error.InvalidArtifact;
2388 return @intCast(value);
2389 }
2390
2391 fn expectedKernelCallArgumentCount(
2392 outline: kernelization_model.KernelOutline,
2393 element_count_argument: ElementCountArgument,
2394 runtime_scalar_argument_count: u32,
2395 static_argument_count: usize,
2396 ) gpu.BackendError!u32 {
2397 var count: usize = 1 + outline.inputCount() + static_argument_count;
2398 count += runtime_scalar_argument_count;
2399 switch (element_count_argument) {
2400 .none => {},
2401 .scalar_u32, .device_buffer_u32 => count += 1,
2402 }
2403 if (count > std.math.maxInt(u32)) return error.InvalidArtifact;
2404 return @intCast(count);
2405 }
2406
2407 fn duplicateCompilePayload(
2408 allocator: std.mem.Allocator,
2409 payload: gpu.CompilePayload,
2410 ) gpu.BackendError!gpu.CompilePayload {
2411 return switch (payload) {
2412 .none => .none,
2413 .bytes => |bytes| .{ .bytes = allocator.dupe(u8, bytes) catch return error.OutOfMemory },
2414 .words_u32 => |words| .{ .words_u32 = allocator.dupe(u32, words) catch return error.OutOfMemory },
2415 .text => |text| .{ .text = allocator.dupe(u8, text) catch return error.OutOfMemory },
2416 };
2417 }
2418
2419 fn compileLaunchPlanForWork(
2420 outline: kernelization_model.KernelOutline,
2421 work: schedule_planning.ScheduleWorkItem,
2422 buffer_plan: *const bufferization.BufferPlanAnalysis,
2423 ) gpu.BackendError!CompileLaunchPlan {
2424 return switch (outline.kind) {
2425 .dot_general => .{ .dot_general = try dotGeneralLaunchPlanForWork(outline, work, buffer_plan) },
2426 .reduction => .{ .reduction = try reductionLaunchPlanForWork(outline, work, buffer_plan) },
2427 .elementwise, .shape, .kernel_call, .row_pipeline, .iterate, .flash_attention, .scan => .generic,
2428 };
2429 }
2430
2431 fn dotGeneralLaunchPlanForWork(
2432 outline: kernelization_model.KernelOutline,
2433 work: schedule_planning.ScheduleWorkItem,
2434 buffer_plan: *const bufferization.BufferPlanAnalysis,
2435 ) gpu.BackendError!DotGeneralLaunchPlan {
2436 const desc = try kernelization.dotGeneralDescriptionForWork(outline, work, buffer_plan);
2437 return .{
2438 .input_dtype = desc.input_dtype,
2439 .output_dtype = desc.output_dtype,
2440 .m = desc.dims.m,
2441 .n = desc.dims.n,
2442 .k = desc.dims.k,
2443 .batch = desc.dims.batch,
2444 };
2445 }
2446
2447 fn reductionLaunchPlanForWork(
2448 outline: kernelization_model.KernelOutline,
2449 work: schedule_planning.ScheduleWorkItem,
2450 buffer_plan: *const bufferization.BufferPlanAnalysis,
2451 ) gpu.BackendError!ReductionLaunchPlan {
2452 const desc = try kernelization.reductionDescriptionForWork(outline, work, buffer_plan);
2453 return .{
2454 .kind = launchReductionKind(desc.kind),
2455 .input_dtype = desc.input_dtype,
2456 .output_dtype = desc.output_dtype,
2457 .input_rank = desc.dims.input_rank,
2458 .axis = desc.dims.axis,
2459 .input_element_count = desc.dims.input_element_count,
2460 .output_element_count = desc.dims.output_element_count,
2461 .reduction_extent = try reductionExtent(desc.dims),
2462 };
2463 }
2464
2465 fn launchReductionKind(kind: kernelization.ReductionKind) LaunchReductionKind {
2466 return switch (kind) {
2467 .sum => .sum,
2468 .max => .max,
2469 .min => .min,
2470 };
2471 }
2472
2473 fn reductionExtent(dims: kernelization.ReductionStaticDims) gpu.BackendError!u32 {
2474 return switch (dims.input_rank) {
2475 1 => dims.input_element_count,
2476 2 => if (dims.axis == 0) dims.rows else dims.cols,
2477 3 => dims.cols,
2478 else => error.InvalidArtifact,
2479 };
2480 }
2481
2482 fn compilePlanFromLoweredKernel(
2483 allocator: std.mem.Allocator,
2484 format: gpu.ArtifactFormat,
2485 lowered: *const kernelization.LoweredKernel,
2486 work: schedule_planning.ScheduleWorkItem,
2487 ) gpu.BackendError!CompilePlan {
2488 const required_features = gpu_codegen.featureRequirementsForModule(lowered.program.kernelModule());
2489 const required_subgroup = gpu_codegen.subgroupRequirementsForModule(lowered.program.kernelModule());
2490 const argument_count = try target_product.abi.argumentCount(format, lowered.argument_count);
2491 const static_arguments = try target_product.abi.staticArguments(
2492 allocator,
2493 format,
2494 work.element_count,
2495 lowered.launchGeometry(),
2496 );
2497 errdefer if (static_arguments.len != 0) allocator.free(static_arguments);
2498 const compilation = try target_product.compileKernelForArtifactFormat(
2499 allocator,
2500 format,
2501 lowered.entry_name,
2502 lowered.program.kernelModule(),
2503 target_product.compileOptionsForArtifactFormat(format, work.element_count),
2504 );
2505 return .{
2506 .entry_name = lowered.entry_name,
2507 .argument_count = argument_count,
2508 .payload = compilation.payload,
2509 .lowered_kernel = lowered,
2510 .required_features = required_features,
2511 .required_subgroup = required_subgroup,
2512 .push_constants = compilation.push_constants,
2513 .runtime_scalar_argument_count = lowered.runtimeScalarArgumentCount(),
2514 .static_arguments = static_arguments,
2515 };
2516 }
2517
2518 fn compilePlanFromCapturedKernel(
2519 allocator: std.mem.Allocator,
2520 format: gpu.ArtifactFormat,
2521 lowered: *const kernelization.LoweredKernel,
2522 work: schedule_planning.ScheduleWorkItem,
2523 target_record: *const accy_choir.record.target.Record,
2524 ) gpu.BackendError!CompilePlan {
2525 for (target_record.kernels) |item| {
2526 if (item.lowered.work_item_id != work.id) continue;
2527 if (item.work_dtype != work.dtype or item.element_count != work.element_count) {
2528 return error.InvalidArtifact;
2529 }
2530 const abi = item.abi orelse return error.InvalidArtifact;
2531 const arguments = try allocator.dupe(choir_abi.ScalarArgument, abi.static_arguments);
2532 errdefer allocator.free(arguments);
2533 const compilation = try target_product.compileKernelForArtifactFormat(
2534 allocator,
2535 format,
2536 lowered.entry_name,
2537 lowered.program.kernelModule(),
2538 abi.compile_options,
2539 );
2540 return .{
2541 .entry_name = lowered.entry_name,
2542 .argument_count = abi.argument_count,
2543 .payload = compilation.payload,
2544 .lowered_kernel = lowered,
2545 .required_dtypes = .{ .bits = item.required_dtype_bits },
2546 .required_features = item.required_features,
2547 .required_subgroup = item.required_subgroup,
2548 .push_constants = compilation.push_constants,
2549 .runtime_scalar_argument_count = item.runtime_scalar_argument_count,
2550 .static_arguments = arguments,
2551 };
2552 }
2553 return error.InvalidArtifact;
2554 }
2555
2556 /// Rebuilds the plans inside the caller's `workspace` from the prepared
2557 /// module's stage records. The call compiles each kernel through `handle` and
2558 /// returns an artifact plan the caller owns. A caller uses this to turn a
2559 /// finished compile into device code for one backend. A `workspace` too small
2560 /// for the rebuilt plans gives `error.WorkExhausted`. The rebuilt plans are
2561 /// freed before the call returns, so the returned plan holds copies of the
2562 /// slots and ids and the compiled code, and nothing that points into
2563 /// `workspace`.
2564 pub fn createBackendArtifactPlanFromPreparedModule(
2565 allocator: std.mem.Allocator,
2566 handle: gpu.BackendHandle,
2567 prepared: *const preparation.pipeline.BackendPreparedModule,
2568 options: ArtifactPlanOptions,
2569 workspace: []u8,
2570 comptime configuration: choir.product.operation.Configuration,
2571 ) !BackendArtifactPlan {
2572 const input = try @import("input.zig").InputJob.create(workspace, prepared, configuration);
2573 defer input.destroy();
2574 return createFromPlans(allocator, handle, input.plans(), options);
2575 }
2576
2577 fn createKernelArtifactFromCompilePlan(
2578 handle: gpu.BackendHandle,
2579 format: gpu.ArtifactFormat,
2580 compile_plan: *const CompilePlan,
2581 diagnostic_id: []const u8,
2582 required_dtypes: gpu.DTypeSet,
2583 required_features: choir_abi.Features,
2584 required_subgroup: choir_abi.SubgroupRequirements,
2585 ) gpu.BackendError!gpu.KernelArtifact {
2586 return try handle.createArtifact(.{
2587 .kernel_name = compile_plan.entry_name,
2588 .requested_format = format,
2589 .argument_count = compile_plan.argument_count,
2590 .scalar_argument_count = compile_plan.runtime_scalar_argument_count + @as(u32, @intCast(compile_plan.static_arguments.len)),
2591 .required_dtypes = required_dtypes,
2592 .required_features = required_features,
2593 .required_subgroup = required_subgroup,
2594 .push_constants = compile_plan.push_constants,
2595 .diagnostic_id = diagnostic_id,
2596 .payload = compile_plan.payload,
2597 });
2598 }
2599
2600 fn deinitCompilePayload(allocator: std.mem.Allocator, payload: gpu.CompilePayload) void {
2601 switch (payload) {
2602 .bytes => |bytes| allocator.free(@constCast(bytes)),
2603 .words_u32 => |words| allocator.free(@constCast(words)),
2604 .text => |text| allocator.free(@constCast(text)),
2605 .none => {},
2606 }
2607 }
2608
2609 fn copyKernelArtifact(allocator: std.mem.Allocator, source: gpu.KernelArtifact) gpu.BackendError!gpu.KernelArtifact {
2610 var artifact = gpu.KernelArtifact.init(allocator, .{
2611 .backend = source.backend,
2612 .format = source.format,
2613 .entry_name = source.entry_name,
2614 .argument_count = source.argument_count,
2615 .scalar_argument_count = source.scalar_argument_count,
2616 .diagnostic_id = source.diagnostic_id,
2617 .interface = source.interface,
2618 }) catch return error.OutOfMemory;
2619 errdefer artifact.deinit();
2620
2621 switch (source.payload) {
2622 .none => {},
2623 .bytes => |bytes| try artifact.setOwnedBytes(bytes),
2624 .words_u32 => |words| try artifact.setOwnedWords(words),
2625 .text => |text| try artifact.setOwnedText(text),
2626 .external => return error.InvalidArtifact,
2627 }
2628
2629 return artifact;
2630 }
2631
2632 fn kernelizationMissError(
2633 outline: kernelization_model.KernelOutline,
2634 work: schedule_planning.ScheduleWorkItem,
2635 buffer_plan: *const bufferization.BufferPlanAnalysis,
2636 ) gpu.BackendError {
2637 return switch (outline.kind) {
2638 .elementwise, .shape, .kernel_call, .row_pipeline, .iterate, .flash_attention, .scan => if (work.dtype == .f32) error.UnsupportedOperation else error.CapabilityMismatch,
2639 .dot_general => {
2640 const dot = kernelization.dotGeneralDescriptionForWork(outline, work, buffer_plan) catch |err| return err;
2641 if (dot.input_dtype != .f32 or dot.output_dtype != .f32) return error.CapabilityMismatch;
2642 return error.UnsupportedOperation;
2643 },
2644 .reduction => {
2645 const desc = kernelization.reductionDescriptionForWork(outline, work, buffer_plan) catch |err| return err;
2646 if (desc.input_dtype != .f32 or desc.output_dtype != .f32) return error.CapabilityMismatch;
2647 return error.UnsupportedOperation;
2648 },
2649 };
2650 }
2651
2652 fn elementCountArgumentForWork(
2653 format: gpu.ArtifactFormat,
2654 work: schedule_planning.ScheduleWorkItem,
2655 compile_plan: CompilePlan,
2656 ) gpu.BackendError!ElementCountArgument {
2657 if (compile_plan.element_count_argument) |argument| return argument;
2658 if (compile_plan.lowered_kernel != null) {
2659 return switch (compile_plan.launch) {
2660 .dot_general => .none,
2661 else => .device_buffer_u32,
2662 };
2663 }
2664 _ = format;
2665 _ = work;
2666 return error.UnsupportedOperation;
2667 }
2668
2669 fn elementCountArgumentValueForWork(
2670 work: schedule_planning.ScheduleWorkItem,
2671 compile_plan: CompilePlan,
2672 ) u64 {
2673 _ = compile_plan;
2674 return work.element_count;
2675 }
2676
2677 fn fixedThreadgroupSizeFor(format: gpu.ArtifactFormat) ?u32 {
2678 return switch (format) {
2679 .vulkan_spirv => 64,
2680 .cpu_machine_code, .cpu_object, .webassembly_module => 1,
2681 else => null,
2682 };
2683 }
2684
2685 fn fixedThreadgroupSize(
2686 caps: gpu.BackendCapabilities,
2687 format: gpu.ArtifactFormat,
2688 requested: u32,
2689 ) gpu.BackendError!u32 {
2690 switch (format) {
2691 .vulkan_spirv, .cpu_machine_code, .cpu_object, .webassembly_module => {},
2692 else => return error.UnsupportedArtifactFormat,
2693 }
2694 const max_threads = if (caps.threadgroup.max_threads == 0) 1 else caps.threadgroup.max_threads;
2695 const max_x = if (caps.threadgroup.max_threads_per_dim[0] == 0) 1 else caps.threadgroup.max_threads_per_dim[0];
2696 if (requested > max_threads or requested > max_x) return error.CapabilityMismatch;
2697 return requested;
2698 }
2699
2700 const LaunchCandidateBuffer = struct {
2701 items: [max_launch_resource_candidates]LaunchResourceCandidate = @as([max_launch_resource_candidates]LaunchResourceCandidate, @splat(.{})),
2702 count: usize = 0,
2703 };
2704
2705 fn fixedLaunchResourceCandidates(
2706 caps: gpu.BackendCapabilities,
2707 format: gpu.ArtifactFormat,
2708 resources: schedule_planning.ScheduleResourceEstimate,
2709 resource_class: LaunchResourceClass,
2710 fixed_threadgroup: u32,
2711 ) gpu.BackendError!LaunchCandidateBuffer {
2712 const threadgroup_size = try fixedThreadgroupSize(caps, format, fixed_threadgroup);
2713 var candidates = LaunchCandidateBuffer{};
2714 try appendLaunchCandidate(&candidates, caps, resources, resource_class, threadgroup_size, threadgroup_size);
2715 return candidates;
2716 }
2717
2718 fn dynamicLaunchResourceCandidates(
2719 caps: gpu.BackendCapabilities,
2720 format: gpu.ArtifactFormat,
2721 resources: schedule_planning.ScheduleResourceEstimate,
2722 resource_class: LaunchResourceClass,
2723 subgroup_size: ?u32,
2724 ) gpu.BackendError!LaunchCandidateBuffer {
2725 const limit = try dynamicThreadgroupLimit(caps, format, resource_class);
2726 const target = dynamicThreadgroupTarget(resources, subgroup_size, limit);
2727 var candidates = LaunchCandidateBuffer{};
2728 try appendLaunchCandidate(&candidates, caps, resources, resource_class, target, target);
2729
2730 const canonical = [_]u32{ 32, 64, 128, 256, 512, 1024, 16, 8, 4, 2, 1 };
2731 for (canonical) |candidate| {
2732 if (!candidateThreadgroupAllowed(candidate, subgroup_size, target)) continue;
2733 try appendLaunchCandidate(&candidates, caps, resources, resource_class, candidate, target);
2734 }
2735 try appendLaunchCandidate(&candidates, caps, resources, resource_class, limit, target);
2736
2737 if (subgroup_size) |size| {
2738 if (size != 0 and limit >= size) {
2739 try appendLaunchCandidate(&candidates, caps, resources, resource_class, size, target);
2740 try appendLaunchCandidate(&candidates, caps, resources, resource_class, largestMultipleAtMost(limit, size), target);
2741 }
2742 }
2743
2744 return candidates;
2745 }
2746
2747 fn dynamicThreadgroupLimit(
2748 caps: gpu.BackendCapabilities,
2749 format: gpu.ArtifactFormat,
2750 resource_class: LaunchResourceClass,
2751 ) gpu.BackendError!u32 {
2752 const preferred = try preferredDynamicThreadgroupSize(format, resource_class);
2753 const max_threads = if (caps.threadgroup.max_threads == 0) 1 else caps.threadgroup.max_threads;
2754 const max_x = if (caps.threadgroup.max_threads_per_dim[0] == 0) 1 else caps.threadgroup.max_threads_per_dim[0];
2755 const limit = @min(@min(preferred, max_threads), max_x);
2756 if (limit == 0) return error.CapabilityMismatch;
2757
2758 return limit;
2759 }
2760
2761 fn dynamicThreadgroupTarget(
2762 resources: schedule_planning.ScheduleResourceEstimate,
2763 subgroup_size: ?u32,
2764 limit: u32,
2765 ) u32 {
2766 var target = elementThreadTarget(resources.element_count, limit);
2767 if (subgroup_size) |size| {
2768 target = alignThreadgroupToSubgroup(target, size, limit);
2769 }
2770 return @max(target, 1);
2771 }
2772
2773 fn appendLaunchCandidate(
2774 candidates: *LaunchCandidateBuffer,
2775 caps: gpu.BackendCapabilities,
2776 resources: schedule_planning.ScheduleResourceEstimate,
2777 resource_class: LaunchResourceClass,
2778 threadgroup_size: u32,
2779 target_threadgroup_size: u32,
2780 ) gpu.BackendError!void {
2781 if (threadgroup_size == 0) return;
2782 const geometry = choir_abi.LaunchGeometry{
2783 .grid = .{ try gridSizeFor(resources.element_count, threadgroup_size), 1, 1 },
2784 .threadgroup = .{ threadgroup_size, 1, 1 },
2785 };
2786 const candidate = LaunchResourceCandidate{
2787 .geometry = geometry,
2788 .score = candidateScore(threadgroup_size, target_threadgroup_size),
2789 .estimated_static_bytes_per_threadgroup = scaleResourceToThreadgroup(
2790 resources.static_total_bytes,
2791 resources.element_count,
2792 threadgroup_size,
2793 ),
2794 .estimated_element_ops_per_threadgroup = scaleResourceToThreadgroup(
2795 resources.estimated_element_ops,
2796 resources.element_count,
2797 threadgroup_size,
2798 ),
2799 };
2800 caps.validateLaunchGeometry(candidate.geometry) catch return;
2801 for (candidates.items[0..candidates.count]) |existing| {
2802 if (sameLaunchGeometry(existing.geometry, candidate.geometry)) return;
2803 }
2804
2805 var insert_index: usize = 0;
2806 while (insert_index < candidates.count and !candidateSortsBefore(candidate, candidates.items[insert_index], resource_class)) {
2807 insert_index += 1;
2808 }
2809 if (insert_index >= max_launch_resource_candidates) return;
2810
2811 if (candidates.count < max_launch_resource_candidates) candidates.count += 1;
2812 var index = candidates.count - 1;
2813 while (index > insert_index) : (index -= 1) {
2814 candidates.items[index] = candidates.items[index - 1];
2815 }
2816 candidates.items[insert_index] = candidate;
2817 }
2818
2819 fn sameLaunchGeometry(lhs: choir_abi.LaunchGeometry, rhs: choir_abi.LaunchGeometry) bool {
2820 return lhs.grid[0] == rhs.grid[0] and
2821 lhs.grid[1] == rhs.grid[1] and
2822 lhs.grid[2] == rhs.grid[2] and
2823 lhs.threadgroup[0] == rhs.threadgroup[0] and
2824 lhs.threadgroup[1] == rhs.threadgroup[1] and
2825 lhs.threadgroup[2] == rhs.threadgroup[2] and
2826 lhs.dynamic_shared_memory_bytes == rhs.dynamic_shared_memory_bytes;
2827 }
2828
2829 fn candidateThreadgroupAllowed(
2830 candidate: u32,
2831 subgroup_size: ?u32,
2832 target: u32,
2833 ) bool {
2834 if (candidate == 0) return false;
2835 if (subgroup_size) |size| {
2836 if (size <= 1) return true;
2837 if (target < size) return true;
2838 return candidate >= size and candidate % size == 0;
2839 }
2840 return true;
2841 }
2842
2843 fn candidateSortsBefore(
2844 lhs: LaunchResourceCandidate,
2845 rhs: LaunchResourceCandidate,
2846 resource_class: LaunchResourceClass,
2847 ) bool {
2848 const lhs_threadgroup = lhs.geometry.threadgroup[0];
2849 const rhs_threadgroup = rhs.geometry.threadgroup[0];
2850 if (lhs.score != rhs.score) return lhs.score < rhs.score;
2851 return switch (resource_class) {
2852 .memory_bound => lhs_threadgroup > rhs_threadgroup,
2853 .compute_weighted => lhs_threadgroup < rhs_threadgroup,
2854 else => lhs_threadgroup < rhs_threadgroup,
2855 };
2856 }
2857
2858 fn candidateScore(candidate: u32, target: u32) u32 {
2859 return if (candidate >= target) candidate - target else target - candidate;
2860 }
2861
2862 fn preferredDynamicThreadgroupSize(
2863 format: gpu.ArtifactFormat,
2864 resource_class: LaunchResourceClass,
2865 ) gpu.BackendError!u32 {
2866 return switch (format) {
2867 .cuda_ptx, .metal_msl, .webgpu_wgsl => switch (resource_class) {
2868 .compute_weighted => 128,
2869 else => 256,
2870 },
2871 else => error.UnsupportedArtifactFormat,
2872 };
2873 }
2874
2875 fn classifyLaunchResources(resources: schedule_planning.ScheduleResourceEstimate) LaunchResourceClass {
2876 if (!resources.static_bytes_complete) return .unknown;
2877 if (resources.static_total_bytes == 0) return .unknown;
2878 if (resources.estimated_element_ops == 0) return .unknown;
2879 const ops_per_kib = resources.elementOpsPerKiB();
2880 if (ops_per_kib <= 256) return .memory_bound;
2881 if (ops_per_kib <= 1024) return .balanced;
2882 return .compute_weighted;
2883 }
2884
2885 fn subgroupSizeFor(caps: gpu.BackendCapabilities) ?u32 {
2886 if (!caps.subgroup.supported) return null;
2887 if (caps.subgroup.size_max != 0) return caps.subgroup.size_max;
2888 if (caps.subgroup.size_min != 0) return caps.subgroup.size_min;
2889 return null;
2890 }
2891
2892 fn elementThreadTarget(element_count: u64, limit: u32) u32 {
2893 if (element_count == 0) return 1;
2894 if (element_count >= limit) return limit;
2895 const count: u32 = @intCast(element_count);
2896 return std.math.ceilPowerOfTwo(u32, count) catch limit;
2897 }
2898
2899 fn alignThreadgroupToSubgroup(target: u32, subgroup_size: u32, limit: u32) u32 {
2900 if (subgroup_size <= 1 or limit < subgroup_size) return target;
2901 const rounded = roundUpToMultiple(target, subgroup_size);
2902 if (rounded <= limit) return rounded;
2903 return largestMultipleAtMost(limit, subgroup_size);
2904 }
2905
2906 fn roundUpToMultiple(value: u32, multiple: u32) u32 {
2907 const remainder = value % multiple;
2908 if (remainder == 0) return value;
2909 return value + (multiple - remainder);
2910 }
2911
2912 fn largestMultipleAtMost(value: u32, multiple: u32) u32 {
2913 return value - (value % multiple);
2914 }
2915
2916 fn gridSizeFor(element_count: u64, threads: u32) gpu.BackendError!u32 {
2917 if (threads == 0) return error.CapabilityMismatch;
2918 const grid = if (element_count == 0)
2919 1
2920 else
2921 ((element_count - 1) / @as(u64, threads)) + 1;
2922 if (grid > std.math.maxInt(u32)) return error.LaunchArgumentMismatch;
2923 return @intCast(grid);
2924 }
2925
2926 fn scaleResourceToThreadgroup(total: u64, element_count: u64, threads: u32) u64 {
2927 if (total == 0 or element_count == 0 or threads == 0) return 0;
2928 return ceilDivSaturated(saturatedMul(total, threads), element_count);
2929 }
2930
2931 fn ceilDivSaturated(numerator: u64, denominator: u64) u64 {
2932 if (denominator == 0) return 0;
2933 return (numerator / denominator) + @intFromBool(numerator % denominator != 0);
2934 }
2935
2936 fn saturatedMul(lhs: u64, rhs: u64) u64 {
2937 return std.math.mul(u64, lhs, rhs) catch std.math.maxInt(u64);
2938 }
2939
2940 const testing = std.testing;
2941
2942 const OwnedSemanticChoirModule = struct {
2943 module: *semantic.SemanticModule,
2944 choir_module: *ir.Operation,
2945 };
2946
2947 const OwnedSemanticChoirModuleWithPayload = struct {
2948 module: *semantic.SemanticModule,
2949 choir_module: *ir.Operation,
2950 payload: [32]u8,
2951 };
2952
2953 fn operationTreeContainsName(op: *ir.Operation, name: []const u8) bool {
2954 if (std.mem.eql(u8, op.name.name, name)) return true;
2955 for (op.regions.items) |*region| {
2956 var block_iter = region.getBlocks();
2957 while (block_iter.next()) |block| {
2958 var op_node = block.operations.head;
2959 while (op_node) |node| {
2960 const child: *ir.Operation = @ptrCast(@alignCast(node));
2961 if (operationTreeContainsName(child, name)) return true;
2962 op_node = child.next_op;
2963 }
2964 }
2965 }
2966 return false;
2967 }
2968
2969 fn addChoirModule(allocator: std.mem.Allocator, name: []const u8) !OwnedSemanticChoirModule {
2970 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
2971 errdefer builder.deinit();
2972 const f32_8 = try builder.tensor(.f32, &.{8});
2973 var fb = try builder.beginFunction(name, &.{ f32_8, f32_8 }, &.{f32_8});
2974 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
2975 try fb.return_(&.{sum});
2976 try fb.finish();
2977 const module = try builder.finish();
2978
2979 return .{
2980 .module = module,
2981 .choir_module = module.choir_module,
2982 };
2983 }
2984
2985 fn kernelCallChoirModule(allocator: std.mem.Allocator, name: []const u8) !OwnedSemanticChoirModule {
2986 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
2987 errdefer builder.deinit();
2988 const f32_8 = try builder.tensor(.f32, &.{8});
2989 var fb = try builder.beginFunction(name, &.{ f32_8, f32_8 }, &.{f32_8});
2990 const call = try fb.kernelCall(
2991 &.{ fb.parameter(0), fb.parameter(1) },
2992 &.{f32_8},
2993 .{
2994 .target = "accy.custom.scale",
2995 .operand_effects = &.{ .read, .write },
2996 .result_aliases = &.{null},
2997 },
2998 );
2999 try fb.return_(&.{call.getFirstResult()});
3000 try fb.finish();
3001 const module = try builder.finish();
3002
3003 return .{
3004 .module = module,
3005 .choir_module = module.choir_module,
3006 };
3007 }
3008
3009 fn reshapeChoirModule(allocator: std.mem.Allocator, name: []const u8) !OwnedSemanticChoirModule {
3010 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
3011 errdefer builder.deinit();
3012 const f32_2x4 = try builder.tensor(.f32, &.{ 2, 4 });
3013 const f32_8 = try builder.tensor(.f32, &.{8});
3014 var fb = try builder.beginFunction(name, &.{f32_2x4}, &.{f32_8});
3015 const reshaped = try fb.reshape(fb.parameter(0), f32_8, &.{8});
3016 try fb.return_(&.{reshaped});
3017 try fb.finish();
3018 const module = try builder.finish();
3019
3020 return .{
3021 .module = module,
3022 .choir_module = module.choir_module,
3023 };
3024 }
3025
3026 fn fusedAddMulChoirModule(allocator: std.mem.Allocator, name: []const u8) !OwnedSemanticChoirModule {
3027 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
3028 errdefer builder.deinit();
3029 const f32_8 = try builder.tensor(.f32, &.{8});
3030 var fb = try builder.beginFunction(name, &.{ f32_8, f32_8, f32_8 }, &.{f32_8});
3031 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
3032 const product = try fb.mul(sum, fb.parameter(2));
3033 try fb.return_(&.{product});
3034 try fb.finish();
3035 const module = try builder.finish();
3036
3037 return .{
3038 .module = module,
3039 .choir_module = module.choir_module,
3040 };
3041 }
3042
3043 fn dotGeneralChoirModule(allocator: std.mem.Allocator, name: []const u8) !OwnedSemanticChoirModule {
3044 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
3045 errdefer builder.deinit();
3046 const f32_16x16 = try builder.tensor(.f32, &.{ 16, 16 });
3047 var fb = try builder.beginFunction(name, &.{ f32_16x16, f32_16x16 }, &.{f32_16x16});
3048 const product = try fb.dotGeneral(
3049 fb.parameter(0),
3050 fb.parameter(1),
3051 f32_16x16,
3052 &.{1},
3053 &.{0},
3054 &.{},
3055 &.{},
3056 );
3057 try fb.return_(&.{product});
3058 try fb.finish();
3059 const module = try builder.finish();
3060
3061 return .{
3062 .module = module,
3063 .choir_module = module.choir_module,
3064 };
3065 }
3066
3067 fn dotGeneralF16ChoirModule(allocator: std.mem.Allocator, name: []const u8) !OwnedSemanticChoirModule {
3068 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
3069 errdefer builder.deinit();
3070 const f16_16x16 = try builder.tensor(.f16, &.{ 16, 16 });
3071 const f32_16x16 = try builder.tensor(.f32, &.{ 16, 16 });
3072 var fb = try builder.beginFunction(name, &.{ f16_16x16, f16_16x16 }, &.{f32_16x16});
3073 const product = try fb.dotGeneral(
3074 fb.parameter(0),
3075 fb.parameter(1),
3076 f32_16x16,
3077 &.{1},
3078 &.{0},
3079 &.{},
3080 &.{},
3081 );
3082 try fb.return_(&.{product});
3083 try fb.finish();
3084 const module = try builder.finish();
3085
3086 return .{
3087 .module = module,
3088 .choir_module = module.choir_module,
3089 };
3090 }
3091
3092 fn reduceChoirModule(allocator: std.mem.Allocator, name: []const u8) !OwnedSemanticChoirModule {
3093 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
3094 errdefer builder.deinit();
3095 const f32_256 = try builder.tensor(.f32, &.{256});
3096 const f32_scalar = try builder.tensor(.f32, &.{});
3097 var fb = try builder.beginFunction(name, &.{f32_256}, &.{f32_scalar});
3098 const zero_value: f32 = 0.0;
3099 const zero = try fb.constant(f32_scalar, std.mem.asBytes(&zero_value));
3100 const reduced = try fb.reduce(fb.parameter(0), zero, f32_scalar, "sum", &.{0});
3101 try fb.return_(&.{reduced});
3102 try fb.finish();
3103 const module = try builder.finish();
3104
3105 return .{
3106 .module = module,
3107 .choir_module = module.choir_module,
3108 };
3109 }
3110
3111 fn reduceI32ChoirModule(allocator: std.mem.Allocator, name: []const u8) !OwnedSemanticChoirModule {
3112 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
3113 errdefer builder.deinit();
3114 const i32_256 = try builder.tensor(.i32, &.{256});
3115 const i32_scalar = try builder.tensor(.i32, &.{});
3116 var fb = try builder.beginFunction(name, &.{i32_256}, &.{i32_scalar});
3117 const zero_value: i32 = 0;
3118 const zero = try fb.constant(i32_scalar, std.mem.asBytes(&zero_value));
3119 const reduced = try fb.reduce(fb.parameter(0), zero, i32_scalar, "sum", &.{0});
3120 try fb.return_(&.{reduced});
3121 try fb.finish();
3122 const module = try builder.finish();
3123
3124 return .{
3125 .module = module,
3126 .choir_module = module.choir_module,
3127 };
3128 }
3129
3130 fn reduceRank2Axis1ChoirModule(allocator: std.mem.Allocator, name: []const u8) !OwnedSemanticChoirModule {
3131 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
3132 errdefer builder.deinit();
3133 const f32_16x16 = try builder.tensor(.f32, &.{ 16, 16 });
3134 const f32_16 = try builder.tensor(.f32, &.{16});
3135 const f32_scalar = try builder.tensor(.f32, &.{});
3136 var fb = try builder.beginFunction(name, &.{f32_16x16}, &.{f32_16});
3137 const zero_value: f32 = 0.0;
3138 const zero = try fb.constant(f32_scalar, std.mem.asBytes(&zero_value));
3139 const reduced = try fb.reduce(fb.parameter(0), zero, f32_16, "sum", &.{1});
3140 try fb.return_(&.{reduced});
3141 try fb.finish();
3142 const module = try builder.finish();
3143
3144 return .{
3145 .module = module,
3146 .choir_module = module.choir_module,
3147 };
3148 }
3149
3150 fn addConstantChoirModule(allocator: std.mem.Allocator, name: []const u8) !OwnedSemanticChoirModuleWithPayload {
3151 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
3152 errdefer builder.deinit();
3153 const f32_8 = try builder.tensor(.f32, &.{8});
3154 var fb = try builder.beginFunction(name, &.{f32_8}, &.{f32_8});
3155 const values = [_]f32{ 1, 2, 3, 4, 5, 6, 7, 8 };
3156 var payload: [32]u8 = undefined;
3157 @memcpy(payload[0..], std.mem.sliceAsBytes(values[0..]));
3158 const c = try fb.constant(f32_8, &payload);
3159 const sum = try fb.add(fb.parameter(0), c);
3160 try fb.return_(&.{sum});
3161 try fb.finish();
3162 const module = try builder.finish();
3163
3164 return .{
3165 .module = module,
3166 .choir_module = module.choir_module,
3167 .payload = payload,
3168 };
3169 }
3170
3171 fn expectScalarU32(actual: choir_abi.ScalarArgument, expected: u32) !void {
3172 switch (actual) {
3173 .u32 => |value| try testing.expectEqual(expected, value),
3174 else => return error.TestExpectedScalarU32,
3175 }
3176 }
3177
3178 fn expectKernelLayoutFingerprints(plan: *const BackendArtifactPlan, kernel: PlannedKernel) !void {
3179 const output_slot = plan.slotById(kernel.output_slot_id) orelse return error.MissingSlot;
3180 try testing.expectEqual(output_slot.layout_fingerprint, kernel.output_layout_fingerprint);
3181 try testing.expect(kernel.output_layout_fingerprint != 0);
3182 try testing.expectEqual(try plan.layoutFingerprintForInputs(kernel.input_slot_ids), kernel.input_layout_fingerprint);
3183 try testing.expect(kernel.input_layout_fingerprint != 0);
3184 }
3185
3186 fn expectReductionTile(
3187 tile: LaunchTilePlan,
3188 output_tile_elements: u32,
3189 reduction_extent_value: u32,
3190 rank: u32,
3191 axis: u32,
3192 input_tile_bytes: u32,
3193 output_tile_bytes: u32,
3194 ) !void {
3195 try testing.expect(tile.active());
3196 try testing.expectEqual(LaunchTileKind.reduction, tile.kind);
3197 try testing.expectEqual(output_tile_elements, tile.m);
3198 try testing.expectEqual(reduction_extent_value, tile.n);
3199 try testing.expectEqual(axis, tile.k);
3200 try testing.expectEqual(@as(u32, 1), tile.batch);
3201 try testing.expectEqual(choir_abi.DType.f32, tile.input_dtype.?);
3202 try testing.expectEqual(choir_abi.DType.f32, tile.output_dtype.?);
3203 try testing.expectEqual(input_tile_bytes, tile.input_tile_bytes);
3204 try testing.expectEqual(output_tile_bytes, tile.output_tile_bytes);
3205 try testing.expectEqual(@as(u32, 0), tile.scratch_memory_bytes);
3206 try testing.expectEqual(LaunchReductionKind.sum, tile.reduction_kind);
3207 try testing.expectEqual(rank, tile.reduction_rank);
3208 try testing.expectEqual(axis, tile.reduction_axis);
3209 try testing.expectEqual(reduction_extent_value, tile.reduction_extent);
3210 }
3211
3212 fn expectDotGeneralTile(
3213 tile: LaunchTilePlan,
3214 m: u32,
3215 n: u32,
3216 k: u32,
3217 input_tile_bytes: u32,
3218 output_tile_bytes: u32,
3219 ) !void {
3220 try testing.expect(tile.active());
3221 try testing.expectEqual(LaunchTileKind.dot_general, tile.kind);
3222 try testing.expectEqual(m, tile.m);
3223 try testing.expectEqual(n, tile.n);
3224 try testing.expectEqual(k, tile.k);
3225 try testing.expectEqual(@as(u32, 1), tile.batch);
3226 try testing.expectEqual(choir_abi.DType.f32, tile.input_dtype.?);
3227 try testing.expectEqual(choir_abi.DType.f32, tile.output_dtype.?);
3228 try testing.expectEqual(input_tile_bytes, tile.input_tile_bytes);
3229 try testing.expectEqual(output_tile_bytes, tile.output_tile_bytes);
3230 try testing.expectEqual(@as(u32, 0), tile.scratch_memory_bytes);
3231 }
3232
3233 fn expectDotGeneralCandidate(
3234 candidate: LaunchResourceCandidate,
3235 grid_xy: [2]u32,
3236 threadgroup_xy: [2]u32,
3237 score: u32,
3238 estimated_static_bytes: u64,
3239 estimated_element_ops: u64,
3240 m: u32,
3241 n: u32,
3242 k: u32,
3243 input_tile_bytes: u32,
3244 output_tile_bytes: u32,
3245 ) !void {
3246 try testing.expectEqual(grid_xy[0], candidate.geometry.grid[0]);
3247 try testing.expectEqual(grid_xy[1], candidate.geometry.grid[1]);
3248 try testing.expectEqual(@as(u32, 1), candidate.geometry.grid[2]);
3249 try testing.expectEqual(threadgroup_xy[0], candidate.geometry.threadgroup[0]);
3250 try testing.expectEqual(threadgroup_xy[1], candidate.geometry.threadgroup[1]);
3251 try testing.expectEqual(@as(u32, 1), candidate.geometry.threadgroup[2]);
3252 try testing.expectEqual(score, candidate.score);
3253 try testing.expectEqual(estimated_static_bytes, candidate.estimated_static_bytes_per_threadgroup);
3254 try testing.expectEqual(estimated_element_ops, candidate.estimated_element_ops_per_threadgroup);
3255 try expectDotGeneralTile(candidate.tile, m, n, k, input_tile_bytes, output_tile_bytes);
3256 }
3257
3258 fn deadTemporaryChoirModule(allocator: std.mem.Allocator, name: []const u8) !OwnedSemanticChoirModule {
3259 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
3260 errdefer builder.deinit();
3261 const f32_8 = try builder.tensor(.f32, &.{8});
3262 const f32_16 = try builder.tensor(.f32, &.{16});
3263 var fb = try builder.beginFunction(name, &.{ f32_8, f32_8, f32_8, f32_8, f32_8, f32_8 }, &.{ f32_8, f32_8 });
3264 const shared = try fb.add(fb.parameter(0), fb.parameter(1));
3265 const product = try fb.mul(shared, fb.parameter(2));
3266 _ = try fb.concatenate(&.{ shared, fb.parameter(3) }, f32_16, 0);
3267 const sum = try fb.add(fb.parameter(4), fb.parameter(5));
3268 try fb.return_(&.{ product, sum });
3269 try fb.finish();
3270 const module = try builder.finish();
3271
3272 return .{
3273 .module = module,
3274 .choir_module = module.choir_module,
3275 };
3276 }
3277
3278 fn launchResourceCaps(
3279 kind: gpu.BackendKind,
3280 format: gpu.ArtifactFormat,
3281 max_threads: u32,
3282 max_x: u32,
3283 subgroup_size: ?u32,
3284 ) gpu.BackendCapabilities {
3285 return .{
3286 .identity = .{
3287 .backend = kind,
3288 .family = gpu.familyForBackendKind(kind),
3289 .name = "resource-test",
3290 },
3291 .subgroup = if (subgroup_size) |size| .{
3292 .supported = true,
3293 .size_min = size,
3294 .size_max = size,
3295 .shuffle = true,
3296 .ballot = true,
3297 .vote = true,
3298 .arithmetic = true,
3299 } else .{},
3300 .threadgroup = .{
3301 .max_threads = max_threads,
3302 .max_blocks = .{ 65_535, 65_535, 65_535 },
3303 .max_threads_per_dim = .{ max_x, 1, 1 },
3304 .max_grid_per_dim = .{ 65_535, 65_535, 65_535 },
3305 },
3306 .dtypes = gpu.DTypeSet.init(&.{.f32}),
3307 .artifact_formats = gpu.ArtifactFormatSet.init(&.{format}),
3308 };
3309 }
3310
3311 fn launchResources(
3312 element_count: u64,
3313 op_count: usize,
3314 static_total_bytes: u64,
3315 ) schedule_planning.ScheduleResourceEstimate {
3316 return .{
3317 .element_count = element_count,
3318 .element_size = 4,
3319 .op_count = op_count,
3320 .static_read_bytes = static_total_bytes,
3321 .static_total_bytes = static_total_bytes,
3322 .estimated_element_ops = element_count * @as(u64, @intCast(op_count)),
3323 };
3324 }
3325
3326 const ArtifactPlanningBackendState = struct {
3327 allocator: std.mem.Allocator,
3328 kind: gpu.BackendKind,
3329 format: gpu.ArtifactFormat,
3330
3331 fn init(allocator: std.mem.Allocator, kind: gpu.BackendKind) ArtifactPlanningBackendState {
3332 return .{
3333 .allocator = allocator,
3334 .kind = kind,
3335 .format = defaultArtifactFormat(kind).?,
3336 };
3337 }
3338
3339 fn deinit(self: *ArtifactPlanningBackendState) void {
3340 self.* = undefined;
3341 }
3342
3343 fn handle(self: *ArtifactPlanningBackendState) gpu.BackendHandle {
3344 return .{
3345 .ptr = self,
3346 .vtable = &artifact_planning_backend_vtable,
3347 .kind = self.kind,
3348 };
3349 }
3350 };
3351
3352 fn artifactPlanningQueryCapabilities(ptr: *anyopaque) gpu.BackendError!gpu.BackendCapabilities {
3353 const state: *ArtifactPlanningBackendState = @ptrCast(@alignCast(ptr));
3354 const dtypes = switch (state.kind) {
3355 .webgpu => gpu.DTypeSet.init(&.{ .i1, .i32, .u32, .f32 }),
3356 .wasm => gpu.DTypeSet.init(&.{ .i1, .i32, .u32, .i64, .u64, .f32, .f64 }),
3357 .cpu => gpu.DTypeSet.init(&.{ .i1, .i32, .u32, .i64, .u64, .f32, .f64 }),
3358 else => gpu.DTypeSet.init(&.{ .i1, .i32, .u32, .f16, .f32 }),
3359 };
3360 const subgroup_size: u32 = switch (state.kind) {
3361 .vulkan, .webgpu, .cpu, .wasm => 0,
3362 else => 32,
3363 };
3364 return .{
3365 .identity = .{
3366 .backend = state.kind,
3367 .family = gpu.familyForBackendKind(state.kind),
3368 .name = "artifact-planning",
3369 },
3370 .memory = switch (state.kind) {
3371 .cuda => .{
3372 .shared_memory_per_threadgroup_bytes = 48 * 1024,
3373 .constant_memory_bytes = 64 * 1024,
3374 .min_buffer_alignment = 256,
3375 },
3376 .metal => .{
3377 .min_buffer_alignment = 256,
3378 .unified_memory = true,
3379 .host_visible_device_memory = true,
3380 },
3381 .webgpu => .{
3382 .shared_memory_per_threadgroup_bytes = 16 * 1024,
3383 .min_buffer_alignment = 4,
3384 },
3385 .wasm => .{
3386 .min_buffer_alignment = 1,
3387 .host_visible_device_memory = true,
3388 },
3389 else => .{
3390 .min_buffer_alignment = 16,
3391 .host_visible_device_memory = true,
3392 },
3393 },
3394 .subgroup = if (subgroup_size == 0) .{} else .{
3395 .supported = true,
3396 .size_min = subgroup_size,
3397 .size_max = subgroup_size,
3398 .shuffle = true,
3399 .ballot = true,
3400 .vote = true,
3401 .arithmetic = true,
3402 .scan = state.kind == .cuda or state.kind == .metal,
3403 },
3404 .threadgroup = .{
3405 .max_threads = if (state.kind == .webgpu) 256 else 1024,
3406 .max_blocks = .{ 65_535, 65_535, 65_535 },
3407 .max_threads_per_dim = if (state.kind == .webgpu) .{ 256, 256, 64 } else .{ 1024, 1024, 64 },
3408 .max_grid_per_dim = .{ 65_535, 65_535, 65_535 },
3409 .shared_memory_bytes = if (state.kind == .cuda) 48 * 1024 else if (state.kind == .webgpu) 16 * 1024 else 0,
3410 },
3411 .dtypes = dtypes,
3412 .layouts = .{
3413 .row_major = true,
3414 .compact_strides = true,
3415 .broadcast_strides = true,
3416 .tiled = true,
3417 .vectorized = state.kind == .cuda,
3418 .opaque_backend_layouts = true,
3419 },
3420 .runtime = .{
3421 .driver_loaded = false,
3422 .device_context = false,
3423 .streams = true,
3424 .events = true,
3425 .timeline_events = state.kind == .vulkan,
3426 },
3427 .features = .{
3428 .atomic_i32 = state.kind != .webgpu and state.kind != .wasm,
3429 .atomic_u32 = state.kind != .webgpu and state.kind != .wasm,
3430 .atomic_index = state.kind != .webgpu and state.kind != .wasm,
3431 .atomic_f32_add_device = state.kind == .cuda or state.kind == .metal,
3432 .atomic_f32_add_shared = state.kind == .cuda,
3433 .async_copy = state.kind == .metal,
3434 .dynamic_shared_memory = state.kind == .cuda,
3435 },
3436 .artifact_formats = gpu.ArtifactFormatSet.init(&.{state.format}),
3437 };
3438 }
3439
3440 fn artifactPlanningCreateArtifact(
3441 ptr: *anyopaque,
3442 request: gpu.CompileRequest,
3443 ) gpu.BackendError!gpu.KernelArtifact {
3444 const state: *ArtifactPlanningBackendState = @ptrCast(@alignCast(ptr));
3445 if (request.requested_format != state.format) return error.UnsupportedOperation;
3446 var artifact = gpu.KernelArtifact.init(state.allocator, .{
3447 .backend = state.kind,
3448 .format = state.format,
3449 .entry_name = request.kernel_name,
3450 .argument_count = request.argument_count,
3451 .scalar_argument_count = request.scalar_argument_count,
3452 .diagnostic_id = request.diagnostic_id,
3453 }) catch return error.OutOfMemory;
3454 errdefer artifact.deinit();
3455 switch (state.format) {
3456 .cuda_ptx, .metal_msl, .webgpu_wgsl => switch (request.payload) {
3457 .text => |text| try artifact.setOwnedText(text),
3458 .bytes => |bytes| try artifact.setOwnedText(bytes),
3459 else => return error.InvalidArtifact,
3460 },
3461 .vulkan_spirv => switch (request.payload) {
3462 .words_u32 => |words| try artifact.setOwnedWords(words),
3463 else => return error.InvalidArtifact,
3464 },
3465 .cpu_machine_code, .cpu_object, .webassembly_module => switch (request.payload) {
3466 .bytes => |bytes| try artifact.setOwnedBytes(bytes),
3467 else => return error.InvalidArtifact,
3468 },
3469 else => return error.UnsupportedArtifactFormat,
3470 }
3471 return artifact;
3472 }
3473
3474 const artifact_planning_backend_vtable = gpu.BackendVTable{
3475 .query_capabilities = artifactPlanningQueryCapabilities,
3476 .create_artifact = artifactPlanningCreateArtifact,
3477 };
3478
3479 fn createTestBackendArtifactPlan(
3480 allocator: std.mem.Allocator,
3481 handle: gpu.BackendHandle,
3482 pass_ctx: *passes.PassContext,
3483 choir_module: *ir.Operation,
3484 options: ArtifactPlanOptions,
3485 ) !BackendArtifactPlan {
3486 const lowered_kernels = try kernelization.getKernelizationAnalysis(pass_ctx, choir_module);
3487 return try createBackendArtifactPlan(
3488 allocator,
3489 handle,
3490 .{
3491 .pass_ctx = pass_ctx,
3492 .choir_module = choir_module,
3493 .lowered_kernels = lowered_kernels,
3494 },
3495 options,
3496 );
3497 }
3498
3499 test "Choir launch resource plan uses subgroup-aligned dynamic threadgroups" {
3500 const caps = launchResourceCaps(.cuda, .cuda_ptx, 96, 96, 32);
3501
3502 const small = try createLaunchResourcePlan(caps, .cuda_ptx, launchResources(17, 1, 204));
3503 try testing.expectEqual(@as(u32, 32), small.geometry.threadgroup[0]);
3504 try testing.expectEqual(@as(u32, 1), small.geometry.grid[0]);
3505 try testing.expectEqual(@as(?u32, 32), small.subgroup_size);
3506 try testing.expect(small.subgroup_aligned);
3507 try testing.expect(!small.fixed_threadgroup);
3508 try testing.expectEqual(LaunchResourceClass.memory_bound, small.resource_class);
3509 try testing.expectEqual(@as(u64, 384), small.estimated_static_bytes_per_threadgroup);
3510 try testing.expectEqual(@as(u64, 32), small.estimated_element_ops_per_threadgroup);
3511 try testing.expectEqual(LaunchTileKind.none, small.tile.kind);
3512 try testing.expect(!small.tile.active());
3513 try testing.expectEqual(@as(usize, 3), small.candidate_count);
3514 try testing.expectEqual(@as(u32, 32), small.selectedCandidate().?.geometry.threadgroup[0]);
3515 try testing.expectEqual(LaunchTileKind.none, small.selectedCandidate().?.tile.kind);
3516 try testing.expectEqual(@as(u32, 64), small.candidates[1].geometry.threadgroup[0]);
3517 try testing.expectEqual(@as(u32, 96), small.candidates[2].geometry.threadgroup[0]);
3518
3519 const capped = try createLaunchResourcePlan(caps, .cuda_ptx, launchResources(1000, 1, 12000));
3520 try testing.expectEqual(@as(u32, 96), capped.geometry.threadgroup[0]);
3521 try testing.expectEqual(@as(u32, 11), capped.geometry.grid[0]);
3522 try testing.expect(capped.subgroup_aligned);
3523 try testing.expectEqual(@as(usize, 3), capped.candidate_count);
3524 try testing.expectEqual(@as(u32, 96), capped.candidates[0].geometry.threadgroup[0]);
3525 try testing.expectEqual(@as(u32, 64), capped.candidates[1].geometry.threadgroup[0]);
3526 try testing.expectEqual(@as(u32, 32), capped.candidates[2].geometry.threadgroup[0]);
3527 }
3528
3529 test "Choir launch resource plan lowers dynamic threadgroups for compute weighted work" {
3530 const caps = launchResourceCaps(.cuda, .cuda_ptx, 256, 256, 32);
3531
3532 const plan = try createLaunchResourcePlan(caps, .cuda_ptx, launchResources(4096, 64, 65536));
3533 try testing.expectEqual(LaunchResourceClass.compute_weighted, plan.resource_class);
3534 try testing.expectEqual(@as(u32, 128), plan.geometry.threadgroup[0]);
3535 try testing.expectEqual(@as(u32, 32), plan.geometry.grid[0]);
3536 try testing.expectEqual(@as(u64, 4096), plan.element_ops_per_kib);
3537 try testing.expectEqual(@as(u64, 2048), plan.estimated_static_bytes_per_threadgroup);
3538 try testing.expectEqual(@as(u64, 8192), plan.estimated_element_ops_per_threadgroup);
3539 try testing.expectEqual(@as(usize, 4), plan.candidate_count);
3540 try testing.expectEqual(@as(u32, 128), plan.candidates[0].geometry.threadgroup[0]);
3541 try testing.expectEqual(@as(u32, 64), plan.candidates[1].geometry.threadgroup[0]);
3542 try testing.expectEqual(@as(u32, 32), plan.candidates[2].geometry.threadgroup[0]);
3543 try testing.expectEqual(@as(u32, 256), plan.candidates[3].geometry.threadgroup[0]);
3544 }
3545
3546 test "Choir launch resource plan preserves fixed Vulkan local size" {
3547 const caps = launchResourceCaps(.vulkan, .vulkan_spirv, 128, 128, 32);
3548
3549 const plan = try createLaunchResourcePlan(caps, .vulkan_spirv, launchResources(65, 1, 780));
3550 try testing.expectEqual(@as(u32, 64), plan.geometry.threadgroup[0]);
3551 try testing.expectEqual(@as(u32, 2), plan.geometry.grid[0]);
3552 try testing.expect(plan.fixed_threadgroup);
3553 try testing.expectEqual(@as(usize, 1), plan.candidate_count);
3554 try testing.expectEqual(@as(u32, 64), plan.candidates[0].geometry.threadgroup[0]);
3555
3556 const too_small = launchResourceCaps(.vulkan, .vulkan_spirv, 63, 63, 32);
3557 try testing.expectError(error.CapabilityMismatch, createLaunchResourcePlan(too_small, .vulkan_spirv, launchResources(65, 1, 780)));
3558 }
3559
3560 test "Choir launch resource plan exposes rank-2 elementwise 2D candidates" {
3561 var caps = launchResourceCaps(.cuda, .cuda_ptx, 256, 256, 32);
3562 caps.threadgroup.max_threads_per_dim = .{ 256, 256, 64 };
3563 const rank2 = kernelization.product.ElementwiseRank2Plan{
3564 .rows = 2048,
3565 .cols = 2048,
3566 .threads_x = 32,
3567 .threads_y = 8,
3568 };
3569
3570 const plan = try createElementwiseRank2LaunchResourcePlan(caps, .cuda_ptx, launchResources(2048 * 2048, 1, 2048 * 2048 * 12), rank2);
3571 try testing.expect(!plan.fixed_threadgroup);
3572 try testing.expectEqual(@as(usize, 8), plan.candidate_count);
3573 try testing.expectEqual(LaunchTileKind.elementwise_rank2, plan.tile.kind);
3574 try testing.expectEqual(@as(u32, 2048), plan.tile.m);
3575 try testing.expectEqual(@as(u32, 2048), plan.tile.n);
3576 try testing.expectEqual(@as(u32, 32), plan.geometry.threadgroup[0]);
3577 try testing.expectEqual(@as(u32, 8), plan.geometry.threadgroup[1]);
3578 try testing.expectEqual(@as(u32, 16), plan.candidates[1].geometry.threadgroup[0]);
3579 try testing.expectEqual(@as(u32, 16), plan.candidates[1].geometry.threadgroup[1]);
3580 try testing.expectEqual(@as(u32, 8), plan.candidates[2].geometry.threadgroup[0]);
3581 try testing.expectEqual(@as(u32, 32), plan.candidates[2].geometry.threadgroup[1]);
3582 }
3583
3584 test "Choir generated launch resource plan carries dynamic shared memory to candidates" {
3585 const allocator = testing.allocator;
3586
3587 var caps = launchResourceCaps(.cuda, .cuda_ptx, 256, 256, 32);
3588 caps.features.dynamic_shared_memory = true;
3589 caps.threadgroup.shared_memory_bytes = 4096;
3590
3591 const plan = try createLaunchResourcePlan(caps, .cuda_ptx, launchResources(4096, 64, 65536));
3592 try testing.expect(plan.candidate_count > 1);
3593
3594 var builder = try kernel_program.Builder.init(
3595 allocator,
3596 kernel_program.Builder.Limits.testing,
3597 "x",
3598 &.{},
3599 );
3600 errdefer builder.deinit();
3601 try builder.return_();
3602 var program = try builder.finish();
3603 defer program.deinit();
3604
3605 const entry_name = try allocator.dupe(u8, "x");
3606 defer allocator.free(entry_name);
3607
3608 const lowered = kernelization.LoweredKernel{
3609 .work_item_id = 0,
3610 .entry_name = entry_name,
3611 .program = program,
3612 .argument_count = 0,
3613 .body_fingerprint = 0,
3614 .dynamic_shared_memory_bytes = 2048,
3615 .schedule = .{
3616 .kind = .flat,
3617 .threads = .{ .x = 64 },
3618 },
3619 };
3620 const compile_plan = CompilePlan{
3621 .entry_name = "x",
3622 .lowered_kernel = &lowered,
3623 };
3624
3625 const updated = try attachGeneratedDynamicSharedMemory(caps, plan, compile_plan);
3626
3627 try testing.expectEqual(@as(u32, 2048), updated.geometry.dynamic_shared_memory_bytes);
3628 var candidate_index: usize = 0;
3629 while (candidate_index < updated.candidate_count) : (candidate_index += 1) {
3630 try testing.expectEqual(@as(u32, 2048), updated.candidates[candidate_index].geometry.dynamic_shared_memory_bytes);
3631 }
3632
3633 var unsupported = caps;
3634 unsupported.features.dynamic_shared_memory = false;
3635 try testing.expectError(error.CapabilityMismatch, attachGeneratedDynamicSharedMemory(unsupported, plan, compile_plan));
3636 }
3637
3638 test "Choir artifact plan creates CUDA artifacts from legal single-op kernels" {
3639 const allocator = testing.allocator;
3640
3641 var owned = try addChoirModule(allocator, "choir_cuda_add_artifact");
3642 defer owned.module.deinit();
3643
3644 var cache = passes.AnalysisCache.init(allocator, null);
3645 defer cache.deinit();
3646 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
3647 defer pass_ctx.deinit();
3648
3649 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
3650 defer state.deinit();
3651
3652 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
3653 defer plan.deinit();
3654
3655 try testing.expectEqual(gpu.BackendKind.cuda, plan.backend_kind);
3656 try testing.expectEqual(gpu.ArtifactFormat.cuda_ptx, plan.format);
3657 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
3658 const kernel = plan.kernels.items[0];
3659 try testing.expectEqual(PlannedKernelSource.tensor, kernel.compile.source);
3660 try testing.expectEqual(PlannedKernelCompileLaunch.generic, kernel.compile.launch);
3661 try testing.expectEqual(gpu.ArtifactFormat.cuda_ptx, kernel.compile.format);
3662 try testing.expectEqualStrings("accy_choir_elementwise_add_0", kernel.compile.entry_name);
3663 try testing.expectEqual(@as(u32, 4), kernel.compile.argument_count);
3664 try testing.expect(kernel.compile.required_dtypes.contains(.f32));
3665 try testing.expectEqual(PlannedKernelCompilePayload.text, kernel.compile.payload);
3666 try testing.expectEqualStrings("accy_choir_elementwise_add_0", kernel.artifact.entry_name);
3667 try testing.expectEqual(@as(u32, 4), kernel.artifact.argument_count);
3668 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
3669 try testing.expectEqual(@as(u64, 8), kernel.element_count);
3670 try testing.expectEqual(@as(u64, 8), kernel.element_count_argument_value);
3671 try testing.expectEqual(@as(usize, 1), kernel.resources.op_count);
3672 try testing.expectEqual(@as(usize, 2), kernel.resources.external_input_value_count);
3673 try testing.expectEqual(@as(u64, 64), kernel.resources.static_read_bytes);
3674 try testing.expectEqual(@as(u64, 32), kernel.resources.static_write_bytes);
3675 try testing.expectEqual(@as(u64, 96), kernel.resources.static_total_bytes);
3676 try testing.expectEqual(@as(u64, 8), kernel.resources.estimated_element_ops);
3677 try testing.expectEqual(LaunchResourceClass.memory_bound, kernel.launch_resources.resource_class);
3678 try testing.expectEqual(@as(u64, 384), kernel.launch_resources.estimated_static_bytes_per_threadgroup);
3679 try testing.expectEqual(@as(u64, 32), kernel.launch_resources.estimated_element_ops_per_threadgroup);
3680 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[0]);
3681 try testing.expectEqual(@as(u32, 32), kernel.launch_resources.geometry.threadgroup[0]);
3682 try expectKernelLayoutFingerprints(&plan, kernel);
3683 try testing.expect(std.mem.indexOf(u8, kernel.artifact.payload.text, ".visible .entry accy_choir_elementwise_add_0") != null);
3684 try testing.expectEqual(kernel.artifact.payload.text.len, kernel.compile.payload_byte_count);
3685 }
3686
3687 test "Choir artifact plan creates native CPU object artifacts from legal single-op kernels" {
3688 if (@import("builtin").cpu.arch != .x86_64) return error.SkipZigTest;
3689 const allocator = testing.allocator;
3690
3691 var owned = try addChoirModule(allocator, "choir_cpu_add_artifact");
3692 defer owned.module.deinit();
3693
3694 var cache = passes.AnalysisCache.init(allocator, null);
3695 defer cache.deinit();
3696 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
3697 defer pass_ctx.deinit();
3698
3699 var state = ArtifactPlanningBackendState.init(allocator, .cpu);
3700 state.format = .cpu_object;
3701 defer state.deinit();
3702
3703 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{ .format = .cpu_object });
3704 defer plan.deinit();
3705
3706 try testing.expectEqual(gpu.BackendKind.cpu, plan.backend_kind);
3707 try testing.expectEqual(gpu.ArtifactFormat.cpu_object, plan.format);
3708 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
3709 const kernel = plan.kernels.items[0];
3710 try testing.expectEqual(PlannedKernelSource.tensor, kernel.compile.source);
3711 try testing.expectEqual(PlannedKernelCompileLaunch.generic, kernel.compile.launch);
3712 try testing.expectEqual(gpu.ArtifactFormat.cpu_object, kernel.compile.format);
3713 try testing.expectEqualStrings("accy_choir_elementwise_add_0", kernel.compile.entry_name);
3714 try testing.expectEqual(@as(u32, 11), kernel.compile.argument_count);
3715 try testing.expectEqual(PlannedKernelCompilePayload.bytes, kernel.compile.payload);
3716 try testing.expectEqualStrings("accy_choir_elementwise_add_0", kernel.artifact.entry_name);
3717 try testing.expectEqual(@as(u32, 11), kernel.artifact.argument_count);
3718 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
3719 try testing.expectEqual(@as(u64, 8), kernel.element_count_argument_value);
3720 try testing.expectEqual(@as(usize, 7), kernel.static_arguments.len);
3721 try expectScalarU32(kernel.static_arguments[0], 8);
3722 try expectScalarU32(kernel.static_arguments[1], 1);
3723 try expectScalarU32(kernel.static_arguments[2], 1);
3724 try expectScalarU32(kernel.static_arguments[3], 1);
3725 try expectScalarU32(kernel.static_arguments[4], 8);
3726 try expectScalarU32(kernel.static_arguments[5], 1);
3727 try expectScalarU32(kernel.static_arguments[6], 1);
3728 try testing.expect(kernel.launch_resources.fixed_threadgroup);
3729 try testing.expectEqual(@as(usize, 1), kernel.launch_resources.candidate_count);
3730 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[0]);
3731 try testing.expectEqual(@as(u32, 8), kernel.launch_resources.geometry.threadgroup[0]);
3732 try testing.expect(kernel.artifact.payload.bytes.len >= 4);
3733 try testing.expectEqualSlices(u8, &.{ 0x7f, 'E', 'L', 'F' }, kernel.artifact.payload.bytes[0..4]);
3734 try testing.expectEqual(kernel.artifact.payload.bytes.len, kernel.compile.payload_byte_count);
3735 }
3736
3737 test "Choir artifact plan creates webassembly module artifacts from legal single-op kernels" {
3738 const allocator = testing.allocator;
3739
3740 var owned = try addChoirModule(allocator, "choir_wasm_add_artifact");
3741 defer owned.module.deinit();
3742
3743 var cache = passes.AnalysisCache.init(allocator, null);
3744 defer cache.deinit();
3745 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
3746 defer pass_ctx.deinit();
3747
3748 var state = ArtifactPlanningBackendState.init(allocator, .wasm);
3749 defer state.deinit();
3750
3751 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
3752 defer plan.deinit();
3753
3754 try testing.expectEqual(gpu.BackendKind.wasm, plan.backend_kind);
3755 try testing.expectEqual(gpu.ArtifactFormat.webassembly_module, plan.format);
3756 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
3757 const kernel = plan.kernels.items[0];
3758 try testing.expectEqual(PlannedKernelSource.tensor, kernel.compile.source);
3759 try testing.expectEqual(PlannedKernelCompileLaunch.generic, kernel.compile.launch);
3760 try testing.expectEqual(gpu.ArtifactFormat.webassembly_module, kernel.compile.format);
3761 try testing.expectEqualStrings("accy_choir_elementwise_add_0", kernel.compile.entry_name);
3762 try testing.expectEqual(@as(u32, 11), kernel.compile.argument_count);
3763 try testing.expectEqual(PlannedKernelCompilePayload.bytes, kernel.compile.payload);
3764 try testing.expectEqualStrings("accy_choir_elementwise_add_0", kernel.artifact.entry_name);
3765 try testing.expectEqual(@as(u32, 11), kernel.artifact.argument_count);
3766 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
3767 try testing.expectEqual(@as(u64, 8), kernel.element_count_argument_value);
3768 try testing.expectEqual(@as(usize, 7), kernel.static_arguments.len);
3769 try expectScalarU32(kernel.static_arguments[0], 8);
3770 try expectScalarU32(kernel.static_arguments[1], 1);
3771 try expectScalarU32(kernel.static_arguments[2], 1);
3772 try expectScalarU32(kernel.static_arguments[3], 1);
3773 try expectScalarU32(kernel.static_arguments[4], 8);
3774 try expectScalarU32(kernel.static_arguments[5], 1);
3775 try expectScalarU32(kernel.static_arguments[6], 1);
3776 try testing.expect(kernel.launch_resources.fixed_threadgroup);
3777 try testing.expectEqual(@as(usize, 1), kernel.launch_resources.candidate_count);
3778 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[0]);
3779 try testing.expectEqual(@as(u32, 8), kernel.launch_resources.geometry.threadgroup[0]);
3780 try testing.expect(kernel.artifact.payload.bytes.len >= 8);
3781 try testing.expectEqualSlices(u8, &.{ 0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00 }, kernel.artifact.payload.bytes[0..8]);
3782 try testing.expectEqual(kernel.artifact.payload.bytes.len, kernel.compile.payload_byte_count);
3783 }
3784
3785 test "Choir artifact plan uses prepared target module lowered kernels" {
3786 const allocator = testing.allocator;
3787
3788 var owned = try addChoirModule(allocator, "choir_cuda_target_artifact_source");
3789 var module_owned = true;
3790 errdefer if (module_owned) owned.module.deinit();
3791
3792 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
3793 defer state.deinit();
3794
3795 const handle = state.handle();
3796 const caps = try handle.queryCapabilities();
3797 const profile = try target_profile.BackendTargetProfile.init(caps, .cuda, .cuda_ptx);
3798 var prepared = try preparation.prepareBackendJobFromSemanticModule(allocator, owned.module, .{ .target_profile = profile });
3799 module_owned = false;
3800 defer prepared.deinit();
3801
3802 const target_module = try prepared.targetModule();
3803 const product = target_module.kernelizationProduct();
3804 try testing.expectEqual(@as(usize, 1), product.kernelCount());
3805 const generated_summary = try product.kernelSummary(0);
3806 const generated_program = try product.kernelProgram(0);
3807 try testing.expectEqual(generated_summary.body_fingerprint, try generated_program.bodyFingerprint(allocator));
3808
3809 var plan = try createBackendArtifactPlanFromTargetJob(allocator, handle, target_module, .{});
3810 defer plan.deinit();
3811
3812 const generated_after_plan = try product.kernelProgramForWork(generated_summary.work_item_id);
3813 try testing.expectEqual(generated_summary.body_fingerprint, try generated_after_plan.bodyFingerprint(allocator));
3814 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
3815 try testing.expectEqualStrings("accy_choir_elementwise_add_0_v4", plan.kernels.items[0].compile.entry_name);
3816 try testing.expectEqual(PlannedKernelSource.tensor, plan.kernels.items[0].compile.source);
3817 try testing.expectEqual(@as(usize, 1), try prepared.generatedKernelCount());
3818 }
3819
3820 test "Choir artifact plan rejects semantic kernel_call until external payload lowering exists" {
3821 const allocator = testing.allocator;
3822
3823 var owned = try kernelCallChoirModule(allocator, "choir_cuda_kernel_call_artifact");
3824 defer owned.module.deinit();
3825
3826 var cache = passes.AnalysisCache.init(allocator, null);
3827 defer cache.deinit();
3828 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
3829 defer pass_ctx.deinit();
3830
3831 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
3832 defer state.deinit();
3833
3834 try testing.expectError(
3835 error.UnsupportedOperation,
3836 createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{}),
3837 );
3838 }
3839
3840 test "Choir artifact plan resolves semantic kernel_call through registry" {
3841 const allocator = testing.allocator;
3842
3843 var owned = try kernelCallChoirModule(allocator, "choir_cuda_registered_kernel_call_artifact");
3844 defer owned.module.deinit();
3845
3846 var cache = passes.AnalysisCache.init(allocator, null);
3847 defer cache.deinit();
3848 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
3849 defer pass_ctx.deinit();
3850
3851 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
3852 defer state.deinit();
3853
3854 const source = ".visible .entry accy_custom_scale() { ret; }";
3855 const registry = KernelCallRegistry{ .entries = &.{.{
3856 .target = "accy.custom.scale",
3857 .version = 1,
3858 .format = .cuda_ptx,
3859 .entry_name = "accy_custom_scale",
3860 .argument_count = 5,
3861 .required_dtypes = gpu.DTypeSet.init(&.{.f32}),
3862 .payload = .{ .text = source },
3863 .launch = .{ .fixed = .{
3864 .grid = .{ 1, 1, 1 },
3865 .threadgroup = .{ 8, 1, 1 },
3866 .dynamic_shared_memory_bytes = 2048,
3867 } },
3868 .element_count_argument = .scalar_u32,
3869 .static_arguments = &.{.{ .u32 = 7 }},
3870 }} };
3871
3872 var plan = try createTestBackendArtifactPlan(
3873 allocator,
3874 state.handle(),
3875 &pass_ctx,
3876 owned.choir_module,
3877 .{ .kernel_call_registry = ®istry },
3878 );
3879 defer plan.deinit();
3880
3881 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
3882 const kernel = plan.kernels.items[0];
3883 try testing.expectEqual(PlannedKernelSource.kernel_call, kernel.compile.source);
3884 try testing.expectEqual(PlannedKernelCompileLaunch.kernel_call, kernel.compile.launch);
3885 try testing.expectEqualStrings("accy_custom_scale", kernel.compile.entry_name);
3886 try testing.expectEqual(@as(u32, 5), kernel.compile.argument_count);
3887 try testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, kernel.compile.required_dtypes.bits);
3888 try testing.expectEqual(PlannedKernelCompilePayload.text, kernel.compile.payload);
3889 try testing.expectEqual(source.len, kernel.compile.payload_byte_count);
3890 try testing.expectEqualStrings("accy_custom_scale", kernel.artifact.entry_name);
3891 try testing.expectEqualStrings(source, kernel.artifact.payload.text);
3892 try testing.expectEqual(@as(usize, 2), kernel.input_slot_ids.len);
3893 try testing.expectEqual(@as(u64, 8), kernel.element_count);
3894 try testing.expectEqual(ElementCountArgument.scalar_u32, kernel.element_count_argument);
3895 try testing.expectEqual(@as(u64, 8), kernel.element_count_argument_value);
3896 try testing.expectEqual(@as(usize, 1), kernel.static_arguments.len);
3897 try testing.expectEqual(choir_abi.ScalarArgument{ .u32 = 7 }, kernel.static_arguments[0]);
3898 try testing.expect(kernel.launch_resources.fixed_threadgroup);
3899 try testing.expectEqual(@as(usize, 1), kernel.launch_resources.candidate_count);
3900 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[0]);
3901 try testing.expectEqual(@as(u32, 8), kernel.launch_resources.geometry.threadgroup[0]);
3902 try testing.expectEqual(@as(u32, 2048), kernel.launch_resources.geometry.dynamic_shared_memory_bytes);
3903 try testing.expectEqual(@as(u32, 8), kernel.launch_resources.candidates[0].geometry.threadgroup[0]);
3904 try testing.expectEqual(@as(u32, 2048), kernel.launch_resources.candidates[0].geometry.dynamic_shared_memory_bytes);
3905 }
3906
3907 test "Choir artifact plan carries prepared backend target profile" {
3908 const allocator = testing.allocator;
3909
3910 var owned = try addChoirModule(allocator, "choir_cuda_target_profile_artifact");
3911 defer owned.module.deinit();
3912
3913 var cache = passes.AnalysisCache.init(allocator, null);
3914 defer cache.deinit();
3915 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
3916 defer pass_ctx.deinit();
3917
3918 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
3919 defer state.deinit();
3920
3921 const profile = target_profile.BackendTargetProfile{
3922 .backend_kind = .cuda,
3923 .artifact_format = .cuda_ptx,
3924 .dtype_bits = gpu.DTypeSet.init(&.{.f32}).bits,
3925 };
3926 try target_profile.setBackendTargetProfile(owned.module.context(), owned.choir_module, profile);
3927
3928 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
3929 defer plan.deinit();
3930
3931 try testing.expect(plan.target_profile.eql(profile));
3932 try testing.expectEqual(gpu.BackendKind.cuda, plan.backend_kind);
3933 try testing.expectEqual(gpu.ArtifactFormat.cuda_ptx, plan.format);
3934 }
3935
3936 test "Choir CUDA artifact planning preserves generated kernel source modules" {
3937 const allocator = testing.allocator;
3938
3939 var owned = try addChoirModule(allocator, "choir_cuda_target_lowering_artifact");
3940 defer owned.module.deinit();
3941
3942 var cache = passes.AnalysisCache.init(allocator, null);
3943 defer cache.deinit();
3944 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
3945 defer pass_ctx.deinit();
3946
3947 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
3948 defer state.deinit();
3949
3950 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
3951 defer plan.deinit();
3952
3953 const planned = plan.kernels.items[0];
3954 const kernelization_plan = try kernelization.getKernelizationAnalysis(&pass_ctx, owned.choir_module);
3955 const lowered = kernelization_plan.getForWork(planned.work_item_id) orelse return error.MissingKernelization;
3956 const module = lowered.program.kernelModule();
3957
3958 try testing.expect(operationTreeContainsName(module, GpuDialect.GlobalIdxOp.operation_name));
3959 try testing.expect(operationTreeContainsName(module, choir.dialects.MemrefDialect.LoadOp.operation_name));
3960 try testing.expect(operationTreeContainsName(module, choir.dialects.MemrefDialect.StoreOp.operation_name));
3961 try testing.expect(!operationTreeContainsName(module, gpu_codegen.nvptx.NvptxDialect.ThreadIdxOp.operation_name));
3962 try testing.expect(!operationTreeContainsName(module, gpu_codegen.nvptx.NvptxDialect.BlockIdxOp.operation_name));
3963 try testing.expect(!operationTreeContainsName(module, gpu_codegen.nvptx.NvptxDialect.LoadGlobalOp.operation_name));
3964 try testing.expect(!operationTreeContainsName(module, gpu_codegen.nvptx.NvptxDialect.StoreGlobalOp.operation_name));
3965 try testing.expect(std.mem.indexOf(u8, planned.artifact.payload.text, planned.compile.entry_name) != null);
3966 }
3967
3968 test "Choir artifact plan creates Vulkan artifacts from legal single-op kernels" {
3969 const allocator = testing.allocator;
3970
3971 var owned = try addChoirModule(allocator, "choir_vulkan_add_artifact");
3972 defer owned.module.deinit();
3973
3974 var cache = passes.AnalysisCache.init(allocator, null);
3975 defer cache.deinit();
3976 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
3977 defer pass_ctx.deinit();
3978
3979 var state = ArtifactPlanningBackendState.init(allocator, .vulkan);
3980 defer state.deinit();
3981
3982 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
3983 defer plan.deinit();
3984
3985 try testing.expectEqual(gpu.BackendKind.vulkan, plan.backend_kind);
3986 try testing.expectEqual(gpu.ArtifactFormat.vulkan_spirv, plan.format);
3987 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
3988 const kernel = plan.kernels.items[0];
3989 try testing.expectEqual(PlannedKernelSource.tensor, kernel.compile.source);
3990 try testing.expectEqual(PlannedKernelCompileLaunch.generic, kernel.compile.launch);
3991 try testing.expectEqual(gpu.ArtifactFormat.vulkan_spirv, kernel.compile.format);
3992 try testing.expectEqualStrings("accy_choir_elementwise_add_0", kernel.compile.entry_name);
3993 try testing.expectEqual(@as(u32, 4), kernel.compile.argument_count);
3994 try testing.expectEqual(PlannedKernelCompilePayload.words_u32, kernel.compile.payload);
3995 const kernelization_plan = try kernelization.getKernelizationAnalysis(&pass_ctx, owned.choir_module);
3996 const lowered = kernelization_plan.getForWork(kernel.work_item_id) orelse return error.MissingKernelization;
3997 const generated_launch = try lowered.program.launch();
3998 try testing.expectEqual(@as(u32, 1), generated_launch.grid[0]);
3999 try testing.expectEqual(@as(u32, 8), generated_launch.block[0]);
4000 try testing.expectEqualStrings("accy_choir_elementwise_add_0", kernel.artifact.entry_name);
4001 try testing.expectEqual(@as(u32, 4), kernel.artifact.argument_count);
4002 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
4003 try testing.expectEqual(@as(u64, 8), kernel.element_count_argument_value);
4004 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[0]);
4005 try testing.expectEqual(@as(u32, 8), kernel.launch_resources.geometry.threadgroup[0]);
4006 try expectKernelLayoutFingerprints(&plan, kernel);
4007 try testing.expectEqual(@as(u32, 0x07230203), kernel.artifact.payload.words_u32[0]);
4008 try testing.expectEqual(kernel.artifact.payload.words_u32.len * @sizeOf(u32), kernel.compile.payload_byte_count);
4009 }
4010
4011 test "Choir artifact planning capabilities reflect implemented backend features" {
4012 var cuda_state = ArtifactPlanningBackendState.init(testing.allocator, .cuda);
4013 defer cuda_state.deinit();
4014 const cuda_caps = try cuda_state.handle().queryCapabilities();
4015 try testing.expect(cuda_caps.features.atomic_i32);
4016 try testing.expect(cuda_caps.features.atomic_u32);
4017 try testing.expect(cuda_caps.features.atomic_index);
4018 try testing.expect(cuda_caps.features.atomic_f32_add_device);
4019 try testing.expect(cuda_caps.features.atomic_f32_add_shared);
4020 try testing.expect(cuda_caps.features.dynamic_shared_memory);
4021
4022 var metal_state = ArtifactPlanningBackendState.init(testing.allocator, .metal);
4023 defer metal_state.deinit();
4024 const metal_caps = try metal_state.handle().queryCapabilities();
4025 try testing.expect(metal_caps.features.atomic_i32);
4026 try testing.expect(metal_caps.features.atomic_u32);
4027 try testing.expect(metal_caps.features.atomic_index);
4028 try testing.expect(metal_caps.features.atomic_f32_add_device);
4029 try testing.expect(!metal_caps.features.atomic_f32_add_shared);
4030 try testing.expect(!metal_caps.features.dynamic_shared_memory);
4031
4032 var vulkan_state = ArtifactPlanningBackendState.init(testing.allocator, .vulkan);
4033 defer vulkan_state.deinit();
4034 const vulkan_caps = try vulkan_state.handle().queryCapabilities();
4035 try testing.expect(vulkan_caps.features.atomic_i32);
4036 try testing.expect(vulkan_caps.features.atomic_u32);
4037 try testing.expect(vulkan_caps.features.atomic_index);
4038 try testing.expect(!vulkan_caps.features.atomic_f32_add_device);
4039 try testing.expect(!vulkan_caps.features.atomic_f32_add_shared);
4040 try testing.expect(!vulkan_caps.features.dynamic_shared_memory);
4041 }
4042
4043 test "Choir Vulkan artifact planning preserves generated kernel source modules" {
4044 const allocator = testing.allocator;
4045
4046 var owned = try addChoirModule(allocator, "choir_vulkan_target_lowering_artifact");
4047 defer owned.module.deinit();
4048
4049 var cache = passes.AnalysisCache.init(allocator, null);
4050 defer cache.deinit();
4051 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4052 defer pass_ctx.deinit();
4053
4054 var state = ArtifactPlanningBackendState.init(allocator, .vulkan);
4055 defer state.deinit();
4056
4057 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4058 defer plan.deinit();
4059
4060 const planned = plan.kernels.items[0];
4061 const kernelization_plan = try kernelization.getKernelizationAnalysis(&pass_ctx, owned.choir_module);
4062 const lowered = kernelization_plan.getForWork(planned.work_item_id) orelse return error.MissingKernelization;
4063 const module = lowered.program.kernelModule();
4064
4065 try testing.expect(operationTreeContainsName(module, GpuDialect.GlobalIdxOp.operation_name));
4066 try testing.expect(operationTreeContainsName(module, choir.dialects.ArithDialect.AddOp.operation_name));
4067 try testing.expect(!operationTreeContainsName(module, gpu_codegen.spirv.SpirvDialect.GlobalInvocationIdOp.operation_name));
4068 try testing.expect(!operationTreeContainsName(module, gpu_codegen.spirv.SpirvDialect.FAddOp.operation_name));
4069 try testing.expectEqual(@as(u32, 0x07230203), planned.artifact.payload.words_u32[0]);
4070 }
4071
4072 test "Choir artifact plan creates Vulkan artifacts from kernel-language reshape kernels" {
4073 const allocator = testing.allocator;
4074
4075 var owned = try reshapeChoirModule(allocator, "choir_vulkan_reshape_artifact");
4076 defer owned.module.deinit();
4077
4078 var cache = passes.AnalysisCache.init(allocator, null);
4079 defer cache.deinit();
4080 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4081 defer pass_ctx.deinit();
4082
4083 var state = ArtifactPlanningBackendState.init(allocator, .vulkan);
4084 defer state.deinit();
4085
4086 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4087 defer plan.deinit();
4088
4089 try testing.expectEqual(gpu.BackendKind.vulkan, plan.backend_kind);
4090 try testing.expectEqual(gpu.ArtifactFormat.vulkan_spirv, plan.format);
4091 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4092 const kernel = plan.kernels.items[0];
4093 const kernelization_plan = try kernelization.getKernelizationAnalysis(&pass_ctx, owned.choir_module);
4094 const lowered = kernelization_plan.getForWork(kernel.work_item_id) orelse return error.MissingKernelization;
4095 const generated_launch = try lowered.program.launch();
4096 try testing.expectEqual(@as(u32, 1), generated_launch.grid[0]);
4097 try testing.expectEqual(@as(u32, 8), generated_launch.block[0]);
4098 try testing.expectEqualStrings("accy_choir_shape_reshape_0", kernel.artifact.entry_name);
4099 try testing.expectEqual(@as(u32, 3), kernel.artifact.argument_count);
4100 try testing.expectEqual(@as(usize, 1), kernel.input_slot_ids.len);
4101 try testing.expectEqual(@as(u64, 8), kernel.element_count);
4102 try testing.expectEqual(@as(u64, 8), kernel.element_count_argument_value);
4103 try testing.expectEqual(@as(u64, 32), kernel.resources.static_read_bytes);
4104 try testing.expectEqual(@as(u64, 32), kernel.resources.static_write_bytes);
4105 try testing.expectEqual(@as(u64, 64), kernel.resources.static_total_bytes);
4106 try testing.expectEqual(@as(u64, 8), kernel.resources.estimated_element_ops);
4107 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
4108 try testing.expectEqual(LaunchResourceClass.memory_bound, kernel.launch_resources.resource_class);
4109 try testing.expect(kernel.launch_resources.fixed_threadgroup);
4110 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[0]);
4111 try testing.expectEqual(@as(u32, 8), kernel.launch_resources.geometry.threadgroup[0]);
4112 try testing.expectEqual(@as(u64, 64), kernel.launch_resources.estimated_static_bytes_per_threadgroup);
4113 try testing.expectEqual(@as(u64, 8), kernel.launch_resources.estimated_element_ops_per_threadgroup);
4114 try expectKernelLayoutFingerprints(&plan, kernel);
4115 try testing.expectEqual(@as(u32, 0x07230203), kernel.artifact.payload.words_u32[0]);
4116 }
4117
4118 test "Choir artifact plan creates native CPU dot_general artifacts" {
4119 if (@import("builtin").cpu.arch != .x86_64) return error.SkipZigTest;
4120 const allocator = testing.allocator;
4121
4122 var owned = try dotGeneralChoirModule(allocator, "choir_cpu_dot_general_artifact");
4123 defer owned.module.deinit();
4124
4125 var cache = passes.AnalysisCache.init(allocator, null);
4126 defer cache.deinit();
4127 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4128 defer pass_ctx.deinit();
4129
4130 var state = ArtifactPlanningBackendState.init(allocator, .cpu);
4131 state.format = .cpu_object;
4132 defer state.deinit();
4133
4134 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{ .format = .cpu_object });
4135 defer plan.deinit();
4136
4137 try testing.expectEqual(gpu.BackendKind.cpu, plan.backend_kind);
4138 try testing.expectEqual(gpu.ArtifactFormat.cpu_object, plan.format);
4139 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4140 const kernel = plan.kernels.items[0];
4141 try testing.expectEqual(PlannedKernelSource.tensor, kernel.compile.source);
4142 try testing.expectEqual(PlannedKernelCompileLaunch.dot_general, kernel.compile.launch);
4143 try testing.expectEqual(gpu.ArtifactFormat.cpu_object, kernel.compile.format);
4144 try testing.expectEqualStrings("accy_choir_dot_general_f32_16x16x16_0", kernel.compile.entry_name);
4145 try testing.expectEqual(@as(u32, 10), kernel.compile.argument_count);
4146 try testing.expect(kernel.compile.required_dtypes.contains(.f32));
4147 try testing.expectEqual(PlannedKernelCompilePayload.bytes, kernel.compile.payload);
4148 try testing.expectEqualStrings("accy_choir_dot_general_f32_16x16x16_0", kernel.artifact.entry_name);
4149 try testing.expectEqual(@as(u32, 10), kernel.artifact.argument_count);
4150 try testing.expectEqual(@as(usize, 2), kernel.input_slot_ids.len);
4151 try testing.expectEqual(@as(u64, 256), kernel.element_count);
4152 try testing.expectEqual(ElementCountArgument.none, kernel.element_count_argument);
4153 try testing.expectEqual(@as(usize, 7), kernel.static_arguments.len);
4154 try expectScalarU32(kernel.static_arguments[0], 256);
4155 try expectScalarU32(kernel.static_arguments[1], 256);
4156 try expectScalarU32(kernel.static_arguments[2], 1);
4157 try expectScalarU32(kernel.static_arguments[3], 1);
4158 try expectScalarU32(kernel.static_arguments[4], 1);
4159 try expectScalarU32(kernel.static_arguments[5], 1);
4160 try expectScalarU32(kernel.static_arguments[6], 1);
4161 try testing.expect(kernel.launch_resources.fixed_threadgroup);
4162 try testing.expectEqual(@as(usize, 1), kernel.launch_resources.candidate_count);
4163 try testing.expectEqual(@as(u32, 256), kernel.launch_resources.geometry.grid[0]);
4164 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[1]);
4165 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[2]);
4166 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.threadgroup[0]);
4167 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.threadgroup[1]);
4168 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.threadgroup[2]);
4169 try expectDotGeneralTile(kernel.launch_resources.tile, 1, 1, 16, 128, 4);
4170 try testing.expect(kernel.artifact.payload.bytes.len >= 4);
4171 try testing.expectEqualSlices(u8, &.{ 0x7f, 'E', 'L', 'F' }, kernel.artifact.payload.bytes[0..4]);
4172 }
4173
4174 test "Choir artifact plan creates CUDA artifacts from choir dot_general kernels" {
4175 const allocator = testing.allocator;
4176
4177 var owned = try dotGeneralChoirModule(allocator, "choir_cuda_dot_general_artifact");
4178 defer owned.module.deinit();
4179
4180 var cache = passes.AnalysisCache.init(allocator, null);
4181 defer cache.deinit();
4182 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4183 defer pass_ctx.deinit();
4184
4185 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
4186 defer state.deinit();
4187
4188 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4189 defer plan.deinit();
4190
4191 try testing.expectEqual(gpu.BackendKind.cuda, plan.backend_kind);
4192 try testing.expectEqual(gpu.ArtifactFormat.cuda_ptx, plan.format);
4193 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4194 const kernel = plan.kernels.items[0];
4195 try testing.expectEqual(PlannedKernelSource.tensor, kernel.compile.source);
4196 try testing.expectEqual(PlannedKernelCompileLaunch.dot_general, kernel.compile.launch);
4197 try testing.expectEqual(gpu.ArtifactFormat.cuda_ptx, kernel.compile.format);
4198 try testing.expectEqualStrings("accy_choir_dot_general_f32_16x16x16_0", kernel.compile.entry_name);
4199 try testing.expectEqual(@as(u32, 3), kernel.compile.argument_count);
4200 try testing.expect(kernel.compile.required_dtypes.contains(.f32));
4201 try testing.expectEqual(PlannedKernelCompilePayload.text, kernel.compile.payload);
4202 try testing.expectEqualStrings("accy_choir_dot_general_f32_16x16x16_0", kernel.artifact.entry_name);
4203 try testing.expectEqual(@as(u32, 3), kernel.artifact.argument_count);
4204 try testing.expectEqual(@as(usize, 2), kernel.input_slot_ids.len);
4205 try testing.expectEqual(@as(u64, 256), kernel.element_count);
4206 try testing.expectEqual(@as(u64, 256), kernel.element_count_argument_value);
4207 try testing.expectEqual(@as(u64, 2048), kernel.resources.static_read_bytes);
4208 try testing.expectEqual(@as(u64, 1024), kernel.resources.static_write_bytes);
4209 try testing.expectEqual(@as(u64, 3072), kernel.resources.static_total_bytes);
4210 try testing.expectEqual(@as(u64, 8192), kernel.resources.estimated_element_ops);
4211 try testing.expectEqual(ElementCountArgument.none, kernel.element_count_argument);
4212 try testing.expectEqual(@as(usize, 0), kernel.static_arguments.len);
4213 try testing.expectEqual(LaunchResourceClass.compute_weighted, kernel.launch_resources.resource_class);
4214 try testing.expect(kernel.launch_resources.fixed_threadgroup);
4215 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[0]);
4216 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[1]);
4217 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[2]);
4218 try testing.expectEqual(@as(u32, 16), kernel.launch_resources.geometry.threadgroup[0]);
4219 try testing.expectEqual(@as(u32, 16), kernel.launch_resources.geometry.threadgroup[1]);
4220 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.threadgroup[2]);
4221 try testing.expectEqual(@as(u64, 3072), kernel.launch_resources.estimated_static_bytes_per_threadgroup);
4222 try testing.expectEqual(@as(u64, 8192), kernel.launch_resources.estimated_element_ops_per_threadgroup);
4223 try testing.expectEqual(@as(usize, 4), kernel.launch_resources.candidate_count);
4224 const tile = kernel.launch_resources.tile;
4225 try expectDotGeneralTile(tile, 16, 16, 16, 2048, 1024);
4226 try testing.expectEqual(tile, kernel.launch_resources.candidates[0].tile);
4227 try expectDotGeneralCandidate(kernel.launch_resources.candidates[0], .{ 1, 1 }, .{ 16, 16 }, 0, 3072, 8192, 16, 16, 16, 2048, 1024);
4228 try expectDotGeneralCandidate(kernel.launch_resources.candidates[1], .{ 1, 2 }, .{ 32, 8 }, 1, 3072, 8192, 8, 32, 16, 2560, 1024);
4229 try expectDotGeneralCandidate(kernel.launch_resources.candidates[2], .{ 2, 1 }, .{ 8, 32 }, 2, 3072, 8192, 32, 8, 16, 2560, 1024);
4230 try expectDotGeneralCandidate(kernel.launch_resources.candidates[3], .{ 2, 2 }, .{ 8, 8 }, 3, 768, 2048, 8, 8, 16, 1024, 256);
4231 try expectKernelLayoutFingerprints(&plan, kernel);
4232 try testing.expect(std.mem.indexOf(u8, kernel.artifact.payload.text, ".visible .entry accy_choir_dot_general_f32_16x16x16_0") != null);
4233 try testing.expectEqual(kernel.artifact.payload.text.len, kernel.compile.payload_byte_count);
4234 }
4235
4236 test "Choir artifact plan requires CUDA f16 choir lowering for f16 dot_general kernels" {
4237 const allocator = testing.allocator;
4238
4239 var owned = try dotGeneralF16ChoirModule(allocator, "choir_cuda_dot_general_f16_artifact");
4240 defer owned.module.deinit();
4241
4242 var cache = passes.AnalysisCache.init(allocator, null);
4243 defer cache.deinit();
4244 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4245 defer pass_ctx.deinit();
4246
4247 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
4248 defer state.deinit();
4249
4250 try testing.expectError(
4251 error.CapabilityMismatch,
4252 createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{}),
4253 );
4254 }
4255
4256 test "Choir artifact plan creates Vulkan artifacts from kernel-language dot_general kernels" {
4257 const allocator = testing.allocator;
4258
4259 var owned = try dotGeneralChoirModule(allocator, "choir_vulkan_dot_general_artifact");
4260 defer owned.module.deinit();
4261
4262 var cache = passes.AnalysisCache.init(allocator, null);
4263 defer cache.deinit();
4264 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4265 defer pass_ctx.deinit();
4266
4267 var state = ArtifactPlanningBackendState.init(allocator, .vulkan);
4268 defer state.deinit();
4269
4270 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4271 defer plan.deinit();
4272
4273 try testing.expectEqual(gpu.BackendKind.vulkan, plan.backend_kind);
4274 try testing.expectEqual(gpu.ArtifactFormat.vulkan_spirv, plan.format);
4275 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4276 const kernel = plan.kernels.items[0];
4277 try testing.expectEqual(PlannedKernelSource.tensor, kernel.compile.source);
4278 try testing.expectEqual(PlannedKernelCompileLaunch.dot_general, kernel.compile.launch);
4279 try testing.expectEqual(gpu.ArtifactFormat.vulkan_spirv, kernel.compile.format);
4280 try testing.expectEqualStrings("accy_choir_dot_general_f32_16x16x16_0", kernel.compile.entry_name);
4281 try testing.expectEqual(@as(u32, 3), kernel.compile.argument_count);
4282 try testing.expect(kernel.compile.required_dtypes.contains(.f32));
4283 try testing.expectEqual(PlannedKernelCompilePayload.words_u32, kernel.compile.payload);
4284 try testing.expectEqualStrings("accy_choir_dot_general_f32_16x16x16_0", kernel.artifact.entry_name);
4285 try testing.expectEqual(@as(u32, 3), kernel.artifact.argument_count);
4286 try testing.expectEqual(@as(usize, 2), kernel.input_slot_ids.len);
4287 try testing.expectEqual(@as(u64, 256), kernel.element_count);
4288 try testing.expectEqual(@as(u64, 256), kernel.element_count_argument_value);
4289 try testing.expectEqual(@as(u64, 8192), kernel.resources.estimated_element_ops);
4290 try testing.expectEqual(ElementCountArgument.none, kernel.element_count_argument);
4291 try testing.expectEqual(@as(usize, 0), kernel.static_arguments.len);
4292 try testing.expectEqual(LaunchResourceClass.compute_weighted, kernel.launch_resources.resource_class);
4293 try testing.expect(kernel.launch_resources.fixed_threadgroup);
4294 try testing.expectEqual(@as(u32, 2), kernel.launch_resources.geometry.grid[0]);
4295 try testing.expectEqual(@as(u32, 2), kernel.launch_resources.geometry.grid[1]);
4296 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[2]);
4297 try testing.expectEqual(@as(u32, 8), kernel.launch_resources.geometry.threadgroup[0]);
4298 try testing.expectEqual(@as(u32, 8), kernel.launch_resources.geometry.threadgroup[1]);
4299 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.threadgroup[2]);
4300 try testing.expectEqual(@as(u64, 768), kernel.launch_resources.estimated_static_bytes_per_threadgroup);
4301 try testing.expectEqual(@as(u64, 2048), kernel.launch_resources.estimated_element_ops_per_threadgroup);
4302 try testing.expectEqual(@as(usize, 4), kernel.launch_resources.candidate_count);
4303 const tile = kernel.launch_resources.tile;
4304 try expectDotGeneralTile(tile, 8, 8, 16, 1024, 256);
4305 try testing.expectEqual(tile, kernel.launch_resources.candidates[0].tile);
4306 try expectDotGeneralCandidate(kernel.launch_resources.candidates[0], .{ 2, 2 }, .{ 8, 8 }, 0, 768, 2048, 8, 8, 16, 1024, 256);
4307 try expectDotGeneralCandidate(kernel.launch_resources.candidates[1], .{ 1, 2 }, .{ 16, 8 }, 1, 1536, 4096, 8, 16, 16, 1536, 512);
4308 try expectDotGeneralCandidate(kernel.launch_resources.candidates[2], .{ 2, 1 }, .{ 8, 16 }, 2, 1536, 4096, 16, 8, 16, 1536, 512);
4309 try expectDotGeneralCandidate(kernel.launch_resources.candidates[3], .{ 1, 1 }, .{ 16, 16 }, 3, 3072, 8192, 16, 16, 16, 2048, 1024);
4310 try expectKernelLayoutFingerprints(&plan, kernel);
4311 try testing.expectEqual(@as(u32, 0x07230203), kernel.artifact.payload.words_u32[0]);
4312 try testing.expectEqual(kernel.artifact.payload.words_u32.len * @sizeOf(u32), kernel.compile.payload_byte_count);
4313 }
4314
4315 test "Choir artifact plan requires Vulkan f16 capability for f16 dot_general kernels" {
4316 const allocator = testing.allocator;
4317
4318 var owned = try dotGeneralF16ChoirModule(allocator, "choir_vulkan_dot_general_f16_artifact");
4319 defer owned.module.deinit();
4320
4321 var cache = passes.AnalysisCache.init(allocator, null);
4322 defer cache.deinit();
4323 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4324 defer pass_ctx.deinit();
4325
4326 var state = ArtifactPlanningBackendState.init(allocator, .vulkan);
4327 defer state.deinit();
4328
4329 try testing.expectError(
4330 error.CapabilityMismatch,
4331 createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{}),
4332 );
4333 }
4334
4335 test "Choir artifact plan creates CUDA artifacts from structured reduction kernels" {
4336 const allocator = testing.allocator;
4337
4338 var owned = try reduceChoirModule(allocator, "choir_cuda_reduce_artifact");
4339 defer owned.module.deinit();
4340
4341 var cache = passes.AnalysisCache.init(allocator, null);
4342 defer cache.deinit();
4343 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4344 defer pass_ctx.deinit();
4345
4346 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
4347 defer state.deinit();
4348
4349 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4350 defer plan.deinit();
4351
4352 try testing.expectEqual(gpu.BackendKind.cuda, plan.backend_kind);
4353 try testing.expectEqual(gpu.ArtifactFormat.cuda_ptx, plan.format);
4354 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4355 const kernel = plan.kernels.items[0];
4356 try testing.expectEqualStrings("accy_choir_reduction_sum_rank1_axis0_f32_0", kernel.artifact.entry_name);
4357 try testing.expectEqual(@as(u32, 3), kernel.artifact.argument_count);
4358 try testing.expectEqual(@as(usize, 1), kernel.input_slot_ids.len);
4359 try testing.expectEqual(@as(u64, 1), kernel.element_count);
4360 try testing.expectEqual(@as(u64, 1), kernel.element_count_argument_value);
4361 try testing.expectEqual(@as(u64, 1024), kernel.resources.static_read_bytes);
4362 try testing.expectEqual(@as(u64, 4), kernel.resources.static_write_bytes);
4363 try testing.expectEqual(@as(u64, 1028), kernel.resources.static_total_bytes);
4364 try testing.expectEqual(@as(u64, 256), kernel.resources.estimated_element_ops);
4365 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
4366 try testing.expectEqual(@as(usize, 0), kernel.static_arguments.len);
4367 try testing.expectEqual(LaunchResourceClass.memory_bound, kernel.launch_resources.resource_class);
4368 try testing.expect(!kernel.launch_resources.fixed_threadgroup);
4369 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[0]);
4370 try testing.expectEqual(@as(u32, 32), kernel.launch_resources.geometry.threadgroup[0]);
4371 try testing.expectEqual(@as(u64, 32896), kernel.launch_resources.estimated_static_bytes_per_threadgroup);
4372 try testing.expectEqual(@as(u64, 8192), kernel.launch_resources.estimated_element_ops_per_threadgroup);
4373 const tile = kernel.launch_resources.tile;
4374 try expectReductionTile(tile, 32, 256, 1, 0, 32768, 128);
4375 try testing.expectEqual(tile, kernel.launch_resources.candidates[0].tile);
4376 try testing.expect(std.mem.indexOf(u8, kernel.artifact.payload.text, ".visible .entry accy_choir_reduction_sum_rank1_axis0_f32_0") != null);
4377 }
4378
4379 test "Choir artifact plan creates CUDA artifacts from i32 structured reductions" {
4380 const allocator = testing.allocator;
4381
4382 var owned = try reduceI32ChoirModule(allocator, "choir_cuda_reduce_i32_artifact");
4383 defer owned.module.deinit();
4384
4385 var cache = passes.AnalysisCache.init(allocator, null);
4386 defer cache.deinit();
4387 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4388 defer pass_ctx.deinit();
4389
4390 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
4391 defer state.deinit();
4392
4393 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4394 defer plan.deinit();
4395
4396 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4397 const kernel = plan.kernels.items[0];
4398 try testing.expectEqualStrings("accy_choir_reduction_sum_rank1_axis0_i32_0", kernel.artifact.entry_name);
4399 try testing.expect(kernel.compile.required_dtypes.contains(.i32));
4400 try testing.expect(std.mem.indexOf(u8, kernel.artifact.payload.text, ".visible .entry accy_choir_reduction_sum_rank1_axis0_i32_0") != null);
4401 }
4402
4403 test "Choir artifact plan creates Vulkan artifacts from i32 structured reductions" {
4404 const allocator = testing.allocator;
4405
4406 var owned = try reduceI32ChoirModule(allocator, "choir_vulkan_reduce_i32_artifact");
4407 defer owned.module.deinit();
4408
4409 var cache = passes.AnalysisCache.init(allocator, null);
4410 defer cache.deinit();
4411 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4412 defer pass_ctx.deinit();
4413
4414 var state = ArtifactPlanningBackendState.init(allocator, .vulkan);
4415 defer state.deinit();
4416
4417 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4418 defer plan.deinit();
4419
4420 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4421 const kernel = plan.kernels.items[0];
4422 try testing.expect(kernel.compile.required_dtypes.contains(.i32));
4423 }
4424
4425 test "Choir artifact plan creates CUDA artifacts from structured rank-2 reduction kernels" {
4426 const allocator = testing.allocator;
4427
4428 var owned = try reduceRank2Axis1ChoirModule(allocator, "choir_cuda_reduce_rank2_artifact");
4429 defer owned.module.deinit();
4430
4431 var cache = passes.AnalysisCache.init(allocator, null);
4432 defer cache.deinit();
4433 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4434 defer pass_ctx.deinit();
4435
4436 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
4437 defer state.deinit();
4438
4439 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4440 defer plan.deinit();
4441
4442 try testing.expectEqual(gpu.BackendKind.cuda, plan.backend_kind);
4443 try testing.expectEqual(gpu.ArtifactFormat.cuda_ptx, plan.format);
4444 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4445 const kernel = plan.kernels.items[0];
4446 try testing.expectEqualStrings("accy_choir_reduction_sum_rank2_axis1_f32_0", kernel.artifact.entry_name);
4447 try testing.expectEqual(@as(u32, 3), kernel.artifact.argument_count);
4448 try testing.expectEqual(@as(usize, 1), kernel.input_slot_ids.len);
4449 try testing.expectEqual(@as(u64, 16), kernel.element_count);
4450 try testing.expectEqual(@as(u64, 16), kernel.element_count_argument_value);
4451 try testing.expectEqual(@as(u64, 1024), kernel.resources.static_read_bytes);
4452 try testing.expectEqual(@as(u64, 64), kernel.resources.static_write_bytes);
4453 try testing.expectEqual(@as(u64, 1088), kernel.resources.static_total_bytes);
4454 try testing.expectEqual(@as(u64, 256), kernel.resources.estimated_element_ops);
4455 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
4456 try testing.expectEqual(@as(usize, 0), kernel.static_arguments.len);
4457 try testing.expectEqual(LaunchResourceClass.memory_bound, kernel.launch_resources.resource_class);
4458 try testing.expect(!kernel.launch_resources.fixed_threadgroup);
4459 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[0]);
4460 try testing.expectEqual(@as(u32, 32), kernel.launch_resources.geometry.threadgroup[0]);
4461 try testing.expectEqual(@as(u64, 2176), kernel.launch_resources.estimated_static_bytes_per_threadgroup);
4462 try testing.expectEqual(@as(u64, 512), kernel.launch_resources.estimated_element_ops_per_threadgroup);
4463 const tile = kernel.launch_resources.tile;
4464 try expectReductionTile(tile, 32, 16, 2, 1, 2048, 128);
4465 try testing.expectEqual(tile, kernel.launch_resources.candidates[0].tile);
4466 try testing.expect(std.mem.indexOf(u8, kernel.artifact.payload.text, ".visible .entry accy_choir_reduction_sum_rank2_axis1_f32_0") != null);
4467 }
4468
4469 test "Choir artifact plan creates Vulkan artifacts from kernel-language reduction kernels" {
4470 const allocator = testing.allocator;
4471
4472 var owned = try reduceChoirModule(allocator, "choir_vulkan_reduce_artifact");
4473 defer owned.module.deinit();
4474
4475 var cache = passes.AnalysisCache.init(allocator, null);
4476 defer cache.deinit();
4477 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4478 defer pass_ctx.deinit();
4479
4480 var state = ArtifactPlanningBackendState.init(allocator, .vulkan);
4481 defer state.deinit();
4482
4483 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4484 defer plan.deinit();
4485
4486 try testing.expectEqual(gpu.BackendKind.vulkan, plan.backend_kind);
4487 try testing.expectEqual(gpu.ArtifactFormat.vulkan_spirv, plan.format);
4488 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4489 const kernel = plan.kernels.items[0];
4490 try testing.expectEqualStrings("accy_choir_reduction_sum_rank1_axis0_f32_0", kernel.artifact.entry_name);
4491 try testing.expectEqual(@as(u32, 3), kernel.artifact.argument_count);
4492 try testing.expectEqual(@as(usize, 1), kernel.input_slot_ids.len);
4493 try testing.expectEqual(@as(u64, 1), kernel.element_count);
4494 try testing.expectEqual(@as(u64, 1), kernel.element_count_argument_value);
4495 try testing.expectEqual(@as(u64, 256), kernel.resources.estimated_element_ops);
4496 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
4497 try testing.expectEqual(@as(usize, 0), kernel.static_arguments.len);
4498 try testing.expectEqual(LaunchResourceClass.memory_bound, kernel.launch_resources.resource_class);
4499 try testing.expect(kernel.launch_resources.fixed_threadgroup);
4500 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[0]);
4501 try testing.expectEqual(@as(u32, 64), kernel.launch_resources.geometry.threadgroup[0]);
4502 try testing.expectEqual(@as(u64, 65792), kernel.launch_resources.estimated_static_bytes_per_threadgroup);
4503 try testing.expectEqual(@as(u64, 16384), kernel.launch_resources.estimated_element_ops_per_threadgroup);
4504 const tile = kernel.launch_resources.tile;
4505 try expectReductionTile(tile, 64, 256, 1, 0, 65536, 256);
4506 try testing.expectEqual(tile, kernel.launch_resources.candidates[0].tile);
4507 try testing.expectEqual(@as(u32, 0x07230203), kernel.artifact.payload.words_u32[0]);
4508 }
4509
4510 test "Choir artifact plan creates Vulkan artifacts from rank-2 kernel-language reductions" {
4511 const allocator = testing.allocator;
4512
4513 var owned = try reduceRank2Axis1ChoirModule(allocator, "choir_vulkan_reduce_rank2_artifact");
4514 defer owned.module.deinit();
4515
4516 var cache = passes.AnalysisCache.init(allocator, null);
4517 defer cache.deinit();
4518 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4519 defer pass_ctx.deinit();
4520
4521 var state = ArtifactPlanningBackendState.init(allocator, .vulkan);
4522 defer state.deinit();
4523
4524 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4525 defer plan.deinit();
4526
4527 try testing.expectEqual(gpu.BackendKind.vulkan, plan.backend_kind);
4528 try testing.expectEqual(gpu.ArtifactFormat.vulkan_spirv, plan.format);
4529 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4530 const kernel = plan.kernels.items[0];
4531 try testing.expectEqualStrings("accy_choir_reduction_sum_rank2_axis1_f32_0", kernel.artifact.entry_name);
4532 try testing.expectEqual(@as(u32, 3), kernel.artifact.argument_count);
4533 try testing.expectEqual(@as(usize, 1), kernel.input_slot_ids.len);
4534 try testing.expectEqual(@as(u64, 16), kernel.element_count);
4535 try testing.expectEqual(@as(u64, 16), kernel.element_count_argument_value);
4536 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
4537 try testing.expectEqual(@as(usize, 0), kernel.static_arguments.len);
4538 try testing.expect(kernel.launch_resources.fixed_threadgroup);
4539 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[0]);
4540 try testing.expectEqual(@as(u32, 64), kernel.launch_resources.geometry.threadgroup[0]);
4541 try testing.expectEqual(@as(u64, 4352), kernel.launch_resources.estimated_static_bytes_per_threadgroup);
4542 try testing.expectEqual(@as(u64, 1024), kernel.launch_resources.estimated_element_ops_per_threadgroup);
4543 const tile = kernel.launch_resources.tile;
4544 try expectReductionTile(tile, 64, 16, 2, 1, 4096, 256);
4545 try testing.expectEqual(tile, kernel.launch_resources.candidates[0].tile);
4546 try testing.expectEqual(@as(u32, 0x07230203), kernel.artifact.payload.words_u32[0]);
4547 }
4548
4549 test "Choir artifact plan creates Metal artifacts from legal single-op kernels" {
4550 const allocator = testing.allocator;
4551
4552 var owned = try addChoirModule(allocator, "choir_metal_add_artifact");
4553 defer owned.module.deinit();
4554
4555 var cache = passes.AnalysisCache.init(allocator, null);
4556 defer cache.deinit();
4557 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4558 defer pass_ctx.deinit();
4559
4560 var state = ArtifactPlanningBackendState.init(allocator, .metal);
4561 defer state.deinit();
4562
4563 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4564 defer plan.deinit();
4565
4566 try testing.expectEqual(gpu.BackendKind.metal, plan.backend_kind);
4567 try testing.expectEqual(gpu.ArtifactFormat.metal_msl, plan.format);
4568 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4569 const kernel = plan.kernels.items[0];
4570 try testing.expectEqualStrings("accy_choir_elementwise_add_0", kernel.artifact.entry_name);
4571 try testing.expectEqual(@as(u32, 4), kernel.artifact.argument_count);
4572 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
4573 try testing.expectEqual(@as(u64, 8), kernel.element_count_argument_value);
4574 try testing.expectEqual(@as(u32, 1), kernel.launch_resources.geometry.grid[0]);
4575 try testing.expectEqual(@as(u32, 32), kernel.launch_resources.geometry.threadgroup[0]);
4576 try testing.expect(std.mem.indexOf(u8, kernel.artifact.payload.text, "kernel void accy_choir_elementwise_add_0") != null);
4577 try testing.expect(std.mem.indexOf(u8, kernel.artifact.payload.text, "thread_position_in_grid") != null);
4578 }
4579
4580 test "Choir artifact plan creates CUDA artifacts from legal fused kernels" {
4581 const allocator = testing.allocator;
4582
4583 var owned = try fusedAddMulChoirModule(allocator, "choir_cuda_fused_add_mul_artifact");
4584 defer owned.module.deinit();
4585
4586 var cache = passes.AnalysisCache.init(allocator, null);
4587 defer cache.deinit();
4588 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4589 defer pass_ctx.deinit();
4590
4591 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
4592 defer state.deinit();
4593
4594 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4595 defer plan.deinit();
4596
4597 try testing.expectEqual(gpu.BackendKind.cuda, plan.backend_kind);
4598 try testing.expectEqual(gpu.ArtifactFormat.cuda_ptx, plan.format);
4599 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4600 const kernel = plan.kernels.items[0];
4601 try testing.expectEqualStrings("accy_choir_elementwise_add_mul_0", kernel.artifact.entry_name);
4602 try testing.expectEqual(@as(u32, 5), kernel.artifact.argument_count);
4603 try testing.expectEqual(@as(usize, 3), kernel.input_slot_ids.len);
4604 try testing.expectEqual(@as(usize, 2), kernel.op_count);
4605 try testing.expectEqual(@as(usize, 3), kernel.resources.external_input_value_count);
4606 try testing.expectEqual(@as(usize, 3), kernel.resources.external_operand_count);
4607 try testing.expectEqual(@as(usize, 1), kernel.resources.chain_operand_count);
4608 try testing.expectEqual(@as(u64, 96), kernel.resources.static_read_bytes);
4609 try testing.expectEqual(@as(u64, 32), kernel.resources.static_write_bytes);
4610 try testing.expectEqual(@as(u64, 128), kernel.resources.static_total_bytes);
4611 try testing.expectEqual(@as(u64, 16), kernel.resources.estimated_element_ops);
4612 try testing.expectEqual(@as(u64, 8), kernel.element_count_argument_value);
4613 try testing.expectEqual(@as(u64, 128), kernel.resources.elementOpsPerKiB());
4614 try testing.expectEqual(LaunchResourceClass.memory_bound, kernel.launch_resources.resource_class);
4615 try testing.expectEqual(@as(u64, 512), kernel.launch_resources.estimated_static_bytes_per_threadgroup);
4616 try testing.expectEqual(@as(u64, 64), kernel.launch_resources.estimated_element_ops_per_threadgroup);
4617 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
4618 try testing.expect(std.mem.indexOf(u8, kernel.artifact.payload.text, ".visible .entry accy_choir_elementwise_add_mul_0") != null);
4619 try testing.expect(std.mem.indexOf(u8, kernel.artifact.payload.text, "mul.f32") != null);
4620 }
4621
4622 test "Choir artifact plan carries memory spaces and layouts into planned slots" {
4623 const allocator = testing.allocator;
4624
4625 var owned = try addConstantChoirModule(allocator, "choir_cuda_layout_metadata_artifact");
4626 defer owned.module.deinit();
4627
4628 var cache = passes.AnalysisCache.init(allocator, null);
4629 defer cache.deinit();
4630 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4631 defer pass_ctx.deinit();
4632
4633 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
4634 defer state.deinit();
4635
4636 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4637 defer plan.deinit();
4638
4639 try testing.expectEqual(@as(usize, 3), plan.slotCount());
4640 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4641
4642 var constant_slot: ?PlannedSlot = null;
4643 var output_slot: ?PlannedSlot = null;
4644 for (plan.slots) |slot| {
4645 if (slot.role.constant) constant_slot = slot;
4646 if (slot.role.output) output_slot = slot;
4647 try testing.expect(slot.layout_fingerprint != 0);
4648 try testing.expectEqual(layout_planning.LayoutKind.row_major, slot.layout_kind);
4649 try testing.expectEqualSlices(i64, &.{8}, slot.dims);
4650 try testing.expectEqualSlices(u64, &.{1}, slot.element_strides.?);
4651 try testing.expectEqualSlices(usize, &.{0}, slot.minor_to_major);
4652 try testing.expectEqual(@as(u64, 4), slot.alignment);
4653 try testing.expect(slot.contiguous);
4654 try testing.expect(slot.static_layout);
4655 }
4656
4657 const constant = constant_slot orelse return error.TestExpectedConstant;
4658 try testing.expectEqual(memory_space.MemorySpace.device_constant, constant.memory_space);
4659 try testing.expectEqual(memory_space.MemoryAccess.read_only, constant.memory_access);
4660 try testing.expectEqual(memory_space.BoundaryTransfer.none, constant.boundary_transfer);
4661 try testing.expectEqualSlices(u8, owned.payload[0..], constant.constantBytes().?);
4662
4663 const output = output_slot orelse return error.TestExpectedSlot;
4664 try testing.expectEqual(memory_space.MemorySpace.device_global, output.memory_space);
4665 try testing.expectEqual(memory_space.MemoryAccess.write_only, output.memory_access);
4666 try testing.expectEqual(memory_space.BoundaryTransfer.device_to_host, output.boundary_transfer);
4667 }
4668
4669 test "Choir backend memory plan keeps layout-incompatible allocations separate" {
4670 const allocation = [_]BackendBufferAllocationPlan{.{
4671 .allocation_id = 0,
4672 .byte_size = 32,
4673 .dtype = .f32,
4674 .memory_space = .device_global,
4675 .layout_kind = .row_major,
4676 .element_count = 8,
4677 .alignment = 4,
4678 .layout_fingerprint = 0x1111,
4679 .last_kernel_index = 0,
4680 }};
4681 const matching = BackendSlotLifetime{
4682 .slot_id = 1,
4683 .role = .{ .temporary = true },
4684 .dtype = .f32,
4685 .memory_space = .device_global,
4686 .layout_kind = .row_major,
4687 .element_count = 8,
4688 .byte_size = 32,
4689 .alignment = 4,
4690 .layout_fingerprint = 0x1111,
4691 .first_kernel_index = 1,
4692 .last_kernel_index = 1,
4693 };
4694 const mismatched = BackendSlotLifetime{
4695 .slot_id = 2,
4696 .role = .{ .temporary = true },
4697 .dtype = .f32,
4698 .memory_space = .device_global,
4699 .layout_kind = .row_major,
4700 .element_count = 8,
4701 .byte_size = 32,
4702 .alignment = 4,
4703 .layout_fingerprint = 0x2222,
4704 .first_kernel_index = 1,
4705 .last_kernel_index = 1,
4706 };
4707
4708 try testing.expectEqual(@as(?usize, 0), reusableAllocationIndex(allocation[0..], matching));
4709 try testing.expectEqual(@as(?usize, null), reusableAllocationIndex(allocation[0..], mismatched));
4710 }
4711
4712 test "Choir backend memory plan reuses dead static tensor slots" {
4713 const allocator = testing.allocator;
4714
4715 var owned = try deadTemporaryChoirModule(allocator, "choir_cuda_memory_reuse_artifact");
4716 defer owned.module.deinit();
4717
4718 var cache = passes.AnalysisCache.init(allocator, null);
4719 defer cache.deinit();
4720 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4721 defer pass_ctx.deinit();
4722
4723 var state = ArtifactPlanningBackendState.init(allocator, .cuda);
4724 defer state.deinit();
4725
4726 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4727 defer plan.deinit();
4728
4729 try testing.expectEqual(@as(usize, 10), plan.slotCount());
4730 try testing.expectEqual(@as(usize, 4), plan.kernelCount());
4731
4732 var memory = try createBackendMemoryPlan(allocator, &plan);
4733 defer memory.deinit();
4734
4735 try testing.expectEqual(@as(usize, 4), memory.assignments.len);
4736 try testing.expectEqual(@as(usize, 3), memory.allocations.len);
4737 try testing.expectEqual(@as(u64, 160), memory.total_static_slot_bytes);
4738 try testing.expectEqual(@as(u64, 128), memory.allocated_static_bytes);
4739 try testing.expectEqual(@as(u64, 128), memory.peak_static_live_bytes);
4740 try testing.expect(memory.assignmentForSlot(plan.output_slot_ids[0]) != null);
4741 try testing.expect(memory.assignmentForSlot(plan.output_slot_ids[1]) != null);
4742 }
4743
4744 test "Choir artifact plan creates Vulkan artifacts from legal fused kernels" {
4745 const allocator = testing.allocator;
4746
4747 var owned = try fusedAddMulChoirModule(allocator, "choir_vulkan_fused_add_mul_artifact");
4748 defer owned.module.deinit();
4749
4750 var cache = passes.AnalysisCache.init(allocator, null);
4751 defer cache.deinit();
4752 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4753 defer pass_ctx.deinit();
4754
4755 var state = ArtifactPlanningBackendState.init(allocator, .vulkan);
4756 defer state.deinit();
4757
4758 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4759 defer plan.deinit();
4760
4761 try testing.expectEqual(gpu.BackendKind.vulkan, plan.backend_kind);
4762 try testing.expectEqual(gpu.ArtifactFormat.vulkan_spirv, plan.format);
4763 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4764 const kernel = plan.kernels.items[0];
4765 try testing.expectEqualStrings("accy_choir_elementwise_add_mul_0", kernel.artifact.entry_name);
4766 try testing.expectEqual(@as(u32, 5), kernel.artifact.argument_count);
4767 try testing.expectEqual(@as(usize, 3), kernel.input_slot_ids.len);
4768 try testing.expectEqual(@as(usize, 2), kernel.op_count);
4769 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
4770 try testing.expectEqual(@as(u64, 8), kernel.element_count_argument_value);
4771 try testing.expectEqual(@as(u32, 0x07230203), kernel.artifact.payload.words_u32[0]);
4772 }
4773
4774 test "Choir artifact plan creates Metal artifacts from legal fused kernels" {
4775 const allocator = testing.allocator;
4776
4777 var owned = try fusedAddMulChoirModule(allocator, "choir_metal_fused_add_mul_artifact");
4778 defer owned.module.deinit();
4779
4780 var cache = passes.AnalysisCache.init(allocator, null);
4781 defer cache.deinit();
4782 var pass_ctx = passes.PassContext.init(owned.choir_module, owned.module.context(), allocator, &cache);
4783 defer pass_ctx.deinit();
4784
4785 var state = ArtifactPlanningBackendState.init(allocator, .metal);
4786 defer state.deinit();
4787
4788 var plan = try createTestBackendArtifactPlan(allocator, state.handle(), &pass_ctx, owned.choir_module, .{});
4789 defer plan.deinit();
4790
4791 try testing.expectEqual(gpu.BackendKind.metal, plan.backend_kind);
4792 try testing.expectEqual(gpu.ArtifactFormat.metal_msl, plan.format);
4793 try testing.expectEqual(@as(usize, 1), plan.kernelCount());
4794 const kernel = plan.kernels.items[0];
4795 try testing.expectEqualStrings("accy_choir_elementwise_add_mul_0", kernel.artifact.entry_name);
4796 try testing.expectEqual(@as(u32, 5), kernel.artifact.argument_count);
4797 try testing.expectEqual(@as(usize, 3), kernel.input_slot_ids.len);
4798 try testing.expectEqual(@as(usize, 2), kernel.op_count);
4799 try testing.expectEqual(ElementCountArgument.device_buffer_u32, kernel.element_count_argument);
4800 try testing.expectEqual(@as(u64, 8), kernel.element_count_argument_value);
4801 try testing.expect(std.mem.indexOf(u8, kernel.artifact.payload.text, "kernel void accy_choir_elementwise_add_mul_0") != null);
4802 try testing.expect(std.mem.indexOf(u8, kernel.artifact.payload.text, " * ") != null);
4803 }