lib/choir/src/egraph/extract.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const ir = @import("../core/root.zig");
3 const rewrite = ir.rewrite;
4 const graph_mod = @import("graph.zig");
5
6 const ClassId = graph_mod.ClassId;
7 const Node = graph_mod.Node;
8 const Graph = graph_mod.Graph;
9 const ValueEntry = graph_mod.ValueEntry;
10
11 pub const infinite_cost = std.math.maxInt(u64);
12
13 pub const OpCost = struct {
14 name: []const u8,
15 cost: u32,
16 };
17
18 pub const CostModel = struct {
19 operation: u32 = 4,
20 constant: u32 = 1,
21 constant_op_name: ?[]const u8 = null,
22 overrides: []const OpCost = &.{},
23
24 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 };
34
35 const Choice = struct {
36 cost: u64 = infinite_cost,
37 node_cost: u64 = infinite_cost,
38 node_index: ?usize = null,
39 };
40
41 pub const Extraction = struct {
42 allocator: std.mem.Allocator,
43 graph: *Graph,
44 model: CostModel,
45 choices: []Choice,
46
47 /// 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 }
52
53 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 }
63
64 pub fn deinit(self: *Extraction) void {
65 self.allocator.free(self.choices);
66 }
67
68 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 }
77
78 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;
85
86 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 }
103
104 pub fn classCost(self: *Extraction, class: ClassId) u64 {
105 const root = self.graph.find(class);
106 return self.choices[@intCast(root.index)].cost;
107 }
108
109 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 }
113
114 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 }
121
122 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 }
129
130 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 }
146
147 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 }
159
160 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);
170
171 if (self.usableValue(root, order_limit, exclude)) |value| return value;
172
173 if (visiting.contains(root.index)) return null;
174 try visiting.put(self.allocator, root.index, {});
175 defer _ = visiting.remove(root.index);
176
177 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;
180
181 const node = self.graph.nodes(root)[node_index];
182 if (node.kind != .operation) return null;
183 if (node.result_types.len != 1) return null;
184
185 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 }
199
200 var state = ir.Operation.State.init(node.op_name, before_op.location);
201 state.addOperands(operand_values.items);
202 state.addTypes(node.result_types);
203
204 if (node.raw_attributes.len > 0) state.addRawAttributes(node.raw_attributes);
205 if (node.properties) |properties| try state.setPropertiesAttr(properties);
206
207 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 };