lib/accy/src/preparation/bufferization/pass.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3 const choir = @import("choir");
4 const accy_root = @import("../../root.zig");
5 const accy_choir = @import("../../choir/root.zig");
6 const dialect_mod = accy_choir.dialect;
7 const fusion = @import("../fusion/root.zig");
8 const shape_analysis = @import("../shape/root.zig");
9
10 const ir = choir.ir;
11 const passes = choir.passes;
12 const accounting = passes.pass.work;
13
14 pub const buffer_plan_analysis_name = "accy-choir-buffer-plan";
15 pub const bufferization_planning_pass_name = "accy-choir-plan-buffers";
16 pub const bufferization_planning_pass_description =
17 "Plan Accy Choir tensor buffer slots before lowering";
18
19 pub const BufferRole = accy_choir.record.memory.BufferRole;
20
21 pub const BufferSlot = struct {
22 id: usize,
23 value: *ir.Value,
24 producer: ?*ir.Operation,
25 function: ?*ir.Operation,
26 role: BufferRole,
27 dtype: choir_abi.DType,
28 dims: []i64,
29 element_count: ?u64,
30 row_major_strides: ?[]u64,
31 byte_size: ?u64,
32
33 fn init(
34 allocator: std.mem.Allocator,
35 id: usize,
36 value: *ir.Value,
37 producer: ?*ir.Operation,
38 function: ?*ir.Operation,
39 role: BufferRole,
40 info: shape_analysis.TensorInfo,
41 ) !BufferSlot {
42 const dims = try allocator.dupe(i64, info.dims);
43 errdefer allocator.free(dims);
44
45 var strides: ?[]u64 = null;
46 if (info.row_major_strides) |existing| {
47 strides = try allocator.dupe(u64, existing);
48 errdefer if (strides) |owned| allocator.free(owned);
49 }
50
51 return .{
52 .id = id,
53 .value = value,
54 .producer = producer,
55 .function = function,
56 .role = role,
57 .dtype = info.dtype,
58 .dims = dims,
59 .element_count = info.element_count,
60 .row_major_strides = strides,
61 .byte_size = byteSize(info),
62 };
63 }
64
65 pub fn hasStaticSize(self: BufferSlot) bool {
66 return self.byte_size != null;
67 }
68
69 fn deinit(self: *BufferSlot, allocator: std.mem.Allocator) void {
70 allocator.free(self.dims);
71 if (self.row_major_strides) |strides| allocator.free(strides);
72 self.* = undefined;
73 }
74 };
75
76 pub const FusionElision = struct {
77 value: *ir.Value,
78 producer: *ir.Operation,
79 root: *ir.Operation,
80 cluster_index: usize,
81 };
82
83 pub const BufferPlanAnalysis = struct {
84 allocator: std.mem.Allocator,
85 slots: std.ArrayListUnmanaged(BufferSlot),
86 elisions: std.ArrayListUnmanaged(FusionElision),
87 value_to_slot: std.AutoHashMap(*ir.Value, usize),
88 value_to_elision: std.AutoHashMap(*ir.Value, usize),
89 input_slot_count: usize = 0,
90 output_slot_count: usize = 0,
91 temporary_slot_count: usize = 0,
92 constant_slot_count: usize = 0,
93 dynamic_slot_count: usize = 0,
94 total_static_bytes: u64 = 0,
95
96 pub fn init(allocator: std.mem.Allocator) BufferPlanAnalysis {
97 return .{
98 .allocator = allocator,
99 .slots = .empty,
100 .elisions = .empty,
101 .value_to_slot = std.AutoHashMap(*ir.Value, usize).init(allocator),
102 .value_to_elision = std.AutoHashMap(*ir.Value, usize).init(allocator),
103 };
104 }
105
106 pub fn deinit(self: *BufferPlanAnalysis) void {
107 for (self.slots.items) |*slot| {
108 slot.deinit(self.allocator);
109 }
110 self.slots.deinit(self.allocator);
111 self.elisions.deinit(self.allocator);
112 self.value_to_slot.deinit();
113 self.value_to_elision.deinit();
114 self.* = undefined;
115 }
116
117 pub fn slotCount(self: BufferPlanAnalysis) usize {
118 return self.slots.items.len;
119 }
120
121 pub fn elisionCount(self: BufferPlanAnalysis) usize {
122 return self.elisions.items.len;
123 }
124
125 pub fn getSlot(self: *const BufferPlanAnalysis, value: *ir.Value) ?*const BufferSlot {
126 const index = self.value_to_slot.get(value) orelse return null;
127 return &self.slots.items[index];
128 }
129
130 pub fn getElision(self: *const BufferPlanAnalysis, value: *ir.Value) ?*const FusionElision {
131 const index = self.value_to_elision.get(value) orelse return null;
132 return &self.elisions.items[index];
133 }
134
135 fn addSlot(
136 self: *BufferPlanAnalysis,
137 value: *ir.Value,
138 producer: ?*ir.Operation,
139 function: ?*ir.Operation,
140 role: BufferRole,
141 info: shape_analysis.TensorInfo,
142 ) !usize {
143 if (self.value_to_slot.get(value)) |existing| return existing;
144
145 const total_bytes = try accounting.add(self.total_static_bytes, byteSize(info) orelse 0);
146 const id = self.slots.items.len;
147 var slot = try BufferSlot.init(
148 self.allocator,
149 id,
150 value,
151 producer,
152 function,
153 role,
154 info,
155 );
156 errdefer slot.deinit(self.allocator);
157
158 try self.value_to_slot.put(value, id);
159 errdefer _ = self.value_to_slot.remove(value);
160 try self.slots.append(self.allocator, slot);
161
162 if (role.input) self.input_slot_count += 1;
163 if (role.output) self.output_slot_count += 1;
164 if (role.temporary) self.temporary_slot_count += 1;
165 if (role.constant) self.constant_slot_count += 1;
166 self.total_static_bytes = total_bytes;
167 if (slot.byte_size == null) self.dynamic_slot_count += 1;
168 return id;
169 }
170
171 fn addAlias(
172 self: *BufferPlanAnalysis,
173 value: *ir.Value,
174 source: *ir.Value,
175 role: BufferRole,
176 ) !usize {
177 if (self.value_to_slot.get(value)) |existing| return existing;
178
179 const id = self.value_to_slot.get(source) orelse return error.MissingAliasedBufferSlot;
180 try self.value_to_slot.put(value, id);
181 self.mergeSlotRole(id, role);
182 return id;
183 }
184
185 fn mergeSlotRole(self: *BufferPlanAnalysis, id: usize, role: BufferRole) void {
186 var slot = &self.slots.items[id];
187 if (role.input and !slot.role.input) {
188 slot.role.input = true;
189 self.input_slot_count += 1;
190 }
191 if (role.output and !slot.role.output) {
192 slot.role.output = true;
193 self.output_slot_count += 1;
194 }
195 if (role.constant and !slot.role.constant) {
196 slot.role.constant = true;
197 self.constant_slot_count += 1;
198 }
199 if (role.temporary and !slot.role.temporary and !slot.role.input and !slot.role.output and !slot.role.constant) {
200 slot.role.temporary = true;
201 self.temporary_slot_count += 1;
202 }
203 if ((slot.role.input or slot.role.output or slot.role.constant) and slot.role.temporary) {
204 slot.role.temporary = false;
205 self.temporary_slot_count -= 1;
206 }
207 }
208
209 fn addElision(self: *BufferPlanAnalysis, elision: FusionElision) !void {
210 if (self.value_to_elision.contains(elision.value)) return;
211 const index = self.elisions.items.len;
212 try self.value_to_elision.put(elision.value, index);
213 errdefer _ = self.value_to_elision.remove(elision.value);
214 try self.elisions.append(self.allocator, elision);
215 }
216 };
217
218 const BufferWork = struct {
219 input: accounting.Census,
220
221 fn storage(self: BufferWork) !u64 {
222 const count = self.input.values;
223 var bytes: u64 = @sizeOf(BufferPlanAnalysis) + @alignOf(BufferPlanAnalysis);
224 bytes = try accounting.add(bytes, try accounting.arrayListGrowth(BufferSlot, count));
225 bytes = try accounting.add(bytes, try accounting.arrayListGrowth(FusionElision, count));
226 const index = try accounting.hashMapGrowth(*ir.Value, usize, count);
227 bytes = try accounting.add(bytes, try accounting.multiply(2, index));
228 bytes = try accounting.add(
229 bytes,
230 try accounting.hashMapGrowth(*ir.Value, void, self.input.operands),
231 );
232 bytes = try accounting.add(
233 bytes,
234 try accounting.hashMapGrowth(*ir.Value, FusionElision, self.input.operations),
235 );
236 const copies = try accounting.multiply(self.input.input_bytes, 2 * @sizeOf(u64));
237 bytes = try accounting.add(bytes, copies);
238 bytes = try accounting.add(bytes, try accounting.multiply(count, 2 * @alignOf(u64)));
239 if (bytes > std.math.maxInt(usize)) return error.WorkOverflow;
240 return bytes;
241 }
242
243 fn bounds(self: BufferWork) !accounting.Bounds {
244 const bytes = try self.storage();
245 const input_units = try accounting.add(self.input.atoms, self.input.input_bytes);
246 const units = try accounting.add(input_units, 1);
247 const operations = self.input.operations;
248 const copies = try accounting.multiply(7, try accounting.multiply(operations, operations));
249 const scans = try accounting.add(try accounting.add(copies, operations), units);
250 const entries = try accounting.add(
251 operations,
252 try accounting.add(self.input.values, self.input.operands),
253 );
254 const inner = try accounting.add(try accounting.hashMapCapacity(entries), units);
255 const visits = try accounting.multiply(64, try accounting.multiply(scans, inner));
256 return .{
257 .work = .{
258 .input_bytes = self.input.input_bytes,
259 .structural_visits = visits,
260 .analysis_computations = 1,
261 .allocation_capacity = bytes,
262 },
263 .workspace = bytes,
264 .retained_storage = bytes,
265 };
266 }
267 };
268
269 fn bufferAnalysisWork(input: accounting.Input) !accounting.Bounds {
270 return (BufferWork{ .input = try accounting.Census.inspect(input.operation) }).bounds();
271 }
272
273 fn bufferPassWork(_: accounting.Input) !accounting.Bounds {
274 return .{ .work = .{ .structural_visits = 1 } };
275 }
276
277 pub const buffer_plan_analysis_descriptor = passes.AnalysisDescriptor{
278 .id = passes.analysisId(buffer_plan_analysis_name),
279 .name = buffer_plan_analysis_name,
280 .work_contract = .{
281 .identity = .{ .name = buffer_plan_analysis_name, .version = 1 },
282 .estimate = bufferAnalysisWork,
283 },
284 };
285
286 pub fn getBufferPlanAnalysis(
287 pass_ctx: *passes.PassContext,
288 op: *ir.Operation,
289 ) !*BufferPlanAnalysis {
290 const ptr = try pass_ctx.getAnalysis(
291 op,
292 &buffer_plan_analysis_descriptor,
293 computeBufferPlanAnalysis,
294 cleanupBufferPlanAnalysis,
295 );
296 return @ptrCast(@alignCast(ptr));
297 }
298
299 pub fn bufferizationPlanningPass() passes.Pass {
300 return .{
301 .name = bufferization_planning_pass_name,
302 .description = bufferization_planning_pass_description,
303 .run_fn = runBufferizationPlanningPass,
304 .work_contract = .{
305 .identity = .{ .name = bufferization_planning_pass_name, .version = 1 },
306 .estimate = bufferPassWork,
307 },
308 };
309 }
310
311 fn runBufferizationPlanningPass(pass_ctx: *passes.PassContext) passes.PassResult {
312 _ = getBufferPlanAnalysis(pass_ctx, pass_ctx.op) catch return .failure;
313 pass_ctx.preserveAllAnalyses();
314 return .success;
315 }
316
317 fn computeBufferPlanAnalysis(
318 pass_ctx: *passes.PassContext,
319 op: *ir.Operation,
320 ) anyerror!*anyopaque {
321 const shapes = try shape_analysis.getShapeLayoutAnalysis(pass_ctx, op);
322 const fusion_plan = try fusion.getFusionPlanAnalysis(pass_ctx, op);
323
324 const analysis = try pass_ctx.allocator.create(BufferPlanAnalysis);
325 analysis.* = BufferPlanAnalysis.init(pass_ctx.allocator);
326 errdefer {
327 analysis.deinit();
328 pass_ctx.allocator.destroy(analysis);
329 }
330
331 var returned = std.AutoHashMap(*ir.Value, void).init(pass_ctx.allocator);
332 defer returned.deinit();
333 try collectReturnedValues(op, &returned);
334
335 var elided_values = std.AutoHashMap(*ir.Value, FusionElision).init(pass_ctx.allocator);
336 defer elided_values.deinit();
337 try collectFusionElisions(fusion_plan, &elided_values);
338
339 try collectBufferSlots(
340 op,
341 null,
342 shapes,
343 &returned,
344 &elided_values,
345 analysis,
346 );
347
348 return @ptrCast(analysis);
349 }
350
351 fn cleanupBufferPlanAnalysis(ptr: *anyopaque, allocator: std.mem.Allocator) void {
352 const analysis: *BufferPlanAnalysis = @ptrCast(@alignCast(ptr));
353 analysis.deinit();
354 allocator.destroy(analysis);
355 }
356
357 fn collectReturnedValues(
358 op: *ir.Operation,
359 returned: *std.AutoHashMap(*ir.Value, void),
360 ) !void {
361 if (isReturnOp(op)) {
362 for (op.getOperandValues()) |value| {
363 try returned.put(value, {});
364 }
365 }
366
367 for (op.regions.items) |*region| {
368 var block_iter = region.getBlocks();
369 while (block_iter.next()) |block| {
370 var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
371 while (current) |current_op| {
372 try collectReturnedValues(current_op, returned);
373 current = current_op.next_op;
374 }
375 }
376 }
377 }
378
379 fn collectFusionElisions(
380 fusion_plan: *const fusion.FusionPlanAnalysis,
381 elided_values: *std.AutoHashMap(*ir.Value, FusionElision),
382 ) !void {
383 for (fusion_plan.clusters.items, 0..) |cluster, cluster_index| {
384 const root = cluster.root() orelse continue;
385 if (cluster.ops.len < 2) continue;
386 for (cluster.ops[0 .. cluster.ops.len - 1]) |producer| {
387 const value = producer.getResult(0) orelse continue;
388 try elided_values.put(value, .{
389 .value = value,
390 .producer = producer,
391 .root = root,
392 .cluster_index = cluster_index,
393 });
394 }
395 }
396
397 for (fusion_plan.elided.items) |producer| {
398 const value = producer.getResult(0) orelse continue;
399 try elided_values.put(value, .{
400 .value = value,
401 .producer = producer,
402 .root = producer,
403 .cluster_index = fusion_plan.clusters.items.len,
404 });
405 }
406 }
407
408 fn collectBufferSlots(
409 op: *ir.Operation,
410 current_function: ?*ir.Operation,
411 shapes: *const shape_analysis.ShapeLayoutAnalysis,
412 returned: *const std.AutoHashMap(*ir.Value, void),
413 elided_values: *const std.AutoHashMap(*ir.Value, FusionElision),
414 analysis: *BufferPlanAnalysis,
415 ) !void {
416 const function = if (isFuncOp(op)) op else current_function;
417
418 if (!isFuncOp(op)) {
419 for (op.results.items, 0..) |*result, result_index| {
420 try recordTensorValue(
421 result,
422 op,
423 function,
424 false,
425 result_index,
426 shapes,
427 returned,
428 elided_values,
429 analysis,
430 );
431 }
432 }
433
434 for (op.regions.items) |*region| {
435 var block_iter = region.getBlocks();
436 while (block_iter.next()) |block| {
437 for (block.arguments.items) |arg| {
438 try recordTensorValue(
439 arg,
440 null,
441 function,
442 isFunctionEntryBlock(function, block),
443 null,
444 shapes,
445 returned,
446 elided_values,
447 analysis,
448 );
449 }
450 var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
451 while (current) |current_op| {
452 try collectBufferSlots(
453 current_op,
454 function,
455 shapes,
456 returned,
457 elided_values,
458 analysis,
459 );
460 current = current_op.next_op;
461 }
462 }
463 }
464 }
465
466 fn recordTensorValue(
467 value: *ir.Value,
468 producer: ?*ir.Operation,
469 function: ?*ir.Operation,
470 is_function_argument: bool,
471 producer_result_index: ?usize,
472 shapes: *const shape_analysis.ShapeLayoutAnalysis,
473 returned: *const std.AutoHashMap(*ir.Value, void),
474 elided_values: *const std.AutoHashMap(*ir.Value, FusionElision),
475 analysis: *BufferPlanAnalysis,
476 ) !void {
477 const info = shapes.get(value) orelse return;
478 if (producer) |op| {
479 if (producer_result_index) |result_index| {
480 if (try aliasedOperandValueForResult(op, result_index)) |source| {
481 _ = try analysis.addAlias(
482 value,
483 source,
484 roleForAliasedResult(value, returned),
485 );
486 return;
487 }
488 }
489 }
490
491 if (!returned.contains(value)) {
492 if (elided_values.get(value)) |elision| {
493 try analysis.addElision(elision);
494 return;
495 }
496 }
497
498 _ = try analysis.addSlot(
499 value,
500 producer,
501 function,
502 roleForValue(value, producer, returned, is_function_argument),
503 info,
504 );
505 }
506
507 fn roleForAliasedResult(
508 value: *ir.Value,
509 returned: *const std.AutoHashMap(*ir.Value, void),
510 ) BufferRole {
511 var role = BufferRole{};
512 if (returned.contains(value)) role.output = true;
513 return role;
514 }
515
516 fn aliasedOperandValueForResult(producer: *ir.Operation, result_index: usize) !?*ir.Value {
517 if (!isName(producer.name.name, dialect_mod.AccyDialect.KernelCallOp.operation_name)) return null;
518 const attr = producer.getAttr("result_aliases") orelse return error.InvalidKernelContract;
519 if (!std.mem.eql(u8, attr.abstract.name, dialect_mod.AccyDialect.KernelCallOp.result_aliases_attr_name)) {
520 return error.InvalidKernelContract;
521 }
522 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return error.InvalidKernelContract;
523 if (dialect_attr.payload.len % @sizeOf(i64) != 0) return error.InvalidKernelContract;
524
525 const alias_count = dialect_attr.payload.len / @sizeOf(i64);
526 if (alias_count != producer.getNumResults()) return error.InvalidKernelContract;
527 if (result_index >= alias_count) return error.InvalidKernelContract;
528
529 var alias_value: i64 = undefined;
530 const start = result_index * @sizeOf(i64);
531 @memcpy(std.mem.asBytes(&alias_value), dialect_attr.payload[start..][0..@sizeOf(i64)]);
532 if (alias_value == -1) return null;
533 if (alias_value < 0) return error.InvalidKernelContract;
534
535 const operand_index: usize = @intCast(alias_value);
536 return producer.getOperand(operand_index) orelse error.InvalidKernelContract;
537 }
538
539 fn roleForValue(
540 value: *ir.Value,
541 producer: ?*ir.Operation,
542 returned: *const std.AutoHashMap(*ir.Value, void),
543 is_function_argument: bool,
544 ) BufferRole {
545 var role = BufferRole{};
546 if (is_function_argument) role.input = true;
547 if (returned.contains(value)) role.output = true;
548 if (producer) |op| {
549 if (isName(op.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name)) {
550 role.constant = true;
551 }
552 }
553 if (!role.input and !role.output and !role.constant) {
554 role.temporary = true;
555 }
556 return role;
557 }
558
559 fn isFunctionEntryBlock(function: ?*ir.Operation, block: *ir.Block) bool {
560 const func = function orelse return false;
561 if (!isFuncOp(func)) return false;
562 const region = func.getRegion(0) orelse return false;
563 const entry = region.getEntryBlock() orelse return false;
564 return entry == block;
565 }
566
567 fn byteSize(info: shape_analysis.TensorInfo) ?u64 {
568 const elements = info.element_count orelse return null;
569 return std.math.mul(u64, elements, @as(u64, info.dtype.sizeOf())) catch null;
570 }
571
572 fn isFuncOp(op: *ir.Operation) bool {
573 return isName(op.name.name, "func.func");
574 }
575
576 fn isReturnOp(op: *ir.Operation) bool {
577 return isName(op.name.name, "func.return") or
578 isName(op.name.name, dialect_mod.AccyDialect.ReturnOp.operation_name);
579 }
580
581 fn isName(actual: []const u8, expected: []const u8) bool {
582 return std.mem.eql(u8, actual, expected);
583 }
584
585 const testing = std.testing;
586 const semantic = accy_choir.semantic;
587
588 fn findOpNamedInBlock(block: *ir.Block, name: []const u8) ?*ir.Operation {
589 var iter = block.operations.head;
590 while (iter) |op_ptr| {
591 const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
592 if (isName(op.name.name, name)) return op;
593 iter = op.next_op;
594 }
595 return null;
596 }
597
598 test "bufferization planning records function boundary slots" {
599 const allocator = testing.allocator;
600
601 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
602 defer builder.deinit();
603 const f32_4 = try builder.tensor(.f32, &.{4});
604 var fb = try builder.beginFunction("buffer_add4", &.{ f32_4, f32_4 }, &.{f32_4});
605 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
606 try fb.return_(&.{sum});
607 try fb.finish();
608 const module = try builder.finish();
609 defer module.deinit();
610
611 const choir_mod = module.choir_module;
612 const ctx = module.context();
613 const ledger = try bufferTestLedger();
614 defer ledger.destroy();
615 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
616 defer cache.deinit();
617 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
618 defer pass_ctx.deinit();
619
620 const analysis = try getBufferPlanAnalysis(&pass_ctx, choir_mod);
621 try ledger.producersComplete();
622 try checkBufferStorage(choir_mod, ctx, &cache, analysis);
623 try testing.expectEqual(@as(usize, 3), analysis.slotCount());
624 try testing.expectEqual(@as(usize, 2), analysis.input_slot_count);
625 try testing.expectEqual(@as(usize, 1), analysis.output_slot_count);
626 try testing.expectEqual(@as(usize, 0), analysis.temporary_slot_count);
627 try testing.expectEqual(@as(usize, 0), analysis.elisionCount());
628 try testing.expectEqual(@as(u64, 48), analysis.total_static_bytes);
629
630 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
631 const func = ir.inspection.functionByNameInBlock(body, "buffer_add4") orelse return error.TestExpectedFunc;
632 const entry = func.getRegion(0).?.getEntryBlock().?;
633 const add = findOpNamedInBlock(entry, dialect_mod.AccyDialect.AddOp.operation_name) orelse return error.TestExpectedAdd;
634 const add_slot = analysis.getSlot(add.getResult(0).?) orelse return error.TestExpectedSlot;
635 try testing.expect(add_slot.role.output);
636 try testing.expect(!add_slot.role.temporary);
637 try testing.expectEqual(@as(?u64, 16), add_slot.byte_size);
638 try testing.expectEqualSlices(i64, &.{4}, add_slot.dims);
639 try testing.expectEqualSlices(u64, &.{1}, add_slot.row_major_strides.?);
640 }
641
642 test "bufferization planning elides fusion-internal tensor values" {
643 const allocator = testing.allocator;
644
645 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
646 defer builder.deinit();
647 const f32_4 = try builder.tensor(.f32, &.{4});
648 var fb = try builder.beginFunction("buffer_fused_add_mul", &.{ f32_4, f32_4, f32_4 }, &.{f32_4});
649 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
650 const product = try fb.mul(sum, fb.parameter(2));
651 try fb.return_(&.{product});
652 try fb.finish();
653 const module = try builder.finish();
654 defer module.deinit();
655
656 const choir_mod = module.choir_module;
657 const ctx = module.context();
658 const ledger = try bufferTestLedger();
659 defer ledger.destroy();
660 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
661 defer cache.deinit();
662 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
663 defer pass_ctx.deinit();
664
665 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
666 const func = ir.inspection.functionByNameInBlock(body, "buffer_fused_add_mul") orelse return error.TestExpectedFunc;
667 const entry = func.getRegion(0).?.getEntryBlock().?;
668 const add = findOpNamedInBlock(entry, dialect_mod.AccyDialect.AddOp.operation_name) orelse return error.TestExpectedAdd;
669 const mul = findOpNamedInBlock(entry, dialect_mod.AccyDialect.MulOp.operation_name) orelse return error.TestExpectedMul;
670
671 const analysis = try getBufferPlanAnalysis(&pass_ctx, choir_mod);
672 try ledger.producersComplete();
673 try checkBufferStorage(choir_mod, ctx, &cache, analysis);
674 try testing.expectEqual(@as(usize, 4), analysis.slotCount());
675 try testing.expectEqual(@as(usize, 3), analysis.input_slot_count);
676 try testing.expectEqual(@as(usize, 1), analysis.output_slot_count);
677 try testing.expectEqual(@as(usize, 1), analysis.elisionCount());
678 try testing.expectEqual(@as(u64, 64), analysis.total_static_bytes);
679 try testing.expect(analysis.getSlot(add.getResult(0).?) == null);
680
681 const elision = analysis.getElision(add.getResult(0).?) orelse return error.TestExpectedElision;
682 try testing.expectEqual(add, elision.producer);
683 try testing.expectEqual(mul, elision.root);
684
685 const root_slot = analysis.getSlot(mul.getResult(0).?) orelse return error.TestExpectedRootSlot;
686 try testing.expect(root_slot.role.output);
687 }
688
689 test "bufferization planning coalesces aliased kernel call results" {
690 const allocator = testing.allocator;
691
692 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
693 defer builder.deinit();
694 const f32_4 = try builder.tensor(.f32, &.{4});
695 var fb = try builder.beginFunction("buffer_kernel_alias", &.{f32_4}, &.{f32_4});
696 const call = try fb.kernelCall(&.{fb.parameter(0)}, &.{f32_4}, .{
697 .target = "update_f32",
698 .operand_effects = &.{.read_write},
699 .result_aliases = &.{0},
700 });
701 try fb.return_(&.{call.getFirstResult()});
702 try fb.finish();
703 const module = try builder.finish();
704 defer module.deinit();
705
706 const choir_mod = module.choir_module;
707 const ctx = module.context();
708 const ledger = try bufferTestLedger();
709 defer ledger.destroy();
710 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
711 defer cache.deinit();
712 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
713 defer pass_ctx.deinit();
714
715 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
716 const func = ir.inspection.functionByNameInBlock(body, "buffer_kernel_alias") orelse return error.TestExpectedFunc;
717 const entry = func.getRegion(0).?.getEntryBlock().?;
718 const call_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse return error.TestExpectedKernelCall;
719
720 const analysis = try getBufferPlanAnalysis(&pass_ctx, choir_mod);
721 try ledger.producersComplete();
722 try checkBufferStorage(choir_mod, ctx, &cache, analysis);
723 try testing.expectEqual(@as(usize, 1), analysis.slotCount());
724 try testing.expectEqual(@as(usize, 1), analysis.input_slot_count);
725 try testing.expectEqual(@as(usize, 1), analysis.output_slot_count);
726 try testing.expectEqual(@as(usize, 0), analysis.temporary_slot_count);
727 try testing.expectEqual(@as(u64, 16), analysis.total_static_bytes);
728
729 const input_value = entry.getArgument(0).?;
730 const input_slot = analysis.getSlot(input_value) orelse return error.TestExpectedInputSlot;
731 const result_slot = analysis.getSlot(call_op.getResult(0).?) orelse return error.TestExpectedResultSlot;
732 try testing.expectEqual(input_slot.id, result_slot.id);
733 try testing.expect(input_slot.role.input);
734 try testing.expect(input_slot.role.output);
735 try testing.expect(!input_slot.role.temporary);
736 }
737
738 test "bufferization planning keeps non-returned aliased results internal" {
739 const allocator = testing.allocator;
740
741 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
742 defer builder.deinit();
743 const f32_4 = try builder.tensor(.f32, &.{4});
744 var fb = try builder.beginFunction("buffer_kernel_alias_internal", &.{ f32_4, f32_4 }, &.{f32_4});
745 const call = try fb.kernelCall(&.{fb.parameter(0)}, &.{f32_4}, .{
746 .target = "update_f32",
747 .operand_effects = &.{.read_write},
748 .result_aliases = &.{0},
749 });
750 const sum = try fb.add(call.getFirstResult(), fb.parameter(1));
751 try fb.return_(&.{sum});
752 try fb.finish();
753 const module = try builder.finish();
754 defer module.deinit();
755
756 const choir_mod = module.choir_module;
757 const ctx = module.context();
758 const ledger = try bufferTestLedger();
759 defer ledger.destroy();
760 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
761 defer cache.deinit();
762 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
763 defer pass_ctx.deinit();
764
765 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
766 const func = ir.inspection.functionByNameInBlock(body, "buffer_kernel_alias_internal") orelse return error.TestExpectedFunc;
767 const entry = func.getRegion(0).?.getEntryBlock().?;
768 const call_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse return error.TestExpectedKernelCall;
769 const add_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.AddOp.operation_name) orelse return error.TestExpectedAdd;
770
771 const analysis = try getBufferPlanAnalysis(&pass_ctx, choir_mod);
772 try ledger.producersComplete();
773 try checkBufferStorage(choir_mod, ctx, &cache, analysis);
774 try testing.expectEqual(@as(usize, 3), analysis.slotCount());
775 try testing.expectEqual(@as(usize, 2), analysis.input_slot_count);
776 try testing.expectEqual(@as(usize, 1), analysis.output_slot_count);
777 try testing.expectEqual(@as(usize, 0), analysis.temporary_slot_count);
778 try testing.expectEqual(@as(u64, 48), analysis.total_static_bytes);
779
780 const input_value = entry.getArgument(0).?;
781 const input_slot = analysis.getSlot(input_value) orelse return error.TestExpectedInputSlot;
782 const result_slot = analysis.getSlot(call_op.getResult(0).?) orelse return error.TestExpectedResultSlot;
783 const add_slot = analysis.getSlot(add_op.getResult(0).?) orelse return error.TestExpectedAddSlot;
784 try testing.expectEqual(input_slot.id, result_slot.id);
785 try testing.expect(input_slot.role.input);
786 try testing.expect(!input_slot.role.output);
787 try testing.expect(!input_slot.role.temporary);
788 try testing.expect(add_slot.role.output);
789 }
790
791 test "bufferization planning records returned constants" {
792 const allocator = testing.allocator;
793
794 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
795 defer builder.deinit();
796 const i32_4 = try builder.tensor(.i32, &.{4});
797 var fb = try builder.beginFunction("buffer_const", &.{}, &.{i32_4});
798 const values = [_]i32{ 1, 2, 3, 4 };
799 const c = try fb.constant(i32_4, std.mem.sliceAsBytes(values[0..]));
800 try fb.return_(&.{c});
801 try fb.finish();
802 const module = try builder.finish();
803 defer module.deinit();
804
805 const choir_mod = module.choir_module;
806 const ctx = module.context();
807 const ledger = try bufferTestLedger();
808 defer ledger.destroy();
809 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
810 defer cache.deinit();
811 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
812 defer pass_ctx.deinit();
813
814 const analysis = try getBufferPlanAnalysis(&pass_ctx, choir_mod);
815 try ledger.producersComplete();
816 try checkBufferStorage(choir_mod, ctx, &cache, analysis);
817 try testing.expectEqual(@as(usize, 1), analysis.slotCount());
818 try testing.expectEqual(@as(usize, 1), analysis.output_slot_count);
819 try testing.expectEqual(@as(usize, 1), analysis.constant_slot_count);
820 try testing.expectEqual(@as(u64, 16), analysis.total_static_bytes);
821
822 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
823 const func = ir.inspection.functionByNameInBlock(body, "buffer_const") orelse return error.TestExpectedFunc;
824 const entry = func.getRegion(0).?.getEntryBlock().?;
825 const constant = findOpNamedInBlock(entry, dialect_mod.AccyDialect.ConstantOp.operation_name) orelse return error.TestExpectedConstant;
826 const slot = analysis.getSlot(constant.getResult(0).?) orelse return error.TestExpectedSlot;
827 try testing.expect(slot.role.constant);
828 try testing.expect(slot.role.output);
829 }
830
831 test "bufferization planning pass preserves IR" {
832 const allocator = testing.allocator;
833
834 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
835 defer builder.deinit();
836 const f32_4 = try builder.tensor(.f32, &.{4});
837 var fb = try builder.beginFunction("buffer_pass_add4", &.{ f32_4, f32_4 }, &.{f32_4});
838 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
839 try fb.return_(&.{sum});
840 try fb.finish();
841 const module = try builder.finish();
842 defer module.deinit();
843
844 const choir_mod = module.choir_module;
845 const ctx = module.context();
846 var pm = passes.PassManager.init(allocator);
847 defer pm.deinit();
848 try pm.addPass(bufferizationPlanningPass());
849
850 const ledger = try choir.product.revision.AccountingV1.create(allocator, .{
851 .allowance = choir.product.revision.WorkVector.uniform(std.math.maxInt(u64)),
852 .workspace = std.math.maxInt(u64),
853 .events = 8,
854 }, &.{.{ .name = bufferization_planning_pass_name, .version = 1 }});
855 defer ledger.destroy();
856 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
857 defer cache.deinit();
858 try testing.expectEqual(
859 passes.PassResult.success,
860 pm.runWithAnalysisCache(choir_mod, ctx, &cache, .{}),
861 );
862 try ledger.producersComplete();
863 try testing.expectEqual(@as(u64, 1), pm.stats.pass_runs);
864 try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
865 }
866
867 fn bufferTestLedger() !*choir.product.revision.AccountingV1 {
868 const revision = choir.product.revision;
869 return revision.AccountingV1.create(testing.allocator, .{
870 .allowance = revision.WorkVector.uniform(std.math.maxInt(u64)),
871 .workspace = std.math.maxInt(u64),
872 .events = 8,
873 }, &.{});
874 }
875
876 fn checkBufferStorage(
877 op: *ir.Operation,
878 ctx: *ir.Context,
879 cache: *passes.AnalysisCache,
880 expected: *const BufferPlanAnalysis,
881 ) !void {
882 const bounds = try bufferAnalysisWork(.{ .operation = op });
883 const fixed = @import("alloc_fixed");
884 const bytes = try testing.allocator.alignedAlloc(u8, .@"64", @intCast(bounds.workspace));
885 defer testing.allocator.free(bytes);
886 var storage = fixed.Tracked.init(bytes);
887 var retained = fixed.Monotonic.init(storage.allocator(), @max(1, bytes.len));
888 const allocator = retained.allocator();
889 var pass_ctx = passes.PassContext.init(op, ctx, allocator, cache);
890 defer pass_ctx.deinit();
891 const ptr = try computeBufferPlanAnalysis(&pass_ctx, op);
892 defer cleanupBufferPlanAnalysis(ptr, allocator);
893 const actual: *BufferPlanAnalysis = @ptrCast(@alignCast(ptr));
894 inline for (.{
895 "input_slot_count", "output_slot_count", "temporary_slot_count",
896 "constant_slot_count", "dynamic_slot_count", "total_static_bytes",
897 }) |field| try testing.expectEqual(@field(expected, field), @field(actual, field));
898 try testing.expectEqual(expected.slotCount(), actual.slotCount());
899 try testing.expectEqual(expected.elisionCount(), actual.elisionCount());
900 for (expected.slots.items, actual.slots.items) |left, right| {
901 inline for (.{
902 "id", "value", "producer", "function",
903 "dtype", "element_count", "byte_size",
904 }) |field| {
905 try testing.expectEqual(@field(left, field), @field(right, field));
906 }
907 try testing.expectEqualDeep(left.role, right.role);
908 try testing.expectEqualSlices(i64, left.dims, right.dims);
909 if (left.row_major_strides) |strides| {
910 try testing.expectEqualSlices(u64, strides, right.row_major_strides.?);
911 } else try testing.expect(right.row_major_strides == null);
912 }
913 for (expected.elisions.items, actual.elisions.items) |left, right| {
914 inline for (.{ "value", "producer", "root", "cluster_index" }) |field| {
915 try testing.expectEqual(@field(left, field), @field(right, field));
916 }
917 }
918 try testing.expectEqual(expected.value_to_slot.count(), actual.value_to_slot.count());
919 var slots = expected.value_to_slot.iterator();
920 while (slots.next()) |entry| {
921 try testing.expectEqual(entry.value_ptr.*, actual.getSlot(entry.key_ptr.*).?.id);
922 }
923 try testing.expectEqual(expected.value_to_elision.count(), actual.value_to_elision.count());
924 var elisions = expected.value_to_elision.iterator();
925 while (elisions.next()) |entry| {
926 try testing.expectEqual(entry.value_ptr.*, actual.value_to_elision.get(entry.key_ptr.*).?);
927 }
928 try testing.expect(!storage.exhausted);
929 const used = if (retained.current) |*current| fixed.used(current) else 0;
930 try testing.expect(used <= bounds.workspace);
931 try testing.expect(used >= @sizeOf(BufferPlanAnalysis));
932 }
933
934 test "bufferization planning work contract covers boundary growth and layouts" {
935 for ([_][]const i64{ &.{}, &.{ 2, 3 }, &.{ 1, 1, 1, 1, 1, 1, 1, 1 }, &.{ -1, 3 } }) |dims| {
936 for ([_]usize{ 0, 1, 6, 7, 16, 64 }) |count| try checkBufferBoundary(count, dims);
937 }
938 try testing.expectError(error.WorkOverflow, (BufferWork{
939 .input = .{ .values = std.math.maxInt(u64) },
940 }).bounds());
941 try testing.expectError(error.WorkOverflow, (BufferWork{
942 .input = .{ .operations = std.math.maxInt(u64) },
943 }).bounds());
944 try testing.expectError(error.WorkOverflow, (BufferWork{
945 .input = .{ .input_bytes = std.math.maxInt(u64) },
946 }).bounds());
947 }
948
949 fn checkBufferBoundary(count: usize, dims: []const i64) !void {
950 std.debug.assert(count <= 64);
951 var builder = try semantic.Builder.init(
952 testing.allocator,
953 semantic.Builder.ContextLimits.standard,
954 );
955 defer builder.deinit();
956 const typ = try builder.tensor(.f32, dims);
957 var types: [64]@TypeOf(typ) = @splat(typ);
958 var function = try builder.beginFunction("buffer_boundary", types[0..count], types[0..count]);
959 var values: [64]@TypeOf(function.parameter(0)) = undefined;
960 for (values[0..count], 0..) |*value, index| value.* = function.parameter(index);
961 try function.return_(values[0..count]);
962 try function.finish();
963 const module = try builder.finish();
964 defer module.deinit();
965 const ledger = try bufferTestLedger();
966 defer ledger.destroy();
967 var cache = try passes.AnalysisCache.initAccounted(testing.allocator, null, ledger, .{}, 3);
968 defer cache.deinit();
969 var pass_ctx = passes.PassContext.init(
970 module.choir_module,
971 module.context(),
972 testing.allocator,
973 &cache,
974 );
975 defer pass_ctx.deinit();
976 const analysis = try getBufferPlanAnalysis(&pass_ctx, module.choir_module);
977 try ledger.producersComplete();
978 try testing.expectEqual(count, analysis.slotCount());
979 try testing.expectEqual(count, analysis.input_slot_count);
980 try testing.expectEqual(count, analysis.output_slot_count);
981 try testing.expectEqual(@as(usize, 0), analysis.elisionCount());
982 for (analysis.slots.items) |slot| try testing.expectEqualSlices(i64, dims, slot.dims);
983 try checkBufferStorage(module.choir_module, module.context(), &cache, analysis);
984 try checkBufferAdmission(module.choir_module, module.context());
985 }
986
987 fn checkBufferAdmission(op: *ir.Operation, ctx: *ir.Context) !void {
988 const revision = choir.product.revision;
989 var charge: u64 = 1;
990 for ([_]passes.AnalysisDescriptor{
991 shape_analysis.shape_layout_analysis_descriptor,
992 fusion.fusion_plan_analysis_descriptor,
993 buffer_plan_analysis_descriptor,
994 }) |descriptor| {
995 const bounds = try descriptor.work_contract.?.estimate(.{ .operation = op });
996 charge = try accounting.add(charge, bounds.work.structural_visits);
997 }
998 for ([_]i8{ -1, 0, 1 }) |offset| {
999 var allowance = revision.WorkVector.uniform(std.math.maxInt(u64));
1000 allowance.structural_visits = @intCast(@as(i128, charge) + offset);
1001 const ledger = try revision.AccountingV1.create(testing.allocator, .{
1002 .allowance = allowance,
1003 .workspace = std.math.maxInt(u64),
1004 .events = 8,
1005 }, &.{.{ .name = bufferization_planning_pass_name, .version = 1 }});
1006 defer ledger.destroy();
1007 var cache = try passes.AnalysisCache.initAccounted(
1008 testing.allocator,
1009 null,
1010 ledger,
1011 .{},
1012 3,
1013 );
1014 defer cache.deinit();
1015 var manager = passes.PassManager.init(testing.allocator);
1016 defer manager.deinit();
1017 try manager.addPass(bufferizationPlanningPass());
1018 const result = manager.runWithAnalysisCache(op, ctx, &cache, .{});
1019 if (offset < 0) {
1020 try testing.expectEqual(passes.PassResult.failure, result);
1021 try testing.expectEqual(revision.receipt.Outcome.exhausted, ledger.view().outcome);
1022 try testing.expectEqual(@as(usize, 1), cache.entries.count());
1023 } else {
1024 try testing.expectEqual(passes.PassResult.success, result);
1025 try ledger.producersComplete();
1026 try testing.expectEqual(@as(usize, 3), cache.entries.count());
1027 }
1028 }
1029 }
1030
1031 test "bufferization planning rejects aggregate byte overflow before publishing slots" {
1032 var builder = try semantic.Builder.init(
1033 testing.allocator,
1034 semantic.Builder.ContextLimits.standard,
1035 );
1036 defer builder.deinit();
1037 const extent = std.math.maxInt(i64) / 4;
1038 const typ = try builder.tensor(.f32, &.{extent});
1039 var function = try builder.beginFunction("buffer_overflow", &.{ typ, typ, typ }, &.{typ});
1040 try function.return_(&.{function.parameter(0)});
1041 try function.finish();
1042 const module = try builder.finish();
1043 defer module.deinit();
1044 const ledger = try bufferTestLedger();
1045 defer ledger.destroy();
1046 var cache = try passes.AnalysisCache.initAccounted(testing.allocator, null, ledger, .{}, 3);
1047 defer cache.deinit();
1048 var pass_ctx = passes.PassContext.init(
1049 module.choir_module,
1050 module.context(),
1051 testing.allocator,
1052 &cache,
1053 );
1054 defer pass_ctx.deinit();
1055 try testing.expectError(
1056 error.WorkOverflow,
1057 getBufferPlanAnalysis(&pass_ctx, module.choir_module),
1058 );
1059 try testing.expectEqual(
1060 choir.product.revision.receipt.Outcome.exhausted,
1061 ledger.view().outcome,
1062 );
1063 try testing.expectEqual(@as(usize, 2), cache.entries.count());
1064 try testing.expectError(
1065 error.TerminalWorkOutcome,
1066 getBufferPlanAnalysis(&pass_ctx, module.choir_module),
1067 );
1068 }
1069
1070 test "bufferization planning work contract covers alias and elision growth" {
1071 for ([_]usize{ 1, 2, 6, 7, 16, 64 }) |count| {
1072 try checkBufferChain(count, true);
1073 try checkBufferChain(count, false);
1074 }
1075 }
1076
1077 fn checkBufferChain(count: usize, alias: bool) !void {
1078 std.debug.assert(count > 0);
1079 std.debug.assert(count <= 64);
1080 var builder = try semantic.Builder.init(
1081 testing.allocator,
1082 semantic.Builder.ContextLimits.standard,
1083 );
1084 defer builder.deinit();
1085 const typ = try builder.tensor(.f32, &.{ 2, 3 });
1086 var function = try builder.beginFunction("buffer_chain", &.{ typ, typ }, &.{typ});
1087 var value = function.parameter(0);
1088 for (0..count) |_| {
1089 if (alias) {
1090 const call = try function.kernelCall(&.{value}, &.{typ}, .{
1091 .target = "update_f32",
1092 .operand_effects = &.{.read_write},
1093 .result_aliases = &.{0},
1094 });
1095 value = call.getFirstResult();
1096 } else value = try function.add(value, function.parameter(1));
1097 }
1098 try function.return_(&.{value});
1099 try function.finish();
1100 const module = try builder.finish();
1101 defer module.deinit();
1102 const ledger = try bufferTestLedger();
1103 defer ledger.destroy();
1104 var cache = try passes.AnalysisCache.initAccounted(testing.allocator, null, ledger, .{}, 3);
1105 defer cache.deinit();
1106 var pass_ctx = passes.PassContext.init(
1107 module.choir_module,
1108 module.context(),
1109 testing.allocator,
1110 &cache,
1111 );
1112 defer pass_ctx.deinit();
1113 const analysis = try getBufferPlanAnalysis(&pass_ctx, module.choir_module);
1114 try ledger.producersComplete();
1115 try testing.expectEqual(@as(usize, if (alias) 2 else 3), analysis.slotCount());
1116 try testing.expectEqual(if (alias) 0 else count - 1, analysis.elisionCount());
1117 try testing.expectEqual(if (alias) count + 2 else 3, analysis.value_to_slot.count());
1118 try testing.expectEqual(@as(usize, 1), analysis.output_slot_count);
1119 try testing.expectEqual(@as(usize, 0), analysis.temporary_slot_count);
1120 try checkBufferStorage(module.choir_module, module.context(), &cache, analysis);
1121 }