lib/accy/src/preparation/outlining/pass.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const alloc_fixed = @import("alloc_fixed");
3 const choir = @import("choir");
4 const bufferization = @import("../bufferization/root.zig");
5 const accy_choir = @import("../../choir/root.zig");
6 const dialect_mod = accy_choir.dialect;
7 const kernelization_model = @import("../kernelization/model/root.zig");
8 const schedule_planning = @import("../schedule/root.zig");
9
10 const ir = choir.ir;
11 const passes = choir.passes;
12 const accounting = passes.pass.work;
13
14 pub const kernel_outlining_planning_pass_name = "accy-choir-plan-kernel-outlines";
15 pub const kernel_outlining_planning_pass_description =
16 "Plan Accy Choir kernel outline candidates";
17
18 const KernelOutlineKind = kernelization_model.KernelOutlineKind;
19 const KernelOutline = kernelization_model.KernelOutline;
20 const KernelOutlinePlanAnalysis = kernelization_model.KernelOutlinePlanAnalysis;
21
22 const OutlineWork = struct {
23 input: accounting.Census,
24
25 fn inspect(input: accounting.Input) !OutlineWork {
26 if (input.options.max_threads != 1 or input.options.worker_allocator != null) {
27 return error.MissingWorkContract;
28 }
29 var scope = input.operation;
30 _ = try scope.walk(.{ .order = .pre_order }, &scope, checkScope);
31 return .{ .input = try accounting.Census.inspect(scope) };
32 }
33
34 fn checkScope(root: **ir.Operation, op: *ir.Operation) !ir.Operation.WalkResult {
35 ir.traits.verifyOperandsWithin(root.*, op) catch return error.UnboundProductInput;
36 return .advance;
37 }
38
39 fn storage(self: OutlineWork) !u64 {
40 const count = self.input.operations;
41 const values = self.input.values;
42 var per_work = try accounting.hashMapGrowth(*ir.Operation, void, count);
43 per_work = try accounting.add(per_work, try accounting.hashMapGrowth(usize, void, values));
44 const memo = try accounting.hashMapGrowth(*ir.Value, u64, values);
45 per_work = try accounting.add(per_work, memo);
46 per_work = try accounting.add(per_work, try accounting.arrayListGrowth(usize, values));
47 const owned_inputs = try accounting.add(
48 try accounting.multiply(values, @sizeOf(usize)),
49 @alignOf(usize),
50 );
51 per_work = try accounting.add(per_work, owned_inputs);
52 const name_bytes = "accy_choir_kernel_".len + 3 * @sizeOf(usize);
53 per_work = try accounting.add(per_work, name_bytes);
54 var bytes: u64 = @sizeOf(KernelOutlinePlanAnalysis) + @alignOf(KernelOutlinePlanAnalysis);
55 bytes = try accounting.add(bytes, try accounting.multiply(count, per_work));
56 bytes = try accounting.add(bytes, try accounting.arrayListGrowth(KernelOutline, count));
57 bytes = try accounting.add(bytes, try accounting.hashMapGrowth(usize, usize, count));
58 bytes = try accounting.add(bytes, try accounting.hashMapGrowth(usize, void, values));
59 if (bytes > std.math.maxInt(usize)) return error.WorkOverflow;
60 return bytes;
61 }
62
63 fn bounds(self: OutlineWork) !accounting.Bounds {
64 const bytes = try self.storage();
65 const input_units = try accounting.add(self.input.atoms, self.input.input_bytes);
66 const units = try accounting.add(input_units, 1);
67 const entries = try accounting.add(self.input.operations, self.input.values);
68 const probes = try accounting.hashMapCapacity(entries);
69 const leaf_states = try accounting.multiply(
70 elided_leaf_walk_depth,
71 try accounting.add(self.input.values, 1),
72 );
73 const edges = try accounting.multiply(
74 leaf_states,
75 try accounting.add(self.input.operands, 1),
76 );
77 const work_visits = try accounting.multiply(
78 try accounting.add(edges, units),
79 try accounting.add(try accounting.add(units, probes), 1),
80 );
81 const visits = try accounting.multiply(
82 try accounting.add(self.input.operations, 1),
83 work_visits,
84 );
85 return .{
86 .work = .{
87 .input_bytes = self.input.input_bytes,
88 .structural_visits = try accounting.multiply(64, visits),
89 .analysis_computations = 1,
90 .allocation_capacity = bytes,
91 },
92 .workspace = bytes,
93 .retained_storage = bytes,
94 };
95 }
96 };
97
98 fn outlineAnalysisWork(input: accounting.Input) !accounting.Bounds {
99 return (try OutlineWork.inspect(input)).bounds();
100 }
101
102 fn outlinePassWork(_: accounting.Input) !accounting.Bounds {
103 return .{ .work = .{ .structural_visits = 1 } };
104 }
105
106 pub const kernel_outline_plan_analysis_descriptor = passes.AnalysisDescriptor{
107 .id = passes.analysisId(kernelization_model.kernel_outline_plan_analysis_name),
108 .name = kernelization_model.kernel_outline_plan_analysis_name,
109 .work_contract = .{
110 .identity = .{
111 .name = kernelization_model.kernel_outline_plan_analysis_name,
112 .version = 1,
113 },
114 .estimate = outlineAnalysisWork,
115 },
116 };
117
118 pub fn getKernelOutlinePlanAnalysis(
119 pass_ctx: *passes.PassContext,
120 op: *ir.Operation,
121 ) !*KernelOutlinePlanAnalysis {
122 const ptr = try pass_ctx.getAnalysis(
123 op,
124 &kernel_outline_plan_analysis_descriptor,
125 computeKernelOutlinePlanAnalysis,
126 cleanupKernelOutlinePlanAnalysis,
127 );
128 return @ptrCast(@alignCast(ptr));
129 }
130
131 pub fn kernelOutliningPlanningPass() passes.Pass {
132 return .{
133 .name = kernel_outlining_planning_pass_name,
134 .description = kernel_outlining_planning_pass_description,
135 .run_fn = runKernelOutliningPlanningPass,
136 .work_contract = .{
137 .identity = .{ .name = kernel_outlining_planning_pass_name, .version = 1 },
138 .estimate = outlinePassWork,
139 },
140 };
141 }
142
143 fn runKernelOutliningPlanningPass(pass_ctx: *passes.PassContext) passes.PassResult {
144 _ = getKernelOutlinePlanAnalysis(pass_ctx, pass_ctx.op) catch return .failure;
145 pass_ctx.preserveAllAnalyses();
146 return .success;
147 }
148
149 fn computeKernelOutlinePlanAnalysis(
150 pass_ctx: *passes.PassContext,
151 op: *ir.Operation,
152 ) anyerror!*anyopaque {
153 const schedule_plan = try schedule_planning.getSchedulePlanAnalysis(pass_ctx, op);
154 const buffer_plan = try bufferization.getBufferPlanAnalysis(pass_ctx, op);
155
156 const analysis = try pass_ctx.allocator.create(KernelOutlinePlanAnalysis);
157 analysis.* = KernelOutlinePlanAnalysis.init(pass_ctx.allocator);
158 errdefer {
159 analysis.deinit();
160 pass_ctx.allocator.destroy(analysis);
161 }
162
163 try addKernelOutlines(pass_ctx, analysis, schedule_plan, buffer_plan);
164 try validateOutlineSlotDataflow(pass_ctx.allocator, analysis, buffer_plan);
165
166 return @ptrCast(analysis);
167 }
168
169 fn validateOutlineSlotDataflow(
170 allocator: std.mem.Allocator,
171 analysis: *const KernelOutlinePlanAnalysis,
172 buffer_plan: *const bufferization.BufferPlanAnalysis,
173 ) !void {
174 var written = std.AutoHashMap(usize, void).init(allocator);
175 defer written.deinit();
176 for (analysis.kernels.items) |outline| {
177 try written.put(outline.output_slot_id, {});
178 if (outline.kind == .kernel_call) {
179 for (outline.input_slot_ids) |slot_id| {
180 try written.put(slot_id, {});
181 }
182 }
183 if (outline.kind == .iterate) {
184 var result_index: usize = 0;
185 while (outline.root.getResult(result_index)) |result| : (result_index += 1) {
186 const slot = buffer_plan.getSlot(result) orelse continue;
187 try written.put(slot.id, {});
188 }
189 }
190 if (outline.kind == .scan) {
191 if (outline.root.getOperand(1)) |scratch| {
192 if (buffer_plan.getSlot(scratch)) |slot| {
193 try written.put(slot.id, {});
194 }
195 }
196 }
197 }
198
199 for (analysis.kernels.items) |outline| {
200 if (outline.kind == .kernel_call) continue;
201 for (outline.input_slot_ids) |slot_id| {
202 if (written.contains(slot_id)) continue;
203 const slot = bufferSlotById(buffer_plan, slot_id) orelse return error.UnwrittenKernelInput;
204 if (slot.role.input or slot.role.constant) continue;
205 return error.UnwrittenKernelInput;
206 }
207 }
208 }
209
210 fn bufferSlotById(
211 buffer_plan: *const bufferization.BufferPlanAnalysis,
212 slot_id: usize,
213 ) ?*const bufferization.BufferSlot {
214 for (buffer_plan.slots.items) |*slot| {
215 if (slot.id == slot_id) return slot;
216 }
217 return null;
218 }
219
220 fn cleanupKernelOutlinePlanAnalysis(ptr: *anyopaque, allocator: std.mem.Allocator) void {
221 const analysis: *KernelOutlinePlanAnalysis = @ptrCast(@alignCast(ptr));
222 analysis.deinit();
223 allocator.destroy(analysis);
224 }
225
226 fn collectInputSlots(
227 allocator: std.mem.Allocator,
228 work: schedule_planning.ScheduleWorkItem,
229 buffer_plan: *const bufferization.BufferPlanAnalysis,
230 ) ![]usize {
231 var inputs: std.ArrayListUnmanaged(usize) = .empty;
232 errdefer inputs.deinit(allocator);
233 var seen = std.AutoHashMap(usize, void).init(allocator);
234 defer seen.deinit();
235 var leaves = LeafWalk.init(allocator, buffer_plan, &seen, &inputs);
236 defer leaves.deinit();
237 var op_set = std.AutoHashMap(*ir.Operation, void).init(allocator);
238 defer op_set.deinit();
239 for (work.ops) |op| try op_set.put(op, {});
240
241 for (work.ops) |op| {
242 for (op.getOperandValues(), 0..) |operand, operand_index| {
243 if ((work.kind == .reduction or work.kind == .row_pipeline) and isConstantReduceInitOperand(buffer_plan, op, operand, operand_index)) continue;
244 if ((work.kind == .row_pipeline or work.kind == .flash_attention) and operandIsConstant(operand)) continue;
245 if (work.kind == .flash_attention and operandIsBroadcastOfConstant(operand)) continue;
246 if (operand.getDefiningOp()) |def_any| {
247 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
248 if (op_set.contains(def_op)) continue;
249 }
250 const slot_value = schedule_planning.kernelInputValueForOperand(work, operand);
251 if (buffer_plan.getSlot(slot_value)) |slot| {
252 if (seen.contains(slot.id)) continue;
253 try seen.put(slot.id, {});
254 try inputs.append(allocator, slot.id);
255 continue;
256 }
257 switch (work.kind) {
258 .elementwise_single, .elementwise_fusion, .reduction, .row_pipeline, .iterate, .flash_attention => {
259 try leaves.collect(operand, elided_leaf_walk_depth);
260 },
261 else => {},
262 }
263 }
264 }
265
266 if (work.kind == .iterate and work.ops.len == 1) {
267 try collectIterateCaptureSlots(&leaves, work.ops[0]);
268 var result_index: usize = 0;
269 while (work.ops[0].getResult(result_index)) |result| : (result_index += 1) {
270 if (result.first_use == null) continue;
271 const slot = buffer_plan.getSlot(result) orelse return error.MissingOutputSlot;
272 if (slot.id == buffer_plan.getSlot(work.output_value).?.id) continue;
273 if (seen.contains(slot.id)) continue;
274 try seen.put(slot.id, {});
275 try inputs.append(allocator, slot.id);
276 }
277 }
278
279 return try inputs.toOwnedSlice(allocator);
280 }
281
282 fn collectIterateCaptureSlots(leaves: *LeafWalk, iterate: *ir.Operation) anyerror!void {
283 if (iterate.regions.items.len != 1) return;
284 const region = &iterate.regions.items[0];
285 var block_iter = region.getBlocks();
286 while (block_iter.next()) |block| {
287 var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
288 while (current) |op| : (current = op.next_op) {
289 for (op.getOperandValues()) |operand| {
290 if (operandIsInsideRegion(operand, region)) continue;
291 if (operandIsSplatMaterial(operand)) continue;
292 const slot = leaves.buffer_plan.getSlot(operand) orelse {
293 try leaves.collect(operand, elided_leaf_walk_depth);
294 continue;
295 };
296 if (leaves.seen.contains(slot.id)) continue;
297 try leaves.seen.put(slot.id, {});
298 try leaves.inputs.append(leaves.allocator, slot.id);
299 }
300 }
301 }
302 }
303
304 fn operandIsSplatMaterial(operand: *ir.Value) bool {
305 const def_any = operand.getDefiningOp() orelse return false;
306 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
307 if (constantOpIsSplat(def_op)) return true;
308 if (!std.mem.eql(u8, def_op.name.name, dialect_mod.AccyDialect.BroadcastOp.operation_name) and
309 !std.mem.eql(u8, def_op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name))
310 {
311 return false;
312 }
313 const operands = def_op.getOperandValues();
314 if (operands.len != 1) return false;
315 const source_any = operands[0].getDefiningOp() orelse return false;
316 const source_op: *ir.Operation = @ptrCast(@alignCast(source_any));
317 return constantOpIsSplat(source_op);
318 }
319
320 fn constantOpIsSplat(op: *ir.Operation) bool {
321 if (!std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name)) return false;
322 const constant = dialect_mod.AccyDialect.ConstantOp{ .op = op };
323 const payload = constant.getPayload() orelse return false;
324 const result = op.getResult(0) orelse return false;
325 var dtype_arena_buffer: [160]u8 = undefined;
326 var dtype_arena = alloc_fixed.FixedBuffer.init(dtype_arena_buffer[0..]);
327 const result_type = dialect_mod.decodeTensorType(dtype_arena.allocator(), result.type) catch return false;
328 switch (result_type.dtype) {
329 .f32, .i32 => {
330 if (payload.len < 4 or payload.len % 4 != 0) return false;
331 var first: u32 = undefined;
332 @memcpy(std.mem.asBytes(&first), payload[0..4]);
333 var offset: usize = 4;
334 while (offset < payload.len) : (offset += 4) {
335 var value: u32 = undefined;
336 @memcpy(std.mem.asBytes(&value), payload[offset..][0..4]);
337 if (value != first) return false;
338 }
339 return true;
340 },
341 .i1 => {
342 if (payload.len == 0) return false;
343 const first = payload[0];
344 for (payload[1..]) |value| {
345 if (value != first) return false;
346 }
347 return true;
348 },
349 else => return false,
350 }
351 }
352
353 fn operandIsInsideRegion(operand: *ir.Value, region: *const ir.Region) bool {
354 if (operand.getDefiningOp()) |def_any| {
355 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
356 const block = def_op.getBlock() orelse return false;
357 return blockInRegion(block, region);
358 }
359 const owner_any = operand.getOwnerBlock() orelse return false;
360 const block: *ir.Block = @ptrCast(@alignCast(owner_any));
361 return blockInRegion(block, region);
362 }
363
364 fn blockInRegion(block: *ir.Block, region: *const ir.Region) bool {
365 var block_iter = @constCast(region).getBlocks();
366 while (block_iter.next()) |candidate| {
367 if (candidate == block) return true;
368 }
369 return false;
370 }
371
372 const elided_leaf_walk_depth = 48;
373
374 const LeafWalk = struct {
375 allocator: std.mem.Allocator,
376 buffer_plan: *const bufferization.BufferPlanAnalysis,
377 seen: *std.AutoHashMap(usize, void),
378 inputs: *std.ArrayListUnmanaged(usize),
379 completed: std.AutoHashMap(*ir.Value, u64),
380 expansions: u64 = 0,
381
382 fn init(
383 allocator: std.mem.Allocator,
384 buffer_plan: *const bufferization.BufferPlanAnalysis,
385 seen: *std.AutoHashMap(usize, void),
386 inputs: *std.ArrayListUnmanaged(usize),
387 ) LeafWalk {
388 return .{
389 .allocator = allocator,
390 .buffer_plan = buffer_plan,
391 .seen = seen,
392 .inputs = inputs,
393 .completed = std.AutoHashMap(*ir.Value, u64).init(allocator),
394 };
395 }
396
397 fn deinit(self: *LeafWalk) void {
398 self.completed.deinit();
399 self.* = undefined;
400 }
401
402 fn collect(self: *LeafWalk, value: *ir.Value, depth: usize) anyerror!void {
403 if (depth == 0) return error.WorkExhausted;
404 std.debug.assert(depth <= elided_leaf_walk_depth);
405 const bit = @as(u64, 1) << @intCast(depth - 1);
406 if (self.completed.get(value)) |mask| {
407 if (mask & bit != 0) return;
408 }
409 self.expansions = try passes.pass.work.add(self.expansions, 1);
410 try self.expand(value, depth);
411 const entry = try self.completed.getOrPut(value);
412 if (!entry.found_existing) entry.value_ptr.* = 0;
413 entry.value_ptr.* |= bit;
414 }
415
416 fn expand(self: *LeafWalk, value: *ir.Value, depth: usize) anyerror!void {
417 if (value.getDefiningOp()) |def_any| {
418 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
419 if (constantOpIsSplat(def_op)) return;
420 const broadcast_name = dialect_mod.AccyDialect.BroadcastInDimOp.operation_name;
421 if (std.mem.eql(u8, def_op.name.name, broadcast_name)) {
422 const operands = def_op.getOperandValues();
423 if (operands.len != 1) return error.UnsupportedOperation;
424 if (operands[0].getDefiningOp()) |source_any| {
425 const source_op: *ir.Operation = @ptrCast(@alignCast(source_any));
426 if (constantOpIsSplat(source_op)) return;
427 }
428 const source_slot = self.buffer_plan.getSlot(operands[0]) orelse
429 return error.UnsupportedOperation;
430 if (!self.seen.contains(source_slot.id)) {
431 try self.seen.put(source_slot.id, {});
432 try self.inputs.append(self.allocator, source_slot.id);
433 }
434 return;
435 }
436 }
437 if (self.buffer_plan.getSlot(value)) |slot| {
438 if (!self.seen.contains(slot.id)) {
439 try self.seen.put(slot.id, {});
440 try self.inputs.append(self.allocator, slot.id);
441 }
442 return;
443 }
444 const def_any = value.getDefiningOp() orelse return error.UnsupportedOperation;
445 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
446 for (def_op.getOperandValues()) |operand| {
447 try self.collect(operand, depth - 1);
448 }
449 }
450 };
451
452 const KernelOutlineSlots = struct {
453 input_slot_ids: []usize,
454 output_slot_id: usize,
455
456 fn deinit(self: KernelOutlineSlots, allocator: std.mem.Allocator) void {
457 allocator.free(self.input_slot_ids);
458 }
459 };
460
461 fn collectKernelOutlineSlots(
462 allocator: std.mem.Allocator,
463 work: schedule_planning.ScheduleWorkItem,
464 buffer_plan: *const bufferization.BufferPlanAnalysis,
465 ) !KernelOutlineSlots {
466 const input_slots = try collectInputSlots(allocator, work, buffer_plan);
467 errdefer allocator.free(input_slots);
468 const output_slot = buffer_plan.getSlot(work.output_value) orelse return error.MissingOutputSlot;
469 return .{
470 .input_slot_ids = input_slots,
471 .output_slot_id = output_slot.id,
472 };
473 }
474
475 fn addKernelOutlines(
476 pass_ctx: *passes.PassContext,
477 analysis: *KernelOutlinePlanAnalysis,
478 schedule_plan: *const schedule_planning.SchedulePlanAnalysis,
479 buffer_plan: *const bufferization.BufferPlanAnalysis,
480 ) !void {
481 const work_items = schedule_plan.work_items.items;
482 if (pass_ctx.workerCount(work_items.len) > 1) {
483 return try addKernelOutlinesParallel(pass_ctx, analysis, work_items, buffer_plan);
484 }
485 for (work_items) |work| {
486 const slots = try collectKernelOutlineSlots(pass_ctx.allocator, work, buffer_plan);
487 try addKernelWithOwnedSlots(pass_ctx.allocator, analysis, work, slots);
488 }
489 }
490
491 const KernelOutlineInputResult = union(enum) {
492 pending,
493 ok: KernelOutlineSlots,
494 err: anyerror,
495 };
496
497 const KernelOutlineInputSlot = struct {
498 result: KernelOutlineInputResult = .pending,
499 };
500
501 const KernelOutlineInputBatch = struct {
502 allocator: std.mem.Allocator,
503 work_items: []const schedule_planning.ScheduleWorkItem,
504 buffer_plan: *const bufferization.BufferPlanAnalysis,
505 slots: []KernelOutlineInputSlot,
506 };
507
508 fn addKernelOutlinesParallel(
509 pass_ctx: *passes.PassContext,
510 analysis: *KernelOutlinePlanAnalysis,
511 work_items: []const schedule_planning.ScheduleWorkItem,
512 buffer_plan: *const bufferization.BufferPlanAnalysis,
513 ) !void {
514 const slots = try pass_ctx.allocator.alloc(KernelOutlineInputSlot, work_items.len);
515 defer pass_ctx.allocator.free(slots);
516 for (slots) |*slot| slot.* = .{};
517
518 const worker_allocator = pass_ctx.workerAllocator();
519 defer deinitKernelOutlineInputSlots(worker_allocator, slots);
520
521 var batch = KernelOutlineInputBatch{
522 .allocator = worker_allocator,
523 .work_items = work_items,
524 .buffer_plan = buffer_plan,
525 .slots = slots,
526 };
527
528 try ir.threading.parallelForEachIndex(
529 pass_ctx.allocator,
530 pass_ctx.run_options,
531 work_items.len,
532 &batch,
533 collectKernelOutlineAt,
534 );
535
536 if (firstKernelOutlineInputError(slots)) |err| return err;
537
538 for (work_items, slots) |work, slot| {
539 const worker_input_slots = switch (slot.result) {
540 .ok => |input_slots| input_slots,
541 .pending => unreachable,
542 .err => unreachable,
543 };
544 try addKernelWithClonedSlots(pass_ctx.allocator, analysis, work, worker_input_slots);
545 }
546 }
547
548 fn collectKernelOutlineAt(batch: *KernelOutlineInputBatch, index: usize) void {
549 batch.slots[index].result = if (collectKernelOutlineSlots(
550 batch.allocator,
551 batch.work_items[index],
552 batch.buffer_plan,
553 )) |input_slots|
554 .{ .ok = input_slots }
555 else |err|
556 .{ .err = err };
557 }
558
559 fn firstKernelOutlineInputError(slots: []const KernelOutlineInputSlot) ?anyerror {
560 for (slots) |slot| {
561 switch (slot.result) {
562 .err => |err| return err,
563 else => {},
564 }
565 }
566 return null;
567 }
568
569 fn deinitKernelOutlineInputSlots(
570 allocator: std.mem.Allocator,
571 slots: []KernelOutlineInputSlot,
572 ) void {
573 for (slots) |*slot| {
574 switch (slot.result) {
575 .ok => |kernel_slots| kernel_slots.deinit(allocator),
576 else => {},
577 }
578 }
579 }
580
581 fn addKernelWithClonedSlots(
582 allocator: std.mem.Allocator,
583 analysis: *KernelOutlinePlanAnalysis,
584 work: schedule_planning.ScheduleWorkItem,
585 worker_slots: KernelOutlineSlots,
586 ) !void {
587 const input_slots = try allocator.dupe(usize, worker_slots.input_slot_ids);
588 try addKernelWithOwnedSlots(allocator, analysis, work, .{
589 .input_slot_ids = input_slots,
590 .output_slot_id = worker_slots.output_slot_id,
591 });
592 }
593
594 fn addKernelWithOwnedSlots(
595 allocator: std.mem.Allocator,
596 analysis: *KernelOutlinePlanAnalysis,
597 work: schedule_planning.ScheduleWorkItem,
598 slots: KernelOutlineSlots,
599 ) !void {
600 errdefer slots.deinit(allocator);
601 try addKernel(analysis, work, slots.input_slot_ids, slots.output_slot_id);
602 }
603
604 fn addKernel(
605 analysis: *KernelOutlinePlanAnalysis,
606 work: schedule_planning.ScheduleWorkItem,
607 input_slot_ids: []usize,
608 output_slot_id: usize,
609 ) !void {
610 if (analysis.work_to_kernel.contains(work.id)) return;
611 const id = analysis.kernels.items.len;
612 const name = try std.fmt.allocPrint(analysis.allocator, "accy_choir_kernel_{d}", .{id});
613 errdefer analysis.allocator.free(name);
614
615 try analysis.work_to_kernel.put(work.id, id);
616 errdefer _ = analysis.work_to_kernel.remove(work.id);
617 try analysis.kernels.append(analysis.allocator, .{
618 .id = id,
619 .name = name,
620 .kind = kernelKindForWork(work.kind),
621 .work_item_id = work.id,
622 .root = work.root,
623 .input_slot_ids = input_slot_ids,
624 .output_slot_id = output_slot_id,
625 .element_count = work.element_count,
626 .op_count = work.opCount(),
627 });
628
629 analysis.total_input_slots += input_slot_ids.len;
630 analysis.total_scheduled_ops += work.opCount();
631 }
632
633 fn kernelKindForWork(kind: schedule_planning.ScheduleWorkKind) KernelOutlineKind {
634 return switch (kind) {
635 .elementwise_single, .elementwise_fusion => .elementwise,
636 .shape => .shape,
637 .dot_general => .dot_general,
638 .reduction => .reduction,
639 .kernel_call => .kernel_call,
640 .row_pipeline => .row_pipeline,
641 .iterate => .iterate,
642 .flash_attention => .flash_attention,
643 .scan => .scan,
644 };
645 }
646
647 fn operandIsBroadcastOfConstant(operand: *ir.Value) bool {
648 const def_any = operand.getDefiningOp() orelse return false;
649 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
650 if (!std.mem.eql(u8, def_op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name) and
651 !std.mem.eql(u8, def_op.name.name, dialect_mod.AccyDialect.BroadcastOp.operation_name))
652 {
653 return false;
654 }
655 const operands = def_op.getOperandValues();
656 if (operands.len != 1) return false;
657 return operandIsConstant(operands[0]);
658 }
659
660 fn operandIsConstant(operand: *ir.Value) bool {
661 const def_any = operand.getDefiningOp() orelse return false;
662 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
663 return std.mem.eql(u8, def_op.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name);
664 }
665
666 fn isReduceInitOperand(op: *ir.Operation, operand_index: usize) bool {
667 return std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ReduceOp.operation_name) and operand_index == 1;
668 }
669
670 fn isConstantReduceInitOperand(
671 buffer_plan: *const bufferization.BufferPlanAnalysis,
672 op: *ir.Operation,
673 operand: *ir.Value,
674 operand_index: usize,
675 ) bool {
676 if (!isReduceInitOperand(op, operand_index)) return false;
677 const slot = buffer_plan.getSlot(operand) orelse return false;
678 return slot.role.constant;
679 }
680
681 const testing = std.testing;
682 const semantic = accy_choir.semantic;
683
684 fn findOpNamedInBlock(block: *ir.Block, name: []const u8) ?*ir.Operation {
685 var iter = block.operations.head;
686 while (iter) |op_ptr| {
687 const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
688 if (std.mem.eql(u8, op.name.name, name)) return op;
689 iter = op.next_op;
690 }
691 return null;
692 }
693
694 test "kernel outlining plans one candidate for an elementwise fusion cluster" {
695 const allocator = testing.allocator;
696
697 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
698 defer builder.deinit();
699 const f32_4 = try builder.tensor(.f32, &.{4});
700 var fb = try builder.beginFunction("kernel_fused_add_mul", &.{ f32_4, f32_4, f32_4 }, &.{f32_4});
701 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
702 const product = try fb.mul(sum, fb.parameter(2));
703 try fb.return_(&.{product});
704 try fb.finish();
705 const module = try builder.finish();
706 defer module.deinit();
707
708 const choir_mod = module.choir_module;
709 const ctx = module.context();
710 const ledger = try outlineTestLedger();
711 defer ledger.destroy();
712 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);
713 defer cache.deinit();
714 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
715 defer pass_ctx.deinit();
716
717 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
718 const func = ir.inspection.functionByNameInBlock(body, "kernel_fused_add_mul") orelse return error.TestExpectedFunc;
719 const entry = func.getRegion(0).?.getEntryBlock().?;
720 const mul = findOpNamedInBlock(entry, dialect_mod.AccyDialect.MulOp.operation_name) orelse return error.TestExpectedMul;
721
722 const schedule_plan = try schedule_planning.getSchedulePlanAnalysis(&pass_ctx, choir_mod);
723 const work = schedule_plan.getWorkForRoot(mul) orelse return error.TestExpectedWorkItem;
724 const buffer_plan = try bufferization.getBufferPlanAnalysis(&pass_ctx, choir_mod);
725 const output_slot = buffer_plan.getSlot(mul.getResult(0).?) orelse return error.TestExpectedOutputSlot;
726
727 const outline_plan = try getKernelOutlinePlanAnalysis(&pass_ctx, choir_mod);
728 try ledger.producersComplete();
729 try checkOutlineStorage(choir_mod, ctx, &cache, outline_plan);
730 try testing.expectEqual(@as(usize, 1), outline_plan.kernelCount());
731 try testing.expectEqual(@as(usize, 3), outline_plan.total_input_slots);
732 try testing.expectEqual(@as(usize, 2), outline_plan.total_scheduled_ops);
733
734 const kernel = outline_plan.getKernelForWork(work.id) orelse return error.TestExpectedKernel;
735 try testing.expectEqualStrings("accy_choir_kernel_0", kernel.name);
736 try testing.expectEqual(KernelOutlineKind.elementwise, kernel.kind);
737 try testing.expectEqual(@as(usize, 3), kernel.inputCount());
738 try testing.expectEqual(output_slot.id, kernel.output_slot_id);
739 try testing.expectEqual(@as(u64, 4), kernel.element_count);
740 try testing.expectEqual(@as(usize, 2), kernel.op_count);
741 }
742
743 test "kernel outlining records standalone elementwise dependencies" {
744 const allocator = testing.allocator;
745
746 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
747 defer builder.deinit();
748 const f32_4 = try builder.tensor(.f32, &.{4});
749 var fb = try builder.beginFunction("kernel_escape", &.{ f32_4, f32_4, f32_4 }, &.{ f32_4, f32_4 });
750 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
751 const product = try fb.mul(sum, fb.parameter(2));
752 try fb.return_(&.{ sum, product });
753 try fb.finish();
754 const module = try builder.finish();
755 defer module.deinit();
756
757 const choir_mod = module.choir_module;
758 const ctx = module.context();
759 const ledger = try outlineTestLedger();
760 defer ledger.destroy();
761 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);
762 defer cache.deinit();
763 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
764 defer pass_ctx.deinit();
765
766 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
767 const func = ir.inspection.functionByNameInBlock(body, "kernel_escape") orelse return error.TestExpectedFunc;
768 const entry = func.getRegion(0).?.getEntryBlock().?;
769 const add = findOpNamedInBlock(entry, dialect_mod.AccyDialect.AddOp.operation_name) orelse return error.TestExpectedAdd;
770 const mul = findOpNamedInBlock(entry, dialect_mod.AccyDialect.MulOp.operation_name) orelse return error.TestExpectedMul;
771
772 const schedule_plan = try schedule_planning.getSchedulePlanAnalysis(&pass_ctx, choir_mod);
773 const add_work = schedule_plan.getWorkForRoot(add) orelse return error.TestExpectedAddWork;
774 const mul_work = schedule_plan.getWorkForRoot(mul) orelse return error.TestExpectedMulWork;
775
776 const outline_plan = try getKernelOutlinePlanAnalysis(&pass_ctx, choir_mod);
777 try ledger.producersComplete();
778 try checkOutlineStorage(choir_mod, ctx, &cache, outline_plan);
779 try testing.expectEqual(@as(usize, 2), outline_plan.kernelCount());
780 try testing.expectEqual(@as(usize, 4), outline_plan.total_input_slots);
781 try testing.expectEqual(@as(usize, 2), outline_plan.total_scheduled_ops);
782
783 const add_kernel = outline_plan.getKernelForWork(add_work.id) orelse return error.TestExpectedAddKernel;
784 const mul_kernel = outline_plan.getKernelForWork(mul_work.id) orelse return error.TestExpectedMulKernel;
785 try testing.expectEqual(@as(usize, 2), add_kernel.inputCount());
786 try testing.expectEqual(@as(usize, 2), mul_kernel.inputCount());
787
788 const buffer_plan = try bufferization.getBufferPlanAnalysis(&pass_ctx, choir_mod);
789 const sum_slot = buffer_plan.getSlot(add.getResult(0).?) orelse return error.TestExpectedSumSlot;
790 try testing.expectEqual(sum_slot.id, mul_kernel.input_slot_ids[0]);
791 }
792
793 test "kernel outlining records reduction data input without init constant" {
794 const allocator = testing.allocator;
795
796 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
797 defer builder.deinit();
798 const f32_256 = try builder.tensor(.f32, &.{256});
799 const f32_scalar = try builder.tensor(.f32, &.{});
800 var fb = try builder.beginFunction("kernel_reduce_sum_rank1", &.{f32_256}, &.{f32_scalar});
801 const zero_value: f32 = 0.0;
802 const zero = try fb.constant(f32_scalar, std.mem.asBytes(&zero_value));
803 const reduced = try fb.reduce(fb.parameter(0), zero, f32_scalar, "sum", &.{0});
804 try fb.return_(&.{reduced});
805 try fb.finish();
806 const module = try builder.finish();
807 defer module.deinit();
808
809 const choir_mod = module.choir_module;
810 const ctx = module.context();
811 const ledger = try outlineTestLedger();
812 defer ledger.destroy();
813 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);
814 defer cache.deinit();
815 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
816 defer pass_ctx.deinit();
817
818 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
819 const func = ir.inspection.functionByNameInBlock(body, "kernel_reduce_sum_rank1") orelse return error.TestExpectedFunc;
820 const entry = func.getRegion(0).?.getEntryBlock().?;
821 const reduce_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.ReduceOp.operation_name) orelse return error.TestExpectedReduce;
822
823 const schedule_plan = try schedule_planning.getSchedulePlanAnalysis(&pass_ctx, choir_mod);
824 const work = schedule_plan.getWorkForRoot(reduce_op) orelse return error.TestExpectedReduceWork;
825 const buffer_plan = try bufferization.getBufferPlanAnalysis(&pass_ctx, choir_mod);
826 const input_slot = buffer_plan.getSlot(reduce_op.getOperand(0).?) orelse return error.TestExpectedInputSlot;
827 const output_slot = buffer_plan.getSlot(reduce_op.getResult(0).?) orelse return error.TestExpectedOutputSlot;
828
829 const outline_plan = try getKernelOutlinePlanAnalysis(&pass_ctx, choir_mod);
830 try ledger.producersComplete();
831 try checkOutlineStorage(choir_mod, ctx, &cache, outline_plan);
832 try testing.expectEqual(@as(usize, 1), outline_plan.kernelCount());
833 try testing.expectEqual(@as(usize, 1), outline_plan.total_input_slots);
834 try testing.expectEqual(@as(usize, 1), outline_plan.total_scheduled_ops);
835
836 const kernel = outline_plan.getKernelForWork(work.id) orelse return error.TestExpectedKernel;
837 try testing.expectEqual(KernelOutlineKind.reduction, kernel.kind);
838 try testing.expectEqual(@as(usize, 1), kernel.inputCount());
839 try testing.expectEqual(input_slot.id, kernel.input_slot_ids[0]);
840 try testing.expectEqual(output_slot.id, kernel.output_slot_id);
841 try testing.expectEqual(@as(u64, 1), kernel.element_count);
842 try testing.expectEqual(@as(usize, 1), kernel.op_count);
843 }
844
845 test "kernel outlining records parameter reduce init as input" {
846 const allocator = testing.allocator;
847
848 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
849 defer builder.deinit();
850 const f32_256 = try builder.tensor(.f32, &.{256});
851 const f32_scalar = try builder.tensor(.f32, &.{});
852 var fb = try builder.beginFunction("kernel_reduce_sum_parameter_init", &.{ f32_256, f32_scalar }, &.{f32_scalar});
853 const reduced = try fb.reduce(fb.parameter(0), fb.parameter(1), f32_scalar, "sum", &.{0});
854 try fb.return_(&.{reduced});
855 try fb.finish();
856 const module = try builder.finish();
857 defer module.deinit();
858
859 const choir_mod = module.choir_module;
860 const ctx = module.context();
861 const ledger = try outlineTestLedger();
862 defer ledger.destroy();
863 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);
864 defer cache.deinit();
865 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
866 defer pass_ctx.deinit();
867
868 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
869 const func = ir.inspection.functionByNameInBlock(body, "kernel_reduce_sum_parameter_init") orelse return error.TestExpectedFunc;
870 const entry = func.getRegion(0).?.getEntryBlock().?;
871 const reduce_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.ReduceOp.operation_name) orelse return error.TestExpectedReduce;
872
873 const schedule_plan = try schedule_planning.getSchedulePlanAnalysis(&pass_ctx, choir_mod);
874 const work = schedule_plan.getWorkForRoot(reduce_op) orelse return error.TestExpectedReduceWork;
875 const buffer_plan = try bufferization.getBufferPlanAnalysis(&pass_ctx, choir_mod);
876 const input_slot = buffer_plan.getSlot(fb.parameter(0)) orelse return error.TestExpectedInputSlot;
877 const init_slot = buffer_plan.getSlot(fb.parameter(1)) orelse return error.TestExpectedInputSlot;
878 const output_slot = buffer_plan.getSlot(reduce_op.getResult(0).?) orelse return error.TestExpectedOutputSlot;
879
880 const outline_plan = try getKernelOutlinePlanAnalysis(&pass_ctx, choir_mod);
881 try ledger.producersComplete();
882 try checkOutlineStorage(choir_mod, ctx, &cache, outline_plan);
883 try testing.expectEqual(@as(usize, 1), outline_plan.kernelCount());
884 try testing.expectEqual(@as(usize, 2), outline_plan.total_input_slots);
885
886 const kernel = outline_plan.getKernelForWork(work.id) orelse return error.TestExpectedKernel;
887 try testing.expectEqual(KernelOutlineKind.reduction, kernel.kind);
888 try testing.expectEqual(@as(usize, 2), kernel.inputCount());
889 try testing.expectEqual(input_slot.id, kernel.input_slot_ids[0]);
890 try testing.expectEqual(init_slot.id, kernel.input_slot_ids[1]);
891 try testing.expectEqual(output_slot.id, kernel.output_slot_id);
892 }
893
894 test "kernel outlining records semantic kernel_call work" {
895 const allocator = testing.allocator;
896
897 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
898 defer builder.deinit();
899 const f32_4 = try builder.tensor(.f32, &.{4});
900 var fb = try builder.beginFunction("kernel_call_outline", &.{ f32_4, f32_4 }, &.{f32_4});
901 const call = try fb.kernelCall(
902 &.{ fb.parameter(0), fb.parameter(1) },
903 &.{f32_4},
904 .{
905 .target = "accy.custom.scale",
906 .operand_effects = &.{ .read, .write },
907 .result_aliases = &.{null},
908 },
909 );
910 try fb.return_(&.{call.getFirstResult()});
911 try fb.finish();
912 const module = try builder.finish();
913 defer module.deinit();
914
915 const choir_mod = module.choir_module;
916 const ctx = module.context();
917 const ledger = try outlineTestLedger();
918 defer ledger.destroy();
919 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);
920 defer cache.deinit();
921 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
922 defer pass_ctx.deinit();
923
924 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
925 const func = ir.inspection.functionByNameInBlock(body, "kernel_call_outline") orelse return error.TestExpectedFunc;
926 const entry = func.getRegion(0).?.getEntryBlock().?;
927 const call_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse return error.TestExpectedKernelCall;
928
929 const schedule_plan = try schedule_planning.getSchedulePlanAnalysis(&pass_ctx, choir_mod);
930 const work = schedule_plan.getWorkForRoot(call_op) orelse return error.TestExpectedWork;
931 const buffer_plan = try bufferization.getBufferPlanAnalysis(&pass_ctx, choir_mod);
932 const input_slot = buffer_plan.getSlot(fb.parameter(0)) orelse return error.TestExpectedInputSlot;
933 const output_slot = buffer_plan.getSlot(call.getFirstResult()) orelse return error.TestExpectedOutputSlot;
934
935 const outline_plan = try getKernelOutlinePlanAnalysis(&pass_ctx, choir_mod);
936 try ledger.producersComplete();
937 try checkOutlineStorage(choir_mod, ctx, &cache, outline_plan);
938 try testing.expectEqual(@as(usize, 1), outline_plan.kernelCount());
939 try testing.expectEqual(@as(usize, 2), outline_plan.total_input_slots);
940 try testing.expectEqual(@as(usize, 1), outline_plan.total_scheduled_ops);
941
942 const kernel = outline_plan.getKernelForWork(work.id) orelse return error.TestExpectedKernel;
943 try testing.expectEqual(KernelOutlineKind.kernel_call, kernel.kind);
944 try testing.expectEqual(@as(usize, 2), kernel.inputCount());
945 try testing.expectEqual(input_slot.id, kernel.input_slot_ids[0]);
946 try testing.expectEqual(output_slot.id, kernel.output_slot_id);
947 try testing.expectEqual(@as(u64, 4), kernel.element_count);
948 try testing.expectEqual(@as(usize, 1), kernel.op_count);
949 }
950
951 test "kernel outlining threads match serial analysis" {
952 const allocator = testing.allocator;
953
954 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
955 defer builder.deinit();
956 const f32_4 = try builder.tensor(.f32, &.{4});
957 var fb = try builder.beginFunction("kernel_threaded", &.{ f32_4, f32_4, f32_4 }, &.{ f32_4, f32_4 });
958 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
959 const product = try fb.mul(sum, fb.parameter(2));
960 try fb.return_(&.{ sum, product });
961 try fb.finish();
962 const module = try builder.finish();
963 defer module.deinit();
964
965 const choir_mod = module.choir_module;
966 const ctx = module.context();
967
968 var serial_cache = passes.AnalysisCache.init(allocator, null);
969 defer serial_cache.deinit();
970 var serial_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &serial_cache);
971 defer serial_ctx.deinit();
972 const serial = try getKernelOutlinePlanAnalysis(&serial_ctx, choir_mod);
973
974 var worker_gpa = std.heap.DebugAllocator(.{}){};
975 defer {
976 const status = worker_gpa.deinit();
977 testing.expect(status == .ok) catch @panic("kernel outlining worker allocator leaked allocations");
978 }
979
980 var threaded_cache = passes.AnalysisCache.init(allocator, null);
981 defer threaded_cache.deinit();
982 var threaded_ctx = passes.PassContext.initWithOptions(choir_mod, ctx, allocator, &threaded_cache, .{
983 .max_threads = 2,
984 .worker_allocator = worker_gpa.allocator(),
985 });
986 defer threaded_ctx.deinit();
987 const threaded = try getKernelOutlinePlanAnalysis(&threaded_ctx, choir_mod);
988
989 try testing.expectEqual(serial.kernelCount(), threaded.kernelCount());
990 try testing.expectEqual(serial.total_input_slots, threaded.total_input_slots);
991 try testing.expectEqual(serial.total_scheduled_ops, threaded.total_scheduled_ops);
992 for (serial.kernels.items, threaded.kernels.items) |expected, actual| {
993 try testing.expectEqual(expected.id, actual.id);
994 try testing.expectEqualStrings(expected.name, actual.name);
995 try testing.expectEqual(expected.kind, actual.kind);
996 try testing.expectEqual(expected.work_item_id, actual.work_item_id);
997 try testing.expectEqual(expected.root, actual.root);
998 try testing.expectEqualSlices(usize, expected.input_slot_ids, actual.input_slot_ids);
999 try testing.expectEqual(expected.output_slot_id, actual.output_slot_id);
1000 try testing.expectEqual(expected.element_count, actual.element_count);
1001 try testing.expectEqual(expected.op_count, actual.op_count);
1002 }
1003 }
1004
1005 test "kernel outlining planning pass preserves IR" {
1006 const allocator = testing.allocator;
1007
1008 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1009 defer builder.deinit();
1010 const f32_4 = try builder.tensor(.f32, &.{4});
1011 var fb = try builder.beginFunction("kernel_pass_add4", &.{ f32_4, f32_4 }, &.{f32_4});
1012 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
1013 try fb.return_(&.{sum});
1014 try fb.finish();
1015 const module = try builder.finish();
1016 defer module.deinit();
1017
1018 const choir_mod = module.choir_module;
1019 const ctx = module.context();
1020 var pm = passes.PassManager.init(allocator);
1021 defer pm.deinit();
1022 try pm.addPass(kernelOutliningPlanningPass());
1023
1024 const ledger = try outlinePassLedger(std.math.maxInt(u64));
1025 defer ledger.destroy();
1026 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);
1027 defer cache.deinit();
1028 try testing.expectEqual(
1029 passes.PassResult.success,
1030 pm.runWithAnalysisCache(choir_mod, ctx, &cache, .{}),
1031 );
1032 try ledger.producersComplete();
1033 try testing.expectEqual(@as(u64, 1), pm.stats.pass_runs);
1034 try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
1035 }
1036
1037 test "kernel outlining plans an iterate that captures a function-scope splat" {
1038 const allocator = testing.allocator;
1039
1040 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1041 defer builder.deinit();
1042 const f32_4 = try builder.tensor(.f32, &.{4});
1043 const f32_scalar = try builder.tensor(.f32, &.{});
1044 const pred_ty = try builder.tensor(.i1, &.{4});
1045 var fb = try builder.beginFunction("kernel_iterate_splat_capture", &.{f32_4}, &.{f32_4});
1046 const step_value: f32 = 1.2;
1047 const step_scalar = try fb.constant(f32_scalar, std.mem.asBytes(&step_value));
1048 const step = try fb.broadcast(step_scalar, f32_4, &.{4});
1049 const limit_value: f32 = 40.0;
1050 const limit_scalar = try fb.constant(f32_scalar, std.mem.asBytes(&limit_value));
1051 const limit = try fb.broadcast(limit_scalar, f32_4, &.{4});
1052
1053 var it = try fb.beginIterate(&.{fb.parameter(0)}, 8);
1054 const body = it.inner();
1055 const advanced = try body.add(it.carry(0), step);
1056 const active = try body.compare(advanced, limit, pred_ty, .lt);
1057 try it.yield_(active, &.{advanced});
1058 try fb.return_(&.{it.result(0)});
1059 try fb.finish();
1060 const module = try builder.finish();
1061 defer module.deinit();
1062
1063 const choir_mod = module.choir_module;
1064 const ctx = module.context();
1065 const ledger = try outlineTestLedger();
1066 defer ledger.destroy();
1067 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);
1068 defer cache.deinit();
1069 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1070 defer pass_ctx.deinit();
1071
1072 const body_block = choir_mod.getRegion(0).?.getEntryBlock().?;
1073 const func = ir.inspection.functionByNameInBlock(body_block, "kernel_iterate_splat_capture") orelse return error.TestExpectedFunc;
1074 const entry = func.getRegion(0).?.getEntryBlock().?;
1075 const iterate_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.IterateOp.operation_name) orelse return error.TestExpectedIterate;
1076
1077 const schedule_plan = try schedule_planning.getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1078 const work = schedule_plan.getWorkForRoot(iterate_op) orelse return error.TestExpectedWorkItem;
1079 const buffer_plan = try bufferization.getBufferPlanAnalysis(&pass_ctx, choir_mod);
1080
1081 const outline_plan = try getKernelOutlinePlanAnalysis(&pass_ctx, choir_mod);
1082 try ledger.producersComplete();
1083 try checkOutlineStorage(choir_mod, ctx, &cache, outline_plan);
1084 const kernel = outline_plan.getKernelForWork(work.id) orelse return error.TestExpectedKernel;
1085 try testing.expectEqual(KernelOutlineKind.iterate, kernel.kind);
1086
1087 const splat_slot = buffer_plan.getSlot(step);
1088 for (kernel.input_slot_ids) |slot_id| {
1089 if (splat_slot) |slot| try testing.expect(slot_id != slot.id);
1090 }
1091 }
1092
1093 test "kernel outlining distinguishes leaf walk exhaustion from unsupported inputs" {
1094 try checkLeafWalkBoundary(elided_leaf_walk_depth - 1, false);
1095 try checkLeafWalkBoundary(elided_leaf_walk_depth, true);
1096 }
1097
1098 fn checkLeafWalkBoundary(count: usize, exhausted: bool) !void {
1099 var builder = try semantic.Builder.init(
1100 testing.allocator,
1101 semantic.Builder.ContextLimits.standard,
1102 );
1103 defer builder.deinit();
1104 const typ = try builder.tensor(.f32, &.{4});
1105 var function = try builder.beginFunction("outline_leaf_depth", &.{ typ, typ }, &.{typ});
1106 const left = function.parameter(0);
1107 const right = function.parameter(1);
1108 var leaf = left;
1109 for (0..count) |_| leaf = try function.add(leaf, right);
1110 const result = try function.add(leaf, right);
1111 try function.return_(&.{result});
1112 try function.finish();
1113 const module = try builder.finish();
1114 defer module.deinit();
1115 var cache = passes.AnalysisCache.init(testing.allocator, null);
1116 defer cache.deinit();
1117 var pass_ctx = passes.PassContext.init(
1118 module.choir_module,
1119 module.context(),
1120 testing.allocator,
1121 &cache,
1122 );
1123 defer pass_ctx.deinit();
1124 const buffers = try bufferization.getBufferPlanAnalysis(&pass_ctx, module.choir_module);
1125 try testing.expect(buffers.getElision(leaf) != null);
1126 var seen = std.AutoHashMap(usize, void).init(testing.allocator);
1127 defer seen.deinit();
1128 var inputs: std.ArrayListUnmanaged(usize) = .empty;
1129 defer inputs.deinit(testing.allocator);
1130 var leaves = LeafWalk.init(testing.allocator, buffers, &seen, &inputs);
1131 defer leaves.deinit();
1132 const collected = leaves.collect(leaf, elided_leaf_walk_depth);
1133 if (exhausted) {
1134 try testing.expectError(error.WorkExhausted, collected);
1135 } else {
1136 try collected;
1137 try testing.expectEqualSlices(
1138 usize,
1139 &.{ buffers.getSlot(left).?.id, buffers.getSlot(right).?.id },
1140 inputs.items,
1141 );
1142 try testing.expectError(
1143 error.WorkExhausted,
1144 leaves.collect(leaf, elided_leaf_walk_depth - 1),
1145 );
1146 }
1147 var empty = bufferization.BufferPlanAnalysis.init(testing.allocator);
1148 defer empty.deinit();
1149 var unsupported = LeafWalk.init(testing.allocator, &empty, &seen, &inputs);
1150 defer unsupported.deinit();
1151 try testing.expectError(
1152 error.UnsupportedOperation,
1153 unsupported.collect(left, elided_leaf_walk_depth),
1154 );
1155 }
1156
1157 test "kernel outlining shares repeated leaf expansion at each remaining depth" {
1158 for ([_]usize{ 1, 2, 4, 8, 12, 16 }) |count| try checkSharedLeafWalk(count);
1159 }
1160
1161 fn checkSharedLeafWalk(count: usize) !void {
1162 std.debug.assert(count < elided_leaf_walk_depth);
1163 var builder = try semantic.Builder.init(
1164 testing.allocator,
1165 semantic.Builder.ContextLimits.standard,
1166 );
1167 defer builder.deinit();
1168 const typ = try builder.tensor(.f32, &.{4});
1169 var function = try builder.beginFunction("shared_leaf_walk", &.{typ}, &.{typ});
1170 const input = function.parameter(0);
1171 var leaf = input;
1172 for (0..count) |_| leaf = try function.add(leaf, leaf);
1173 const result = try function.add(leaf, input);
1174 try function.return_(&.{result});
1175 try function.finish();
1176 const module = try builder.finish();
1177 defer module.deinit();
1178 var cache = passes.AnalysisCache.init(testing.allocator, null);
1179 defer cache.deinit();
1180 var pass_ctx = passes.PassContext.init(
1181 module.choir_module,
1182 module.context(),
1183 testing.allocator,
1184 &cache,
1185 );
1186 defer pass_ctx.deinit();
1187 const buffers = try bufferization.getBufferPlanAnalysis(&pass_ctx, module.choir_module);
1188 try testing.expect(buffers.getElision(leaf) != null);
1189 var seen = std.AutoHashMap(usize, void).init(testing.allocator);
1190 defer seen.deinit();
1191 var inputs: std.ArrayListUnmanaged(usize) = .empty;
1192 defer inputs.deinit(testing.allocator);
1193 var leaves = LeafWalk.init(testing.allocator, buffers, &seen, &inputs);
1194 defer leaves.deinit();
1195 try leaves.collect(leaf, elided_leaf_walk_depth);
1196 const slot_id = buffers.getSlot(input).?.id;
1197 try testing.expectEqualSlices(usize, &.{slot_id}, inputs.items);
1198 try testing.expectEqual(count + 1, leaves.expansions);
1199 try testing.expectEqual(count + 1, leaves.completed.count());
1200 try leaves.collect(leaf, elided_leaf_walk_depth);
1201 try testing.expectEqual(count + 1, leaves.expansions);
1202 try testing.expectEqual(@as(usize, 1), inputs.items.len);
1203 try checkLeafWalkStorage(leaf, buffers, count + 1, slot_id);
1204 }
1205
1206 fn checkLeafWalkStorage(
1207 value: *ir.Value,
1208 buffers: *const bufferization.BufferPlanAnalysis,
1209 reached_values: u64,
1210 slot_id: usize,
1211 ) !void {
1212 const memo = try accounting.hashMapGrowth(*ir.Value, u64, reached_values);
1213 const seen_bytes = try accounting.hashMapGrowth(usize, void, 1);
1214 const input_bytes = try accounting.arrayListGrowth(usize, 1);
1215 const allowance = try accounting.add(memo, try accounting.add(seen_bytes, input_bytes));
1216 const bytes = try testing.allocator.alignedAlloc(u8, .@"64", @intCast(allowance));
1217 defer testing.allocator.free(bytes);
1218 var backing = alloc_fixed.Tracked.init(bytes);
1219 var retained = alloc_fixed.Monotonic.init(backing.allocator(), bytes.len);
1220 const allocator = retained.allocator();
1221 var seen = std.AutoHashMap(usize, void).init(allocator);
1222 defer seen.deinit();
1223 var inputs: std.ArrayListUnmanaged(usize) = .empty;
1224 defer inputs.deinit(allocator);
1225 var leaves = LeafWalk.init(allocator, buffers, &seen, &inputs);
1226 defer leaves.deinit();
1227 try leaves.collect(value, elided_leaf_walk_depth);
1228 try testing.expectEqualSlices(usize, &.{slot_id}, inputs.items);
1229 try testing.expectEqual(reached_values, leaves.expansions);
1230 try testing.expectEqual(reached_values, leaves.completed.count());
1231 try testing.expect(!backing.exhausted);
1232 const used = if (retained.current) |*current| alloc_fixed.used(current) else 0;
1233 try testing.expect(used <= allowance);
1234 try testing.expect(used > 0);
1235 }
1236
1237 fn outlineTestLedger() !*choir.product.revision.AccountingV1 {
1238 const revision = choir.product.revision;
1239 return revision.AccountingV1.create(testing.allocator, .{
1240 .allowance = revision.WorkVector.uniform(std.math.maxInt(u64)),
1241 .workspace = std.math.maxInt(u64),
1242 .events = 32,
1243 }, &.{});
1244 }
1245
1246 fn checkOutlineStorage(
1247 op: *ir.Operation,
1248 ctx: *ir.Context,
1249 cache: *passes.AnalysisCache,
1250 expected: *const KernelOutlinePlanAnalysis,
1251 ) !void {
1252 const bounds = try outlineAnalysisWork(.{ .operation = op });
1253 const bytes = try testing.allocator.alignedAlloc(u8, .@"64", @intCast(bounds.workspace));
1254 defer testing.allocator.free(bytes);
1255 var backing = alloc_fixed.Tracked.init(bytes);
1256 var retained = alloc_fixed.Monotonic.init(backing.allocator(), @max(1, bytes.len));
1257 const allocator = retained.allocator();
1258 var pass_ctx = passes.PassContext.init(op, ctx, allocator, cache);
1259 defer pass_ctx.deinit();
1260 const ptr = try computeKernelOutlinePlanAnalysis(&pass_ctx, op);
1261 defer cleanupKernelOutlinePlanAnalysis(ptr, allocator);
1262 const actual: *KernelOutlinePlanAnalysis = @ptrCast(@alignCast(ptr));
1263 try expectOutlinePlan(expected, actual);
1264 const used = if (retained.current) |*current| alloc_fixed.used(current) else 0;
1265 try testing.expect(!backing.exhausted);
1266 try testing.expect(used <= bounds.workspace);
1267 try testing.expect(used >= @sizeOf(KernelOutlinePlanAnalysis));
1268 }
1269
1270 fn expectOutlinePlan(
1271 expected: *const KernelOutlinePlanAnalysis,
1272 actual: *const KernelOutlinePlanAnalysis,
1273 ) !void {
1274 try testing.expectEqual(expected.kernelCount(), actual.kernelCount());
1275 try testing.expectEqual(expected.total_input_slots, actual.total_input_slots);
1276 try testing.expectEqual(expected.total_scheduled_ops, actual.total_scheduled_ops);
1277 for (expected.kernels.items, actual.kernels.items) |left, right| {
1278 inline for (.{
1279 "id", "kind", "work_item_id", "root", "output_slot_id", "element_count", "op_count",
1280 }) |field| try testing.expectEqual(@field(left, field), @field(right, field));
1281 try testing.expectEqualStrings(left.name, right.name);
1282 try testing.expectEqualSlices(usize, left.input_slot_ids, right.input_slot_ids);
1283 try testing.expectEqual(right.id, actual.getKernelForWork(right.work_item_id).?.id);
1284 }
1285 }
1286
1287 fn buildOutlineTestModule(builder: *semantic.Builder, count: usize) !*semantic.SemanticModule {
1288 const typ = try builder.tensor(.f32, &.{4});
1289 const types = try testing.allocator.alloc(ir.Type, count);
1290 defer testing.allocator.free(types);
1291 @memset(types, typ);
1292 const results = try testing.allocator.alloc(*ir.Value, count);
1293 defer testing.allocator.free(results);
1294 var function = try builder.beginFunction("outline_work", &.{ typ, typ }, types);
1295 for (results) |*result| {
1296 result.* = try function.add(function.parameter(0), function.parameter(1));
1297 }
1298 try function.return_(results);
1299 try function.finish();
1300 return builder.finish();
1301 }
1302
1303 test "kernel outlining work contract covers growing independent outputs" {
1304 for ([_]usize{ 0, 1, 2, 6, 7, 16, 64 }) |count| {
1305 var builder = try semantic.Builder.init(
1306 testing.allocator,
1307 semantic.Builder.ContextLimits.standard,
1308 );
1309 defer builder.deinit();
1310 const module = try buildOutlineTestModule(&builder, count);
1311 defer module.deinit();
1312 const ledger = try outlineTestLedger();
1313 defer ledger.destroy();
1314 var cache = try passes.AnalysisCache.initAccounted(
1315 testing.allocator,
1316 null,
1317 ledger,
1318 .{},
1319 8,
1320 );
1321 defer cache.deinit();
1322 var ctx = passes.PassContext.init(
1323 module.choir_module,
1324 module.context(),
1325 testing.allocator,
1326 &cache,
1327 );
1328 defer ctx.deinit();
1329 const outlines = try getKernelOutlinePlanAnalysis(&ctx, module.choir_module);
1330 try testing.expectEqual(count, outlines.kernelCount());
1331 try testing.expectEqual(count * 2, outlines.total_input_slots);
1332 try ledger.producersComplete();
1333 try checkOutlineStorage(module.choir_module, module.context(), &cache, outlines);
1334 }
1335 try testing.expectError(error.WorkOverflow, (OutlineWork{
1336 .input = .{ .operations = std.math.maxInt(u64) },
1337 }).bounds());
1338 try testing.expectError(error.WorkOverflow, (OutlineWork{
1339 .input = .{ .values = std.math.maxInt(u64) },
1340 }).bounds());
1341 try testing.expectError(error.WorkOverflow, (OutlineWork{
1342 .input = .{ .input_bytes = std.math.maxInt(u64) },
1343 }).bounds());
1344 }
1345
1346 test "kernel outlining rejects unmodeled execution options before compute and refresh" {
1347 for ([_]ir.ThreadingOptions{
1348 .{ .max_threads = 0 },
1349 .{ .max_threads = 2 },
1350 .{ .worker_allocator = testing.allocator },
1351 }) |options| {
1352 try checkOutlineOptions(options, false);
1353 try checkOutlineOptions(options, true);
1354 }
1355 }
1356
1357 fn rejectOutlineRefresh(_: *passes.PassContext, _: *ir.Operation, _: *anyopaque) !void {
1358 return error.TestUnexpectedResult;
1359 }
1360
1361 fn checkOutlineOptions(options: ir.ThreadingOptions, refresh: bool) !void {
1362 var builder = try semantic.Builder.init(
1363 testing.allocator,
1364 semantic.Builder.ContextLimits.standard,
1365 );
1366 defer builder.deinit();
1367 const module = try buildOutlineTestModule(&builder, 2);
1368 defer module.deinit();
1369 const ledger = try outlineTestLedger();
1370 defer ledger.destroy();
1371 var cache = try passes.AnalysisCache.initAccounted(testing.allocator, null, ledger, .{}, 8);
1372 defer cache.deinit();
1373 var ctx = passes.PassContext.init(
1374 module.choir_module,
1375 module.context(),
1376 testing.allocator,
1377 &cache,
1378 );
1379 defer ctx.deinit();
1380 const existing = if (refresh)
1381 try getKernelOutlinePlanAnalysis(&ctx, module.choir_module)
1382 else
1383 null;
1384 const before = ledger.view().charged;
1385 const entries = cache.entries.count();
1386 ctx.run_options = options;
1387 if (refresh) {
1388 try testing.expectError(error.MissingWorkContract, ctx.refreshAnalysis(
1389 module.choir_module,
1390 &kernel_outline_plan_analysis_descriptor,
1391 rejectOutlineRefresh,
1392 ));
1393 try testing.expectEqual(@as(usize, 2), existing.?.kernelCount());
1394 } else {
1395 try testing.expectError(
1396 error.MissingWorkContract,
1397 getKernelOutlinePlanAnalysis(&ctx, module.choir_module),
1398 );
1399 }
1400 try testing.expectEqualDeep(before, ledger.view().charged);
1401 try testing.expectEqual(entries, cache.entries.count());
1402 try testing.expectEqual(.rejected, ledger.view().outcome);
1403 ctx.run_options = .{};
1404 try testing.expectError(
1405 error.TerminalWorkOutcome,
1406 getKernelOutlinePlanAnalysis(&ctx, module.choir_module),
1407 );
1408 }
1409
1410 test "kernel outlining accounts closed scopes and rejects external operand owners" {
1411 var builder = try semantic.Builder.init(
1412 testing.allocator,
1413 semantic.Builder.ContextLimits.standard,
1414 );
1415 defer builder.deinit();
1416 const module = try buildOutlineTestModule(&builder, 2);
1417 defer module.deinit();
1418 const body = module.choir_module.getRegion(0).?.getEntryBlock().?;
1419 const function = ir.inspection.functionByNameInBlock(body, "outline_work").?;
1420 const entry = function.getRegion(0).?.getEntryBlock().?;
1421 const add = findOpNamedInBlock(entry, dialect_mod.AccyDialect.AddOp.operation_name).?;
1422 const returned: *ir.Operation = @ptrCast(@alignCast(entry.operations.tail.?));
1423 _ = try outlineAnalysisWork(.{ .operation = module.choir_module });
1424 _ = try outlineAnalysisWork(.{ .operation = function });
1425 for ([_]*ir.Operation{ add, returned }) |op| {
1426 const ledger = try outlineTestLedger();
1427 defer ledger.destroy();
1428 var cache = try passes.AnalysisCache.initAccounted(
1429 testing.allocator,
1430 null,
1431 ledger,
1432 .{},
1433 8,
1434 );
1435 defer cache.deinit();
1436 var ctx = passes.PassContext.init(op, module.context(), testing.allocator, &cache);
1437 defer ctx.deinit();
1438 try testing.expectError(error.UnboundProductInput, getKernelOutlinePlanAnalysis(&ctx, op));
1439 try testing.expectEqual(@as(usize, 0), cache.entries.count());
1440 try testing.expectEqual(@as(u64, 0), ledger.view().charged.analysis_computations);
1441 try testing.expectEqual(.rejected, ledger.view().outcome);
1442 }
1443 }
1444
1445 fn outlinePassLedger(visits: u64) !*choir.product.revision.AccountingV1 {
1446 const revision = choir.product.revision;
1447 var allowance = revision.WorkVector.uniform(std.math.maxInt(u64));
1448 allowance.structural_visits = visits;
1449 return revision.AccountingV1.create(testing.allocator, .{
1450 .allowance = allowance,
1451 .workspace = std.math.maxInt(u64),
1452 .events = 32,
1453 }, &.{.{ .name = kernel_outlining_planning_pass_name, .version = 1 }});
1454 }
1455
1456 test "kernel outlining preserves its exact admission boundary across cache hits" {
1457 var builder = try semantic.Builder.init(
1458 testing.allocator,
1459 semantic.Builder.ContextLimits.standard,
1460 );
1461 defer builder.deinit();
1462 const module = try buildOutlineTestModule(&builder, 2);
1463 defer module.deinit();
1464 const charge = try checkOutlineAllowance(module, std.math.maxInt(u64), true);
1465 try testing.expect(charge > 1);
1466 _ = try checkOutlineAllowance(module, charge - 1, false);
1467 try testing.expectEqual(charge, try checkOutlineAllowance(module, charge, true));
1468 try testing.expectEqual(charge, try checkOutlineAllowance(module, charge + 1, true));
1469 }
1470
1471 fn checkOutlineAllowance(module: *semantic.SemanticModule, visits: u64, success: bool) !u64 {
1472 const ledger = try outlinePassLedger(visits);
1473 defer ledger.destroy();
1474 var cache = try passes.AnalysisCache.initAccounted(testing.allocator, null, ledger, .{}, 8);
1475 defer cache.deinit();
1476 var manager = passes.PassManager.init(testing.allocator);
1477 defer manager.deinit();
1478 try manager.addPass(kernelOutliningPlanningPass());
1479 const result = manager.runWithAnalysisCache(module.choir_module, module.context(), &cache, .{});
1480 var ctx = passes.PassContext.init(
1481 module.choir_module,
1482 module.context(),
1483 testing.allocator,
1484 &cache,
1485 );
1486 defer ctx.deinit();
1487 if (success) {
1488 try testing.expectEqual(passes.PassResult.success, result);
1489 try ledger.producersComplete();
1490 const before = ledger.view().charged;
1491 const outlines = try getKernelOutlinePlanAnalysis(&ctx, module.choir_module);
1492 try testing.expectEqual(@as(usize, 2), outlines.kernelCount());
1493 try testing.expectEqualDeep(before, ledger.view().charged);
1494 try testing.expect(ledger.view().executed.counters.analysis_hits > 0);
1495 } else {
1496 try testing.expectEqual(passes.PassResult.failure, result);
1497 try testing.expectEqual(.exhausted, ledger.view().outcome);
1498 try testing.expectError(
1499 error.TerminalWorkOutcome,
1500 getKernelOutlinePlanAnalysis(&ctx, module.choir_module),
1501 );
1502 }
1503 return ledger.view().charged.structural_visits;
1504 }
1505
1506 test "kernel outlining parallel input failure returns after workers are joined" {
1507 var builder = try semantic.Builder.init(
1508 testing.allocator,
1509 semantic.Builder.ContextLimits.standard,
1510 );
1511 defer builder.deinit();
1512 const module = try buildOutlineTestModule(&builder, 2);
1513 defer module.deinit();
1514 var cache = passes.AnalysisCache.init(testing.allocator, null);
1515 defer cache.deinit();
1516 var ctx = passes.PassContext.initWithOptions(
1517 module.choir_module,
1518 module.context(),
1519 testing.allocator,
1520 &cache,
1521 .{ .max_threads = 2 },
1522 );
1523 defer ctx.deinit();
1524 const schedule = try schedule_planning.getSchedulePlanAnalysis(&ctx, module.choir_module);
1525 try testing.expectEqual(@as(usize, 2), schedule.work_items.items.len);
1526 var missing = bufferization.BufferPlanAnalysis.init(testing.allocator);
1527 defer missing.deinit();
1528 var analysis = KernelOutlinePlanAnalysis.init(testing.allocator);
1529 defer analysis.deinit();
1530 try testing.expectError(error.UnsupportedOperation, addKernelOutlinesParallel(
1531 &ctx,
1532 &analysis,
1533 schedule.work_items.items,
1534 &missing,
1535 ));
1536 try testing.expectEqual(@as(usize, 0), analysis.kernelCount());
1537 }