lib/accy/src/preparation/kernelization/lowering/builder.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const common = @import("common.zig");
5 const product_mod = @import("../model/root.zig");
6
7 const kernel_root = common.kernel_root;
8 const mapKernelBuildError = common.mapKernelBuildError;
9 const mapGeneratedKernelError = common.mapGeneratedKernelError;
10
11 const LoweredKernel = product_mod.LoweredKernel;
12
13 /// The kernel stage calls this once per try at emitting a kernel under one candidate schedule to
14 /// charge one kernel builder for its work bound. A pass declares those costs before it runs so
15 /// compilation can refuse a pass that would exceed the caller's limits. The bound adds up every
16 /// allocation one generated-kernel builder makes over its life, together with the working memory
17 /// for hashing its finished kernel, and the alignment slack of each part. The kernel's entry name
18 /// and the scratch each kernel family uses are owned elsewhere and charged separately, along with
19 /// the shared compiler context, the object that owns the operations and values of the generated
20 /// kernels. The call returns `error.WorkOverflow` when a size does not fit.
21 pub fn storageBound(context_limits: common.ir.Context.Limits) !u64 {
22 const work = choir.passes.pass.work;
23 const limits = kernel_root.Builder.Limits.borrowed_standard;
24 const program = kernel_root.Builder.Capacity.deriveBorrowing(limits) catch
25 return error.WorkOverflow;
26 const context = common.ir.Context.Capacity.derive(context_limits) catch
27 return error.WorkOverflow;
28 const hashing = choir.product.hashing;
29 const value_count = std.math.cast(u32, context.storage_bytes / @sizeOf(common.ir.Value)) orelse
30 return error.WorkOverflow;
31 const hash = hashing.StableHasher.Capacity.derive(.{
32 .root = undefined,
33 .facts = .{
34 .value_count = value_count,
35 .operation_depth = hashing.maximum_operation_depth,
36 .attribute_depth = hashing.maximum_attribute_depth,
37 },
38 }) catch return error.WorkOverflow;
39 var bytes = try work.add(program.construction_bytes, hash.working_bytes);
40 bytes = try work.add(bytes, program.body.storage_alignment.toByteUnits());
41 bytes = try work.add(bytes, program.schedule.storage_alignment.toByteUnits());
42 bytes = try work.add(bytes, @alignOf(hashing.capacity.ValueEntry));
43 bytes = try work.add(bytes, @alignOf(hashing.capacity.OperationFrame));
44 bytes = try work.add(bytes, @alignOf(hashing.capacity.AttributeFrame));
45 if (bytes > std.math.maxInt(usize)) return error.WorkOverflow;
46 return bytes;
47 }
48
49 const Launch = enum {
50 none,
51 generated,
52 };
53
54 pub fn withLaunch(
55 allocator: std.mem.Allocator,
56 ir_ctx: *common.ir.Context,
57 work_item_id: usize,
58 entry_name: []u8,
59 params: []const kernel_root.Param,
60 schedule_plan: anytype,
61 context: anytype,
62 comptime body: anytype,
63 ) common.LoweringError!LoweredKernel {
64 return build(allocator, ir_ctx, work_item_id, entry_name, params, schedule_plan, context, body, true);
65 }
66
67 pub fn withoutLaunch(
68 allocator: std.mem.Allocator,
69 ir_ctx: *common.ir.Context,
70 work_item_id: usize,
71 entry_name: []u8,
72 params: []const kernel_root.Param,
73 schedule_plan: anytype,
74 context: anytype,
75 comptime body: anytype,
76 ) common.LoweringError!LoweredKernel {
77 return build(allocator, ir_ctx, work_item_id, entry_name, params, schedule_plan, context, body, false);
78 }
79
80 fn build(
81 allocator: std.mem.Allocator,
82 ir_ctx: *common.ir.Context,
83 work_item_id: usize,
84 entry_name: []u8,
85 params: []const kernel_root.Param,
86 schedule_plan: anytype,
87 context: anytype,
88 comptime body: anytype,
89 comptime capture_launch: bool,
90 ) common.LoweringError!LoweredKernel {
91 var builder = kernel_root.Builder.initBorrowing(
92 allocator,
93 kernel_root.Builder.Limits.borrowed_standard,
94 ir_ctx,
95 entry_name,
96 params,
97 ) catch |err| return mapKernelBuildError(err);
98 errdefer builder.deinit();
99
100 var logical = kernel_root.logical.wrap(&builder, schedule_plan.policy);
101 body(&logical, context) catch |err| return mapGeneratedKernelError(err);
102
103 if (comptime capture_launch) {
104 return finish(allocator, &builder, work_item_id, entry_name, params.len, schedule_plan.generated, .generated);
105 }
106 return finish(allocator, &builder, work_item_id, entry_name, params.len, schedule_plan.generated, .none);
107 }
108
109 fn finish(
110 allocator: std.mem.Allocator,
111 builder: *kernel_root.Builder,
112 work_item_id: usize,
113 entry_name: []u8,
114 argument_count: usize,
115 schedule: product_mod.GeneratedSchedule,
116 launch: Launch,
117 ) common.LoweringError!LoweredKernel {
118 builder.return_() catch |err| return mapKernelBuildError(err);
119
120 var kernel = builder.finish() catch |err| return mapKernelBuildError(err);
121 errdefer kernel.deinit();
122
123 const body_fingerprint = kernel.bodyFingerprint(allocator) catch |err|
124 return mapKernelBuildError(err);
125 const dynamic_shared_memory_bytes = product_mod.dynamicSharedMemoryBytes(kernel.kernelModule()) catch |err| return mapGeneratedKernelError(err);
126 const lowered_argument_count = std.math.cast(u32, argument_count) orelse return error.InvalidArtifact;
127
128 return switch (launch) {
129 .none => .{
130 .work_item_id = work_item_id,
131 .entry_name = entry_name,
132 .program = kernel,
133 .argument_count = lowered_argument_count,
134 .body_fingerprint = body_fingerprint,
135 .dynamic_shared_memory_bytes = dynamic_shared_memory_bytes,
136 .schedule = schedule,
137 },
138 .generated => .{
139 .work_item_id = work_item_id,
140 .entry_name = entry_name,
141 .program = kernel,
142 .argument_count = lowered_argument_count,
143 .body_fingerprint = body_fingerprint,
144 .dynamic_shared_memory_bytes = dynamic_shared_memory_bytes,
145 .schedule = schedule,
146 .launch = kernel.launch() catch |err| return mapGeneratedKernelError(err),
147 },
148 };
149 }
150
151 const testing = std.testing;
152 const choir = @import("choir");
153 const generated_schedule = @import("schedule.zig");
154
155 fn testContext() !common.ir.Context {
156 var ctx = try common.ir.Context.init(testing.allocator, common.ir.Context.Limits.testing);
157 errdefer ctx.deinit(testing.allocator);
158 try choir.dialects.registerChoirDialect(&ctx);
159 return ctx;
160 }
161
162 fn emitEmpty(_: anytype, _: void) !void {}
163
164 fn emitStorageBody(logical: anytype, count: usize) !void {
165 const builder: *kernel_root.Builder = @ptrCast(@alignCast(logical.token));
166 for (0..count) |index| _ = try builder.constantInt(.i32, @intCast(index));
167 for (0..kernel_root.Builder.Limits.borrowed_standard.schedule.axes) |_| {
168 _ = try builder.axis("x", 1);
169 }
170 }
171
172 test "kernelization generated builder storage covers construction and fingerprint traffic" {
173 for ([_]usize{ 0, 1, 512 }) |count| try checkBuilderStorage(count);
174 var overflow = common.ir.Context.Limits.testing;
175 overflow.operations.storage_bytes = std.math.maxInt(usize);
176 try testing.expectError(error.WorkOverflow, storageBound(overflow));
177 }
178
179 fn checkBuilderStorage(count: usize) !void {
180 const fixed = @import("alloc_fixed");
181 var ctx = try testContext();
182 defer ctx.deinit(testing.allocator);
183 const bound = try storageBound(ctx.capacity.asLimits());
184 const bytes = try testing.allocator.alignedAlloc(u8, .@"64", @intCast(bound));
185 defer testing.allocator.free(bytes);
186 var backing = fixed.Tracked.init(bytes);
187 var retained = fixed.Monotonic.init(backing.allocator(), bytes.len);
188 const allocator = retained.allocator();
189 const name = try testing.allocator.dupe(u8, "storage");
190 defer testing.allocator.free(name);
191 const parameter_count = kernel_root.Builder.Limits.borrowed_standard.parameters;
192 const params: [parameter_count]kernel_root.Param = @splat(kernel_root.scalar(.i32));
193 var result = try withoutLaunch(
194 allocator,
195 &ctx,
196 0,
197 name,
198 ¶ms,
199 generated_schedule.flat(),
200 count,
201 emitStorageBody,
202 );
203 defer result.program.deinit();
204 try result.program.verify();
205 try testing.expectEqual(params.len, result.argument_count);
206 try testing.expectEqual(
207 try result.program.bodyFingerprint(testing.allocator),
208 result.body_fingerprint,
209 );
210 try testing.expect(!backing.exhausted);
211 try testing.expect(backing.status().used_bytes > 0);
212 try testing.expect(backing.status().used_bytes <= bound);
213 }
214
215 test "kernelization generated builder distinguishes name and parameter exhaustion" {
216 const limits = kernel_root.Builder.Limits.borrowed_standard;
217 for ([_]usize{ limits.parameters - 1, limits.parameters, limits.parameters + 1 }) |count| {
218 var ctx = try testContext();
219 defer ctx.deinit(testing.allocator);
220 const params = try testing.allocator.alloc(kernel_root.Param, count);
221 defer testing.allocator.free(params);
222 @memset(params, kernel_root.scalar(.i32));
223 const name = try testing.allocator.dupe(u8, "parameters");
224 const before = ctx.capacityUsage();
225 const result = withoutLaunch(
226 testing.allocator,
227 &ctx,
228 0,
229 name,
230 params,
231 generated_schedule.flat(),
232 {},
233 emitEmpty,
234 );
235 if (count > limits.parameters) {
236 defer testing.allocator.free(name);
237 try testing.expectError(error.WorkExhausted, result);
238 try testing.expectEqualDeep(before, ctx.capacityUsage());
239 } else {
240 var kernel = try result;
241 defer kernel.deinit(testing.allocator);
242 try testing.expectEqual(count, kernel.program.params().len);
243 try testing.expectEqual(count, kernel.argument_count);
244 try testing.expectEqual(@as(?kernel_root.Launch, null), kernel.launch);
245 }
246 }
247 for ([_]usize{
248 limits.kernel_name_bytes - 1,
249 limits.kernel_name_bytes,
250 limits.kernel_name_bytes + 1,
251 }) |count| {
252 var ctx = try testContext();
253 defer ctx.deinit(testing.allocator);
254 const name = try testing.allocator.alloc(u8, count);
255 @memset(name, 'k');
256 const result = withoutLaunch(
257 testing.allocator,
258 &ctx,
259 0,
260 name,
261 &.{},
262 generated_schedule.flat(),
263 {},
264 emitEmpty,
265 );
266 if (count > limits.kernel_name_bytes) {
267 defer testing.allocator.free(name);
268 try testing.expectError(error.WorkExhausted, result);
269 } else {
270 var kernel = try result;
271 defer kernel.deinit(testing.allocator);
272 try testing.expectEqualStrings(name, kernel.entry_name);
273 try testing.expectEqualStrings(name, kernel.program.storage.kernel.func().getName().?);
274 }
275 }
276 }
277
278 fn emitAxes(logical: anytype, count: usize) !void {
279 const builder: *kernel_root.Builder = @ptrCast(@alignCast(logical.token));
280 for (0..count) |_| _ = try builder.axis("x", 1);
281 }
282
283 test "kernelization generated builder preserves schedule capacity failure" {
284 const limit = kernel_root.Builder.Limits.borrowed_standard.schedule.axes;
285 for ([_]usize{ limit - 1, limit, limit + 1 }) |count| {
286 var ctx = try testContext();
287 defer ctx.deinit(testing.allocator);
288 const name = try testing.allocator.dupe(u8, "schedule");
289 const result = withoutLaunch(
290 testing.allocator,
291 &ctx,
292 0,
293 name,
294 &.{},
295 generated_schedule.flat(),
296 count,
297 emitAxes,
298 );
299 if (count > limit) {
300 defer testing.allocator.free(name);
301 try testing.expectError(error.WorkExhausted, result);
302 } else {
303 var kernel = try result;
304 defer kernel.deinit(testing.allocator);
305 var snapshot = try kernel.program.scheduleSnapshot(testing.allocator);
306 defer snapshot.deinit(testing.allocator);
307 try testing.expectEqual(count, snapshot.allAxes().len);
308 }
309 }
310 }
311
312 test "kernelization generated builder distinguishes borrowed budget from host OOM" {
313 var limits = common.ir.Context.Limits.testing;
314 limits.operations.storage_bytes = 1;
315 var ctx = try common.ir.Context.init(testing.allocator, limits);
316 defer ctx.deinit(testing.allocator);
317 try choir.dialects.registerChoirDialect(&ctx);
318 const name = try testing.allocator.dupe(u8, "budget");
319 defer testing.allocator.free(name);
320 var failing = testing.FailingAllocator.init(testing.allocator, .{ .fail_index = 0 });
321 try testing.expectError(
322 error.WorkExhausted,
323 withoutLaunch(
324 failing.allocator(),
325 &ctx,
326 0,
327 name,
328 &.{},
329 generated_schedule.flat(),
330 {},
331 emitEmpty,
332 ),
333 );
334 try testing.expectEqual(@as(usize, 0), failing.alloc_index);
335 try testing.expectEqual(@as(?common.ir.Context.Segment, null), ctx.exhaustedSegment());
336 try testing.checkAllAllocationFailures(testing.allocator, testHostAllocationFailure, .{});
337 }
338
339 fn emitIndex(logical: anytype, _: void) !void {
340 _ = try logical.index1D("x", 5);
341 }
342
343 fn testHostAllocationFailure(allocator: std.mem.Allocator) !void {
344 var ctx = try testContext();
345 defer ctx.deinit(testing.allocator);
346 const name = try allocator.dupe(u8, "host");
347 var kernel = withLaunch(
348 allocator,
349 &ctx,
350 7,
351 name,
352 &.{},
353 generated_schedule.flat(),
354 {},
355 emitIndex,
356 ) catch |err| {
357 allocator.free(name);
358 return err;
359 };
360 defer kernel.deinit(allocator);
361 try testing.expectEqual(@as(usize, 7), kernel.work_item_id);
362 try testing.expectEqual(@as(u32, 1), kernel.launch.?.grid[0]);
363 try testing.expectEqual(@as(u32, 5), kernel.launch.?.block[0]);
364 try testing.expectEqual(try kernel.program.bodyFingerprint(allocator), kernel.body_fingerprint);
365 }
366
367 fn emitTemporaryFailure(logical: anytype, types: bool) !void {
368 const value = try logical.constantInt(.i32, 0);
369 var values: [65]kernel_root.Value = undefined;
370 @memset(&values, value);
371 var result_types: [65]kernel_root.Type = undefined;
372 @memset(&result_types, value.valueType());
373 _ = try logical.forScope(
374 value,
375 value,
376 value,
377 if (types) &.{} else &values,
378 if (types) &result_types else &.{},
379 );
380 }
381
382 test "kernelization generated builder preserves temporary capacity failure" {
383 for ([_]bool{ false, true }) |types| {
384 var ctx = try testContext();
385 defer ctx.deinit(testing.allocator);
386 const name = try testing.allocator.dupe(u8, "temporary");
387 defer testing.allocator.free(name);
388 try testing.expectError(
389 error.WorkExhausted,
390 withoutLaunch(
391 testing.allocator,
392 &ctx,
393 0,
394 name,
395 &.{},
396 generated_schedule.flat(),
397 types,
398 emitTemporaryFailure,
399 ),
400 );
401 }
402 }
403
404 fn testReductionModule(
405 source: anytype,
406 count: usize,
407 ) !*@import("../../../choir/root.zig").semantic.SemanticModule {
408 const typ = try source.tensor(.f32, &.{4});
409 const types = try testing.allocator.alloc(common.ir.Type, count);
410 defer testing.allocator.free(types);
411 @memset(types, typ);
412 var function = try source.beginFunction("kernel_capacity", types, &.{typ});
413 const values = try testing.allocator.alloc(*common.ir.Value, count);
414 defer testing.allocator.free(values);
415 for (values, 0..) |*value, index| value.* = function.parameter(index);
416 var remaining = count;
417 while (remaining > 1) {
418 var next: usize = 0;
419 var index: usize = 0;
420 while (index < remaining) : (index += 2) {
421 values[next] = if (index + 1 < remaining)
422 try function.add(values[index], values[index + 1])
423 else
424 values[index];
425 next += 1;
426 }
427 remaining = next;
428 }
429 try function.return_(values[0..1]);
430 try function.finish();
431 return source.finish();
432 }
433
434 test "kernelization generated builder exhaustion reaches the analysis cache" {
435 const semantic = @import("../../../choir/root.zig").semantic;
436 const lowering = @import("root.zig");
437 const outline_owner = @import("../../outlining/root.zig");
438 for ([_]usize{ 61, 62, 63 }) |count| {
439 var source = try semantic.Builder.init(
440 testing.allocator,
441 semantic.Builder.ContextLimits.standard,
442 );
443 defer source.deinit();
444 const module = try testReductionModule(&source, count);
445 defer module.deinit();
446 const revision = choir.product.revision;
447 const ledger = try revision.AccountingV1.create(testing.allocator, .{
448 .allowance = revision.WorkVector.uniform(std.math.maxInt(u64)),
449 .workspace = std.math.maxInt(u64),
450 .events = 32,
451 }, &.{});
452 defer ledger.destroy();
453 var cache = try choir.passes.AnalysisCache.initAccounted(
454 testing.allocator,
455 null,
456 ledger,
457 .{},
458 8,
459 );
460 defer cache.deinit();
461 var ctx = choir.passes.PassContext.init(
462 module.choir_module,
463 module.context(),
464 testing.allocator,
465 &cache,
466 );
467 defer ctx.deinit();
468 const outline = try outline_owner.getKernelOutlinePlanAnalysis(&ctx, module.choir_module);
469 try testing.expectEqual(@as(usize, 1), outline.kernels.items.len);
470 try testing.expectEqual(count, outline.kernels.items[0].inputCount());
471 const before = cache.entries.count();
472 const result = lowering.getKernelizationAnalysis(&ctx, module.choir_module);
473 if (count == 63) {
474 try testing.expectError(error.WorkExhausted, result);
475 try testing.expectEqual(before, cache.entries.count());
476 try testing.expectEqual(.exhausted, ledger.view().outcome);
477 try testing.expectError(
478 error.TerminalWorkOutcome,
479 lowering.getKernelizationAnalysis(&ctx, module.choir_module),
480 );
481 } else {
482 const analysis = try result;
483 try testing.expectEqual(@as(usize, 1), analysis.kernelCount());
484 try testing.expectEqual(count + 2, analysis.kernels.items[0].argument_count);
485 }
486 }
487 }
488
489 test "kernelization generated builder ABI distinguishes size overflow from host OOM" {
490 const abi = @import("abi.zig");
491 const byte_limit = std.math.maxInt(usize) / @sizeOf(kernel_root.Param);
492 for ([_]usize{
493 byte_limit - 2,
494 byte_limit - 1,
495 std.math.maxInt(usize) - 1,
496 std.math.maxInt(usize),
497 }) |count| {
498 var failing = testing.FailingAllocator.init(testing.allocator, .{ .fail_index = 0 });
499 const expected = if (count == byte_limit - 2) error.OutOfMemory else error.WorkOverflow;
500 try testing.expectError(expected, abi.flat(failing.allocator(), .f32, count));
501 try testing.expectEqual(count == byte_limit - 2, failing.has_induced_failure);
502 }
503 const inputs = [_]choir_abi.DType{ .i32, .f32, .f16 };
504 var mixed = try abi.flatTyped(testing.allocator, .f32, &inputs);
505 defer mixed.deinit(testing.allocator);
506 var uniform = try abi.flat(testing.allocator, .f32, inputs.len);
507 defer uniform.deinit(testing.allocator);
508 try testing.expectEqual(inputs.len + 2, mixed.params().len);
509 try testing.expectEqual(mixed.params().len, uniform.params().len);
510 try testing.expectEqualDeep(uniform.params()[0], mixed.params()[0]);
511 try testing.expectEqualDeep(uniform.params()[4], mixed.params()[4]);
512 for (inputs, 1..) |dtype, index| {
513 try testing.expectEqual(dtype, mixed.params()[index].buffer.dtype);
514 try testing.expectEqual(.f32, uniform.params()[index].buffer.dtype);
515 }
516 }
517
518 fn occupySourceAttributes(ctx: *common.ir.Context) !void {
519 const remaining = ctx.capacityUsage().attribute_payloads.remainingBytes();
520 const reserved = kernel_root.Builder.Limits.borrowed_standard.context.attributes.payload_bytes;
521 const bytes = try testing.allocator.alloc(u8, remaining - reserved / 2);
522 defer testing.allocator.free(bytes);
523 @memset(bytes, 'q');
524 _ = try ctx.getStringAttr(bytes);
525 try testing.expect(ctx.capacityUsage().attribute_payloads.remainingBytes() < reserved);
526 }
527
528 fn testDetachedProgram(
529 cache: *choir.passes.AnalysisCache,
530 occupied: bool,
531 vector: bool,
532 ) !*product_mod.KernelizationAnalysis {
533 const semantic = @import("../../../choir/root.zig").semantic;
534 const lowering = @import("root.zig");
535 var source = try semantic.Builder.init(
536 testing.allocator,
537 semantic.Builder.ContextLimits.standard,
538 );
539 defer source.deinit();
540 const module = try testReductionModule(&source, 2);
541 defer module.deinit();
542 if (vector) {
543 const target = @import("../../root.zig").target;
544 var dtypes = gpu.DTypeSet{};
545 dtypes.insert(.f32);
546 try target.setBackendTargetProfile(module.context(), module.choir_module, .{
547 .backend_kind = .cuda,
548 .artifact_format = .cuda_ptx,
549 .dtype_bits = dtypes.bits,
550 });
551 }
552 var ctx = choir.passes.PassContext.init(
553 module.choir_module,
554 module.context(),
555 testing.allocator,
556 cache,
557 );
558 defer ctx.deinit();
559 errdefer {
560 var preserved = choir.passes.PreservedAnalyses.init(testing.allocator);
561 defer preserved.deinit();
562 cache.invalidate(&preserved);
563 }
564 const outline = @import("../../outlining/root.zig");
565 _ = try outline.getKernelOutlinePlanAnalysis(&ctx, module.choir_module);
566 const before = try choir.bytecode.encodeModule(testing.allocator, module.choir_module);
567 defer testing.allocator.free(before);
568 if (occupied) try occupySourceAttributes(module.context());
569 const generated = try lowering.getKernelizationAnalysis(&ctx, module.choir_module);
570 try testing.expect(generated.context != module.context());
571 try testing.expectEqual(@as(usize, 1), generated.kernelCount());
572 try testing.expect(
573 generated.kernels.items[0].program.kernelModule().context == generated.context,
574 );
575 const expected_body: std.meta.Tag(product_mod.LoweredKernelBody) =
576 if (vector) .elementwise_vector else .generic;
577 try testing.expectEqual(expected_body, std.meta.activeTag(generated.kernels.items[0].body));
578 try testing.expectEqualDeep(
579 module.context().capacity.asLimits(),
580 generated.context.capacity.asLimits(),
581 );
582 const after = try choir.bytecode.encodeModule(testing.allocator, module.choir_module);
583 defer testing.allocator.free(after);
584 try testing.expectEqualSlices(u8, before, after);
585 return generated;
586 }
587
588 test "kernelization generated builder owns programs after occupied origin destruction" {
589 for ([_]bool{ false, true }) |vector| try checkDetachedPrograms(vector);
590 }
591
592 fn checkDetachedPrograms(vector: bool) !void {
593 var cold_cache = choir.passes.AnalysisCache.init(testing.allocator, null);
594 defer cold_cache.deinit();
595 var occupied_cache = choir.passes.AnalysisCache.init(testing.allocator, null);
596 defer occupied_cache.deinit();
597 const cold = try testDetachedProgram(&cold_cache, false, vector);
598 const occupied = try testDetachedProgram(&occupied_cache, true, vector);
599 try testing.expectEqual(@as(usize, 1), cold.kernelCount());
600 try testing.expectEqual(cold.kernelCount(), occupied.kernelCount());
601 for (cold.kernels.items, occupied.kernels.items) |*left, *right| {
602 try left.program.verify();
603 try right.program.verify();
604 const left_bytes = try choir.bytecode.encodeModule(
605 testing.allocator,
606 left.program.kernelModule(),
607 );
608 defer testing.allocator.free(left_bytes);
609 const right_bytes = try choir.bytecode.encodeModule(
610 testing.allocator,
611 right.program.kernelModule(),
612 );
613 defer testing.allocator.free(right_bytes);
614 try testing.expectEqualSlices(u8, left_bytes, right_bytes);
615 try testing.expectEqualDeep(left.program.params(), right.program.params());
616 try testing.expectEqualDeep(left.summary(), right.summary());
617 try testing.expectEqualDeep(left.body, right.body);
618 try testing.expectEqual(left.output_fill_pattern, right.output_fill_pattern);
619 try testing.expectEqual(left.scratch_fill_pattern, right.scratch_fill_pattern);
620 var left_schedule = try left.program.scheduleSnapshot(testing.allocator);
621 defer left_schedule.deinit(testing.allocator);
622 var right_schedule = try right.program.scheduleSnapshot(testing.allocator);
623 defer right_schedule.deinit(testing.allocator);
624 try testing.expectEqualDeep(left_schedule.allAxes(), right_schedule.allAxes());
625 try testing.expectEqualDeep(left_schedule.allSteps(), right_schedule.allSteps());
626 }
627 }
628
629 fn testSplatSource(source: anytype, broadcast: u2) !*common.ir.Value {
630 const typ = try source.tensor(.f32, &.{2});
631 const result_type = if (broadcast == 2) try source.tensor(.f32, &.{ 2, 2 }) else typ;
632 var function = try source.beginFunction("splat_failure", &.{}, &.{result_type});
633 const data = [_]f32{ 7, 7 };
634 var value = try function.constant(typ, std.mem.sliceAsBytes(&data));
635 if (broadcast == 1) {
636 value = try function.broadcastInDim(value, typ, &.{2}, &.{0});
637 } else if (broadcast == 2) {
638 value = try function.broadcast(value, result_type, &.{2});
639 }
640 try function.return_(&.{value});
641 try function.finish();
642 return value;
643 }
644
645 fn testSplatRefusal(broadcast: u2) !void {
646 const semantic = @import("../../../choir/root.zig").semantic;
647 const elementwise = @import("elementwise.zig");
648 const abi_owner = @import("abi.zig");
649 var source = try semantic.Builder.init(
650 testing.allocator,
651 semantic.Builder.ContextLimits.standard,
652 );
653 defer source.deinit();
654 const value = try testSplatSource(&source, broadcast);
655 const root: *common.ir.Operation = @ptrCast(@alignCast(value.getDefiningOp().?));
656 var ctx = try testContext();
657 defer ctx.deinit(testing.allocator);
658 var abi = try abi_owner.flat(testing.allocator, .f32, 0);
659 defer abi.deinit(testing.allocator);
660 var builder = try kernel_root.Builder.initBorrowing(
661 testing.allocator,
662 kernel_root.Builder.Limits.borrowed_standard,
663 &ctx,
664 "splat",
665 abi.params(),
666 );
667 defer builder.deinit();
668 const index = try builder.constantIndex(0);
669 const payload_allocator = common.ir.context.attributePayloadAllocator(&ctx);
670 const unused = try payload_allocator.alloc(
671 u8,
672 ctx.capacityUsage().attribute_payloads.remainingBytes(),
673 );
674 defer payload_allocator.free(unused);
675 var buffers = common.bufferization.BufferPlanAnalysis.init(testing.allocator);
676 defer buffers.deinit();
677 const outline = product_mod.KernelOutline{
678 .id = 0,
679 .name = &.{},
680 .kind = .elementwise,
681 .work_item_id = 0,
682 .root = root,
683 .input_slot_ids = &.{},
684 .output_slot_id = 0,
685 .element_count = 2,
686 .op_count = 1,
687 };
688 try testing.expectError(error.OutOfMemory, elementwise.shapeReadValue(
689 &builder,
690 index,
691 value,
692 outline,
693 &buffers,
694 abi,
695 elementwise.max_recompute_depth,
696 ));
697 try testing.expectEqual(common.ir.Context.Segment.attribute_payloads, ctx.exhaustedSegment().?);
698 }
699
700 test "kernelization generated builder preserves splat emission exhaustion" {
701 for ([_]u2{ 0, 1, 2 }) |broadcast| try testSplatRefusal(broadcast);
702 }
703
704 fn testKernelContextHostFailure(allocator: std.mem.Allocator) !void {
705 const semantic = @import("../../../choir/root.zig").semantic;
706 const lowering = @import("root.zig");
707 var source = try semantic.Builder.init(
708 testing.allocator,
709 semantic.Builder.ContextLimits.standard,
710 );
711 defer source.deinit();
712 const module = try testReductionModule(&source, 2);
713 defer module.deinit();
714 var cache = choir.passes.AnalysisCache.init(allocator, null);
715 defer cache.deinit();
716 var ctx = choir.passes.PassContext.init(
717 module.choir_module,
718 module.context(),
719 allocator,
720 &cache,
721 );
722 defer ctx.deinit();
723 const analysis = try lowering.getKernelizationAnalysis(&ctx, module.choir_module);
724 try testing.expectEqual(@as(usize, 1), analysis.kernelCount());
725 try testing.expectEqual(@as(u32, 4), analysis.kernels.items[0].argument_count);
726 try testing.expect(analysis.context != module.context());
727 }
728
729 test "kernelization generated builder owns cleanup across native allocation failures" {
730 try testing.checkAllAllocationFailures(testing.allocator, testKernelContextHostFailure, .{});
731 }
732
733 test "kernelization generated builder declines nonsplats without losing valid literals" {
734 const semantic = @import("../../../choir/root.zig").semantic;
735 var source = try semantic.Builder.init(
736 testing.allocator,
737 semantic.Builder.ContextLimits.standard,
738 );
739 defer source.deinit();
740 const typ = try source.tensor(.f32, &.{2});
741 var function = try source.beginFunction("constants", &.{}, &.{typ});
742 const uniform_data = [_]f32{ 7, 7 };
743 const varied_data = [_]f32{ 7, 8 };
744 const uniform = try function.constant(typ, std.mem.sliceAsBytes(&uniform_data));
745 const varied = try function.constant(typ, std.mem.sliceAsBytes(&varied_data));
746 try function.return_(&.{uniform});
747 try function.finish();
748 var ctx = try testContext();
749 defer ctx.deinit(testing.allocator);
750 var builder = try kernel_root.Builder.initBorrowing(
751 testing.allocator,
752 kernel_root.Builder.Limits.borrowed_standard,
753 &ctx,
754 "literal",
755 &.{},
756 );
757 defer builder.deinit();
758 const before = ctx.capacityUsage();
759 const varied_op: *common.ir.Operation = @ptrCast(@alignCast(varied.getDefiningOp().?));
760 try testing.expect(try common.optionalSplatConstantValue(&builder, varied_op) == null);
761 try testing.expectEqualDeep(before, ctx.capacityUsage());
762 const uniform_op: *common.ir.Operation = @ptrCast(@alignCast(uniform.getDefiningOp().?));
763 try testing.expect(try common.optionalSplatConstantValue(&builder, uniform_op) != null);
764 try builder.return_();
765 var program = try builder.finish();
766 defer program.deinit();
767 try program.verify();
768 }