lib/choir/src/egraph/graph.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const ir = @import("../core/root.zig");
3 const node_mod = @import("node.zig");
4
5 pub const ClassId = node_mod.ClassId;
6 pub const ValueEntry = node_mod.ValueEntry;
7 pub const Node = node_mod.Node;
8
9 const Class = struct {
10 parent: ClassId,
11 rank: u8 = 0,
12 nodes: std.ArrayListUnmanaged(Node) = .empty,
13 values: std.ArrayListUnmanaged(ValueEntry) = .empty,
14
15 fn init(id: ClassId) Class {
16 return .{ .parent = id };
17 }
18
19 fn deinit(self: *Class, allocator: std.mem.Allocator) void {
20 for (self.nodes.items) |*node| {
21 node.deinit(allocator);
22 }
23 self.nodes.deinit(allocator);
24 self.values.deinit(allocator);
25 }
26 };
27
28 pub const GraphStats = struct {
29 classes_created: usize = 0,
30 nodes_added: usize = 0,
31 unions: usize = 0,
32 rebuilds: usize = 0,
33 };
34
35 pub const Graph = struct {
36 allocator: std.mem.Allocator,
37 classes: std.ArrayListUnmanaged(Class) = .empty,
38 memo: std.AutoHashMapUnmanaged(u64, std.ArrayListUnmanaged(ClassId)) = .empty,
39 stats: GraphStats = .{},
40
41 /// Cumulative allocation traffic for a graph whose rules only merge existing
42 /// classes. Counts cover all blocks together; each class and collision bucket
43 /// may grow to the entire population. Rebuild passes include terminal scans.
44 pub fn mergeStorageBound(
45 node_count: u64,
46 atoms: u64,
47 graphs: u64,
48 rebuild_passes: u64,
49 ) !u64 {
50 if (node_count > std.math.maxInt(u32)) return error.CapacityOverflow;
51 const classes = try mul(graphs, try listStorage(Class, node_count));
52 const contents = try mul(node_count, try sum(
53 try listStorage(Node, node_count),
54 try listStorage(ValueEntry, node_count),
55 ));
56 const payload_item = @sizeOf(ClassId) + @sizeOf(ir.Type) +
57 2 * @sizeOf(ir.NamedAttribute) + 64;
58 const payload = try mul(try mul(2, atoms), payload_item);
59 const table = try mul(graphs, try memoStorage(node_count));
60 const buckets = try mul(node_count, try listStorage(ClassId, 1));
61 const pairs = try listStorage(struct { lhs: ClassId, rhs: ClassId }, node_count);
62 const rebuilds = try mul(rebuild_passes, try sum(buckets, pairs));
63 return sum(try sum(classes, contents), try sum(payload, try sum(
64 table,
65 try sum(buckets, rebuilds),
66 )));
67 }
68
69 pub fn init(allocator: std.mem.Allocator) Graph {
70 return .{ .allocator = allocator };
71 }
72
73 pub fn deinit(self: *Graph) void {
74 self.clearMemo();
75 self.memo.deinit(self.allocator);
76 for (self.classes.items) |*class| {
77 class.deinit(self.allocator);
78 }
79 self.classes.deinit(self.allocator);
80 }
81
82 pub fn find(self: *Graph, id: ClassId) ClassId {
83 const idx: usize = @intCast(id.index);
84 const parent = self.classes.items[idx].parent;
85 if (parent.eql(id)) return id;
86 const root = self.find(parent);
87 self.classes.items[idx].parent = root;
88 return root;
89 }
90
91 pub fn classCount(self: *const Graph) usize {
92 return self.classes.items.len;
93 }
94
95 pub fn addValue(self: *Graph, value: *ir.Value, cost: u32, order: usize) !ClassId {
96 var node = Node.valueNode(value);
97 const id = try self.addNode(&node);
98 try self.attachValue(id, value, cost, order);
99 return id;
100 }
101
102 pub fn addOperation(
103 self: *Graph,
104 op: *ir.Operation,
105 operands: []const ClassId,
106 cost: u32,
107 order: usize,
108 ) !ClassId {
109 var node = try Node.operationNode(self.allocator, op, operands, op.hasTrait("is_commutative"));
110 defer node.deinit(self.allocator);
111 self.canonicalizeNode(&node);
112 const id = try self.addNode(&node);
113 if (op.getResult(0)) |value| {
114 try self.attachValue(id, value, cost, order);
115 }
116 return id;
117 }
118
119 pub fn addNode(self: *Graph, node: *const Node) !ClassId {
120 const hash = node.hash();
121 if (self.memo.get(hash)) |ids| {
122 for (ids.items) |candidate| {
123 const root = self.find(candidate);
124 const class = &self.classes.items[@intCast(root.index)];
125 for (class.nodes.items) |*existing| {
126 if (existing.eql(node)) return root;
127 }
128 }
129 }
130
131 const id = ClassId{ .index = @intCast(self.classes.items.len) };
132 var class = Class.init(id);
133 var class_owned = true;
134 errdefer if (class_owned) class.deinit(self.allocator);
135
136 var owned = try node.clone(self.allocator);
137 var owned_in_class = false;
138 errdefer if (!owned_in_class) owned.deinit(self.allocator);
139 try class.nodes.append(self.allocator, owned);
140 owned_in_class = true;
141
142 try self.classes.append(self.allocator, class);
143 class_owned = false;
144 try self.insertMemo(hash, id);
145 self.stats.classes_created += 1;
146 self.stats.nodes_added += 1;
147 return id;
148 }
149
150 pub fn attachValue(self: *Graph, id: ClassId, value: *ir.Value, cost: u32, order: usize) !void {
151 const root = self.find(id);
152 try self.attachValueToRoot(root, .{ .value = value, .cost = cost, .order = order });
153 }
154
155 pub fn representativeValue(self: *Graph, id: ClassId) ?ValueEntry {
156 const root = self.find(id);
157 const class = &self.classes.items[@intCast(root.index)];
158 if (class.values.items.len == 0) return null;
159
160 var best = class.values.items[0];
161 for (class.values.items[1..]) |candidate| {
162 if (candidate.cost < best.cost or
163 (candidate.cost == best.cost and candidate.order < best.order))
164 {
165 best = candidate;
166 }
167 }
168 return best;
169 }
170
171 pub fn nodes(self: *Graph, id: ClassId) []const Node {
172 const root = self.find(id);
173 return self.classes.items[@intCast(root.index)].nodes.items;
174 }
175
176 pub fn classValues(self: *Graph, id: ClassId) []const ValueEntry {
177 const root = self.find(id);
178 return self.classes.items[@intCast(root.index)].values.items;
179 }
180
181 pub fn merge(self: *Graph, lhs: ClassId, rhs: ClassId) !bool {
182 var lhs_root = self.find(lhs);
183 var rhs_root = self.find(rhs);
184 if (lhs_root.eql(rhs_root)) return false;
185
186 var lhs_class = &self.classes.items[@intCast(lhs_root.index)];
187 var rhs_class = &self.classes.items[@intCast(rhs_root.index)];
188 if (lhs_class.rank < rhs_class.rank) {
189 const tmp_root = lhs_root;
190 lhs_root = rhs_root;
191 rhs_root = tmp_root;
192
193 const tmp_class = lhs_class;
194 lhs_class = rhs_class;
195 rhs_class = tmp_class;
196 }
197
198 try lhs_class.nodes.appendSlice(self.allocator, rhs_class.nodes.items);
199 rhs_class.nodes.clearRetainingCapacity();
200
201 for (rhs_class.values.items) |entry| {
202 try self.attachValueToRoot(lhs_root, entry);
203 }
204 rhs_class.values.clearRetainingCapacity();
205
206 rhs_class.parent = lhs_root;
207 if (lhs_class.rank == rhs_class.rank) {
208 lhs_class.rank += 1;
209 }
210
211 self.stats.unions += 1;
212 return true;
213 }
214
215 pub fn rebuild(self: *Graph) !bool {
216 var any_changed = false;
217 while (try self.rebuildOnce()) {
218 any_changed = true;
219 self.stats.rebuilds += 1;
220 }
221 return any_changed;
222 }
223
224 fn attachValueToRoot(self: *Graph, root: ClassId, entry: ValueEntry) !void {
225 const class = &self.classes.items[@intCast(root.index)];
226 for (class.values.items) |*existing| {
227 if (existing.value == entry.value) {
228 if (entry.cost < existing.cost or
229 (entry.cost == existing.cost and entry.order < existing.order))
230 {
231 existing.cost = entry.cost;
232 existing.order = entry.order;
233 }
234 return;
235 }
236 }
237 try class.values.append(self.allocator, entry);
238 }
239
240 fn canonicalizeNode(self: *Graph, node: *Node) void {
241 if (node.kind != .operation) return;
242 for (node.operands) |*operand| {
243 operand.* = self.find(operand.*);
244 }
245 node.normalizeOperands();
246 }
247
248 fn rebuildOnce(self: *Graph) !bool {
249 self.clearMemo();
250
251 var pairs: std.ArrayListUnmanaged(struct { lhs: ClassId, rhs: ClassId }) = .empty;
252 defer pairs.deinit(self.allocator);
253
254 for (self.classes.items, 0..) |*class, index| {
255 const id = ClassId{ .index = @intCast(index) };
256 const root = self.find(id);
257 if (!root.eql(id)) continue;
258
259 for (class.nodes.items) |*node| {
260 self.canonicalizeNode(node);
261 const hash = node.hash();
262 var found: ?ClassId = null;
263 if (self.memo.get(hash)) |ids| {
264 for (ids.items) |candidate| {
265 const candidate_root = self.find(candidate);
266 const candidate_class = &self.classes.items[@intCast(candidate_root.index)];
267 for (candidate_class.nodes.items) |*existing| {
268 if (existing.eql(node)) {
269 found = candidate_root;
270 break;
271 }
272 }
273 if (found != null) break;
274 }
275 }
276
277 if (found) |existing| {
278 if (!existing.eql(root)) {
279 try pairs.append(self.allocator, .{ .lhs = root, .rhs = existing });
280 }
281 } else {
282 try self.insertMemo(hash, root);
283 }
284 }
285 }
286
287 var changed = false;
288 for (pairs.items) |pair| {
289 if (try self.merge(pair.lhs, pair.rhs)) {
290 changed = true;
291 }
292 }
293 return changed;
294 }
295
296 fn insertMemo(self: *Graph, hash: u64, id: ClassId) !void {
297 var gop = try self.memo.getOrPut(self.allocator, hash);
298 if (!gop.found_existing) {
299 gop.value_ptr.* = .empty;
300 }
301 const root = self.find(id);
302 for (gop.value_ptr.items) |existing| {
303 if (self.find(existing).eql(root)) return;
304 }
305 try gop.value_ptr.append(self.allocator, root);
306 }
307
308 fn clearMemo(self: *Graph) void {
309 var iter = self.memo.valueIterator();
310 while (iter.next()) |ids| {
311 ids.deinit(self.allocator);
312 }
313 self.memo.clearRetainingCapacity();
314 }
315 };
316
317 fn sum(a: u64, b: u64) !u64 {
318 return std.math.add(u64, a, b) catch error.CapacityOverflow;
319 }
320
321 fn mul(a: u64, b: u64) !u64 {
322 return std.math.mul(u64, a, b) catch error.CapacityOverflow;
323 }
324
325 fn listStorage(comptime Item: type, count: u64) !u64 {
326 if (count == 0) return 0;
327 return mul(try mul(4, try sum(try mul(4, count), 64)), @sizeOf(Item) + @alignOf(Item));
328 }
329
330 fn memoStorage(count: u64) !u64 {
331 if (count == 0) return 0;
332 const required = try sum(try mul(count, 100) / 80, 1);
333 if (required > std.math.maxInt(u32)) return error.CapacityOverflow;
334 const capacity = std.math.ceilPowerOfTwo(u32, @intCast(required)) catch
335 return error.CapacityOverflow;
336 if (capacity > std.math.maxInt(u32) / 80) return error.CapacityOverflow;
337 const bytes = 1 + @sizeOf(u64) + @sizeOf(std.ArrayListUnmanaged(ClassId)) +
338 4 * @sizeOf(usize) + 3 * @alignOf(usize);
339 return mul(try mul(2, @max(8, capacity)), bytes);
340 }
341
342 test "saturation graph storage covers collisions merges and repeated rebuilds" {
343 for ([_]u32{ 1, 8, 64 }) |count| try checkMergeStorage(count);
344 try std.testing.expectError(
345 error.CapacityOverflow,
346 Graph.mergeStorageBound(std.math.maxInt(u64), 1, 1, 1),
347 );
348 }
349
350 fn checkMergeStorage(count: u32) !void {
351 const allocator = std.testing.allocator;
352 const fixed = @import("alloc_fixed");
353 var context = try ir.Context.init(allocator, ir.Context.Limits.testing);
354 defer context.deinit(allocator);
355 const bound = try Graph.mergeStorageBound(count, count, 1, count + 2);
356 const storage = try allocator.alignedAlloc(u8, .@"64", @intCast(bound));
357 defer allocator.free(storage);
358 var backing = fixed.Tracked.init(storage);
359 var retained = fixed.Monotonic.init(backing.allocator(), @intCast(bound));
360 var graph = Graph.init(retained.allocator());
361 defer graph.deinit();
362 for (0..count) |index| {
363 var attrs = [_]ir.NamedAttribute{.{
364 .name = "value",
365 .value = try context.getI64Attr(@intCast(index)),
366 }};
367 const node = Node{ .kind = .operation, .op_name = "test.node", .attributes = &attrs };
368 _ = try graph.addNode(&node);
369 }
370 graph.clearMemo();
371 for (0..count) |index| try graph.insertMemo(0, .{ .index = @intCast(index) });
372 try std.testing.expectEqual(count, graph.memo.get(0).?.items.len);
373 _ = try graph.rebuild();
374 for (1..count) |index| {
375 try std.testing.expect(try graph.merge(.{ .index = 0 }, .{ .index = @intCast(index) }));
376 _ = try graph.rebuild();
377 }
378 try std.testing.expectEqual(count, graph.nodes(.{ .index = 0 }).len);
379 try std.testing.expect(!backing.exhausted);
380 const used = if (retained.current) |*current| fixed.used(current) else 0;
381 try std.testing.expect(used <= bound);
382 }