lib/accy/src/preparation/memory.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("choir");
3 const bufferization = @import("bufferization/root.zig");
4 const schedule = @import("schedule/root.zig");
5 const accy_choir = @import("../choir/root.zig");
6 const dialect_mod = accy_choir.dialect;
7
8 const ir = choir.ir;
9 const passes = choir.passes;
10 const accounting = passes.pass.work;
11
12 pub const memory_space_plan_analysis_name = "accy-choir-memory-space-plan";
13 pub const memory_space_planning_pass_name = "accy-choir-plan-memory-spaces";
14 pub const memory_space_planning_pass_description =
15 "Plan Accy Choir buffer memory spaces and boundary transfers";
16
17 pub const MemorySpace = accy_choir.record.memory.MemorySpace;
18
19 pub const MemoryAccess = accy_choir.record.memory.MemoryAccess;
20
21 pub const BoundaryTransfer = accy_choir.record.memory.BoundaryTransfer;
22
23 /// A caller reads this to know whether a function output is written by a kernel or is an existing
24 /// buffer. The value is either the id of the scheduled unit of work that writes the output, or the
25 /// id of an input or constant buffer that the output is. Planning fails with
26 /// `error.MissingOutputWriter` when an output has neither.
27 pub const OutputSource = accy_choir.record.memory.OutputSource;
28
29 pub const MemorySpaceAssignment = struct {
30 slot_id: usize,
31 value: *ir.Value,
32 producer: ?*ir.Operation,
33 role: bufferization.BufferRole,
34 space: MemorySpace,
35 access: MemoryAccess,
36 transfer: BoundaryTransfer,
37 byte_size: ?u64,
38 output_source: ?OutputSource = null,
39
40 pub fn isDynamic(self: MemorySpaceAssignment) bool {
41 return self.byte_size == null;
42 }
43 };
44
45 pub const MemorySpacePlanAnalysis = struct {
46 allocator: std.mem.Allocator,
47 assignments: std.ArrayListUnmanaged(MemorySpaceAssignment),
48 slot_to_assignment: std.AutoHashMap(usize, usize),
49 host_slot_count: usize = 0,
50 device_global_slot_count: usize = 0,
51 device_constant_slot_count: usize = 0,
52 device_shared_slot_count: usize = 0,
53 unified_slot_count: usize = 0,
54 host_input_transfer_count: usize = 0,
55 host_output_transfer_count: usize = 0,
56 dynamic_slot_count: usize = 0,
57 elided_value_count: usize = 0,
58 total_static_bytes: u64 = 0,
59
60 pub fn init(allocator: std.mem.Allocator) MemorySpacePlanAnalysis {
61 return .{
62 .allocator = allocator,
63 .assignments = .empty,
64 .slot_to_assignment = std.AutoHashMap(usize, usize).init(allocator),
65 };
66 }
67
68 pub fn deinit(self: *MemorySpacePlanAnalysis) void {
69 self.assignments.deinit(self.allocator);
70 self.slot_to_assignment.deinit();
71 self.* = undefined;
72 }
73
74 pub fn assignmentCount(self: MemorySpacePlanAnalysis) usize {
75 return self.assignments.items.len;
76 }
77
78 pub fn getAssignmentForSlot(
79 self: *const MemorySpacePlanAnalysis,
80 slot_id: usize,
81 ) ?*const MemorySpaceAssignment {
82 const index = self.slot_to_assignment.get(slot_id) orelse return null;
83 return &self.assignments.items[index];
84 }
85
86 pub fn getAssignmentForValue(
87 self: *const MemorySpacePlanAnalysis,
88 buffers: *const bufferization.BufferPlanAnalysis,
89 value: *ir.Value,
90 ) ?*const MemorySpaceAssignment {
91 const slot = buffers.getSlot(value) orelse return null;
92 return self.getAssignmentForSlot(slot.id);
93 }
94
95 fn addAssignment(self: *MemorySpacePlanAnalysis, assignment: MemorySpaceAssignment) !void {
96 if (self.slot_to_assignment.contains(assignment.slot_id)) return;
97 const index = self.assignments.items.len;
98 try self.slot_to_assignment.put(assignment.slot_id, index);
99 errdefer _ = self.slot_to_assignment.remove(assignment.slot_id);
100 try self.assignments.append(self.allocator, assignment);
101
102 switch (assignment.space) {
103 .host => self.host_slot_count += 1,
104 .device_global => self.device_global_slot_count += 1,
105 .device_constant => self.device_constant_slot_count += 1,
106 .device_shared => self.device_shared_slot_count += 1,
107 .unified => self.unified_slot_count += 1,
108 }
109 if (assignment.transfer.needsHostInput()) self.host_input_transfer_count += 1;
110 if (assignment.transfer.needsHostOutput()) self.host_output_transfer_count += 1;
111 if (assignment.byte_size) |bytes| {
112 self.total_static_bytes += bytes;
113 } else {
114 self.dynamic_slot_count += 1;
115 }
116 }
117 };
118
119 const missing_output_format =
120 "function @{s} result #{d} (%{d}) has no kernel writer or backing alias; " ++
121 "schedule its producer or provide an input/constant alias";
122
123 const MemoryWork = struct {
124 input: accounting.Census,
125 symbols: u64,
126
127 fn inspect(op: *ir.Operation) !MemoryWork {
128 var result: MemoryWork = .{ .input = try accounting.Census.inspect(op), .symbols = 0 };
129 _ = try op.walk(.{ .order = .pre_order }, &result, visit);
130 return result;
131 }
132
133 fn visit(self: *MemoryWork, op: *ir.Operation) !ir.Operation.WalkResult {
134 if (isName(op.name.name, "func.func")) {
135 const function = choir.dialects.FuncDialect.FuncOp{ .op = op };
136 const name = function.getName() orelse "<unknown>";
137 self.symbols = try accounting.add(self.symbols, name.len);
138 }
139 return .advance;
140 }
141
142 fn diagnosticBytes(self: MemoryWork) !u64 {
143 return accounting.add(self.symbols, missing_output_format.len + 2 * 20 + "<unknown>".len);
144 }
145
146 fn storage(self: MemoryWork) !u64 {
147 var bytes: u64 = @sizeOf(MemorySpacePlanAnalysis) + @alignOf(MemorySpacePlanAnalysis);
148 bytes = try accounting.add(
149 bytes,
150 try accounting.arrayListGrowth(MemorySpaceAssignment, self.input.values),
151 );
152 bytes = try accounting.add(
153 bytes,
154 try accounting.hashMapGrowth(usize, usize, self.input.values),
155 );
156 bytes = try accounting.add(bytes, try self.diagnosticBytes());
157 if (bytes > std.math.maxInt(usize)) return error.WorkOverflow;
158 return bytes;
159 }
160
161 fn bounds(self: MemoryWork) !accounting.Bounds {
162 const bytes = try self.storage();
163 const input_bytes = try accounting.add(self.input.input_bytes, self.symbols);
164 const input_units = try accounting.add(self.input.atoms, input_bytes);
165 const units = try accounting.add(input_units, try self.diagnosticBytes());
166 const entries = try accounting.add(
167 self.input.operations,
168 try accounting.add(self.input.values, self.input.operands),
169 );
170 const probes = try accounting.hashMapCapacity(entries);
171 const scans = try accounting.multiply(units, try accounting.add(self.input.operations, 1));
172 const inner = try accounting.add(try accounting.add(self.input.values, probes), units);
173 const visits = try accounting.multiply(64, try accounting.multiply(scans, inner));
174 return .{
175 .work = .{
176 .input_bytes = input_bytes,
177 .structural_visits = visits,
178 .analysis_computations = 1,
179 .allocation_capacity = bytes,
180 },
181 .workspace = bytes,
182 .retained_storage = bytes,
183 };
184 }
185 };
186
187 fn memoryAnalysisWork(input: accounting.Input) !accounting.Bounds {
188 return (try MemoryWork.inspect(input.operation)).bounds();
189 }
190
191 fn memoryPassWork(_: accounting.Input) !accounting.Bounds {
192 return .{ .work = .{ .structural_visits = 1 } };
193 }
194
195 pub const memory_space_plan_analysis_descriptor = passes.AnalysisDescriptor{
196 .id = passes.analysisId(memory_space_plan_analysis_name),
197 .name = memory_space_plan_analysis_name,
198 .work_contract = .{
199 .identity = .{ .name = memory_space_plan_analysis_name, .version = 1 },
200 .estimate = memoryAnalysisWork,
201 },
202 };
203
204 pub fn getMemorySpacePlanAnalysis(
205 pass_ctx: *passes.PassContext,
206 op: *ir.Operation,
207 ) !*MemorySpacePlanAnalysis {
208 const ptr = try pass_ctx.getAnalysis(
209 op,
210 &memory_space_plan_analysis_descriptor,
211 computeMemorySpacePlanAnalysis,
212 cleanupMemorySpacePlanAnalysis,
213 );
214 return @ptrCast(@alignCast(ptr));
215 }
216
217 pub fn memorySpacePlanningPass() passes.Pass {
218 return .{
219 .name = memory_space_planning_pass_name,
220 .description = memory_space_planning_pass_description,
221 .run_fn = runMemorySpacePlanningPass,
222 .work_contract = .{
223 .identity = .{ .name = memory_space_planning_pass_name, .version = 1 },
224 .estimate = memoryPassWork,
225 },
226 };
227 }
228
229 fn runMemorySpacePlanningPass(pass_ctx: *passes.PassContext) passes.PassResult {
230 _ = getMemorySpacePlanAnalysis(pass_ctx, pass_ctx.op) catch return .failure;
231 pass_ctx.preserveAllAnalyses();
232 return .success;
233 }
234
235 fn computeMemorySpacePlanAnalysis(
236 pass_ctx: *passes.PassContext,
237 op: *ir.Operation,
238 ) anyerror!*anyopaque {
239 const buffers = try bufferization.getBufferPlanAnalysis(pass_ctx, op);
240 const work = try schedule.getSchedulePlanAnalysis(pass_ctx, op);
241
242 const analysis = try pass_ctx.allocator.create(MemorySpacePlanAnalysis);
243 analysis.* = MemorySpacePlanAnalysis.init(pass_ctx.allocator);
244 errdefer {
245 analysis.deinit();
246 pass_ctx.allocator.destroy(analysis);
247 }
248
249 analysis.elided_value_count = buffers.elisionCount();
250 for (buffers.slots.items) |slot| {
251 var assignment = classifySlot(slot);
252 assignment.output_source = try classifyOutput(pass_ctx.allocator, slot, buffers, work);
253 try analysis.addAssignment(assignment);
254 }
255 std.debug.assert(analysis.assignmentCount() == buffers.slotCount());
256 std.debug.assert(analysis.total_static_bytes == buffers.total_static_bytes);
257
258 return @ptrCast(analysis);
259 }
260
261 fn classifyOutput(
262 allocator: std.mem.Allocator,
263 slot: bufferization.BufferSlot,
264 buffers: *const bufferization.BufferPlanAnalysis,
265 work: *const schedule.SchedulePlanAnalysis,
266 ) !?OutputSource {
267 if (!slot.role.output) return null;
268 for (work.work_items.items) |item| {
269 if (workWritesSlot(item, buffers, slot.id)) return .{ .kernel_written = item.id };
270 }
271 if (slot.role.input or slot.role.constant) return .{ .aliased = slot.id };
272 try diagnoseMissingOutput(allocator, slot);
273 return error.MissingOutputWriter;
274 }
275
276 fn workWritesSlot(
277 work: schedule.ScheduleWorkItem,
278 buffers: *const bufferization.BufferPlanAnalysis,
279 slot_id: usize,
280 ) bool {
281 if (buffers.getSlot(work.output_value)) |output| {
282 if (output.id == slot_id) return true;
283 }
284 if (work.kind == .iterate) {
285 for (work.root.results.items) |*result| {
286 if (buffers.getSlot(result)) |output| {
287 if (output.id == slot_id) return true;
288 }
289 }
290 }
291 if (work.kind == .scan) {
292 if (work.root.getOperand(1)) |scratch| {
293 if (buffers.getSlot(scratch)) |output| return output.id == slot_id;
294 }
295 }
296 return false;
297 }
298
299 fn diagnoseMissingOutput(allocator: std.mem.Allocator, slot: bufferization.BufferSlot) !void {
300 var uses = slot.value.useIterator();
301 while (uses.next()) |use| {
302 const user: *ir.Operation = @ptrCast(@alignCast(use.owner));
303 if (!isName(user.name.name, "func.return") and
304 !isName(user.name.name, dialect_mod.AccyDialect.ReturnOp.operation_name)) continue;
305 const function_name = if (slot.function) |function|
306 (choir.dialects.FuncDialect.FuncOp{ .op = function }).getName() orelse "<unknown>"
307 else
308 "<unknown>";
309 const message = try std.fmt.allocPrint(
310 allocator,
311 missing_output_format,
312 .{ function_name, use.operand_number, slot.value.id },
313 );
314 defer allocator.free(message);
315 var diagnostic = (slot.producer orelse user).emitError(message);
316 defer diagnostic.deinit();
317 _ = try diagnostic.emit();
318 return;
319 }
320 return error.InvalidOutputSlot;
321 }
322
323 fn cleanupMemorySpacePlanAnalysis(ptr: *anyopaque, allocator: std.mem.Allocator) void {
324 const analysis: *MemorySpacePlanAnalysis = @ptrCast(@alignCast(ptr));
325 analysis.deinit();
326 allocator.destroy(analysis);
327 }
328
329 fn classifySlot(slot: bufferization.BufferSlot) MemorySpaceAssignment {
330 return .{
331 .slot_id = slot.id,
332 .value = slot.value,
333 .producer = slot.producer,
334 .role = slot.role,
335 .space = spaceForSlot(slot),
336 .access = accessForRole(slot.role),
337 .transfer = transferForRole(slot.role),
338 .byte_size = slot.byte_size,
339 };
340 }
341
342 fn spaceForSlot(slot: bufferization.BufferSlot) MemorySpace {
343 if (slot.role.constant and !slot.role.output) return .device_constant;
344 return .device_global;
345 }
346
347 fn accessForRole(role: bufferization.BufferRole) MemoryAccess {
348 if (role.input and role.output) return .read_write;
349 if (role.output) return .write_only;
350 if (role.input or role.constant) return .read_only;
351 return .read_write;
352 }
353
354 fn transferForRole(role: bufferization.BufferRole) BoundaryTransfer {
355 if (role.input and role.output) return .bidirectional;
356 if (role.input) return .host_to_device;
357 if (role.output) return .device_to_host;
358 return .none;
359 }
360
361 const testing = std.testing;
362 const semantic = accy_choir.semantic;
363
364 fn findOpNamedInBlock(block: *ir.Block, name: []const u8) ?*ir.Operation {
365 var iter = block.operations.head;
366 while (iter) |op_ptr| {
367 const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
368 if (isName(op.name.name, name)) return op;
369 iter = op.next_op;
370 }
371 return null;
372 }
373
374 fn isName(actual: []const u8, expected: []const u8) bool {
375 return std.mem.eql(u8, actual, expected);
376 }
377
378 test "memory-space planning classifies function boundary transfers" {
379 const allocator = testing.allocator;
380
381 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
382 defer builder.deinit();
383 const f32_4 = try builder.tensor(.f32, &.{4});
384 var fb = try builder.beginFunction("memory_add4", &.{ f32_4, f32_4 }, &.{f32_4});
385 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
386 try fb.return_(&.{sum});
387 try fb.finish();
388 const module = try builder.finish();
389 defer module.deinit();
390
391 const choir_mod = module.choir_module;
392 const ctx = module.context();
393 const ledger = try memoryTestLedger();
394 defer ledger.destroy();
395 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 5);
396 defer cache.deinit();
397 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
398 defer pass_ctx.deinit();
399
400 const buffers = try bufferization.getBufferPlanAnalysis(&pass_ctx, choir_mod);
401 const analysis = try getMemorySpacePlanAnalysis(&pass_ctx, choir_mod);
402 try ledger.producersComplete();
403 try checkMemoryStorage(choir_mod, ctx, &cache, analysis);
404 try testing.expectEqual(@as(usize, 3), analysis.assignmentCount());
405 try testing.expectEqual(@as(usize, 3), analysis.device_global_slot_count);
406 try testing.expectEqual(@as(usize, 2), analysis.host_input_transfer_count);
407 try testing.expectEqual(@as(usize, 1), analysis.host_output_transfer_count);
408 try testing.expectEqual(@as(usize, 0), analysis.device_constant_slot_count);
409 try testing.expectEqual(@as(usize, 0), analysis.dynamic_slot_count);
410 try testing.expectEqual(@as(u64, 48), analysis.total_static_bytes);
411
412 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
413 const func = ir.inspection.functionByNameInBlock(body, "memory_add4") orelse return error.TestExpectedFunc;
414 const entry = func.getRegion(0).?.getEntryBlock().?;
415
416 const first_arg_slot = buffers.getSlot(entry.arguments.items[0]) orelse return error.TestExpectedSlot;
417 const first_arg = analysis.getAssignmentForSlot(first_arg_slot.id) orelse return error.TestExpectedAssignment;
418 try testing.expectEqual(MemorySpace.device_global, first_arg.space);
419 try testing.expectEqual(MemoryAccess.read_only, first_arg.access);
420 try testing.expectEqual(BoundaryTransfer.host_to_device, first_arg.transfer);
421
422 const add = findOpNamedInBlock(entry, dialect_mod.AccyDialect.AddOp.operation_name) orelse return error.TestExpectedAdd;
423 const output = analysis.getAssignmentForValue(buffers, add.getResult(0).?) orelse return error.TestExpectedAssignment;
424 try testing.expectEqual(MemorySpace.device_global, output.space);
425 try testing.expectEqual(MemoryAccess.write_only, output.access);
426 try testing.expectEqual(BoundaryTransfer.device_to_host, output.transfer);
427 try testing.expect(output.output_source.? == .kernel_written);
428 try testing.expectEqual(@as(?OutputSource, null), first_arg.output_source);
429 }
430
431 test "memory-space planning aliases returned input and constant backing" {
432 const allocator = testing.allocator;
433 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
434 defer builder.deinit();
435 const f32_2 = try builder.tensor(.f32, &.{2});
436 var function = try builder.beginFunction("backed_outputs", &.{f32_2}, &.{ f32_2, f32_2 });
437 const values = [_]f32{ 3, -4 };
438 const constant = try function.constant(f32_2, std.mem.sliceAsBytes(&values));
439 const input = function.parameter(0);
440 try function.return_(&.{ input, constant });
441 try function.finish();
442 const module = try builder.finish();
443 defer module.deinit();
444 const ledger = try memoryTestLedger();
445 defer ledger.destroy();
446 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 5);
447 defer cache.deinit();
448 var pass_ctx = passes.PassContext.init(
449 module.choir_module,
450 module.context(),
451 allocator,
452 &cache,
453 );
454 defer pass_ctx.deinit();
455 const buffers = try bufferization.getBufferPlanAnalysis(&pass_ctx, module.choir_module);
456 const analysis = try getMemorySpacePlanAnalysis(&pass_ctx, module.choir_module);
457 try ledger.producersComplete();
458 try checkMemoryStorage(module.choir_module, module.context(), &cache, analysis);
459 for ([_]*ir.Value{ input, constant }) |value| {
460 const slot = buffers.getSlot(value).?;
461 const output = analysis.getAssignmentForSlot(slot.id).?;
462 try testing.expect(output.output_source.? == .aliased);
463 try testing.expectEqual(slot.id, output.output_source.?.aliased);
464 }
465 }
466
467 test "memory-space planning classifies broadcast and refuses a removed writer at its location" {
468 const allocator = testing.allocator;
469 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
470 defer builder.deinit();
471 const scalar = try builder.tensor(.f32, &.{});
472 const vector = try builder.tensor(.f32, &.{8});
473 var function = try builder.beginFunction("writerless", &.{scalar}, &.{ scalar, vector });
474 const location = ir.Location.getFile("writerless.accy", 12, 7);
475 function.setLocation(location);
476 const broadcast = try function.broadcastInDim(function.parameter(0), vector, &.{8}, &.{});
477 try function.return_(&.{ function.parameter(0), broadcast });
478 try function.finish();
479 const module = try builder.finish();
480 defer module.deinit();
481 const ledger = try memoryTestLedger();
482 defer ledger.destroy();
483 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 5);
484 defer cache.deinit();
485 var pass_ctx = passes.PassContext.init(
486 module.choir_module,
487 module.context(),
488 allocator,
489 &cache,
490 );
491 defer pass_ctx.deinit();
492 const work = try schedule.getSchedulePlanAnalysis(&pass_ctx, module.choir_module);
493 const buffers = try bufferization.getBufferPlanAnalysis(&pass_ctx, module.choir_module);
494 const output_slot = buffers.getSlot(broadcast).?;
495 const source = (try classifyOutput(allocator, output_slot.*, buffers, work)).?;
496 try testing.expect(source == .kernel_written);
497 try testing.expectEqual(@as(usize, 1), work.workItemCount());
498 try testing.expectEqual(work.work_items.items[0].id, source.kernel_written);
499 work.deinit();
500 work.* = schedule.SchedulePlanAnalysis.init(allocator);
501 try checkMemoryFailureStorage(module.choir_module, module.context(), &cache);
502 var captured = choir.diagnostics.CaptureBuffer.init(allocator);
503 defer captured.deinit();
504 var scope = module.context().captureDiagnostics(&captured);
505 var guard = scope.enter();
506 defer guard.deinit();
507 try testing.expectError(
508 error.MissingOutputWriter,
509 getMemorySpacePlanAnalysis(&pass_ctx, module.choir_module),
510 );
511 try testing.expect(!ledger.view().missing_work_contract);
512 try testing.expectEqual(choir.product.revision.receipt.Outcome.rejected, ledger.view().outcome);
513 try testing.expectEqual(@as(usize, 4), cache.entries.count());
514 try testing.expectError(error.TerminalWorkOutcome, ledger.producersComplete());
515 try testing.expectEqual(@as(usize, 1), captured.diagnostics.items.len);
516 const diagnostic = captured.diagnostics.items[0];
517 try testing.expectEqual(choir.diagnostics.Severity.err, diagnostic.severity);
518 try testing.expect(diagnostic.location.eql(location));
519 const expected = try std.fmt.allocPrint(
520 allocator,
521 "function @writerless result #1 (%{d}) has no kernel writer or backing alias; " ++
522 "schedule its producer or provide an input/constant alias",
523 .{broadcast.id},
524 );
525 defer allocator.free(expected);
526 try testing.expectEqualStrings(expected, diagnostic.message);
527 }
528
529 test "memory-space planning places non-returned constants in constant memory" {
530 const allocator = testing.allocator;
531
532 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
533 defer builder.deinit();
534 const i32_4 = try builder.tensor(.i32, &.{4});
535 var fb = try builder.beginFunction("memory_const_add", &.{i32_4}, &.{i32_4});
536 const values = [_]i32{ 1, 2, 3, 4 };
537 const c = try fb.constant(i32_4, std.mem.sliceAsBytes(values[0..]));
538 const sum = try fb.add(fb.parameter(0), c);
539 try fb.return_(&.{sum});
540 try fb.finish();
541 const module = try builder.finish();
542 defer module.deinit();
543
544 const choir_mod = module.choir_module;
545 const ctx = module.context();
546 const ledger = try memoryTestLedger();
547 defer ledger.destroy();
548 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 5);
549 defer cache.deinit();
550 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
551 defer pass_ctx.deinit();
552
553 const buffers = try bufferization.getBufferPlanAnalysis(&pass_ctx, choir_mod);
554 const analysis = try getMemorySpacePlanAnalysis(&pass_ctx, choir_mod);
555 try ledger.producersComplete();
556 try checkMemoryStorage(choir_mod, ctx, &cache, analysis);
557 try testing.expectEqual(@as(usize, 3), analysis.assignmentCount());
558 try testing.expectEqual(@as(usize, 2), analysis.device_global_slot_count);
559 try testing.expectEqual(@as(usize, 1), analysis.device_constant_slot_count);
560 try testing.expectEqual(@as(usize, 1), analysis.host_input_transfer_count);
561 try testing.expectEqual(@as(usize, 1), analysis.host_output_transfer_count);
562 try testing.expectEqual(@as(u64, 48), analysis.total_static_bytes);
563
564 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
565 const func = ir.inspection.functionByNameInBlock(body, "memory_const_add") orelse return error.TestExpectedFunc;
566 const entry = func.getRegion(0).?.getEntryBlock().?;
567 const constant = findOpNamedInBlock(entry, dialect_mod.AccyDialect.ConstantOp.operation_name) orelse return error.TestExpectedConstant;
568 const assignment = analysis.getAssignmentForValue(buffers, constant.getResult(0).?) orelse return error.TestExpectedAssignment;
569 try testing.expectEqual(MemorySpace.device_constant, assignment.space);
570 try testing.expectEqual(MemoryAccess.read_only, assignment.access);
571 try testing.expectEqual(BoundaryTransfer.none, assignment.transfer);
572 }
573
574 test "memory-space planning records fusion elisions without assigning storage" {
575 const allocator = testing.allocator;
576
577 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
578 defer builder.deinit();
579 const f32_4 = try builder.tensor(.f32, &.{4});
580 var fb = try builder.beginFunction("memory_fused_add_mul", &.{ f32_4, f32_4, f32_4 }, &.{f32_4});
581 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
582 const product = try fb.mul(sum, fb.parameter(2));
583 try fb.return_(&.{product});
584 try fb.finish();
585 const module = try builder.finish();
586 defer module.deinit();
587
588 const choir_mod = module.choir_module;
589 const ctx = module.context();
590 const ledger = try memoryTestLedger();
591 defer ledger.destroy();
592 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 5);
593 defer cache.deinit();
594 var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
595 defer pass_ctx.deinit();
596
597 const buffers = try bufferization.getBufferPlanAnalysis(&pass_ctx, choir_mod);
598 const analysis = try getMemorySpacePlanAnalysis(&pass_ctx, choir_mod);
599 try ledger.producersComplete();
600 try checkMemoryStorage(choir_mod, ctx, &cache, analysis);
601 try testing.expectEqual(@as(usize, 4), analysis.assignmentCount());
602 try testing.expectEqual(@as(usize, 1), analysis.elided_value_count);
603 try testing.expectEqual(@as(usize, 4), analysis.device_global_slot_count);
604 try testing.expectEqual(@as(usize, 3), analysis.host_input_transfer_count);
605 try testing.expectEqual(@as(usize, 1), analysis.host_output_transfer_count);
606
607 const body = choir_mod.getRegion(0).?.getEntryBlock().?;
608 const func = ir.inspection.functionByNameInBlock(body, "memory_fused_add_mul") orelse return error.TestExpectedFunc;
609 const entry = func.getRegion(0).?.getEntryBlock().?;
610 const add = findOpNamedInBlock(entry, dialect_mod.AccyDialect.AddOp.operation_name) orelse return error.TestExpectedAdd;
611 try testing.expect(buffers.getSlot(add.getResult(0).?) == null);
612 try testing.expect(analysis.getAssignmentForValue(buffers, add.getResult(0).?) == null);
613 }
614
615 test "memory-space planning pass preserves IR" {
616 const allocator = testing.allocator;
617
618 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
619 defer builder.deinit();
620 const f32_4 = try builder.tensor(.f32, &.{4});
621 var fb = try builder.beginFunction("memory_pass_add4", &.{ f32_4, f32_4 }, &.{f32_4});
622 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
623 try fb.return_(&.{sum});
624 try fb.finish();
625 const module = try builder.finish();
626 defer module.deinit();
627
628 const choir_mod = module.choir_module;
629 const ctx = module.context();
630 var pm = passes.PassManager.init(allocator);
631 defer pm.deinit();
632 try pm.addPass(memorySpacePlanningPass());
633
634 const ledger = try choir.product.revision.AccountingV1.create(allocator, .{
635 .allowance = choir.product.revision.WorkVector.uniform(std.math.maxInt(u64)),
636 .workspace = std.math.maxInt(u64),
637 .events = 12,
638 }, &.{.{ .name = memory_space_planning_pass_name, .version = 1 }});
639 defer ledger.destroy();
640 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 5);
641 defer cache.deinit();
642 try testing.expectEqual(
643 passes.PassResult.success,
644 pm.runWithAnalysisCache(choir_mod, ctx, &cache, .{}),
645 );
646 try ledger.producersComplete();
647 try testing.expectEqual(@as(u64, 1), pm.stats.pass_runs);
648 try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
649 }
650
651 fn memoryTestLedger() !*choir.product.revision.AccountingV1 {
652 const revision = choir.product.revision;
653 return revision.AccountingV1.create(testing.allocator, .{
654 .allowance = revision.WorkVector.uniform(std.math.maxInt(u64)),
655 .workspace = std.math.maxInt(u64),
656 .events = 12,
657 }, &.{});
658 }
659
660 fn checkMemoryStorage(
661 op: *ir.Operation,
662 ctx: *ir.Context,
663 cache: *passes.AnalysisCache,
664 expected: *const MemorySpacePlanAnalysis,
665 ) !void {
666 const bounds = try memoryAnalysisWork(.{ .operation = op });
667 const fixed = @import("alloc_fixed");
668 const bytes = try testing.allocator.alignedAlloc(u8, .@"64", @intCast(bounds.workspace));
669 defer testing.allocator.free(bytes);
670 var storage = fixed.Tracked.init(bytes);
671 var retained = fixed.Monotonic.init(storage.allocator(), @max(1, bytes.len));
672 const allocator = retained.allocator();
673 var pass_ctx = passes.PassContext.init(op, ctx, allocator, cache);
674 defer pass_ctx.deinit();
675 const ptr = try computeMemorySpacePlanAnalysis(&pass_ctx, op);
676 defer cleanupMemorySpacePlanAnalysis(ptr, allocator);
677 const actual: *MemorySpacePlanAnalysis = @ptrCast(@alignCast(ptr));
678 inline for (.{
679 "host_slot_count", "device_global_slot_count", "device_constant_slot_count",
680 "device_shared_slot_count", "unified_slot_count", "host_input_transfer_count",
681 "host_output_transfer_count", "dynamic_slot_count", "elided_value_count",
682 "total_static_bytes",
683 }) |field| try testing.expectEqual(@field(expected, field), @field(actual, field));
684 try testing.expectEqual(expected.assignmentCount(), actual.assignmentCount());
685 for (expected.assignments.items, actual.assignments.items) |left, right| {
686 inline for (.{
687 "slot_id", "value", "producer", "space", "access", "transfer", "byte_size",
688 }) |field| try testing.expectEqual(@field(left, field), @field(right, field));
689 try testing.expectEqualDeep(left.role, right.role);
690 try testing.expectEqualDeep(left.output_source, right.output_source);
691 try testing.expectEqual(right.value, actual.getAssignmentForSlot(right.slot_id).?.value);
692 }
693 try testing.expectEqual(expected.slot_to_assignment.count(), actual.slot_to_assignment.count());
694 try testing.expect(!storage.exhausted);
695 const used = if (retained.current) |*current| fixed.used(current) else 0;
696 try testing.expect(used <= bounds.workspace);
697 try testing.expect(used >= @sizeOf(MemorySpacePlanAnalysis));
698 }
699
700 fn checkMemoryFailureStorage(
701 op: *ir.Operation,
702 ctx: *ir.Context,
703 cache: *passes.AnalysisCache,
704 ) !void {
705 const counts = try MemoryWork.inspect(op);
706 const bounds = try counts.bounds();
707 const fixed = @import("alloc_fixed");
708 const bytes = try testing.allocator.alignedAlloc(u8, .@"64", @intCast(bounds.workspace));
709 defer testing.allocator.free(bytes);
710 var storage = fixed.Tracked.init(bytes);
711 var retained = fixed.Monotonic.init(storage.allocator(), @max(1, bytes.len));
712 var pass_ctx = passes.PassContext.init(op, ctx, retained.allocator(), cache);
713 defer pass_ctx.deinit();
714 var captured = choir.diagnostics.CaptureBuffer.init(testing.allocator);
715 defer captured.deinit();
716 var scope = ctx.captureDiagnostics(&captured);
717 var guard = scope.enter();
718 defer guard.deinit();
719 try testing.expectError(
720 error.MissingOutputWriter,
721 computeMemorySpacePlanAnalysis(&pass_ctx, op),
722 );
723 try testing.expectEqual(@as(usize, 1), captured.diagnostics.items.len);
724 const message = captured.diagnostics.items[0].message;
725 try testing.expect(message.len <= try counts.diagnosticBytes());
726 try testing.expect(!storage.exhausted);
727 const used = if (retained.current) |*current| fixed.used(current) else 0;
728 try testing.expect(used <= bounds.workspace);
729 try testing.expect(used >= @sizeOf(MemorySpacePlanAnalysis) + message.len);
730 }
731
732 test "memory-space planning work contract covers assignment and writer growth" {
733 for ([_]usize{ 0, 1, 6, 7, 16, 64 }) |count| {
734 try checkMemoryBoundary(count, false, &.{ 2, 3 });
735 try checkMemoryBoundary(count, false, &.{ -1, 3 });
736 try checkMemoryBoundary(count, true, &.{ 2, 3 });
737 }
738 try testing.expectError(error.WorkOverflow, (MemoryWork{
739 .input = .{ .values = std.math.maxInt(u64) },
740 .symbols = 0,
741 }).bounds());
742 try testing.expectError(error.WorkOverflow, (MemoryWork{
743 .input = .{ .operations = std.math.maxInt(u64) },
744 .symbols = 0,
745 }).bounds());
746 try testing.expectError(error.WorkOverflow, (MemoryWork{
747 .input = .{},
748 .symbols = std.math.maxInt(u64),
749 }).bounds());
750 }
751
752 fn checkMemoryBoundary(count: usize, written: bool, dims: []const i64) !void {
753 std.debug.assert(count <= 64);
754 var builder = try semantic.Builder.init(
755 testing.allocator,
756 semantic.Builder.ContextLimits.standard,
757 );
758 defer builder.deinit();
759 const typ = try builder.tensor(.f32, dims);
760 var types: [64]@TypeOf(typ) = @splat(typ);
761 var function = try builder.beginFunction("memory_boundary", types[0..count], types[0..count]);
762 var values: [64]@TypeOf(function.parameter(0)) = undefined;
763 for (values[0..count], 0..) |*value, index| {
764 const input = function.parameter(index);
765 value.* = if (written) try function.reshape(input, typ, dims) else input;
766 }
767 try function.return_(values[0..count]);
768 try function.finish();
769 const module = try builder.finish();
770 defer module.deinit();
771 const ledger = try memoryTestLedger();
772 defer ledger.destroy();
773 var cache = try passes.AnalysisCache.initAccounted(testing.allocator, null, ledger, .{}, 5);
774 defer cache.deinit();
775 var pass_ctx = passes.PassContext.init(
776 module.choir_module,
777 module.context(),
778 testing.allocator,
779 &cache,
780 );
781 defer pass_ctx.deinit();
782 const analysis = try getMemorySpacePlanAnalysis(&pass_ctx, module.choir_module);
783 try ledger.producersComplete();
784 try testing.expectEqual(count * @as(usize, if (written) 2 else 1), analysis.assignmentCount());
785 try testing.expectEqual(count, analysis.host_input_transfer_count);
786 try testing.expectEqual(count, analysis.host_output_transfer_count);
787 for (values[0..count], 0..) |value, index| {
788 const slot_id = if (written) count + index else index;
789 const assignment = analysis.getAssignmentForSlot(slot_id).?;
790 try testing.expectEqual(value, assignment.value);
791 const source = assignment.output_source.?;
792 if (written) {
793 try testing.expectEqual(index, source.kernel_written);
794 } else try testing.expectEqual(slot_id, source.aliased);
795 }
796 try checkMemoryStorage(module.choir_module, module.context(), &cache, analysis);
797 try checkMemoryAdmission(module.choir_module, module.context());
798 }
799
800 fn checkMemoryAdmission(op: *ir.Operation, ctx: *ir.Context) !void {
801 const revision = choir.product.revision;
802 var charge: u64 = 1;
803 for ([_]passes.AnalysisDescriptor{
804 @import("shape/root.zig").shape_layout_analysis_descriptor,
805 @import("fusion/root.zig").fusion_plan_analysis_descriptor,
806 schedule.schedule_plan_analysis_descriptor,
807 bufferization.buffer_plan_analysis_descriptor,
808 memory_space_plan_analysis_descriptor,
809 }) |descriptor| {
810 const bounds = try descriptor.work_contract.?.estimate(.{ .operation = op });
811 charge = try accounting.add(charge, bounds.work.structural_visits);
812 }
813 for ([_]i8{ -1, 0, 1 }) |offset| {
814 var allowance = revision.WorkVector.uniform(std.math.maxInt(u64));
815 allowance.structural_visits = @intCast(@as(i128, charge) + offset);
816 const ledger = try revision.AccountingV1.create(testing.allocator, .{
817 .allowance = allowance,
818 .workspace = std.math.maxInt(u64),
819 .events = 12,
820 }, &.{.{ .name = memory_space_planning_pass_name, .version = 1 }});
821 defer ledger.destroy();
822 var cache = try passes.AnalysisCache.initAccounted(
823 testing.allocator,
824 null,
825 ledger,
826 .{},
827 5,
828 );
829 defer cache.deinit();
830 var manager = passes.PassManager.init(testing.allocator);
831 defer manager.deinit();
832 try manager.addPass(memorySpacePlanningPass());
833 const result = manager.runWithAnalysisCache(op, ctx, &cache, .{});
834 if (offset < 0) {
835 try testing.expectEqual(passes.PassResult.failure, result);
836 try testing.expectEqual(revision.receipt.Outcome.exhausted, ledger.view().outcome);
837 try testing.expectEqual(@as(usize, 3), cache.entries.count());
838 } else {
839 try testing.expectEqual(passes.PassResult.success, result);
840 try ledger.producersComplete();
841 try testing.expectEqual(@as(usize, 5), cache.entries.count());
842 }
843 }
844 }
845
846 test "memory-space planning bounds rejected diagnostics with long function symbols" {
847 var name: [4096]u8 = @splat('n');
848 for ([_]usize{ 1, 64, 4096 }) |length| try checkMemoryDiagnostic(name[0..length]);
849 }
850
851 fn checkMemoryDiagnostic(name: []const u8) !void {
852 var builder = try semantic.Builder.init(
853 testing.allocator,
854 semantic.Builder.ContextLimits.standard,
855 );
856 defer builder.deinit();
857 const scalar = try builder.tensor(.f32, &.{});
858 const vector = try builder.tensor(.f32, &.{8});
859 var function = try builder.beginFunction(name, &.{scalar}, &.{vector});
860 const result = try function.broadcastInDim(function.parameter(0), vector, &.{8}, &.{});
861 try function.return_(&.{result});
862 try function.finish();
863 const module = try builder.finish();
864 defer module.deinit();
865 const counts = try MemoryWork.inspect(module.choir_module);
866 try testing.expectEqual(name.len, counts.symbols);
867 const ledger = try memoryTestLedger();
868 defer ledger.destroy();
869 var cache = try passes.AnalysisCache.initAccounted(testing.allocator, null, ledger, .{}, 5);
870 defer cache.deinit();
871 var pass_ctx = passes.PassContext.init(
872 module.choir_module,
873 module.context(),
874 testing.allocator,
875 &cache,
876 );
877 defer pass_ctx.deinit();
878 const work = try schedule.getSchedulePlanAnalysis(&pass_ctx, module.choir_module);
879 _ = try bufferization.getBufferPlanAnalysis(&pass_ctx, module.choir_module);
880 work.deinit();
881 work.* = schedule.SchedulePlanAnalysis.init(testing.allocator);
882 try checkMemoryFailureStorage(module.choir_module, module.context(), &cache);
883 try ledger.producersComplete();
884 }