tiny.accy.preparation.memory

1const std = @import("std");2const choir = @import("choir");3const bufferization = @import("bufferization/root.zig");4const schedule = @import("schedule/root.zig");5const accy_choir = @import("../choir/root.zig");6const dialect_mod = accy_choir.dialect;78const ir = choir.ir;9const passes = choir.passes;10const accounting = passes.pass.work;1112pub const memory_space_plan_analysis_name = "accy-choir-memory-space-plan";13pub const memory_space_planning_pass_name = "accy-choir-plan-memory-spaces";14pub const memory_space_planning_pass_description =15    "Plan Accy Choir buffer memory spaces and boundary transfers";1617pub const MemorySpace = accy_choir.record.memory.MemorySpace;1819pub const MemoryAccess = accy_choir.record.memory.MemoryAccess;2021pub const BoundaryTransfer = accy_choir.record.memory.BoundaryTransfer;2223/// A caller reads this to know whether a function output is written by a kernel or is an existing24/// buffer. The value is either the id of the scheduled unit of work that writes the output, or the25/// id of an input or constant buffer that the output is. Planning fails with26/// `error.MissingOutputWriter` when an output has neither.27pub const OutputSource = accy_choir.record.memory.OutputSource;2829pub 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,3940    pub fn isDynamic(self: MemorySpaceAssignment) bool {41        return self.byte_size == null;42    }43};4445pub 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,5960    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    }6768    pub fn deinit(self: *MemorySpacePlanAnalysis) void {69        self.assignments.deinit(self.allocator);70        self.slot_to_assignment.deinit();71        self.* = undefined;72    }7374    pub fn assignmentCount(self: MemorySpacePlanAnalysis) usize {75        return self.assignments.items.len;76    }7778    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    }8586    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    }9495    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);101102        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};118119const 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";122123const MemoryWork = struct {124    input: accounting.Census,125    symbols: u64,126127    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    }132133    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    }141142    fn diagnosticBytes(self: MemoryWork) !u64 {143        return accounting.add(self.symbols, missing_output_format.len + 2 * 20 + "<unknown>".len);144    }145146    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    }160161    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};186187fn memoryAnalysisWork(input: accounting.Input) !accounting.Bounds {188    return (try MemoryWork.inspect(input.operation)).bounds();189}190191fn memoryPassWork(_: accounting.Input) !accounting.Bounds {192    return .{ .work = .{ .structural_visits = 1 } };193}194195pub 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};203204pub 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}216217pub 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}228229fn 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}234235fn 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);241242    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    }248249    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);257258    return @ptrCast(analysis);259}260261fn 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}275276fn 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}298299fn 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") and304            !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        else308            "<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}322323fn cleanupMemorySpacePlanAnalysis(ptr: *anyopaque, allocator: std.mem.Allocator) void {324    const analysis: *MemorySpacePlanAnalysis = @ptrCast(@alignCast(ptr));325    analysis.deinit();326    allocator.destroy(analysis);327}328329fn 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}341342fn spaceForSlot(slot: bufferization.BufferSlot) MemorySpace {343    if (slot.role.constant and !slot.role.output) return .device_constant;344    return .device_global;345}346347fn 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}353354fn 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}360361const testing = std.testing;362const semantic = accy_choir.semantic;363364fn 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}373374fn isName(actual: []const u8, expected: []const u8) bool {375    return std.mem.eql(u8, actual, expected);376}377378test "memory-space planning classifies function boundary transfers" {379    const allocator = testing.allocator;380381    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();390391    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();399400    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);411412    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().?;415416    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);421422    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}430431test "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}466467test "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}528529test "memory-space planning places non-returned constants in constant memory" {530    const allocator = testing.allocator;531532    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();543544    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();552553    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);563564    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}573574test "memory-space planning records fusion elisions without assigning storage" {575    const allocator = testing.allocator;576577    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();587588    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();596597    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);606607    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}614615test "memory-space planning pass preserves IR" {616    const allocator = testing.allocator;617618    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();627628    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());633634    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}650651fn 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}659660fn 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}699700fn 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}731732test "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}751752fn 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}799800fn 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}845846test "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}850851fn 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}

1pub const memory = @import("memory.zig");