lib/choir/src/egraph/pattern.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const alloc_arena = @import("alloc_arena");
3 const ir = @import("../core/root.zig");
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
10 pub const max_variables = 4;
11
12 pub const Scalar = union(enum) {
13 int: i64,
14 float: f64,
15 boolean: bool,
16
17 pub fn eql(self: Scalar, other: Scalar) bool {
18 if (std.meta.activeTag(self) != std.meta.activeTag(other)) return false;
19 return switch (self) {
20 .int => |value| value == other.int,
21 .float => |value| @as(u64, @bitCast(value)) == @as(u64, @bitCast(other.float)),
22 .boolean => |value| value == other.boolean,
23 };
24 }
25 };
26
27 pub const TypeClass = enum {
28 integer,
29 float,
30 boolean,
31 other,
32 };
33
34 pub const ClassifyFn = *const fn (ty: ir.Type) TypeClass;
35 pub const ReadFn = *const fn (attr: ir.Attribute) ?Scalar;
36 pub const BuildFn = *const fn (ir_ctx: *ir.Context, value: Scalar) anyerror!ir.Attribute;
37
38 pub const ConstantStorage = enum {
39 attribute,
40 properties,
41 };
42
43 pub const ConstantModel = struct {
44 op_name: []const u8,
45 attr_name: []const u8 = "value",
46 classify: ClassifyFn,
47 read: ReadFn = readScalar,
48 build: BuildFn = buildScalar,
49 storage: ConstantStorage = .attribute,
50 };
51
52 pub fn readScalar(attr: ir.Attribute) ?Scalar {
53 if (attr.cast(ir.Attribute.IntegerAttr)) |int_attr| return .{ .int = int_attr.getValue() };
54 if (attr.cast(ir.Attribute.FloatAttr)) |float_attr| return .{ .float = float_attr.getValue() };
55 if (attr.cast(ir.Attribute.BoolAttr)) |bool_attr| return .{ .boolean = bool_attr.getValue() };
56 return null;
57 }
58
59 pub fn buildScalar(ir_ctx: *ir.Context, value: Scalar) anyerror!ir.Attribute {
60 return switch (value) {
61 .int => |int_value| try ir_ctx.getI64Attr(int_value),
62 .float => |float_value| try ir_ctx.getF64Attr(float_value),
63 .boolean => |bool_value| try ir_ctx.getBoolAttr(bool_value),
64 };
65 }
66
67 pub const OpShape = struct {
68 name: []const u8,
69 operands: []const Pattern,
70 };
71
72 pub const Pattern = union(enum) {
73 variable: u8,
74 constant: ?Scalar,
75 operation: OpShape,
76 };
77
78 pub const TemplateShape = struct {
79 name: []const u8,
80 operands: []const Template,
81 };
82
83 pub const Template = union(enum) {
84 variable: u8,
85 constant: Scalar,
86 operation: TemplateShape,
87 };
88
89 pub const Rule = struct {
90 name: []const u8,
91 benefit: u32 = 1,
92 classes: ?[]const TypeClass = null,
93 lhs: OpShape,
94 rhs: Template,
95 };
96
97 pub const Bindings = struct {
98 slots: [max_variables]?ClassId = @as([max_variables]?ClassId, @splat(null)),
99
100 pub fn get(self: *const Bindings, variable: u8) ?ClassId {
101 return self.slots[variable];
102 }
103
104 pub fn set(self: *Bindings, variable: u8, class: ClassId) void {
105 self.slots[variable] = class;
106 }
107 };
108
109 pub fn classConstant(graph: *Graph, model: *const ConstantModel, class: ClassId) ?Scalar {
110 for (graph.nodes(class)) |*node| {
111 if (node.kind != .operation) continue;
112 if (!std.mem.eql(u8, node.op_name, model.op_name)) continue;
113 const attr = node.getAttr(model.attr_name) orelse continue;
114 if (model.read(attr)) |value| return value;
115 }
116 return null;
117 }
118
119 pub fn matchNode(
120 graph: *Graph,
121 model: *const ConstantModel,
122 rule: *const Rule,
123 node: *const Node,
124 bindings: *Bindings,
125 ) bool {
126 if (node.kind != .operation) return false;
127 if (!std.mem.eql(u8, node.op_name, rule.lhs.name)) return false;
128 if (rule.classes) |classes| {
129 if (node.result_types.len == 0) return false;
130 const type_class = model.classify(node.result_types[0]);
131 if (std.mem.indexOfScalar(TypeClass, classes, type_class) == null) return false;
132 }
133 return matchOperands(graph, model, rule.lhs.operands, node, bindings);
134 }
135
136 fn matchOperands(
137 graph: *Graph,
138 model: *const ConstantModel,
139 patterns: []const Pattern,
140 node: *const Node,
141 bindings: *Bindings,
142 ) bool {
143 if (patterns.len != node.operands.len) return false;
144
145 const saved = bindings.*;
146 var swapped: [2]ClassId = undefined;
147 const attempt_count: usize = if (node.commutative and node.operands.len == 2) 2 else 1;
148
149 attempts: for (0..attempt_count) |attempt| {
150 bindings.* = saved;
151 const operands = if (attempt == 0) node.operands else operands: {
152 swapped = .{ node.operands[1], node.operands[0] };
153 break :operands swapped[0..];
154 };
155
156 for (patterns, operands) |pattern, operand| {
157 switch (pattern) {
158 .variable => |variable| {
159 if (bindings.get(variable)) |bound| {
160 if (!graph.find(bound).eql(graph.find(operand))) continue :attempts;
161 } else {
162 bindings.set(variable, graph.find(operand));
163 }
164 },
165 .constant => |expected| {
166 const actual = classConstant(graph, model, operand) orelse continue :attempts;
167 if (expected) |value| {
168 if (!value.eql(actual)) continue :attempts;
169 }
170 },
171 .operation => |shape| {
172 var matched = false;
173 const class_nodes = graph.nodes(operand);
174 for (class_nodes) |*candidate| {
175 if (candidate.kind != .operation) continue;
176 if (!std.mem.eql(u8, candidate.op_name, shape.name)) continue;
177 if (matchOperands(graph, model, shape.operands, candidate, bindings)) {
178 matched = true;
179 break;
180 }
181 }
182 if (!matched) continue :attempts;
183 },
184 }
185 }
186 return true;
187 }
188
189 bindings.* = saved;
190 return false;
191 }
192
193 pub fn instantiate(
194 graph: *Graph,
195 ir_ctx: *ir.Context,
196 model: *const ConstantModel,
197 template: Template,
198 bindings: *const Bindings,
199 result_types: []ir.Type,
200 ) anyerror!ClassId {
201 switch (template) {
202 .variable => |variable| {
203 const bound = bindings.get(variable) orelse return error.UnboundPatternVariable;
204 return graph.find(bound);
205 },
206 .constant => |value| {
207 const attr = try model.build(ir_ctx, value);
208 var attributes = [_]ir.NamedAttribute{.{ .name = model.attr_name, .value = attr }};
209 var node = Node{
210 .kind = .operation,
211 .op_name = model.op_name,
212 .operands = &.{},
213 .result_types = result_types,
214 .attributes = attributes[0..],
215 .raw_attributes = if (model.storage == .attribute) attributes[0..] else &.{},
216 .properties = if (model.storage == .properties) attr else null,
217 .commutative = false,
218 };
219 return graph.addNode(&node);
220 },
221 .operation => |shape| {
222 var operand_classes: [max_variables]ClassId = undefined;
223 if (shape.operands.len > operand_classes.len) return error.PatternTooWide;
224 for (shape.operands, 0..) |operand_template, index| {
225 operand_classes[index] = try instantiate(graph, ir_ctx, model, operand_template, bindings, result_types);
226 }
227
228 const commutative = if (ir_ctx.lookupOperation(shape.name)) |info| info.traits.is_commutative else false;
229 var node = Node{
230 .kind = .operation,
231 .op_name = shape.name,
232 .operands = operand_classes[0..shape.operands.len],
233 .result_types = result_types,
234 .attributes = &.{},
235 .raw_attributes = &.{},
236 .properties = null,
237 .commutative = commutative,
238 };
239 node.normalizeOperands();
240 return graph.addNode(&node);
241 },
242 }
243 }
244
245 pub fn applyRule(
246 graph: *Graph,
247 ir_ctx: *ir.Context,
248 model: *const ConstantModel,
249 rule: *const Rule,
250 class: ClassId,
251 node: *const Node,
252 ) anyerror!bool {
253 var bindings = Bindings{};
254 if (!matchNode(graph, model, rule, node, &bindings)) return false;
255 const rhs_class = try instantiate(graph, ir_ctx, model, rule.rhs, &bindings, node.result_types);
256 return graph.merge(class, rhs_class);
257 }
258
259 const test_dialect_mod = @import("../dialects/fixture/root.zig");
260
261 fn testClassify(ty: ir.Type) TypeClass {
262 const name = ty.getDialectTypeName() orelse return .other;
263 if (std.mem.eql(u8, name, "test.i64")) return .integer;
264 return .other;
265 }
266
267 const test_model = ConstantModel{
268 .op_name = "test.constant",
269 .classify = testClassify,
270 };
271
272 fn testGraphContext(allocator: std.mem.Allocator) !ir.Context {
273 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
274 errdefer ctx.deinit(allocator);
275 try test_dialect_mod.registerTestDialect(&ctx);
276 _ = try ctx.registerOperation("test.constant", .{});
277 _ = try ctx.registerOperation("test.binary", .{ .is_commutative = true });
278 return ctx;
279 }
280
281 fn addTestConstant(graph: *Graph, ir_ctx: *ir.Context, i64_type: ir.Type, value: i64) !ClassId {
282 const attr = try ir_ctx.getI64Attr(value);
283 var attributes = [_]ir.NamedAttribute{.{ .name = "value", .value = attr }};
284 var result_types = [_]ir.Type{i64_type};
285 var node = Node{
286 .kind = .operation,
287 .op_name = "test.constant",
288 .operands = &.{},
289 .result_types = result_types[0..],
290 .attributes = attributes[0..],
291 .raw_attributes = attributes[0..],
292 .properties = null,
293 .commutative = false,
294 };
295 return graph.addNode(&node);
296 }
297
298 fn addTestBinary(graph: *Graph, i64_type: ir.Type, lhs: ClassId, rhs: ClassId) !ClassId {
299 var operands = [_]ClassId{ lhs, rhs };
300 var result_types = [_]ir.Type{i64_type};
301 var node = Node{
302 .kind = .operation,
303 .op_name = "test.binary",
304 .operands = operands[0..],
305 .result_types = result_types[0..],
306 .attributes = &.{},
307 .raw_attributes = &.{},
308 .properties = null,
309 .commutative = true,
310 };
311 node.normalizeOperands();
312 return graph.addNode(&node);
313 }
314
315 test "scalar float equality distinguishes signed zero" {
316 const positive = Scalar{ .float = 0.0 };
317 const negative = Scalar{ .float = -0.0 };
318 try std.testing.expect(!positive.eql(negative));
319 try std.testing.expect(positive.eql(.{ .float = 0.0 }));
320 try std.testing.expect(!positive.eql(.{ .int = 0 }));
321 }
322
323 const add_zero_rule = Rule{
324 .name = "test-binary-zero",
325 .classes = &.{.integer},
326 .lhs = .{ .name = "test.binary", .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } },
327 .rhs = .{ .variable = 0 },
328 };
329
330 const self_to_zero_rule = Rule{
331 .name = "test-binary-self-zero",
332 .classes = &.{.integer},
333 .lhs = .{ .name = "test.binary", .operands = &.{ .{ .variable = 0 }, .{ .variable = 0 } } },
334 .rhs = .{ .constant = .{ .int = 0 } },
335 };
336
337 test "pattern rule merges class with bound variable" {
338 const testing = std.testing;
339 var arena = alloc_arena.Arena.init(testing.allocator);
340 defer arena.deinit();
341 const allocator = arena.allocator();
342
343 var ctx = try testGraphContext(allocator);
344 defer ctx.deinit(allocator);
345 const i64_type = try test_dialect_mod.TestDialect.getI64Type(&ctx);
346
347 var graph = Graph.init(allocator);
348 defer graph.deinit();
349
350 const seven = try addTestConstant(&graph, &ctx, i64_type, 7);
351 const zero = try addTestConstant(&graph, &ctx, i64_type, 0);
352 const sum = try addTestBinary(&graph, i64_type, seven, zero);
353
354 try testing.expect(!graph.find(sum).eql(graph.find(seven)));
355
356 const sum_nodes = graph.nodes(sum);
357 var fired = false;
358 for (sum_nodes) |*node| {
359 if (node.kind != .operation) continue;
360 if (!std.mem.eql(u8, node.op_name, "test.binary")) continue;
361 fired = try applyRule(&graph, &ctx, &test_model, &add_zero_rule, sum, node);
362 break;
363 }
364
365 try testing.expect(fired);
366 try testing.expect(graph.find(sum).eql(graph.find(seven)));
367 }
368
369 test "pattern rule with commutative swap matches constant on the left" {
370 const testing = std.testing;
371 var arena = alloc_arena.Arena.init(testing.allocator);
372 defer arena.deinit();
373 const allocator = arena.allocator();
374
375 var ctx = try testGraphContext(allocator);
376 defer ctx.deinit(allocator);
377 const i64_type = try test_dialect_mod.TestDialect.getI64Type(&ctx);
378
379 var graph = Graph.init(allocator);
380 defer graph.deinit();
381
382 const zero = try addTestConstant(&graph, &ctx, i64_type, 0);
383 const seven = try addTestConstant(&graph, &ctx, i64_type, 7);
384 const sum = try addTestBinary(&graph, i64_type, zero, seven);
385
386 const sum_nodes = graph.nodes(sum);
387 var fired = false;
388 for (sum_nodes) |*node| {
389 if (node.kind != .operation) continue;
390 if (!std.mem.eql(u8, node.op_name, "test.binary")) continue;
391 fired = try applyRule(&graph, &ctx, &test_model, &add_zero_rule, sum, node);
392 break;
393 }
394
395 try testing.expect(fired);
396 try testing.expect(graph.find(sum).eql(graph.find(seven)));
397 }
398
399 test "pattern rule instantiates a fresh constant node" {
400 const testing = std.testing;
401 var arena = alloc_arena.Arena.init(testing.allocator);
402 defer arena.deinit();
403 const allocator = arena.allocator();
404
405 var ctx = try testGraphContext(allocator);
406 defer ctx.deinit(allocator);
407 const i64_type = try test_dialect_mod.TestDialect.getI64Type(&ctx);
408
409 var graph = Graph.init(allocator);
410 defer graph.deinit();
411
412 const seven = try addTestConstant(&graph, &ctx, i64_type, 7);
413 const sum = try addTestBinary(&graph, i64_type, seven, seven);
414
415 const sum_nodes = graph.nodes(sum);
416 var fired = false;
417 for (sum_nodes) |*node| {
418 if (node.kind != .operation) continue;
419 if (!std.mem.eql(u8, node.op_name, "test.binary")) continue;
420 fired = try applyRule(&graph, &ctx, &test_model, &self_to_zero_rule, sum, node);
421 break;
422 }
423
424 try testing.expect(fired);
425 const merged = classConstant(&graph, &test_model, sum) orelse return error.TestExpectedResult;
426 try testing.expect(merged.eql(.{ .int = 0 }));
427 }
428
429 test "pattern rule respects type class guard" {
430 const testing = std.testing;
431 var arena = alloc_arena.Arena.init(testing.allocator);
432 defer arena.deinit();
433 const allocator = arena.allocator();
434
435 var ctx = try testGraphContext(allocator);
436 defer ctx.deinit(allocator);
437 const i64_type = try test_dialect_mod.TestDialect.getI64Type(&ctx);
438
439 var graph = Graph.init(allocator);
440 defer graph.deinit();
441
442 const seven = try addTestConstant(&graph, &ctx, i64_type, 7);
443 const zero = try addTestConstant(&graph, &ctx, i64_type, 0);
444 const sum = try addTestBinary(&graph, i64_type, seven, zero);
445
446 const float_guarded = Rule{
447 .name = "float-only",
448 .classes = &.{.float},
449 .lhs = add_zero_rule.lhs,
450 .rhs = add_zero_rule.rhs,
451 };
452
453 const sum_nodes = graph.nodes(sum);
454 var fired = false;
455 for (sum_nodes) |*node| {
456 if (node.kind != .operation) continue;
457 if (!std.mem.eql(u8, node.op_name, "test.binary")) continue;
458 fired = try applyRule(&graph, &ctx, &test_model, &float_guarded, sum, node);
459 break;
460 }
461
462 try testing.expect(!fired);
463 try testing.expect(!graph.find(sum).eql(graph.find(seven)));
464 }
465
466 test "nested pattern matches through operand classes" {
467 const testing = std.testing;
468 var arena = alloc_arena.Arena.init(testing.allocator);
469 defer arena.deinit();
470 const allocator = arena.allocator();
471
472 var ctx = try testGraphContext(allocator);
473 defer ctx.deinit(allocator);
474 const i64_type = try test_dialect_mod.TestDialect.getI64Type(&ctx);
475
476 var graph = Graph.init(allocator);
477 defer graph.deinit();
478
479 const seven = try addTestConstant(&graph, &ctx, i64_type, 7);
480 const zero = try addTestConstant(&graph, &ctx, i64_type, 0);
481 const inner = try addTestBinary(&graph, i64_type, seven, zero);
482 const outer = try addTestBinary(&graph, i64_type, inner, zero);
483
484 const nested_rule = Rule{
485 .name = "nested-binary-zero",
486 .lhs = .{ .name = "test.binary", .operands = &.{
487 .{ .operation = .{ .name = "test.binary", .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } } },
488 .{ .constant = .{ .int = 0 } },
489 } },
490 .rhs = .{ .variable = 0 },
491 };
492
493 const outer_nodes = graph.nodes(outer);
494 var fired = false;
495 for (outer_nodes) |*node| {
496 if (node.kind != .operation) continue;
497 if (!std.mem.eql(u8, node.op_name, "test.binary")) continue;
498 fired = try applyRule(&graph, &ctx, &test_model, &nested_rule, outer, node);
499 break;
500 }
501
502 try testing.expect(fired);
503 try testing.expect(graph.find(outer).eql(graph.find(seven)));
504 }