tiny.choir.passes.promotion

1const std = @import("std");2const ir = @import("../core/root.zig");3const dialects = @import("../dialects/root.zig");4const arith_types = @import("../dialects/arith/types.zig");5const pass_mod = @import("pass/root.zig");6const registry_mod = @import("pipeline.zig");7const effects = @import("effects.zig");89const arith = dialects.arith.ArithDialect;10const memref = dialects.memref.MemrefDialect;11const scf = dialects.scf.ScfDialect;12const func_dialect = dialects.func.FuncDialect;1314const Pass = pass_mod.Pass;15const PassContext = pass_mod.PassContext;16const PassResult = pass_mod.PassResult;1718pub const memory_promotion_pass_name = "choir-mem2reg";19pub const memory_promotion_pass_description =20    "Promote single-cell allocas to SSA values and scf iter args";2122const PromotionError = error{23    OutOfMemory,24    InvalidPromotion,25};2627/// The operations this pass creates for each operation it is given.28///29/// WHERE AN OPERATION COMES FROM HERE. The pass creates operations in three30/// places and nowhere else. A region a promoted cell is stored inside is31/// rebuilt, because an operation cannot gain a result in place, and a rebuild32/// creates one operation. The rebuilt region's terminators are replaced so33/// they carry the promoted values, which creates one for an `scf.for`, two for34/// an `scf.while`, and one per branch of an `scf.if`. A cell a region carries35/// before anything has stored to it gets a zero constant.36///37/// WHY A CONSTANT AND NOT A COUNT OF CELLS. A region is rebuilt at most once38/// however many cells it carries, so its rebuild and its terminators are39/// charged against the region operation and the terminators the caller handed40/// the pass, and a zero constant is charged against the cell's own41/// `memref.alloca`, which the caller handed it too. No operation of the input42/// is charged by two rebuilds, so three times the operations given holds43/// whatever shape the input has. It is an upper bound and not a measurement:44/// the shapes pinned in this file's tests create fewer.45///46/// WHY A CALLER CARES ABOUT OPERATIONS CREATED AND NOT OPERATIONS KEPT. The47/// segment these operations draw from is a fixed buffer whose free reclaims48/// only the last allocation, so an operation this pass erases holds its bytes49/// for the rest of the compile exactly as a kept one does. A caller sizing50/// that segment multiplies this figure by the operations it will hand over.51pub const operations_created_per_operation: usize = 3;5253/// The operations this pass may create given `operations` operations.54///55/// A caller that knows how many operations it will lower calls this instead of56/// reading the constant, so the shape of the bound can change without every57/// caller changing with it.58pub fn operationsCreated(operations: usize) usize {59    return operations_created_per_operation * operations;60}6162/// What one run of the pass created, for a caller that wants the figure it63/// paid rather than the bound it was promised.64pub const Created = struct {65    /// Operations the run created, erased ones included.66    operations: usize = 0,67    /// Regions the run rebuilt. Each is rebuilt once however many cells it68    /// carries, so this never passes the regions the caller handed over.69    regions_rebuilt: usize = 0,70};7172pub fn createMemoryPromotionPass() Pass {73    return .{74        .name = memory_promotion_pass_name,75        .description = memory_promotion_pass_description,76        .run_fn = runMemoryPromotion,77        .mutation_scope = .whole_module,78    };79}8081pub const memory_promotion_pass_registration = registry_mod.PassRegistration{82    .name = memory_promotion_pass_name,83    .description = memory_promotion_pass_description,84    .pass = createMemoryPromotionPass(),85};8687/// Promotes every local cell under `op`, and answers whether anything moved.88///89/// The pass above is how a pipeline asks for this. This is how a caller that90/// runs no pipeline asks for the same thing, which is what the x86_64 backend91/// does: it lowers one module and the page and the file are both read from92/// that module, so a transform it runs is seen by both or by neither.93pub fn promote(op: *ir.Operation, allocator: std.mem.Allocator) PromotionError!bool {94    var created = Created{};95    return promoteCounting(op, allocator, &created);96}9798/// Promotes as `promote` does and reports what the run created.99///100/// A caller that sizes the context this module is built in reads `created` to101/// see what the run spent, against `operationsCreated` for what it was allowed102/// to spend.103pub fn promoteCounting(104    op: *ir.Operation,105    allocator: std.mem.Allocator,106    created: *Created,107) PromotionError!bool {108    var modified = false;109    try promoteInOp(op, allocator, &modified, created);110    return modified;111}112113fn runMemoryPromotion(ctx: *PassContext) PassResult {114    var modified = false;115    var created = Created{};116    promoteInOp(ctx.op, ctx.allocator, &modified, &created) catch return .failure;117    if (modified) {118        ctx.markModified();119    } else {120        ctx.preserveAllAnalyses();121    }122    return .success;123}124125fn promoteInOp(126    op: *ir.Operation,127    allocator: std.mem.Allocator,128    modified: *bool,129    created: *Created,130) !void {131    if (std.mem.eql(u8, op.name.name, func_dialect.FuncOp.operation_name)) {132        if (op.getRegion(0)) |region| {133            if (region.getEntryBlock()) |entry| {134                try promoteFunction(entry, allocator, modified, created);135            }136        }137        return;138    }139    for (op.regions.items) |*region| {140        var block_iter = region.getBlocks();141        while (block_iter.next()) |block| {142            var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));143            while (current) |current_op| {144                const next = current_op.next_op;145                try promoteInOp(current_op, allocator, modified, created);146                current = next;147            }148        }149    }150}151152/// Replaces the accesses of every local cell with the values they carried.153///154/// THE ALLOCATION IS LEFT WHERE IT STANDS. The permission this asks for covers155/// the reads and the writes and nothing else: an allocation can still fail, and156/// a failure is an event no derivation here discharges. Removing the `alloca`157/// would delete that failure along with the bytes, and a program that cannot158/// run out of stack is not the program the footprint charged. A later pass that159/// can state the failure is dead may remove it.160///161/// THE CELLS OF ONE BLOCK TRAVEL TOGETHER. A rewrite is run once per block that162/// declares cells, carrying every cell that block declares, so an `scf` region163/// is rebuilt at most once however many cells it holds. Cells are grouped in164/// the order they were collected, which is outer before inner, and each cell's165/// block is read when its group is reached, because an outer rewrite moves the166/// operations of an inner block into the block it built.167fn promoteFunction(168    entry: *ir.Block,169    allocator: std.mem.Allocator,170    modified: *bool,171    created: *Created,172) !void {173    var allocas = std.ArrayListUnmanaged(*ir.Operation).empty;174    defer allocas.deinit(allocator);175    try collectAllocas(entry, allocator, &allocas);176177    var cells = std.ArrayListUnmanaged(Cell).empty;178    defer cells.deinit(allocator);179    for (allocas.items) |alloca| {180        switch (effects.permitsDiscardingLocalAccesses(alloca)) {181            .yes => {},182            .no => continue,183            .unaffordable => return error.OutOfMemory,184        }185        if (!allUsesArePromotable(alloca)) continue;186        const element = elementType(alloca) orelse continue;187        const cell = alloca.getResult(0) orelse continue;188        try cells.append(allocator, .{ .alloca = alloca, .value = cell, .element_type = element });189    }190    if (cells.items.len == 0) return;191192    const grouped = try allocator.alloc(bool, cells.items.len);193    defer allocator.free(grouped);194    @memset(grouped, false);195196    var group = std.ArrayListUnmanaged(Cell).empty;197    defer group.deinit(allocator);198199    for (cells.items, 0..) |first, index| {200        if (grouped[index]) continue;201        const block = first.alloca.getBlock() orelse {202            grouped[index] = true;203            continue;204        };205206        group.clearRetainingCapacity();207        var start: ?*ir.Operation = null;208        for (cells.items[index..], index..) |candidate, at| {209            if (grouped[at]) continue;210            if (candidate.alloca.getBlock() != block) continue;211            grouped[at] = true;212            if (!analyzeFrom(candidate.alloca, block, false).ok) continue;213            if (start == null) start = candidate.alloca.next_op;214            try group.append(allocator, candidate);215        }216        if (group.items.len == 0) continue;217218        const values = try allocator.alloc(?*ir.Value, group.items.len);219        defer allocator.free(values);220        @memset(values, null);221222        var promotion = Promotion{223            .allocator = allocator,224            .ctx = first.alloca.context,225            .cells = group.items,226            .created = created,227        };228        try promotion.rewriteBlock(block, start, values);229        if (promotion.changed) modified.* = true;230    }231}232233fn collectAllocas(block: *ir.Block, allocator: std.mem.Allocator, out: *std.ArrayListUnmanaged(*ir.Operation)) !void {234    var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));235    while (current) |op| {236        const next = op.next_op;237        if (std.mem.eql(u8, op.name.name, memref.AllocaOp.operation_name)) {238            try out.append(allocator, op);239        }240        for (op.regions.items) |*region| {241            var block_iter = region.getBlocks();242            while (block_iter.next()) |nested| {243                try collectAllocas(nested, allocator, out);244            }245        }246        current = next;247    }248}249250fn elementType(alloca: *ir.Operation) ?ir.Type {251    const result = alloca.getResult(0) orelse return null;252    const name = result.type.getDialectTypeName() orelse return null;253    const params = result.type.getDialectParamKey() orelse return null;254    _ = name;255    const parsed = memref.parseMemrefParams(params) orelse return null;256    if (parsed.size) |size| {257        if (size != 1) return null;258    } else {259        return null;260    }261    const ctx = alloca.context;262    return ctx.getDialectTypeFromName(parsed.element_type_name) catch null;263}264265/// Whether this pass can MODEL every use, which is a narrower question than266/// whether discarding them is permitted.267///268/// The permission comes from `effects.permitsDiscardingLocalAccesses`, which269/// reads the declarations. What is left here is arithmetic this pass does: one270/// SSA value stands for one cell, so a second cell has nowhere to go, and the271/// value that stands for it is only the right one when every access names the272/// same cell, which a literal index at zero is what this pass knows how to273/// check.274fn allUsesArePromotable(alloca: *ir.Operation) bool {275    const cell = alloca.getResult(0) orelse return false;276    if (elementType(alloca) == null) return false;277    var use = cell.first_use;278    while (use) |operand| {279        const owner: *ir.Operation = @ptrCast(@alignCast(operand.owner));280        if (std.mem.eql(u8, owner.name.name, memref.LoadOp.operation_name)) {281            if (operand.operand_number != 0) return false;282            if (!indexIsZero(owner, 1)) return false;283        } else if (std.mem.eql(u8, owner.name.name, memref.StoreOp.operation_name)) {284            if (operand.operand_number != 1) return false;285            if (!indexIsZero(owner, 2)) return false;286        } else {287            return false;288        }289        use = operand.next_use;290    }291    return true;292}293294fn indexIsZero(op: *ir.Operation, operand_index: usize) bool {295    const index_value = op.getOperand(operand_index) orelse return false;296    const defining = index_value.getDefiningOp() orelse return false;297    const defining_op: *ir.Operation = @ptrCast(@alignCast(defining));298    if (!std.mem.eql(u8, defining_op.name.name, arith.ConstantOp.operation_name)) return false;299    const value = (arith.ConstantOp{ .op = defining_op }).getIntValue() orelse return false;300    return value == 0;301}302303const Analysis = struct {304    ok: bool,305    defined: bool,306    stores: bool,307};308309fn analyzeFrom(alloca: *ir.Operation, block: *ir.Block, defined_at_entry: bool) Analysis {310    return analyzeOps(alloca, if (block == alloca.getBlock()) alloca.next_op else firstOp(block), defined_at_entry);311}312313fn firstOp(block: *ir.Block) ?*ir.Operation {314    return @ptrCast(@alignCast(block.operations.head));315}316317fn analyzeOps(alloca: *ir.Operation, start: ?*ir.Operation, defined_at_entry: bool) Analysis {318    const cell = alloca.getResult(0) orelse return .{ .ok = false, .defined = false, .stores = false };319    var defined = defined_at_entry;320    var stores = false;321    var current = start;322    while (current) |op| {323        const next = op.next_op;324        if (std.mem.eql(u8, op.name.name, memref.LoadOp.operation_name) and op.getOperand(0) == cell) {325            if (!defined) return .{ .ok = false, .defined = defined, .stores = stores };326        } else if (std.mem.eql(u8, op.name.name, memref.StoreOp.operation_name) and op.getOperand(1) == cell) {327            defined = true;328            stores = true;329        } else if (op.regions.items.len != 0) {330            if (!subtreeUsesCell(op, cell)) {331                current = next;332                continue;333            }334            if (std.mem.eql(u8, op.name.name, scf.ForOp.operation_name)) {335                const body = (scf.ForOp{ .op = op }).getBodyBlock();336                const body_result = analyzeOps(alloca, firstOp(body), defined);337                if (!body_result.ok) return .{ .ok = false, .defined = defined, .stores = stores };338                stores = stores or body_result.stores;339            } else if (std.mem.eql(u8, op.name.name, scf.WhileOp.operation_name)) {340                const loop = analyzeWhile(alloca, op, defined);341                if (!loop.ok) return .{ .ok = false, .defined = defined, .stores = stores };342                defined = loop.defined;343                stores = stores or loop.stores;344            } else if (std.mem.eql(u8, op.name.name, scf.IfOp.operation_name)) {345                const if_op = scf.IfOp{ .op = op };346                const then_result = analyzeOps(alloca, firstOp(if_op.getThenBlock()), defined);347                if (!then_result.ok) return .{ .ok = false, .defined = defined, .stores = stores };348                var else_defined = defined;349                var else_stores = false;350                if (if_op.getElseBlock()) |else_block| {351                    const else_result = analyzeOps(alloca, firstOp(else_block), defined);352                    if (!else_result.ok) return .{ .ok = false, .defined = defined, .stores = stores };353                    else_defined = else_result.defined;354                    else_stores = else_result.stores;355                }356                defined = then_result.defined and else_defined;357                stores = stores or then_result.stores or else_stores;358            } else {359                return .{ .ok = false, .defined = defined, .stores = stores };360            }361        }362        current = next;363    }364    return .{ .ok = true, .defined = defined, .stores = stores };365}366367/// A `scf.while` runs its before block at least once and its after block zero368/// or more times, each time after a before block. So the after block starts369/// with whatever the before block defined, and after the loop the cell is370/// defined exactly when the before block defines it.371///372/// A later trip through the before block starts at least as defined as the373/// first, because a trip through the after block only adds definitions, so374/// checking the first trip is enough.375fn analyzeWhile(alloca: *ir.Operation, op: *ir.Operation, defined: bool) Analysis {376    const loop = scf.WhileOp{ .op = op };377    const before = analyzeOps(alloca, firstOp(loop.getBeforeBlock()), defined);378    if (!before.ok) return .{ .ok = false, .defined = defined, .stores = false };379    const after = analyzeOps(alloca, firstOp(loop.getAfterBlock()), before.defined);380    if (!after.ok) return .{ .ok = false, .defined = defined, .stores = false };381    return .{ .ok = true, .defined = before.defined, .stores = before.stores or after.stores };382}383384fn subtreeUsesCell(op: *ir.Operation, cell: *ir.Value) bool {385    for (op.operands.items) |operand| {386        if (operand.value == cell) return true;387    }388    for (op.regions.items) |*region| {389        var block_iter = region.getBlocks();390        while (block_iter.next()) |block| {391            var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));392            while (current) |nested| {393                if (subtreeUsesCell(nested, cell)) return true;394                current = nested.next_op;395            }396        }397    }398    return false;399}400401/// A local cell the pass promotes, and the type of the value it held.402const Cell = struct {403    alloca: *ir.Operation,404    value: *ir.Value,405    element_type: ir.Type,406};407408/// Nested context bytes one region rewrite consumes for its own splice before409/// the cells it carries are counted, not counting the rewrites nested inside410/// it, which reserve for themselves.411///412/// MEASURED AND NOT DERIVED. Over the nested-loop reproducer and two413/// generated nests, one splice took between 992 and 4576 bytes at cell counts414/// of 1, 2, 8 and 12, the widest being a twelve cell block. This pair holds at415/// least 1.5 times the widest measured point at every one of those counts. A416/// rewrite that outgrows the pair exhausts the segment after it has begun417/// mutating, and `refuseOutgrownRewrite` turns that into a refusal at the end418/// of the rewrite rather than leaving a module for a later reader to trip419/// over. A bound derived from the new operations, blocks and carried values420/// would replace this measurement.421pub const rewrite_bytes_base: usize = 1024;422423/// Nested context bytes one region rewrite consumes for each cell it carries.424/// Read `rewrite_bytes_base` above for what the pair was measured against.425pub const rewrite_bytes_per_cell: usize = 512;426427/// The nested context bytes one region rewrite reserves before it mutates428/// anything, carrying `cells` cells.429///430/// A caller sizing `operation_nested` calls this instead of reading the two431/// constants, so the shape of the bound can change without every caller432/// changing with it.433pub fn rewriteBytes(cells: usize) usize {434    return rewrite_bytes_base + rewrite_bytes_per_cell * cells;435}436437/// Rewrites one block and everything below it, carrying every cell at once.438///439/// A REGION IS REBUILT AT MOST ONCE. An `scf` region a promoted cell is stored440/// inside gains a carried value, and the operation holding that region cannot441/// gain a result in place, so it is rebuilt. Rebuilding once per cell would442/// create one operation per cell and per region, and the context segment these443/// operations draw from is a bump allocator, so the storage an erased operation444/// held stays spent for the rest of the compile. Carrying every cell of a block445/// through one walk keeps the operations this pass creates proportional to the446/// operations it was given.447const Promotion = struct {448    allocator: std.mem.Allocator,449    ctx: *ir.Context,450    cells: []const Cell,451    created: *Created,452    changed: bool = false,453454    /// Refuses before the first mutation when `operation_nested` cannot hold455    /// the splice this rewrite is about to perform.456    ///457    /// WHY A RESERVATION AND NOT A ROLLBACK. A rewrite creates the new region,458    /// splices it into the parent, moves the body across, and only then erases459    /// the old one. A shortfall anywhere between those steps leaves a module460    /// Choir's own verifier rejects, and there is no461    /// undo to run. So the one failure point is moved in front of the first462    /// mutation and the steps after it are reached only when the bytes are463    /// already there.464    ///465    /// WHAT IT GUARDS AND WHAT IT DOES NOT. It guards the three region466    /// rewrites, which are the steps that leave a module the verifier rejects467    /// when they stop halfway. It does not guard the walk that reaches them:468    /// `rewriteBlock` replaces a block's own loads and stores with the values469    /// last stored as it goes, so a refusal deeper in the nest leaves those470    /// already replaced. The module is then SMALLER than the one that went in471    /// and it still verifies, because replacing a load with the value last472    /// stored is a complete edit on its own. A pass that refuses leaves a473    /// valid module, not an untouched one.474    ///475    /// IT ASKS THE SEGMENT RATHER THAN READING ITS HEADROOM. A request the476    /// segment cannot serve is how the segment learns which one ran short and477    /// by how much, and that record is what the driver turns into a refusal478    /// naming the capacity and the flag. A reservation that only compared479    /// figures would refuse with nothing to say. The bytes come straight back:480    /// the segment bumps, and this was the last allocation in it.481    fn reserveRewrite(self: *Promotion) PromotionError!void {482        const allocator = ir.context.operationAllocator(self.ctx);483        const held = allocator.alloc(u8, rewriteBytes(self.cells.len)) catch484            return error.OutOfMemory;485        allocator.free(held);486    }487488    /// Refuses when the splice outgrew the reservation it was taken under.489    ///490    /// ON THE PATH THAT FINISHED, AND NOT ON THE PATH THAT REFUSED. A rewrite491    /// nested inside this one refuses by asking the segment for its own492    /// reservation and being told no, and that request is recorded on the493    /// context. A refusal leaves through the error and never reaches here, so494    /// an exhaustion recorded at the end of a rewrite that COMPLETED means495    /// this splice outgrew the bytes it reserved and a shortfall inside it was496    /// swallowed, which is the defect this file is fixing.497    ///498    /// WHY A REFUSAL AND NOT AN ASSERTION. An assertion is compiled out of a499    /// release build, so the one shape that outgrew the bound would emit an500    /// object in release and say nothing, which is the class of defect this501    /// pass exists to close. This is one load and one compare per rebuilt502    /// region. The refusal is safe to take because nothing reads the module503    /// after it: `promote` returns the error, the backend maps it to504    /// `BackendError.OutOfMemory`, and the consumer's `measureModule` forgets505    /// every measurement and writes no object.506    fn refuseOutgrownRewrite(self: *Promotion) PromotionError!void {507        if (self.ctx.exhaustion() != null) return error.OutOfMemory;508    }509510    /// The value each cell holds at the point the walk has reached. A null511    /// entry is a cell nothing has stored to yet, which a load may not read.512    /// Every `values` slice below holds one entry per cell, in cell order.513    fn rewriteBlock(514        self: *Promotion,515        block: *ir.Block,516        start: ?*ir.Operation,517        values: []?*ir.Value,518    ) PromotionError!void {519        std.debug.assert(values.len == self.cells.len);520        var current = start;521        while (current) |op| {522            const next = op.next_op;523            if (self.loadedCell(op)) |index| {524                const loaded = op.getResult(0) orelse return error.InvalidPromotion;525                const replacement = values[index] orelse return error.InvalidPromotion;526                loaded.replaceAllUsesWith(replacement);527                op.erase();528                self.changed = true;529            } else if (self.storedCell(op)) |index| {530                values[index] = op.getOperand(0) orelse return error.InvalidPromotion;531                op.erase();532                self.changed = true;533            } else if (op.regions.items.len != 0 and self.subtreeUsesAnyCell(op)) {534                if (std.mem.eql(u8, op.name.name, scf.ForOp.operation_name)) {535                    try self.rewriteFor(block, op, values);536                } else if (std.mem.eql(u8, op.name.name, scf.WhileOp.operation_name)) {537                    try self.rewriteWhile(block, op, values);538                } else if (std.mem.eql(u8, op.name.name, scf.IfOp.operation_name)) {539                    try self.rewriteIf(block, op, values);540                } else {541                    return error.InvalidPromotion;542                }543            }544            current = next;545        }546    }547548    fn loadedCell(self: *Promotion, op: *ir.Operation) ?usize {549        if (!std.mem.eql(u8, op.name.name, memref.LoadOp.operation_name)) return null;550        return self.cellIndex(op.getOperand(0) orelse return null);551    }552553    fn storedCell(self: *Promotion, op: *ir.Operation) ?usize {554        if (!std.mem.eql(u8, op.name.name, memref.StoreOp.operation_name)) return null;555        return self.cellIndex(op.getOperand(1) orelse return null);556    }557558    fn cellIndex(self: *Promotion, value: *ir.Value) ?usize {559        for (self.cells, 0..) |cell, index| {560            if (cell.value == value) return index;561        }562        return null;563    }564565    fn subtreeUsesAnyCell(self: *Promotion, op: *ir.Operation) bool {566        for (self.cells) |cell| {567            if (subtreeUsesCell(op, cell.value)) return true;568        }569        return false;570    }571572    /// The cells a region stores to, and how many of them there are. A cell the573    /// region only reads keeps the value it had outside, so it is not carried.574    const Carried = struct {575        flags: []bool,576        count: usize,577    };578579    fn carriedBuffer(self: *Promotion) PromotionError![]bool {580        const flags = self.allocator.alloc(bool, self.cells.len) catch return error.OutOfMemory;581        @memset(flags, false);582        return flags;583    }584585    fn valuesBuffer(self: *Promotion, from: []?*ir.Value) PromotionError![]?*ir.Value {586        const copied = self.allocator.alloc(?*ir.Value, self.cells.len) catch return error.OutOfMemory;587        @memcpy(copied, from);588        return copied;589    }590591    /// Collects the carried values a rebuilt region receives, in cell order,592    /// and the type of each, so the inits and the results agree on the order.593    fn collectCarried(594        self: *Promotion,595        carried: Carried,596        values: []?*ir.Value,597        parent: *ir.Block,598        before: *ir.Operation,599        loc: ir.Location,600        inits: *std.ArrayListUnmanaged(*ir.Value),601        result_types: *std.ArrayListUnmanaged(ir.Type),602    ) PromotionError!void {603        for (self.cells, 0..) |cell, index| {604            if (!carried.flags[index]) continue;605            const init_value = values[index] orelse try self.zeroValue(cell, parent, before, loc);606            try inits.append(self.allocator, init_value);607            try result_types.append(self.allocator, cell.element_type);608        }609    }610611    /// Points `values` at the carried arguments a rebuilt block receives.612    fn bindCarried(613        self: *Promotion,614        carried: Carried,615        block: *ir.Block,616        values: []?*ir.Value,617    ) void {618        const base = block.arguments.items.len - carried.count;619        var rank: usize = 0;620        for (self.cells, 0..) |_, index| {621            if (!carried.flags[index]) continue;622            values[index] = block.arguments.items[base + rank];623            rank += 1;624        }625    }626627    fn rewriteFor(628        self: *Promotion,629        parent: *ir.Block,630        old_op: *ir.Operation,631        values: []?*ir.Value,632    ) PromotionError!void {633        try self.reserveRewrite();634        const old_for = scf.ForOp{ .op = old_op };635        const old_body = old_for.getBodyBlock();636637        const flags = try self.carriedBuffer();638        defer self.allocator.free(flags);639        var count: usize = 0;640        for (self.cells, 0..) |cell, index| {641            const body = analyzeOps(cell.alloca, firstOp(old_body), values[index] != null);642            flags[index] = body.stores;643            if (body.stores) count += 1;644        }645        const carried = Carried{ .flags = flags, .count = count };646647        if (count == 0) {648            try self.rewriteBlock(old_body, firstOp(old_body), values);649            return;650        }651652        const loc = old_op.getLoc();653654        var inits = std.ArrayListUnmanaged(*ir.Value).empty;655        defer inits.deinit(self.allocator);656        for (old_op.operands.items[3..]) |operand| {657            try inits.append(self.allocator, operand.value);658        }659660        var result_types = std.ArrayListUnmanaged(ir.Type).empty;661        defer result_types.deinit(self.allocator);662        for (old_op.results.items) |result| {663            try result_types.append(self.allocator, result.type);664        }665666        try self.collectCarried(carried, values, parent, old_op, loc, &inits, &result_types);667668        const new_for = scf.ForOp.create(669            self.ctx,670            loc,671            old_for.getLowerBound(),672            old_for.getUpperBound(),673            old_for.getStep(),674            inits.items,675            result_types.items,676        ) catch return error.OutOfMemory;677        parent.addOperation(new_for.op) catch return error.OutOfMemory;678        new_for.op.moveBefore(old_op) catch return error.OutOfMemory;679680        const new_body = new_for.getBodyBlock();681        try self.moveBlockBody(old_body, new_body);682        for (old_op.results.items, 0..) |*old_result, index| {683            old_result.replaceAllUsesWith(new_for.op.getResult(index) orelse return error.InvalidPromotion);684        }685        old_op.erase();686687        const body_values = try self.valuesBuffer(values);688        defer self.allocator.free(body_values);689        self.bindCarried(carried, new_body, body_values);690691        const entering = try self.valuesBuffer(body_values);692        defer self.allocator.free(entering);693694        try self.rewriteBlock(new_body, firstOp(new_body), body_values);695        try self.extendYield(new_body, loc, carried, body_values, entering);696        try self.takeResults(carried, new_for.op, values);697        self.created.operations += 1;698        self.created.regions_rebuilt += 1;699        self.changed = true;700        try self.refuseOutgrownRewrite();701    }702703    /// Rebuilds a `scf.while` that stores cells with those cells' values as704    /// more carried values. The before block receives them, forwards what it705    /// holds at the condition to both the after block and the loop's results,706    /// and the after block yields them back.707    fn rewriteWhile(708        self: *Promotion,709        parent: *ir.Block,710        old_op: *ir.Operation,711        values: []?*ir.Value,712    ) PromotionError!void {713        try self.reserveRewrite();714        const old_while = scf.WhileOp{ .op = old_op };715        const old_before = old_while.getBeforeBlock();716        const old_after = old_while.getAfterBlock();717718        const flags = try self.carriedBuffer();719        defer self.allocator.free(flags);720        var count: usize = 0;721        for (self.cells, 0..) |cell, index| {722            const loop = analyzeWhile(cell.alloca, old_op, values[index] != null);723            std.debug.assert(loop.ok);724            flags[index] = loop.stores;725            if (loop.stores) count += 1;726        }727        const carried = Carried{ .flags = flags, .count = count };728729        if (count == 0) {730            const before_values = try self.valuesBuffer(values);731            defer self.allocator.free(before_values);732            try self.rewriteBlock(old_before, firstOp(old_before), before_values);733            const after_values = try self.valuesBuffer(values);734            defer self.allocator.free(after_values);735            try self.rewriteBlock(old_after, firstOp(old_after), after_values);736            return;737        }738739        const loc = old_op.getLoc();740741        var inits = std.ArrayListUnmanaged(*ir.Value).empty;742        defer inits.deinit(self.allocator);743        for (old_op.operands.items) |operand| {744            try inits.append(self.allocator, operand.value);745        }746747        var result_types = std.ArrayListUnmanaged(ir.Type).empty;748        defer result_types.deinit(self.allocator);749        for (old_op.results.items) |result| {750            try result_types.append(self.allocator, result.type);751        }752753        try self.collectCarried(carried, values, parent, old_op, loc, &inits, &result_types);754755        const new_while = scf.WhileOp.create(self.ctx, loc, inits.items, result_types.items) catch756            return error.OutOfMemory;757        parent.addOperation(new_while.op) catch return error.OutOfMemory;758        new_while.op.moveBefore(old_op) catch return error.OutOfMemory;759760        const new_before = new_while.getBeforeBlock();761        const new_after = new_while.getAfterBlock();762        try self.moveBlockBody(old_before, new_before);763        try self.moveBlockBody(old_after, new_after);764        for (old_op.results.items, 0..) |*old_result, index| {765            old_result.replaceAllUsesWith(766                new_while.op.getResult(index) orelse return error.InvalidPromotion,767            );768        }769        old_op.erase();770771        const before_values = try self.valuesBuffer(values);772        defer self.allocator.free(before_values);773        self.bindCarried(carried, new_before, before_values);774        const before_entering = try self.valuesBuffer(before_values);775        defer self.allocator.free(before_entering);776777        try self.rewriteBlock(new_before, firstOp(new_before), before_values);778        try self.extendCondition(new_before, loc, carried, before_values, before_entering);779780        const after_values = try self.valuesBuffer(values);781        defer self.allocator.free(after_values);782        self.bindCarried(carried, new_after, after_values);783        const after_entering = try self.valuesBuffer(after_values);784        defer self.allocator.free(after_entering);785786        try self.rewriteBlock(new_after, firstOp(new_after), after_values);787        try self.extendYield(new_after, loc, carried, after_values, after_entering);788789        try self.takeResults(carried, new_while.op, values);790        self.created.operations += 1;791        self.created.regions_rebuilt += 1;792        self.changed = true;793        try self.refuseOutgrownRewrite();794    }795796    fn rewriteIf(797        self: *Promotion,798        parent: *ir.Block,799        old_op: *ir.Operation,800        values: []?*ir.Value,801    ) PromotionError!void {802        try self.reserveRewrite();803        const old_if = scf.IfOp{ .op = old_op };804        const old_then = old_if.getThenBlock();805        const old_else = old_if.getElseBlock();806807        const flags = try self.carriedBuffer();808        defer self.allocator.free(flags);809        var count: usize = 0;810        for (self.cells, 0..) |cell, index| {811            const defined = values[index] != null;812            const then_branch = analyzeOps(cell.alloca, firstOp(old_then), defined);813            const else_stores = if (old_else) |block|814                analyzeOps(cell.alloca, firstOp(block), defined).stores815            else816                false;817            flags[index] = then_branch.stores or else_stores;818            if (flags[index]) count += 1;819        }820        const carried = Carried{ .flags = flags, .count = count };821822        if (count == 0) {823            const then_values = try self.valuesBuffer(values);824            defer self.allocator.free(then_values);825            try self.rewriteBlock(old_then, firstOp(old_then), then_values);826            if (old_else) |block| {827                const else_values = try self.valuesBuffer(values);828                defer self.allocator.free(else_values);829                try self.rewriteBlock(block, firstOp(block), else_values);830            }831            return;832        }833834        const loc = old_op.getLoc();835836        var result_types = std.ArrayListUnmanaged(ir.Type).empty;837        defer result_types.deinit(self.allocator);838        for (old_op.results.items) |result| {839            try result_types.append(self.allocator, result.type);840        }841        for (self.cells, 0..) |cell, index| {842            if (!carried.flags[index]) continue;843            try result_types.append(self.allocator, cell.element_type);844        }845846        const new_if = scf.IfOp.create(847            self.ctx,848            loc,849            old_if.getCondition(),850            result_types.items,851        ) catch return error.OutOfMemory;852        parent.addOperation(new_if.op) catch return error.OutOfMemory;853        new_if.op.moveBefore(old_op) catch return error.OutOfMemory;854855        const new_then = new_if.getThenBlock();856        const new_else = new_if.getElseBlock() orelse return error.InvalidPromotion;857        try self.moveBranch(old_then, new_then);858        if (old_else) |block| {859            try self.moveBranch(block, new_else);860        }861862        for (old_op.results.items, 0..) |*old_result, index| {863            old_result.replaceAllUsesWith(new_if.op.getResult(index) orelse return error.InvalidPromotion);864        }865        old_op.erase();866867        const then_values = try self.valuesBuffer(values);868        defer self.allocator.free(then_values);869        try self.rewriteBlock(new_then, firstOp(new_then), then_values);870        try self.extendBranchYield(new_then, loc, carried, then_values, values);871872        const else_values = try self.valuesBuffer(values);873        defer self.allocator.free(else_values);874        try self.rewriteBlock(new_else, firstOp(new_else), else_values);875        try self.extendBranchYield(new_else, loc, carried, else_values, values);876877        try self.takeResults(carried, new_if.op, values);878        self.created.operations += 1;879        self.created.regions_rebuilt += 1;880        self.changed = true;881        try self.refuseOutgrownRewrite();882    }883884    /// Points `values` at the carried results a rebuilt operation returns.885    fn takeResults(886        self: *Promotion,887        carried: Carried,888        op: *ir.Operation,889        values: []?*ir.Value,890    ) PromotionError!void {891        const base = op.results.items.len - carried.count;892        var rank: usize = 0;893        for (self.cells, 0..) |_, index| {894            if (!carried.flags[index]) continue;895            values[index] = op.getResult(base + rank) orelse return error.InvalidPromotion;896            rank += 1;897        }898    }899900    /// Moves every operation of `from` into `to` and points uses of `from`'s901    /// arguments at the arguments `to` holds in the same positions. `to` may902    /// hold more arguments than `from`.903    fn moveBlockBody(self: *Promotion, from: *ir.Block, to: *ir.Block) PromotionError!void {904        std.debug.assert(from.arguments.items.len <= to.arguments.items.len);905        try self.moveBranch(from, to);906        for (from.arguments.items, 0..) |old_arg, index| {907            old_arg.replaceAllUsesWith(to.arguments.items[index]);908        }909    }910911    fn moveBranch(self: *Promotion, from: *ir.Block, to: *ir.Block) PromotionError!void {912        _ = self;913        var op_iter: ?*ir.Operation = @ptrCast(@alignCast(from.operations.head));914        while (op_iter) |moved| {915            const next = moved.next_op;916            moved.moveToEnd(to) catch return error.OutOfMemory;917            op_iter = next;918        }919    }920921    /// The value each carried cell leaves a region with, in cell order. A cell922    /// the walk left holding nothing leaves with what it entered holding.923    fn finalValues(924        self: *Promotion,925        carried: Carried,926        values: []?*ir.Value,927        entering: []const ?*ir.Value,928        out: *std.ArrayListUnmanaged(*ir.Value),929    ) PromotionError!void {930        for (self.cells, 0..) |_, index| {931            if (!carried.flags[index]) continue;932            const final = values[index] orelse entering[index] orelse return error.InvalidPromotion;933            try out.append(self.allocator, final);934        }935    }936937    fn extendYield(938        self: *Promotion,939        block: *ir.Block,940        loc: ir.Location,941        carried: Carried,942        values: []?*ir.Value,943        entering: []const ?*ir.Value,944    ) PromotionError!void {945        const yield_op = lastOp(block) orelse return error.InvalidPromotion;946        if (!std.mem.eql(u8, yield_op.name.name, scf.YieldOp.operation_name)) return error.InvalidPromotion;947        var yields = std.ArrayListUnmanaged(*ir.Value).empty;948        defer yields.deinit(self.allocator);949        for (yield_op.operands.items) |operand| {950            try yields.append(self.allocator, operand.value);951        }952        try self.finalValues(carried, values, entering, &yields);953        const new_yield = scf.YieldOp.create(self.ctx, loc, yields.items) catch return error.OutOfMemory;954        block.addOperation(new_yield.op) catch return error.OutOfMemory;955        new_yield.op.moveBefore(yield_op) catch return error.OutOfMemory;956        yield_op.erase();957        self.created.operations += 1;958    }959960    /// A branch of an `scf.if` may hold no terminator yet, so this adds one961    /// when the branch has none and extends the one it has otherwise.962    fn extendBranchYield(963        self: *Promotion,964        block: *ir.Block,965        loc: ir.Location,966        carried: Carried,967        values: []?*ir.Value,968        entering: []const ?*ir.Value,969    ) PromotionError!void {970        var yields = std.ArrayListUnmanaged(*ir.Value).empty;971        defer yields.deinit(self.allocator);972        const existing = lastOp(block);973        const extends = existing != null and974            std.mem.eql(u8, existing.?.name.name, scf.YieldOp.operation_name);975        if (extends) {976            for (existing.?.operands.items) |operand| {977                try yields.append(self.allocator, operand.value);978            }979        }980        try self.finalValues(carried, values, entering, &yields);981        const new_yield = scf.YieldOp.create(self.ctx, loc, yields.items) catch return error.OutOfMemory;982        block.addOperation(new_yield.op) catch return error.OutOfMemory;983        if (extends) {984            new_yield.op.moveBefore(existing.?) catch return error.OutOfMemory;985            existing.?.erase();986        }987        self.created.operations += 1;988    }989990    fn extendCondition(991        self: *Promotion,992        block: *ir.Block,993        loc: ir.Location,994        carried: Carried,995        values: []?*ir.Value,996        entering: []const ?*ir.Value,997    ) PromotionError!void {998        const condition_op = lastOp(block) orelse return error.InvalidPromotion;999        if (!std.mem.eql(u8, condition_op.name.name, scf.ConditionOp.operation_name)) {1000            return error.InvalidPromotion;1001        }1002        const old_condition = scf.ConditionOp{ .op = condition_op };1003        var forwarded = std.ArrayListUnmanaged(*ir.Value).empty;1004        defer forwarded.deinit(self.allocator);1005        for (condition_op.operands.items[1..]) |operand| {1006            try forwarded.append(self.allocator, operand.value);1007        }1008        try self.finalValues(carried, values, entering, &forwarded);1009        const new_condition = scf.ConditionOp.create(1010            self.ctx,1011            loc,1012            old_condition.getCondition(),1013            forwarded.items,1014        ) catch return error.OutOfMemory;1015        block.addOperation(new_condition.op) catch return error.OutOfMemory;1016        new_condition.op.moveBefore(condition_op) catch return error.OutOfMemory;1017        condition_op.erase();1018        self.created.operations += 1;1019    }10201021    fn zeroValue(1022        self: *Promotion,1023        cell: Cell,1024        parent: *ir.Block,1025        before: *ir.Operation,1026        loc: ir.Location,1027    ) PromotionError!*ir.Value {1028        const name = cell.element_type.getDialectTypeName() orelse return error.InvalidPromotion;1029        const kind = arith_types.scalarKindFromTypeName(name) orelse return error.InvalidPromotion;1030        const constant = if (arith_types.scalarDescriptor(kind).class == .float)1031            arith.ConstantOp.createFloat(self.ctx, loc, cell.element_type, 0.0) catch return error.OutOfMemory1032        else if (kind == .bool)1033            arith.ConstantOp.createBool(self.ctx, loc, false) catch return error.OutOfMemory1034        else1035            arith.ConstantOp.createInt(self.ctx, loc, cell.element_type, 0) catch return error.OutOfMemory;1036        parent.addOperation(constant.op) catch return error.OutOfMemory;1037        constant.op.moveBefore(before) catch return error.OutOfMemory;1038        self.created.operations += 1;1039        var mutable = constant;1040        return mutable.getResult();1041    }1042};10431044fn lastOp(block: *ir.Block) ?*ir.Operation {1045    return @ptrCast(@alignCast(block.operations.tail));1046}10471048const testing = std.testing;10491050fn runPromotion(allocator: std.mem.Allocator, module: *ir.Operation) !bool {1051    var manager = pass_mod.PassManager.init(allocator);1052    defer manager.deinit();1053    try manager.addPass(createMemoryPromotionPass());1054    try testing.expectEqual(PassResult.success, manager.run(module, module.context));1055    return manager.stats.passes_modified > 0;1056}10571058fn countOps(op: *ir.Operation, name: []const u8) usize {1059    var count: usize = 0;1060    if (std.mem.eql(u8, op.name.name, name)) count += 1;1061    for (op.regions.items) |*region| {1062        var block_iter = region.getBlocks();1063        while (block_iter.next()) |block| {1064            var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));1065            while (current) |nested| {1066                count += countOps(nested, name);1067                current = nested.next_op;1068            }1069        }1070    }1071    return count;1072}10731074const TestFunction = struct {1075    ctx: ir.Context,1076    module: dialects.builtin.BuiltinDialect.ModuleOp,1077    func: func_dialect.FuncOp,1078    alloca: memref.AllocaOp,1079    zero: *ir.Value,10801081    fn init(allocator: std.mem.Allocator) !*TestFunction {1082        const self = try allocator.create(TestFunction);1083        self.ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);1084        errdefer self.ctx.deinit(allocator);1085        try dialects.registerAllDialects(&self.ctx);1086        const loc = ir.Location.getUnknown();1087        self.module = try dialects.builtin.BuiltinDialect.ModuleOp.create(&self.ctx, loc);1088        const i64_type = try arith.getScalarType(&self.ctx, .i64);1089        self.func = try func_dialect.FuncOp.create(&self.ctx, loc, "f", &.{i64_type}, &.{i64_type});1090        try self.module.getBodyBlock().addOperation(self.func.op);1091        const entry = self.func.getEntryBlock();10921093        const index_type = try arith.getIndexType(&self.ctx);1094        var zero = try arith.ConstantOp.createInt(&self.ctx, loc, index_type, 0);1095        try entry.addOperation(zero.op);1096        self.zero = zero.getResult();10971098        const cell_type = try memref.getMemrefType1D(&self.ctx, 1, i64_type, .host);1099        self.alloca = try memref.AllocaOp.createStatic(&self.ctx, loc, cell_type);1100        try entry.addOperation(self.alloca.op);1101        return self;1102    }11031104    fn deinit(self: *TestFunction, allocator: std.mem.Allocator) void {1105        self.ctx.deinit(allocator);1106        allocator.destroy(self);1107    }1108};11091110test "F10a promotion discards the accesses of a local cell in straight line code" {1111    const allocator = testing.allocator;1112    var fixture = try TestFunction.init(allocator);1113    defer fixture.deinit(allocator);1114    const loc = ir.Location.getUnknown();1115    const entry = fixture.func.getEntryBlock();1116    const i64_type = try arith.getScalarType(&fixture.ctx, .i64);11171118    const store = try memref.StoreOp.create(&fixture.ctx, loc, fixture.func.getArgument(0), fixture.alloca.getResult(), fixture.zero);1119    try entry.addOperation(store.op);1120    var load = try memref.LoadOp.create(&fixture.ctx, loc, fixture.alloca.getResult(), fixture.zero, i64_type);1121    try entry.addOperation(load.op);1122    const ret = try func_dialect.ReturnOp.create(&fixture.ctx, loc, &.{load.getResult()});1123    try entry.addOperation(ret.op);11241125    try testing.expect(try runPromotion(allocator, fixture.module.op));1126    try testing.expectEqual(@as(usize, 1), countOps(fixture.module.op, memref.AllocaOp.operation_name));1127    try testing.expectEqual(@as(usize, 0), countOps(fixture.module.op, memref.LoadOp.operation_name));1128    try testing.expectEqual(1129        @as(usize, 0),1130        countOps(fixture.module.op, memref.StoreOp.operation_name),1131    );1132    try testing.expect(ret.op.getOperand(0).? == fixture.func.getArgument(0));1133    try ir.verifyOperation(fixture.module.op, ir.verify.default_options);1134}11351136test "F10a promotion carries a local cell through a loop as a value" {1137    const allocator = testing.allocator;1138    var fixture = try TestFunction.init(allocator);1139    defer fixture.deinit(allocator);1140    const loc = ir.Location.getUnknown();1141    const entry = fixture.func.getEntryBlock();1142    const i64_type = try arith.getScalarType(&fixture.ctx, .i64);1143    const index_type = try arith.getIndexType(&fixture.ctx);11441145    var init_value = try arith.ConstantOp.createInt(&fixture.ctx, loc, i64_type, 0);1146    try entry.addOperation(init_value.op);1147    const init_store = try memref.StoreOp.create(&fixture.ctx, loc, init_value.getResult(), fixture.alloca.getResult(), fixture.zero);1148    try entry.addOperation(init_store.op);11491150    var bound = try arith.ConstantOp.createInt(&fixture.ctx, loc, index_type, 8);1151    try entry.addOperation(bound.op);1152    var step = try arith.ConstantOp.createInt(&fixture.ctx, loc, index_type, 1);1153    try entry.addOperation(step.op);1154    var for_op = try scf.ForOp.create(&fixture.ctx, loc, fixture.zero, bound.getResult(), step.getResult(), &.{}, &.{});1155    try entry.addOperation(for_op.op);1156    const body = for_op.getBodyBlock();1157    {1158        var body_zero = try arith.ConstantOp.createInt(&fixture.ctx, loc, index_type, 0);1159        try body.addOperation(body_zero.op);1160        var current = try memref.LoadOp.create(&fixture.ctx, loc, fixture.alloca.getResult(), body_zero.getResult(), i64_type);1161        try body.addOperation(current.op);1162        var one = try arith.ConstantOp.createInt(&fixture.ctx, loc, i64_type, 1);1163        try body.addOperation(one.op);1164        var next = try arith.AddOp.create(&fixture.ctx, loc, current.getResult(), one.getResult());1165        try body.addOperation(next.op);1166        const body_store = try memref.StoreOp.create(&fixture.ctx, loc, next.getResult(), fixture.alloca.getResult(), body_zero.getResult());1167        try body.addOperation(body_store.op);1168        const yield = try scf.YieldOp.create(&fixture.ctx, loc, &.{});1169        try body.addOperation(yield.op);1170    }11711172    var final = try memref.LoadOp.create(&fixture.ctx, loc, fixture.alloca.getResult(), fixture.zero, i64_type);1173    try entry.addOperation(final.op);1174    const ret = try func_dialect.ReturnOp.create(&fixture.ctx, loc, &.{final.getResult()});1175    try entry.addOperation(ret.op);11761177    try testing.expect(try runPromotion(allocator, fixture.module.op));1178    try testing.expectEqual(@as(usize, 1), countOps(fixture.module.op, memref.AllocaOp.operation_name));1179    try testing.expectEqual(1180        @as(usize, 0),1181        countOps(fixture.module.op, memref.LoadOp.operation_name),1182    );1183    try testing.expectEqual(1184        @as(usize, 0),1185        countOps(fixture.module.op, memref.StoreOp.operation_name),1186    );11871188    var found_for: ?*ir.Operation = null;1189    var current: ?*ir.Operation = @ptrCast(@alignCast(entry.operations.head));1190    while (current) |op| {1191        if (std.mem.eql(u8, op.name.name, scf.ForOp.operation_name)) found_for = op;1192        current = op.next_op;1193    }1194    const promoted_for = found_for orelse return error.TestFailure;1195    try testing.expectEqual(@as(usize, 4), promoted_for.operands.items.len);1196    try testing.expectEqual(@as(usize, 1), promoted_for.results.items.len);1197    try testing.expect(ret.op.getOperand(0).? == &promoted_for.results.items[0]);1198    try ir.verifyOperation(fixture.module.op, ir.verify.default_options);1199}12001201test "leaves read before write and escaping cells alone" {1202    const allocator = testing.allocator;1203    var fixture = try TestFunction.init(allocator);1204    defer fixture.deinit(allocator);1205    const loc = ir.Location.getUnknown();1206    const entry = fixture.func.getEntryBlock();1207    const i64_type = try arith.getScalarType(&fixture.ctx, .i64);12081209    var early = try memref.LoadOp.create(&fixture.ctx, loc, fixture.alloca.getResult(), fixture.zero, i64_type);1210    try entry.addOperation(early.op);1211    const ret = try func_dialect.ReturnOp.create(&fixture.ctx, loc, &.{early.getResult()});1212    try entry.addOperation(ret.op);12131214    try testing.expect(!try runPromotion(allocator, fixture.module.op));1215    try testing.expectEqual(@as(usize, 1), countOps(fixture.module.op, memref.AllocaOp.operation_name));1216    try testing.expectEqual(@as(usize, 1), countOps(fixture.module.op, memref.LoadOp.operation_name));1217}12181219test "F10a promotion retains a cell whose second use is a call" {1220    const allocator = testing.allocator;1221    var fixture = try TestFunction.init(allocator);1222    defer fixture.deinit(allocator);1223    const loc = ir.Location.getUnknown();1224    const entry = fixture.func.getEntryBlock();1225    const i64_type = try arith.getScalarType(&fixture.ctx, .i64);12261227    const store = try memref.StoreOp.create(1228        &fixture.ctx,1229        loc,1230        fixture.func.getArgument(0),1231        fixture.alloca.getResult(),1232        fixture.zero,1233    );1234    try entry.addOperation(store.op);1235    const call = try func_dialect.CallOp.create(1236        &fixture.ctx,1237        loc,1238        "observer",1239        &.{fixture.alloca.getResult()},1240        &.{},1241    );1242    try entry.addOperation(call.op);1243    const load = try memref.LoadOp.create(1244        &fixture.ctx,1245        loc,1246        fixture.alloca.getResult(),1247        fixture.zero,1248        i64_type,1249    );1250    try entry.addOperation(load.op);1251    const ret = try func_dialect.ReturnOp.create(&fixture.ctx, loc, &.{load.getResult()});1252    try entry.addOperation(ret.op);12531254    try testing.expect(!try runPromotion(allocator, fixture.module.op));1255    try testing.expectEqual(1256        @as(usize, 1),1257        countOps(fixture.module.op, memref.LoadOp.operation_name),1258    );1259    try testing.expectEqual(1260        @as(usize, 1),1261        countOps(fixture.module.op, memref.StoreOp.operation_name),1262    );1263    try testing.expect(ret.op.getOperand(0).? == load.getResult());1264}12651266test "F10a promotion retains a cell an index the operands do not settle reaches" {1267    const allocator = testing.allocator;1268    var fixture = try TestFunction.init(allocator);1269    defer fixture.deinit(allocator);1270    const loc = ir.Location.getUnknown();1271    const entry = fixture.func.getEntryBlock();1272    const i64_type = try arith.getScalarType(&fixture.ctx, .i64);12731274    const store = try memref.StoreOp.create(1275        &fixture.ctx,1276        loc,1277        fixture.func.getArgument(0),1278        fixture.alloca.getResult(),1279        fixture.zero,1280    );1281    try entry.addOperation(store.op);1282    var computed = try arith.AddOp.create(&fixture.ctx, loc, fixture.zero, fixture.zero);1283    try entry.addOperation(computed.op);1284    const load = try memref.LoadOp.create(1285        &fixture.ctx,1286        loc,1287        fixture.alloca.getResult(),1288        computed.getResult(),1289        i64_type,1290    );1291    try entry.addOperation(load.op);1292    const ret = try func_dialect.ReturnOp.create(&fixture.ctx, loc, &.{load.getResult()});1293    try entry.addOperation(ret.op);12941295    try testing.expect(!try runPromotion(allocator, fixture.module.op));1296    try testing.expectEqual(1297        @as(usize, 1),1298        countOps(fixture.module.op, memref.LoadOp.operation_name),1299    );1300    try testing.expectEqual(1301        @as(usize, 1),1302        countOps(fixture.module.op, memref.StoreOp.operation_name),1303    );1304}13051306test "F10a promotion retains a cell of two elements" {1307    const allocator = testing.allocator;1308    var fixture = try TestFunction.init(allocator);1309    defer fixture.deinit(allocator);1310    const loc = ir.Location.getUnknown();1311    const entry = fixture.func.getEntryBlock();1312    const i64_type = try arith.getScalarType(&fixture.ctx, .i64);13131314    const pair_type = try memref.getMemrefType1D(&fixture.ctx, 2, i64_type, .host);1315    const pair = try memref.AllocaOp.createStatic(&fixture.ctx, loc, pair_type);1316    try entry.addOperation(pair.op);1317    const store = try memref.StoreOp.create(1318        &fixture.ctx,1319        loc,1320        fixture.func.getArgument(0),1321        pair.getResult(),1322        fixture.zero,1323    );1324    try entry.addOperation(store.op);1325    const load = try memref.LoadOp.create(1326        &fixture.ctx,1327        loc,1328        pair.getResult(),1329        fixture.zero,1330        i64_type,1331    );1332    try entry.addOperation(load.op);1333    const ret = try func_dialect.ReturnOp.create(&fixture.ctx, loc, &.{load.getResult()});1334    try entry.addOperation(ret.op);13351336    try testing.expect(!try runPromotion(allocator, fixture.module.op));1337    try testing.expectEqual(1338        @as(usize, 1),1339        countOps(fixture.module.op, memref.LoadOp.operation_name),1340    );1341    try testing.expectEqual(1342        @as(usize, 1),1343        countOps(fixture.module.op, memref.StoreOp.operation_name),1344    );1345}13461347test "F10a promotion retains a cell a store hands to another cell" {1348    const allocator = testing.allocator;1349    var fixture = try TestFunction.init(allocator);1350    defer fixture.deinit(allocator);1351    const loc = ir.Location.getUnknown();1352    const entry = fixture.func.getEntryBlock();1353    const i64_type = try arith.getScalarType(&fixture.ctx, .i64);13541355    const cell_type = try memref.getMemrefType1D(&fixture.ctx, 1, i64_type, .host);1356    const holder_type = try memref.getMemrefType1D(&fixture.ctx, 1, cell_type, .host);1357    const holder = try memref.AllocaOp.createStatic(&fixture.ctx, loc, holder_type);1358    try entry.addOperation(holder.op);1359    const keeper = try func_dialect.CallOp.create(1360        &fixture.ctx,1361        loc,1362        "observer",1363        &.{holder.getResult()},1364        &.{},1365    );1366    try entry.addOperation(keeper.op);1367    const store = try memref.StoreOp.create(1368        &fixture.ctx,1369        loc,1370        fixture.func.getArgument(0),1371        fixture.alloca.getResult(),1372        fixture.zero,1373    );1374    try entry.addOperation(store.op);1375    const escape = try memref.StoreOp.create(1376        &fixture.ctx,1377        loc,1378        fixture.alloca.getResult(),1379        holder.getResult(),1380        fixture.zero,1381    );1382    try entry.addOperation(escape.op);1383    const load = try memref.LoadOp.create(1384        &fixture.ctx,1385        loc,1386        fixture.alloca.getResult(),1387        fixture.zero,1388        i64_type,1389    );1390    try entry.addOperation(load.op);1391    const ret = try func_dialect.ReturnOp.create(&fixture.ctx, loc, &.{load.getResult()});1392    try entry.addOperation(ret.op);13931394    try testing.expect(!try runPromotion(allocator, fixture.module.op));1395    try testing.expectEqual(1396        @as(usize, 1),1397        countOps(fixture.module.op, memref.LoadOp.operation_name),1398    );1399    try testing.expectEqual(1400        @as(usize, 2),1401        countOps(fixture.module.op, memref.StoreOp.operation_name),1402    );1403    try testing.expect(ret.op.getOperand(0).? == load.getResult());1404}14051406test "F10a promotion carries a local cell through a while as a loop value" {1407    const allocator = testing.allocator;1408    var fixture = try TestFunction.init(allocator);1409    defer fixture.deinit(allocator);1410    const loc = ir.Location.getUnknown();1411    const entry = fixture.func.getEntryBlock();1412    const i64_type = try arith.getScalarType(&fixture.ctx, .i64);14131414    const init_store = try memref.StoreOp.create(1415        &fixture.ctx,1416        loc,1417        fixture.func.getArgument(0),1418        fixture.alloca.getResult(),1419        fixture.zero,1420    );1421    try entry.addOperation(init_store.op);14221423    var loop = try scf.WhileOp.create(&fixture.ctx, loc, &.{}, &.{});1424    try entry.addOperation(loop.op);1425    {1426        const before = loop.getBeforeBlock();1427        var held = try memref.LoadOp.create(1428            &fixture.ctx,1429            loc,1430            fixture.alloca.getResult(),1431            fixture.zero,1432            i64_type,1433        );1434        try before.addOperation(held.op);1435        var limit = try arith.ConstantOp.createInt(&fixture.ctx, loc, i64_type, 10);1436        try before.addOperation(limit.op);1437        var below = try arith.CmpOp.create(&fixture.ctx, loc, .lt, held.getResult(), limit.getResult());1438        try before.addOperation(below.op);1439        const decided = try scf.ConditionOp.create(&fixture.ctx, loc, below.getResult(), &.{});1440        try before.addOperation(decided.op);14411442        const after = loop.getAfterBlock();1443        var current = try memref.LoadOp.create(1444            &fixture.ctx,1445            loc,1446            fixture.alloca.getResult(),1447            fixture.zero,1448            i64_type,1449        );1450        try after.addOperation(current.op);1451        var one = try arith.ConstantOp.createInt(&fixture.ctx, loc, i64_type, 1);1452        try after.addOperation(one.op);1453        var next = try arith.AddOp.create(&fixture.ctx, loc, current.getResult(), one.getResult());1454        try after.addOperation(next.op);1455        const body_store = try memref.StoreOp.create(1456            &fixture.ctx,1457            loc,1458            next.getResult(),1459            fixture.alloca.getResult(),1460            fixture.zero,1461        );1462        try after.addOperation(body_store.op);1463        const yield = try scf.YieldOp.create(&fixture.ctx, loc, &.{});1464        try after.addOperation(yield.op);1465    }14661467    var final = try memref.LoadOp.create(1468        &fixture.ctx,1469        loc,1470        fixture.alloca.getResult(),1471        fixture.zero,1472        i64_type,1473    );1474    try entry.addOperation(final.op);1475    const ret = try func_dialect.ReturnOp.create(&fixture.ctx, loc, &.{final.getResult()});1476    try entry.addOperation(ret.op);14771478    try testing.expect(try runPromotion(allocator, fixture.module.op));1479    try testing.expectEqual(1480        @as(usize, 0),1481        countOps(fixture.module.op, memref.LoadOp.operation_name),1482    );1483    try testing.expectEqual(1484        @as(usize, 0),1485        countOps(fixture.module.op, memref.StoreOp.operation_name),1486    );1487    try testing.expectEqual(1488        @as(usize, 1),1489        countOps(fixture.module.op, memref.AllocaOp.operation_name),1490    );14911492    var found: ?*ir.Operation = null;1493    var current: ?*ir.Operation = @ptrCast(@alignCast(entry.operations.head));1494    while (current) |op| {1495        if (std.mem.eql(u8, op.name.name, scf.WhileOp.operation_name)) found = op;1496        current = op.next_op;1497    }1498    const promoted = found orelse return error.TestFailure;1499    try testing.expectEqual(@as(usize, 1), promoted.operands.items.len);1500    try testing.expectEqual(@as(usize, 1), promoted.results.items.len);1501    try testing.expect(ret.op.getOperand(0).? == &promoted.results.items[0]);1502    try ir.verifyOperation(fixture.module.op, ir.verify.default_options);1503}15041505fn countAllOps(op: *ir.Operation) usize {1506    var count: usize = 1;1507    for (op.regions.items) |*region| {1508        var block_iter = region.getBlocks();1509        while (block_iter.next()) |block| {1510            var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));1511            while (current) |nested| {1512                count += countAllOps(nested);1513                current = nested.next_op;1514            }1515        }1516    }1517    return count;1518}15191520const CellLoop = struct {1521    ctx: ir.Context,1522    module: dialects.builtin.BuiltinDialect.ModuleOp,15231524    fn init(allocator: std.mem.Allocator, cells: usize, depth: usize, seeded: bool) !*CellLoop {1525        return initUnder(allocator, cells, depth, seeded, ir.Context.Limits.testing);1526    }15271528    /// The same program in a context of stated limits, for a test that wants1529    /// the nested segment to run short at a chosen point.1530    fn initUnder(1531        allocator: std.mem.Allocator,1532        cells: usize,1533        depth: usize,1534        seeded: bool,1535        limits: ir.Context.Limits,1536    ) !*CellLoop {1537        std.debug.assert(cells >= 1);1538        std.debug.assert(depth >= 1);1539        const self = try allocator.create(CellLoop);1540        self.ctx = try ir.Context.init(allocator, limits);1541        errdefer self.ctx.deinit(allocator);1542        try dialects.registerAllDialects(&self.ctx);15431544        const loc = ir.Location.getUnknown();1545        self.module = try dialects.builtin.BuiltinDialect.ModuleOp.create(&self.ctx, loc);1546        const i64_type = try arith.getScalarType(&self.ctx, .i64);1547        const index_type = try arith.getIndexType(&self.ctx);1548        var func = try func_dialect.FuncOp.create(&self.ctx, loc, "f", &.{i64_type}, &.{i64_type});1549        try self.module.getBodyBlock().addOperation(func.op);1550        const entry = func.getEntryBlock();15511552        var zero = try arith.ConstantOp.createInt(&self.ctx, loc, index_type, 0);1553        try entry.addOperation(zero.op);1554        const index_zero = zero.getResult();15551556        const cell_type = try memref.getMemrefType1D(&self.ctx, 1, i64_type, .host);1557        const holders = try allocator.alloc(*ir.Value, cells);1558        defer allocator.free(holders);1559        for (holders) |*holder| {1560            var alloca = try memref.AllocaOp.createStatic(&self.ctx, loc, cell_type);1561            try entry.addOperation(alloca.op);1562            holder.* = alloca.getResult();1563            if (!seeded) continue;1564            var seed = try arith.ConstantOp.createInt(&self.ctx, loc, i64_type, 0);1565            try entry.addOperation(seed.op);1566            const store = try memref.StoreOp.create(&self.ctx, loc, seed.getResult(), holder.*, index_zero);1567            try entry.addOperation(store.op);1568        }15691570        var bound = try arith.ConstantOp.createInt(&self.ctx, loc, index_type, 8);1571        try entry.addOperation(bound.op);1572        var step = try arith.ConstantOp.createInt(&self.ctx, loc, index_type, 1);1573        try entry.addOperation(step.op);15741575        var outer = entry;1576        var level: usize = 0;1577        while (level < depth) : (level += 1) {1578            var loop = try scf.ForOp.create(&self.ctx, loc, index_zero, bound.getResult(), step.getResult(), &.{}, &.{});1579            try outer.addOperation(loop.op);1580            outer = loop.getBodyBlock();1581        }15821583        var body_zero = try arith.ConstantOp.createInt(&self.ctx, loc, index_type, 0);1584        try outer.addOperation(body_zero.op);1585        for (holders) |holder| {1586            var one = try arith.ConstantOp.createInt(&self.ctx, loc, i64_type, 1);1587            try outer.addOperation(one.op);1588            if (seeded) {1589                var current = try memref.LoadOp.create(&self.ctx, loc, holder, body_zero.getResult(), i64_type);1590                try outer.addOperation(current.op);1591                var next = try arith.AddOp.create(&self.ctx, loc, current.getResult(), one.getResult());1592                try outer.addOperation(next.op);1593                const store = try memref.StoreOp.create(&self.ctx, loc, next.getResult(), holder, body_zero.getResult());1594                try outer.addOperation(store.op);1595                continue;1596            }1597            const store = try memref.StoreOp.create(&self.ctx, loc, one.getResult(), holder, body_zero.getResult());1598            try outer.addOperation(store.op);1599        }16001601        var closing = outer;1602        var closed: usize = 0;1603        while (closed < depth) : (closed += 1) {1604            const yield = try scf.YieldOp.create(&self.ctx, loc, &.{});1605            try closing.addOperation(yield.op);1606            const holder_op = closing.getParentOperation() orelse break;1607            closing = holder_op.getBlock() orelse break;1608        }16091610        var final = try memref.LoadOp.create(&self.ctx, loc, holders[0], index_zero, i64_type);1611        try entry.addOperation(final.op);1612        const ret = try func_dialect.ReturnOp.create(&self.ctx, loc, &.{final.getResult()});1613        try entry.addOperation(ret.op);1614        return self;1615    }16161617    fn deinit(self: *CellLoop, allocator: std.mem.Allocator) void {1618        self.ctx.deinit(allocator);1619        allocator.destroy(self);1620    }1621};16221623test "the operations promotion creates for one loop do not grow with the cells it carries" {1624    const allocator = testing.allocator;1625    const counts = [_]usize{ 1, 2, 8, 32 };1626    var first: ?Created = null;1627    for (counts) |cells| {1628        var fixture = try CellLoop.init(allocator, cells, 1, true);1629        defer fixture.deinit(allocator);1630        var created = Created{};1631        try testing.expect(try promoteCounting(fixture.module.op, allocator, &created));1632        try testing.expectEqual(@as(usize, 1), created.regions_rebuilt);1633        if (first) |seen| {1634            try testing.expectEqual(seen.operations, created.operations);1635        } else {1636            first = created;1637        }1638        try testing.expectEqual(@as(usize, 0), countOps(fixture.module.op, memref.LoadOp.operation_name));1639        try testing.expectEqual(@as(usize, 0), countOps(fixture.module.op, memref.StoreOp.operation_name));1640        try ir.verifyOperation(fixture.module.op, ir.verify.default_options);1641    }1642}16431644test "promotion rebuilds one region per region and stays inside the figure it declares" {1645    const allocator = testing.allocator;1646    const depths = [_]usize{ 1, 2, 4 };1647    for (depths) |depth| {1648        var fixture = try CellLoop.init(allocator, 4, depth, true);1649        defer fixture.deinit(allocator);1650        const given = countAllOps(fixture.module.op);1651        var created = Created{};1652        try testing.expect(try promoteCounting(fixture.module.op, allocator, &created));1653        try testing.expectEqual(depth, created.regions_rebuilt);1654        try testing.expect(created.operations <= operationsCreated(given));1655        try ir.verifyOperation(fixture.module.op, ir.verify.default_options);1656    }1657}16581659test "a loop carrying a cell no store defined creates one zero for it and refuses a later read" {1660    const allocator = testing.allocator;1661    var fixture = try CellLoop.init(allocator, 4, 1, false);1662    defer fixture.deinit(allocator);1663    const given = countAllOps(fixture.module.op);1664    var created = Created{};1665    try testing.expect(try promoteCounting(fixture.module.op, allocator, &created));16661667    try testing.expectEqual(@as(usize, 1), countOps(fixture.module.op, memref.LoadOp.operation_name));1668    try testing.expectEqual(@as(usize, 1), countOps(fixture.module.op, memref.StoreOp.operation_name));16691670    try testing.expectEqual(@as(usize, 1), created.regions_rebuilt);1671    try testing.expectEqual(@as(usize, 5), created.operations);1672    try testing.expect(created.operations <= operationsCreated(given));1673    try ir.verifyOperation(fixture.module.op, ir.verify.default_options);1674}16751676test "a rewrite that cannot reserve its splice leaves the module verifying" {1677    const allocator = testing.allocator;1678    const cells: usize = 4;1679    const depth: usize = 2;16801681    var built = try CellLoop.init(allocator, cells, depth, true);1682    const construction = built.ctx.capacityUsage().operation_nested.frontier_bytes;1683    built.deinit(allocator);16841685    var limits = ir.Context.Limits.testing;1686    limits.operations.nested_bytes = construction + rewriteBytes(cells) - 1;16871688    var fixture = try CellLoop.initUnder(allocator, cells, depth, true, limits);1689    defer fixture.deinit(allocator);16901691    const ops = countAllOps(fixture.module.op);1692    const loads = countOps(fixture.module.op, memref.LoadOp.operation_name);1693    const stores = countOps(fixture.module.op, memref.StoreOp.operation_name);1694    try testing.expect(loads > 0);1695    try testing.expect(stores > 0);16961697    var created = Created{};1698    try testing.expectError(1699        error.OutOfMemory,1700        promoteCounting(fixture.module.op, allocator, &created),1701    );17021703    try ir.verifyOperation(fixture.module.op, ir.verify.default_options);1704    try testing.expectEqual(@as(usize, 0), created.regions_rebuilt);1705    try testing.expectEqual(@as(usize, 0), created.operations);1706    try testing.expect(countAllOps(fixture.module.op) <= ops);1707    try testing.expect(countOps(fixture.module.op, memref.LoadOp.operation_name) <= loads);1708    try testing.expect(countOps(fixture.module.op, memref.StoreOp.operation_name) <= stores);1709}

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