lib/accy/src/preparation/backend.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const sys = @import("sys");
3 const choir = @import("choir");
4 const bufferization = @import("bufferization/root.zig");
5 const accy_choir = @import("../choir/root.zig");
6 const kernelization_model = @import("kernelization/model/root.zig");
7 const kernel_outlining = @import("outlining/root.zig");
8 const target_profile = @import("target.zig");
9
10 const ir = choir.ir;
11 const passes = choir.passes;
12 const accounting = passes.pass.work;
13
14 pub const backend_legalization_analysis_name = "accy-choir-backend-legalization";
15 pub const backend_legalization_pass_name = "accy-choir-legalize-backend";
16 pub const backend_legalization_pass_description =
17 "Validate Accy Choir kernel candidates for backend artifact planning";
18
19 pub const BackendKernelStatus = enum {
20 legal,
21 missing_input_slot,
22 missing_output_slot,
23 dynamic_input_slot,
24 dynamic_output_slot,
25 illegal_output_role,
26 unsupported_dtype,
27 };
28
29 pub const BackendKernelLegalization = struct {
30 kernel_id: usize,
31 status: BackendKernelStatus,
32 input_count: usize,
33 output_slot_id: usize,
34 static_bytes: u64 = 0,
35
36 pub fn isLegal(self: BackendKernelLegalization) bool {
37 return self.status == .legal;
38 }
39 };
40
41 pub const BackendLegalizationAnalysis = struct {
42 allocator: std.mem.Allocator,
43 kernels: std.ArrayListUnmanaged(BackendKernelLegalization),
44 target: ?target_profile.BackendTargetProfile = null,
45 legal_kernel_count: usize = 0,
46 illegal_kernel_count: usize = 0,
47 total_static_bytes: u64 = 0,
48
49 pub fn init(
50 allocator: std.mem.Allocator,
51 target: ?target_profile.BackendTargetProfile,
52 ) BackendLegalizationAnalysis {
53 return .{
54 .allocator = allocator,
55 .kernels = .empty,
56 .target = target,
57 };
58 }
59
60 pub fn deinit(self: *BackendLegalizationAnalysis) void {
61 self.kernels.deinit(self.allocator);
62 self.* = undefined;
63 }
64
65 pub fn kernelCount(self: BackendLegalizationAnalysis) usize {
66 return self.kernels.items.len;
67 }
68
69 pub fn isLegal(self: BackendLegalizationAnalysis) bool {
70 return self.illegal_kernel_count == 0;
71 }
72
73 pub fn hasPipelineFailure(self: BackendLegalizationAnalysis) bool {
74 for (self.kernels.items) |kernel| {
75 if (kernel.status != .legal and kernel.status != .unsupported_dtype) return true;
76 }
77 return false;
78 }
79
80 fn addKernel(self: *BackendLegalizationAnalysis, result: BackendKernelLegalization) !void {
81 const total = try accounting.add(
82 self.total_static_bytes,
83 if (result.isLegal()) result.static_bytes else 0,
84 );
85 try self.kernels.append(self.allocator, result);
86 if (result.isLegal()) {
87 self.legal_kernel_count += 1;
88 self.total_static_bytes = total;
89 } else {
90 self.illegal_kernel_count += 1;
91 }
92 }
93
94 pub fn firstIllegalStatus(self: BackendLegalizationAnalysis) ?BackendKernelStatus {
95 for (self.kernels.items) |kernel| {
96 if (!kernel.isLegal()) return kernel.status;
97 }
98 return null;
99 }
100 };
101
102 const BackendWork = struct {
103 input: accounting.Census,
104
105 fn bounds(self: BackendWork) !accounting.Bounds {
106 const count = self.input.operations;
107 const storage = try accounting.add(
108 @sizeOf(BackendLegalizationAnalysis) + @alignOf(BackendLegalizationAnalysis),
109 try accounting.arrayListGrowth(BackendKernelLegalization, count),
110 );
111 if (storage > std.math.maxInt(usize)) return error.WorkOverflow;
112 const input_units = try accounting.add(self.input.atoms, self.input.input_bytes);
113 const units = try accounting.add(input_units, 1);
114 const checks = try accounting.multiply(
115 try accounting.add(count, 1),
116 try accounting.add(self.input.values, 1),
117 );
118 const visits = try accounting.multiply(64, try accounting.multiply(units, checks));
119 return .{
120 .work = .{
121 .input_bytes = self.input.input_bytes,
122 .structural_visits = visits,
123 .analysis_computations = 1,
124 .allocation_capacity = storage,
125 },
126 .workspace = storage,
127 .retained_storage = storage,
128 };
129 }
130 };
131
132 fn backendAnalysisWork(input: accounting.Input) !accounting.Bounds {
133 if (input.options.max_threads != 1 or input.options.worker_allocator != null) {
134 return error.MissingWorkContract;
135 }
136 return (BackendWork{ .input = try accounting.Census.inspect(input.operation) }).bounds();
137 }
138
139 fn backendPassWork(input: accounting.Input) !accounting.Bounds {
140 const counts = try accounting.Census.inspect(input.operation);
141 return .{ .work = .{ .structural_visits = try accounting.add(counts.operations, 1) } };
142 }
143
144 pub const backend_legalization_analysis_descriptor = passes.AnalysisDescriptor{
145 .id = passes.analysisId(backend_legalization_analysis_name),
146 .name = backend_legalization_analysis_name,
147 .work_contract = .{
148 .identity = .{ .name = backend_legalization_analysis_name, .version = 1 },
149 .estimate = backendAnalysisWork,
150 },
151 };
152
153 pub fn getBackendLegalizationAnalysis(
154 pass_ctx: *passes.PassContext,
155 op: *ir.Operation,
156 ) !*BackendLegalizationAnalysis {
157 const ptr = try pass_ctx.getAnalysis(
158 op,
159 &backend_legalization_analysis_descriptor,
160 computeBackendLegalizationAnalysis,
161 cleanupBackendLegalizationAnalysis,
162 );
163 return @ptrCast(@alignCast(ptr));
164 }
165
166 pub fn backendLegalizationPass() passes.Pass {
167 return .{
168 .name = backend_legalization_pass_name,
169 .description = backend_legalization_pass_description,
170 .run_fn = runBackendLegalizationPass,
171 .work_contract = .{
172 .identity = .{ .name = backend_legalization_pass_name, .version = 1 },
173 .estimate = backendPassWork,
174 },
175 };
176 }
177
178 fn runBackendLegalizationPass(pass_ctx: *passes.PassContext) passes.PassResult {
179 const analysis = getBackendLegalizationAnalysis(pass_ctx, pass_ctx.op) catch return .failure;
180 if (analysis.hasPipelineFailure()) return .failure;
181 pass_ctx.preserveAllAnalyses();
182 return .success;
183 }
184
185 fn computeBackendLegalizationAnalysis(
186 pass_ctx: *passes.PassContext,
187 op: *ir.Operation,
188 ) anyerror!*anyopaque {
189 const outline_plan = try kernel_outlining.getKernelOutlinePlanAnalysis(pass_ctx, op);
190 const buffer_plan = try bufferization.getBufferPlanAnalysis(pass_ctx, op);
191
192 const analysis = try pass_ctx.allocator.create(BackendLegalizationAnalysis);
193 const target = target_profile.readBackendTargetProfile(op);
194 analysis.* = BackendLegalizationAnalysis.init(pass_ctx.allocator, target);
195 errdefer {
196 analysis.deinit();
197 pass_ctx.allocator.destroy(analysis);
198 }
199
200 try addLegalizedKernels(pass_ctx, analysis, outline_plan, buffer_plan);
201
202 return @ptrCast(analysis);
203 }
204
205 fn cleanupBackendLegalizationAnalysis(ptr: *anyopaque, allocator: std.mem.Allocator) void {
206 const analysis: *BackendLegalizationAnalysis = @ptrCast(@alignCast(ptr));
207 analysis.deinit();
208 allocator.destroy(analysis);
209 }
210
211 pub fn backendKernelStatusError(status: BackendKernelStatus) anyerror {
212 return switch (status) {
213 .unsupported_dtype => error.CapabilityMismatch,
214 else => error.IllegalKernel,
215 };
216 }
217
218 /// A caller uses this to learn whether one outlined kernel can run on the chosen device with the
219 /// buffers the memory plan gives it. The function looks up the kernel's output buffer and each
220 /// input buffer in the memory plan and returns a result with a status, the kernel id, the input
221 /// count, the output buffer id and the total static size in bytes. The status names the first
222 /// problem found: a missing or dynamically sized output buffer, an output buffer in a role a kernel
223 /// may not write, a missing or dynamically sized input buffer, or an element type the device does
224 /// not support. The function checks element types only when given a target profile, the facts about
225 /// the chosen device, such as which element types it supports. An illegal kernel comes back as a
226 /// status and never as an error, and the only error is an overflow while adding up the byte sizes.
227 pub fn legalizeKernel(
228 kernel: kernelization_model.KernelOutline,
229 buffer_plan: *const bufferization.BufferPlanAnalysis,
230 target: ?target_profile.BackendTargetProfile,
231 ) !BackendKernelLegalization {
232 var result = BackendKernelLegalization{
233 .kernel_id = kernel.id,
234 .status = .legal,
235 .input_count = kernel.inputCount(),
236 .output_slot_id = kernel.output_slot_id,
237 };
238
239 const output_slot = slotById(buffer_plan, kernel.output_slot_id) orelse {
240 result.status = .missing_output_slot;
241 return result;
242 };
243 if (!output_slot.hasStaticSize()) {
244 result.status = .dynamic_output_slot;
245 return result;
246 }
247 if (!outputSlotRoleIsLegal(output_slot.role)) {
248 result.status = .illegal_output_role;
249 return result;
250 }
251 if (target) |profile| {
252 if (!profile.supportsDType(output_slot.dtype)) {
253 result.status = .unsupported_dtype;
254 return result;
255 }
256 }
257 result.static_bytes = try accounting.add(result.static_bytes, output_slot.byte_size.?);
258
259 for (kernel.input_slot_ids) |slot_id| {
260 const input_slot = slotById(buffer_plan, slot_id) orelse {
261 result.status = .missing_input_slot;
262 return result;
263 };
264 if (!input_slot.hasStaticSize()) {
265 result.status = .dynamic_input_slot;
266 return result;
267 }
268 if (target) |profile| {
269 if (!profile.supportsDType(input_slot.dtype)) {
270 result.status = .unsupported_dtype;
271 return result;
272 }
273 }
274 result.static_bytes = try accounting.add(result.static_bytes, input_slot.byte_size.?);
275 }
276
277 return result;
278 }
279
280 fn addLegalizedKernels(
281 pass_ctx: *passes.PassContext,
282 analysis: *BackendLegalizationAnalysis,
283 outline_plan: *const kernelization_model.KernelOutlinePlanAnalysis,
284 buffer_plan: *const bufferization.BufferPlanAnalysis,
285 ) !void {
286 const kernels = outline_plan.kernels.items;
287 if (pass_ctx.workerCount(kernels.len) > 1) {
288 return try addLegalizedKernelsParallel(pass_ctx, analysis, kernels, buffer_plan);
289 }
290 for (kernels) |kernel| {
291 try analysis.addKernel(try legalizeKernel(kernel, buffer_plan, analysis.target));
292 }
293 }
294
295 const BackendLegalizationSlot = struct {
296 result: ?(anyerror!BackendKernelLegalization) = null,
297 };
298
299 const BackendLegalizationBatch = struct {
300 kernels: []const kernelization_model.KernelOutline,
301 buffer_plan: *const bufferization.BufferPlanAnalysis,
302 target: ?target_profile.BackendTargetProfile,
303 slots: []BackendLegalizationSlot,
304 };
305
306 fn addLegalizedKernelsParallel(
307 pass_ctx: *passes.PassContext,
308 analysis: *BackendLegalizationAnalysis,
309 kernels: []const kernelization_model.KernelOutline,
310 buffer_plan: *const bufferization.BufferPlanAnalysis,
311 ) !void {
312 const slots = try pass_ctx.allocator.alloc(BackendLegalizationSlot, kernels.len);
313 defer pass_ctx.allocator.free(slots);
314 for (slots) |*slot| slot.* = .{};
315
316 var batch = BackendLegalizationBatch{
317 .kernels = kernels,
318 .buffer_plan = buffer_plan,
319 .target = analysis.target,
320 .slots = slots,
321 };
322
323 try ir.threading.parallelForEachIndex(
324 pass_ctx.allocator,
325 pass_ctx.run_options,
326 kernels.len,
327 &batch,
328 legalizeKernelAt,
329 );
330
331 for (slots) |slot| {
332 try analysis.addKernel(try (slot.result orelse unreachable));
333 }
334 }
335
336 fn legalizeKernelAt(batch: *BackendLegalizationBatch, index: usize) void {
337 batch.slots[index].result = legalizeKernel(
338 batch.kernels[index],
339 batch.buffer_plan,
340 batch.target,
341 );
342 }
343
344 fn slotById(
345 buffer_plan: *const bufferization.BufferPlanAnalysis,
346 slot_id: usize,
347 ) ?*const bufferization.BufferSlot {
348 if (slot_id >= buffer_plan.slots.items.len) return null;
349 const slot = &buffer_plan.slots.items[slot_id];
350 if (slot.id != slot_id) return null;
351 return slot;
352 }
353
354 fn outputSlotRoleIsLegal(role: bufferization.BufferRole) bool {
355 return role.output or role.temporary;
356 }
357
358 const testing = std.testing;
359 const semantic = accy_choir.semantic;
360
361 test "backend legalization accepts static elementwise outlined kernels" {
362 const allocator = testing.allocator;
363
364 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
365 defer builder.deinit();
366 const f32_4 = try builder.tensor(.f32, &.{4});
367 var fb = try builder.beginFunction("backend_legal_fused_add_mul", &.{ f32_4, f32_4, f32_4 }, &.{f32_4});
368 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
369 const product = try fb.mul(sum, fb.parameter(2));
370 try fb.return_(&.{product});
371 try fb.finish();
372 const module = try builder.finish();
373 defer module.deinit();
374
375 const choir_mod = module.choir_module;
376 const ctx = module.context();
377 const ledger = try backendTestLedger(false, std.math.maxInt(u64));
378 defer ledger.destroy();
379 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);
380 defer cache.deinit();
381 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
382 defer pass_ctx.deinit();
383
384 const analysis = try getBackendLegalizationAnalysis(&pass_ctx, choir_mod);
385 try ledger.producersComplete();
386 try checkBackendStorage(choir_mod, ctx, &cache, analysis);
387 try testing.expect(analysis.isLegal());
388 try testing.expectEqual(@as(usize, 1), analysis.kernelCount());
389 try testing.expectEqual(@as(usize, 1), analysis.legal_kernel_count);
390 try testing.expectEqual(@as(usize, 0), analysis.illegal_kernel_count);
391 try testing.expectEqual(@as(u64, 64), analysis.total_static_bytes);
392 try testing.expectEqual(BackendKernelStatus.legal, analysis.kernels.items[0].status);
393 }
394
395 test "backend legalization threads match serial analysis" {
396 const allocator = testing.allocator;
397
398 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
399 defer builder.deinit();
400 const f32_4 = try builder.tensor(.f32, &.{4});
401 var fb = try builder.beginFunction("backend_legal_threaded", &.{ f32_4, f32_4, f32_4 }, &.{ f32_4, f32_4 });
402 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
403 const product = try fb.mul(sum, fb.parameter(2));
404 try fb.return_(&.{ sum, product });
405 try fb.finish();
406 const module = try builder.finish();
407 defer module.deinit();
408
409 const choir_mod = module.choir_module;
410 const ctx = module.context();
411
412 var serial_cache = passes.AnalysisCache.init(allocator, null);
413 defer serial_cache.deinit();
414 var serial_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &serial_cache);
415 defer serial_ctx.deinit();
416 const serial = try getBackendLegalizationAnalysis(&serial_ctx, choir_mod);
417
418 var threaded_cache = passes.AnalysisCache.init(allocator, null);
419 defer threaded_cache.deinit();
420 var threaded_ctx = passes.PassContext.initWithOptions(choir_mod, ctx, allocator, &threaded_cache, .{
421 .max_threads = 2,
422 });
423 defer threaded_ctx.deinit();
424 const threaded = try getBackendLegalizationAnalysis(&threaded_ctx, choir_mod);
425
426 try testing.expectEqual(serial.kernelCount(), threaded.kernelCount());
427 try testing.expectEqual(serial.legal_kernel_count, threaded.legal_kernel_count);
428 try testing.expectEqual(serial.illegal_kernel_count, threaded.illegal_kernel_count);
429 try testing.expectEqual(serial.total_static_bytes, threaded.total_static_bytes);
430 for (serial.kernels.items, threaded.kernels.items) |expected, actual| {
431 try testing.expectEqual(expected.kernel_id, actual.kernel_id);
432 try testing.expectEqual(expected.status, actual.status);
433 try testing.expectEqual(expected.input_count, actual.input_count);
434 try testing.expectEqual(expected.output_slot_id, actual.output_slot_id);
435 try testing.expectEqual(expected.static_bytes, actual.static_bytes);
436 }
437 }
438
439 test "backend legalization pass preserves legal IR" {
440 const allocator = testing.allocator;
441
442 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
443 defer builder.deinit();
444 const f32_4 = try builder.tensor(.f32, &.{4});
445 var fb = try builder.beginFunction("backend_legal_pass_add4", &.{ f32_4, f32_4 }, &.{f32_4});
446 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
447 try fb.return_(&.{sum});
448 try fb.finish();
449 const module = try builder.finish();
450 defer module.deinit();
451
452 const choir_mod = module.choir_module;
453 const ctx = module.context();
454 var pm = passes.PassManager.init(allocator);
455 defer pm.deinit();
456 try pm.addPass(backendLegalizationPass());
457
458 const ledger = try backendTestLedger(true, std.math.maxInt(u64));
459 defer ledger.destroy();
460 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);
461 defer cache.deinit();
462 try testing.expectEqual(
463 passes.PassResult.success,
464 pm.runWithAnalysisCache(choir_mod, ctx, &cache, .{}),
465 );
466 try ledger.producersComplete();
467 try testing.expectEqual(@as(u64, 1), pm.stats.pass_runs);
468 try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
469 }
470
471 fn backendTestLedger(pass: bool, visits: u64) !*choir.product.revision.AccountingV1 {
472 const revision = choir.product.revision;
473 var allowance = revision.WorkVector.uniform(std.math.maxInt(u64));
474 allowance.structural_visits = visits;
475 const pipeline = [_]revision.record.Version{
476 .{ .name = backend_legalization_pass_name, .version = 1 },
477 };
478 return revision.AccountingV1.create(testing.allocator, .{
479 .allowance = allowance,
480 .workspace = std.math.maxInt(u64),
481 .events = 32,
482 }, if (pass) &pipeline else &.{});
483 }
484
485 fn buildBackendModule(
486 builder: *semantic.Builder,
487 count: usize,
488 extent: i64,
489 alias: bool,
490 ) !*semantic.SemanticModule {
491 const typ = try builder.tensor(.f32, &.{extent});
492 const types = try testing.allocator.alloc(ir.Type, count);
493 defer testing.allocator.free(types);
494 @memset(types, typ);
495 const results = try testing.allocator.alloc(*ir.Value, count);
496 defer testing.allocator.free(results);
497 var function = try builder.beginFunction(
498 "backend_work",
499 if (alias) &.{typ} else &.{ typ, typ },
500 types,
501 );
502 for (results) |*result| {
503 if (alias) {
504 const call = try function.kernelCall(&.{function.parameter(0)}, &.{typ}, .{
505 .target = "update_f32",
506 .operand_effects = &.{.read_write},
507 .result_aliases = &.{0},
508 });
509 result.* = call.getFirstResult();
510 } else {
511 result.* = try function.add(function.parameter(0), function.parameter(1));
512 }
513 }
514 try function.return_(results);
515 try function.finish();
516 return builder.finish();
517 }
518
519 fn checkBackendStorage(
520 op: *ir.Operation,
521 ctx: *ir.Context,
522 cache: *passes.AnalysisCache,
523 expected: *const BackendLegalizationAnalysis,
524 ) !void {
525 const bounds = try backendAnalysisWork(.{ .operation = op });
526 const fixed = @import("alloc_fixed");
527 const bytes = try testing.allocator.alignedAlloc(u8, .@"64", @intCast(bounds.workspace));
528 defer testing.allocator.free(bytes);
529 var backing = fixed.Tracked.init(bytes);
530 var retained = fixed.Monotonic.init(backing.allocator(), @max(1, bytes.len));
531 const allocator = retained.allocator();
532 var ctx_fixed = passes.PassContext.init(op, ctx, allocator, cache);
533 defer ctx_fixed.deinit();
534 const ptr = try computeBackendLegalizationAnalysis(&ctx_fixed, op);
535 defer cleanupBackendLegalizationAnalysis(ptr, allocator);
536 const actual: *BackendLegalizationAnalysis = @ptrCast(@alignCast(ptr));
537 try testing.expectEqual(expected.legal_kernel_count, actual.legal_kernel_count);
538 try testing.expectEqual(expected.illegal_kernel_count, actual.illegal_kernel_count);
539 try testing.expectEqual(expected.total_static_bytes, actual.total_static_bytes);
540 try testing.expectEqualDeep(expected.target, actual.target);
541 try testing.expectEqualDeep(expected.kernels.items, actual.kernels.items);
542 const used = if (retained.current) |*current| fixed.used(current) else 0;
543 try testing.expect(!backing.exhausted);
544 try testing.expect(used >= @sizeOf(BackendLegalizationAnalysis));
545 try testing.expect(used <= bounds.workspace);
546 }
547
548 test "backend legalization work contract covers growing output lists" {
549 for ([_]usize{ 0, 1, 2, 6, 7, 16, 64 }) |count| {
550 var builder = try semantic.Builder.init(
551 testing.allocator,
552 semantic.Builder.ContextLimits.standard,
553 );
554 defer builder.deinit();
555 const module = try buildBackendModule(&builder, count, 4, false);
556 defer module.deinit();
557 const ledger = try backendTestLedger(false, std.math.maxInt(u64));
558 defer ledger.destroy();
559 var cache = try passes.AnalysisCache.initAccounted(
560 testing.allocator,
561 null,
562 ledger,
563 .{},
564 8,
565 );
566 defer cache.deinit();
567 var ctx = passes.PassContext.init(
568 module.choir_module,
569 module.context(),
570 testing.allocator,
571 &cache,
572 );
573 defer ctx.deinit();
574 const analysis = try getBackendLegalizationAnalysis(&ctx, module.choir_module);
575 try testing.expectEqual(count, analysis.kernelCount());
576 try testing.expectEqual(@as(u64, count * 48), analysis.total_static_bytes);
577 try ledger.producersComplete();
578 try checkBackendStorage(module.choir_module, module.context(), &cache, analysis);
579 }
580 try testing.expectError(error.WorkOverflow, (BackendWork{
581 .input = .{ .operations = std.math.maxInt(u64) },
582 }).bounds());
583 try testing.expectError(error.WorkOverflow, (BackendWork{
584 .input = .{ .values = std.math.maxInt(u64) },
585 }).bounds());
586 try testing.expectError(error.WorkOverflow, (BackendWork{
587 .input = .{ .input_bytes = std.math.maxInt(u64) },
588 }).bounds());
589 }
590
591 test "backend legalization checks individual and repeated-buffer byte totals" {
592 const aggregate_edge: i64 = @intCast(std.math.maxInt(u64) / 36);
593 const alias_edge: i64 = @intCast(std.math.maxInt(u64) / 8);
594 for ([_]i64{ -1, 0, 1 }) |offset| {
595 try checkBackendBytes(3, aggregate_edge + offset, false, false);
596 try checkBackendBytes(3, aggregate_edge + offset, false, true);
597 try checkBackendBytes(1, alias_edge + offset, true, false);
598 }
599 try checkBackendBytes(2, alias_edge + 1, true, true);
600 }
601
602 fn checkBackendBytes(count: usize, extent: i64, alias: bool, parallel: bool) !void {
603 if (parallel and !sys.thread.threadsSupported()) return error.SkipZigTest;
604 var builder = try semantic.Builder.init(
605 testing.allocator,
606 semantic.Builder.ContextLimits.standard,
607 );
608 defer builder.deinit();
609 const module = try buildBackendModule(&builder, count, extent, alias);
610 defer module.deinit();
611 const ledger = if (parallel) null else try backendTestLedger(false, std.math.maxInt(u64));
612 defer if (ledger) |value| value.destroy();
613 var cache = if (ledger) |value|
614 try passes.AnalysisCache.initAccounted(testing.allocator, null, value, .{}, 8)
615 else
616 passes.AnalysisCache.init(testing.allocator, null);
617 defer cache.deinit();
618 var ctx = passes.PassContext.initWithOptions(
619 module.choir_module,
620 module.context(),
621 testing.allocator,
622 &cache,
623 .{ .max_threads = if (parallel) 2 else 1 },
624 );
625 defer ctx.deinit();
626 const outlines = try kernel_outlining.getKernelOutlinePlanAnalysis(&ctx, module.choir_module);
627 const buffers = try bufferization.getBufferPlanAnalysis(&ctx, module.choir_module);
628 try testing.expectEqual(count, outlines.kernelCount());
629 try testing.expect(buffers.total_static_bytes > 0);
630 const entries = cache.entries.count();
631 const total: u128 = @as(u128, @intCast(extent)) * 4 * count * @as(u128, if (alias) 2 else 3);
632 if (total <= std.math.maxInt(u64)) {
633 const result = try getBackendLegalizationAnalysis(&ctx, module.choir_module);
634 try testing.expectEqual(@as(u64, @intCast(total)), result.total_static_bytes);
635 try testing.expectEqual(count, result.legal_kernel_count);
636 if (ledger) |value| try value.producersComplete();
637 } else {
638 try testing.expectError(
639 error.WorkOverflow,
640 getBackendLegalizationAnalysis(&ctx, module.choir_module),
641 );
642 try testing.expectEqual(entries, cache.entries.count());
643 if (ledger) |value| {
644 try testing.expectEqual(.exhausted, value.view().outcome);
645 try testing.expectError(
646 error.TerminalWorkOutcome,
647 getBackendLegalizationAnalysis(&ctx, module.choir_module),
648 );
649 }
650 }
651 }
652
653 test "backend legalization overflow preserves the admitted result prefix" {
654 var analysis = BackendLegalizationAnalysis.init(testing.allocator, null);
655 defer analysis.deinit();
656 const first = BackendKernelLegalization{
657 .kernel_id = 0,
658 .status = .legal,
659 .input_count = 0,
660 .output_slot_id = 0,
661 .static_bytes = std.math.maxInt(u64),
662 };
663 try analysis.addKernel(first);
664 try testing.expectError(error.WorkOverflow, analysis.addKernel(.{
665 .kernel_id = 1,
666 .status = .legal,
667 .input_count = 0,
668 .output_slot_id = 1,
669 .static_bytes = 1,
670 }));
671 try testing.expectEqual(@as(usize, 1), analysis.kernelCount());
672 try testing.expectEqual(@as(usize, 1), analysis.legal_kernel_count);
673 try testing.expectEqual(std.math.maxInt(u64), analysis.total_static_bytes);
674 try testing.expectEqualDeep(first, analysis.kernels.items[0]);
675 }
676
677 test "backend legalization rejects unmodeled options with cached dependencies" {
678 for ([_]ir.ThreadingOptions{
679 .{ .max_threads = 0 },
680 .{ .max_threads = 2 },
681 .{ .worker_allocator = testing.allocator },
682 }) |options| {
683 var builder = try semantic.Builder.init(
684 testing.allocator,
685 semantic.Builder.ContextLimits.standard,
686 );
687 defer builder.deinit();
688 const module = try buildBackendModule(&builder, 2, 4, false);
689 defer module.deinit();
690 const ledger = try backendTestLedger(false, std.math.maxInt(u64));
691 defer ledger.destroy();
692 var cache = try passes.AnalysisCache.initAccounted(
693 testing.allocator,
694 null,
695 ledger,
696 .{},
697 8,
698 );
699 defer cache.deinit();
700 var ctx = passes.PassContext.init(
701 module.choir_module,
702 module.context(),
703 testing.allocator,
704 &cache,
705 );
706 defer ctx.deinit();
707 _ = try kernel_outlining.getKernelOutlinePlanAnalysis(&ctx, module.choir_module);
708 const before = ledger.view().charged;
709 const entries = cache.entries.count();
710 ctx.run_options = options;
711 try testing.expectError(
712 error.MissingWorkContract,
713 getBackendLegalizationAnalysis(&ctx, module.choir_module),
714 );
715 try testing.expectEqualDeep(before, ledger.view().charged);
716 try testing.expectEqual(entries, cache.entries.count());
717 try testing.expectEqual(.rejected, ledger.view().outcome);
718 ctx.run_options = .{};
719 try testing.expectError(
720 error.TerminalWorkOutcome,
721 getBackendLegalizationAnalysis(&ctx, module.choir_module),
722 );
723 }
724 }
725
726 test "backend legalization preserves target capability failure as a distinct result" {
727 var builder = try semantic.Builder.init(
728 testing.allocator,
729 semantic.Builder.ContextLimits.standard,
730 );
731 defer builder.deinit();
732 const module = try buildBackendModule(&builder, 2, 4, false);
733 defer module.deinit();
734 const profile = target_profile.BackendTargetProfile{
735 .backend_kind = .cuda,
736 .artifact_format = .cuda_ptx,
737 .math_tier = .exact,
738 .dtype_bits = 0,
739 .feature_bits = 0,
740 };
741 try target_profile.setBackendTargetProfile(module.context(), module.choir_module, profile);
742 const ledger = try backendTestLedger(true, std.math.maxInt(u64));
743 defer ledger.destroy();
744 var cache = try passes.AnalysisCache.initAccounted(
745 testing.allocator,
746 null,
747 ledger,
748 .{},
749 8,
750 );
751 defer cache.deinit();
752 var manager = passes.PassManager.init(testing.allocator);
753 defer manager.deinit();
754 try manager.addPass(backendLegalizationPass());
755 try testing.expectEqual(
756 passes.PassResult.success,
757 manager.runWithAnalysisCache(module.choir_module, module.context(), &cache, .{}),
758 );
759 try ledger.producersComplete();
760 var ctx = passes.PassContext.init(
761 module.choir_module,
762 module.context(),
763 testing.allocator,
764 &cache,
765 );
766 defer ctx.deinit();
767 const analysis = try getBackendLegalizationAnalysis(&ctx, module.choir_module);
768 try testing.expect(!analysis.isLegal());
769 try testing.expect(!analysis.hasPipelineFailure());
770 try testing.expectEqual(@as(usize, 2), analysis.illegal_kernel_count);
771 try testing.expectEqual(@as(u64, 0), analysis.total_static_bytes);
772 try testing.expectEqualDeep(profile, analysis.target.?);
773 try testing.expectEqual(.unsupported_dtype, analysis.firstIllegalStatus().?);
774 try testing.expectEqual(error.CapabilityMismatch, backendKernelStatusError(.unsupported_dtype));
775 try checkBackendStorage(module.choir_module, module.context(), &cache, analysis);
776 }
777
778 test "backend legalization retains slot failure distinctions" {
779 var builder = try semantic.Builder.init(
780 testing.allocator,
781 semantic.Builder.ContextLimits.standard,
782 );
783 defer builder.deinit();
784 const module = try buildBackendModule(&builder, 1, 4, false);
785 defer module.deinit();
786 var cache = passes.AnalysisCache.init(testing.allocator, null);
787 defer cache.deinit();
788 var ctx = passes.PassContext.init(
789 module.choir_module,
790 module.context(),
791 testing.allocator,
792 &cache,
793 );
794 defer ctx.deinit();
795 const outlines = try kernel_outlining.getKernelOutlinePlanAnalysis(&ctx, module.choir_module);
796 const buffers = try bufferization.getBufferPlanAnalysis(&ctx, module.choir_module);
797 const kernel = outlines.kernels.items[0];
798 var bad_output = kernel;
799 bad_output.output_slot_id = buffers.slots.items.len;
800 try testing.expectEqual(
801 .missing_output_slot,
802 (try legalizeKernel(bad_output, buffers, null)).status,
803 );
804 var bad_inputs = kernel;
805 var missing = [_]usize{buffers.slots.items.len};
806 bad_inputs.input_slot_ids = &missing;
807 try testing.expectEqual(
808 .missing_input_slot,
809 (try legalizeKernel(bad_inputs, buffers, null)).status,
810 );
811 const output = &buffers.slots.items[kernel.output_slot_id];
812 const saved_output = output.*;
813 defer output.* = saved_output;
814 output.byte_size = null;
815 try testing.expectEqual(
816 .dynamic_output_slot,
817 (try legalizeKernel(kernel, buffers, null)).status,
818 );
819 output.* = saved_output;
820 output.role = .{};
821 try testing.expectEqual(
822 .illegal_output_role,
823 (try legalizeKernel(kernel, buffers, null)).status,
824 );
825 output.* = saved_output;
826 const input = &buffers.slots.items[kernel.input_slot_ids[0]];
827 const saved_input = input.*;
828 defer input.* = saved_input;
829 input.byte_size = null;
830 try testing.expectEqual(
831 .dynamic_input_slot,
832 (try legalizeKernel(kernel, buffers, null)).status,
833 );
834 }
835
836 test "backend legalization enforces its pass allowance and terminal retry" {
837 var builder = try semantic.Builder.init(
838 testing.allocator,
839 semantic.Builder.ContextLimits.standard,
840 );
841 defer builder.deinit();
842 const module = try buildBackendModule(&builder, 2, 4, false);
843 defer module.deinit();
844 const charge = try checkBackendAllowance(module, std.math.maxInt(u64), true);
845 try testing.expect(charge > 1);
846 _ = try checkBackendAllowance(module, charge - 1, false);
847 try testing.expectEqual(charge, try checkBackendAllowance(module, charge, true));
848 try testing.expectEqual(charge, try checkBackendAllowance(module, charge + 1, true));
849 }
850
851 fn checkBackendAllowance(module: *semantic.SemanticModule, visits: u64, success: bool) !u64 {
852 const ledger = try backendTestLedger(true, visits);
853 defer ledger.destroy();
854 var cache = try passes.AnalysisCache.initAccounted(
855 testing.allocator,
856 null,
857 ledger,
858 .{},
859 8,
860 );
861 defer cache.deinit();
862 var manager = passes.PassManager.init(testing.allocator);
863 defer manager.deinit();
864 try manager.addPass(backendLegalizationPass());
865 const result = manager.runWithAnalysisCache(module.choir_module, module.context(), &cache, .{});
866 var ctx = passes.PassContext.init(
867 module.choir_module,
868 module.context(),
869 testing.allocator,
870 &cache,
871 );
872 defer ctx.deinit();
873 if (success) {
874 try testing.expectEqual(passes.PassResult.success, result);
875 try ledger.producersComplete();
876 const before = ledger.view().charged;
877 const analysis = try getBackendLegalizationAnalysis(&ctx, module.choir_module);
878 try testing.expectEqual(@as(usize, 2), analysis.kernelCount());
879 try testing.expectEqualDeep(before, ledger.view().charged);
880 } else {
881 try testing.expectEqual(passes.PassResult.failure, result);
882 try testing.expectEqual(.exhausted, ledger.view().outcome);
883 try testing.expectError(
884 error.TerminalWorkOutcome,
885 getBackendLegalizationAnalysis(&ctx, module.choir_module),
886 );
887 }
888 return ledger.view().charged.structural_visits;
889 }