tiny.choir.egraph.extract

1const std = @import("std");2const ir = @import("../core/root.zig");3const rewrite = ir.rewrite;4const graph_mod = @import("graph.zig");56const ClassId = graph_mod.ClassId;7const Node = graph_mod.Node;8const Graph = graph_mod.Graph;9const ValueEntry = graph_mod.ValueEntry;1011pub const infinite_cost = std.math.maxInt(u64);1213pub const OpCost = struct {14    name: []const u8,15    cost: u32,16};1718pub const CostModel = struct {19    operation: u32 = 4,20    constant: u32 = 1,21    constant_op_name: ?[]const u8 = null,22    overrides: []const OpCost = &.{},2324    pub fn operationCost(self: *const CostModel, op_name: []const u8) u32 {25        for (self.overrides) |entry| {26            if (std.mem.eql(u8, entry.name, op_name)) return entry.cost;27        }28        if (self.constant_op_name) |constant_name| {29            if (std.mem.eql(u8, constant_name, op_name)) return self.constant;30        }31        return self.operation;32    }33};3435const Choice = struct {36    cost: u64 = infinite_cost,37    node_cost: u64 = infinite_cost,38    node_index: ?usize = null,39};4041pub const Extraction = struct {42    allocator: std.mem.Allocator,43    graph: *Graph,44    model: CostModel,45    choices: []Choice,4647    /// Choice-table storage. Temporary materialization paths are charged separately.48    pub fn choiceStorageBound(classes: u64, blocks: u64) !u64 {49        const bytes = try std.math.mul(u64, classes, @sizeOf(Choice));50        return std.math.add(u64, bytes, try std.math.mul(u64, blocks, @alignOf(Choice)));51    }5253    pub fn init(allocator: std.mem.Allocator, graph: *Graph, model: CostModel) !Extraction {54        const choices = try allocator.alloc(Choice, graph.classCount());55        @memset(choices, .{});56        return .{57            .allocator = allocator,58            .graph = graph,59            .model = model,60            .choices = choices,61        };62    }6364    pub fn deinit(self: *Extraction) void {65        self.allocator.free(self.choices);66    }6768    pub fn analyze(self: *Extraction) void {69        for (0..self.graph.classCount()) |index| {70            const id = ClassId{ .index = @intCast(index) };71            const root = self.graph.find(id);72            if (!root.eql(id)) continue;73            if (self.graph.classValues(root).len > 0) {74                self.choices[index].cost = 0;75            }76        }7778        var changed = true;79        while (changed) {80            changed = false;81            for (0..self.graph.classCount()) |index| {82                const id = ClassId{ .index = @intCast(index) };83                const root = self.graph.find(id);84                if (!root.eql(id)) continue;8586                const slot = &self.choices[index];87                for (self.graph.nodes(root), 0..) |*node, node_index| {88                    if (node.kind != .operation) continue;89                    const cost = self.termCostOfNode(node);90                    if (cost < slot.node_cost) {91                        slot.node_cost = cost;92                        slot.node_index = node_index;93                        changed = true;94                    }95                }96                if (slot.node_cost < slot.cost) {97                    slot.cost = slot.node_cost;98                    changed = true;99                }100            }101        }102    }103104    pub fn classCost(self: *Extraction, class: ClassId) u64 {105        const root = self.graph.find(class);106        return self.choices[@intCast(root.index)].cost;107    }108109    pub fn nodeCost(self: *Extraction, class: ClassId) u64 {110        const root = self.graph.find(class);111        return self.choices[@intCast(root.index)].node_cost;112    }113114    pub fn termCost(self: *Extraction, op_name: []const u8, operand_classes: []const ClassId) u64 {115        var total: u64 = self.model.operationCost(op_name);116        for (operand_classes) |operand| {117            total +|= self.classCost(operand);118        }119        return @min(total, infinite_cost);120    }121122    fn termCostOfNode(self: *Extraction, node: *const Node) u64 {123        var total: u64 = self.model.operationCost(node.op_name);124        for (node.operands) |operand| {125            total +|= self.classCost(operand);126        }127        return total;128    }129130    pub fn usableValue(131        self: *Extraction,132        class: ClassId,133        order_limit: usize,134        exclude: ?*ir.Value,135    ) ?*ir.Value {136        const root = self.graph.find(class);137        var best: ?ValueEntry = null;138        for (self.graph.classValues(root)) |entry| {139            if (entry.order > order_limit) continue;140            if (exclude != null and entry.value == exclude.?) continue;141            if (best == null or entry.order < best.?.order) best = entry;142        }143        if (best) |entry| return entry.value;144        return null;145    }146147    pub fn materialize(148        self: *Extraction,149        rewriter: *rewrite.PatternRewriter,150        class: ClassId,151        before_op: *ir.Operation,152        order_limit: usize,153        exclude: ?*ir.Value,154    ) anyerror!?*ir.Value {155        var visiting = std.AutoHashMapUnmanaged(u32, void).empty;156        defer visiting.deinit(self.allocator);157        return self.materializeClass(rewriter, class, before_op, order_limit, exclude, &visiting);158    }159160    fn materializeClass(161        self: *Extraction,162        rewriter: *rewrite.PatternRewriter,163        class: ClassId,164        before_op: *ir.Operation,165        order_limit: usize,166        exclude: ?*ir.Value,167        visiting: *std.AutoHashMapUnmanaged(u32, void),168    ) anyerror!?*ir.Value {169        const root = self.graph.find(class);170171        if (self.usableValue(root, order_limit, exclude)) |value| return value;172173        if (visiting.contains(root.index)) return null;174        try visiting.put(self.allocator, root.index, {});175        defer _ = visiting.remove(root.index);176177        const slot = self.choices[@intCast(root.index)];178        if (slot.node_cost == infinite_cost) return null;179        const node_index = slot.node_index orelse return null;180181        const node = self.graph.nodes(root)[node_index];182        if (node.kind != .operation) return null;183        if (node.result_types.len != 1) return null;184185        var operand_values: std.ArrayListUnmanaged(*ir.Value) = .empty;186        defer operand_values.deinit(self.allocator);187        try operand_values.ensureTotalCapacity(self.allocator, node.operands.len);188        for (node.operands) |operand| {189            const operand_value = try self.materializeClass(190                rewriter,191                operand,192                before_op,193                order_limit,194                exclude,195                visiting,196            ) orelse return null;197            operand_values.appendAssumeCapacity(operand_value);198        }199200        var state = ir.Operation.State.init(node.op_name, before_op.location);201        state.addOperands(operand_values.items);202        state.addTypes(node.result_types);203204        if (node.raw_attributes.len > 0) state.addRawAttributes(node.raw_attributes);205        if (node.properties) |properties| try state.setPropertiesAttr(properties);206207        rewriter.setInsertionPointBefore(before_op);208        const new_op = try rewriter.ir_ctx.createOperation(state);209        var attached = false;210        defer if (!attached) new_op.erase();211        var declaration = try ir.interfaces.effects.inspect(self.allocator, new_op);212        defer declaration.deinit(self.allocator);213        if (new_op.hasTrait("is_terminator")) return null;214        if (new_op.hasInterface(ir.interfaces.SymbolOpInterface)) return null;215        if (!ir.interfaces.effects.duplicate(declaration.facts, .{})) return null;216        _ = try rewriter.insert(new_op);217        attached = true;218        const result = new_op.getResult(0) orelse return null;219        try self.graph.attachValue(root, result, 0, order_limit);220        return result;221    }222};

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