lib/choir/src/passes/test.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const passes = @import("root.zig");
3 const cse = @import("cse/test.zig");
4 const pass_tests = @import("pass/test.zig");
5 const observations = @import("observations/test.zig");
6
7 test {
8 _ = cse;
9 _ = pass_tests;
10 _ = observations;
11 std.testing.refAllDecls(passes);
12 }
13
14 const alloc_arena = @import("alloc_arena");
15 const egraph = @import("../egraph/root.zig");
16 const ir = @import("../core/root.zig");
17 const dialects = @import("../dialects/root.zig");
18 const test_support = @import("../dialects/fixture/root.zig");
19 const pass_mod = passes.pass;
20 const arith = dialects.ArithDialect;
21 const ClassId = egraph.ClassId;
22 const EGraphPass = passes.EGraphPass;
23 const Graph = egraph.Graph;
24 const Node = egraph.Node;
25 const RewriteContext = egraph.RewriteContext;
26 const RewriteSet = egraph.RewriteSet;
27 const Scalar = egraph.Scalar;
28 const TypeClass = egraph.TypeClass;
29 fn noRules(_: *RewriteSet) anyerror!void {}
30
31 fn testCandidate(_: ?*anyopaque, op: *ir.Operation) anyerror!bool {
32 return std.mem.eql(u8, op.name.name, "arith.constant") or
33 std.mem.eql(u8, op.name.name, "arith.add");
34 }
35
36 const EmptyTestPass = EGraphPass(.{
37 .name = "choir-egraph-empty-test",
38 .description = "test egraph pass",
39 .populate_rules = noRules,
40 .options = .{ .candidate = testCandidate },
41 });
42
43 fn isZeroClass(ctx: *RewriteContext, id: ClassId) bool {
44 for (ctx.nodes(id)) |*entry| {
45 if (entry.kind != .operation) continue;
46 if (!std.mem.eql(u8, entry.op_name, "arith.constant")) continue;
47 const attr = entry.getAttr("value") orelse continue;
48 const value = arith.getIntValue(attr) orelse continue;
49 if (value == 0) return true;
50 }
51 return false;
52 }
53
54 fn foldBinaryZero(ctx: *RewriteContext, class: ClassId, entry: *const Node) anyerror!bool {
55 if (entry.kind != .operation) return false;
56 if (!std.mem.eql(u8, entry.op_name, "arith.add")) return false;
57 if (entry.operands.len != 2) return false;
58 if (isZeroClass(ctx, entry.operands[0])) return try ctx.merge(class, entry.operands[1]);
59 if (isZeroClass(ctx, entry.operands[1])) return try ctx.merge(class, entry.operands[0]);
60 return false;
61 }
62
63 fn zeroRules(rule_set: *RewriteSet) anyerror!void {
64 try rule_set.add(.{
65 .name = "test-binary-zero",
66 .benefit = 10,
67 .apply = foldBinaryZero,
68 });
69 }
70
71 const ZeroTestPass = EGraphPass(.{
72 .name = "choir-egraph-zero-test",
73 .description = "test egraph zero rewrite pass",
74 .populate_rules = zeroRules,
75 .options = .{ .candidate = testCandidate },
76 });
77
78 const BudgetTestPass = EGraphPass(.{
79 .name = "choir-egraph-budget-test",
80 .description = "test egraph rewrite budget pass",
81 .populate_rules = zeroRules,
82 .options = .{ .candidate = testCandidate, .max_rewrites = 1 },
83 });
84
85 fn buildTestContext(allocator: std.mem.Allocator) !ir.Context {
86 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
87 errdefer ctx.deinit(allocator);
88 try test_support.registerTestDialect(&ctx);
89 try ir.dialects.loadDialectSpec(&ctx, dialects.arith.spec);
90 _ = try ctx.registerOperation("arith.constant", .{});
91 _ = try ctx.registerOperation("arith.add", .{ .is_commutative = true });
92 return ctx;
93 }
94
95 test "egraph default candidate uses effect summaries" {
96 const testing = std.testing;
97
98 var arena = alloc_arena.Arena.init(testing.allocator);
99 defer arena.deinit();
100 const allocator = arena.allocator();
101
102 var ctx = try buildTestContext(allocator);
103 defer ctx.deinit(allocator);
104 try ctx.registerOperationInterface(
105 "test.effect_free_value",
106 ir.interfaces.EffectOpInterface.entryFor(.{}),
107 );
108 try ctx.registerOperationInterface(
109 "test.read_value",
110 ir.interfaces.EffectOpInterface.entryFor(.{}),
111 );
112
113 const loc = ir.Location.getUnknown();
114 const i64_type = try arith.getScalarType(&ctx, .i64);
115 var producer_state = ir.Operation.State.init("arith.constant", loc);
116 producer_state.addTypes(&.{i64_type});
117 const producer = try ctx.createOperation(producer_state);
118
119 var effect_free_state = ir.Operation.State.init("test.effect_free_value", loc);
120 effect_free_state.addTypes(&.{i64_type});
121 const effect_free = try ctx.createOperation(effect_free_state);
122
123 var read_state = ir.Operation.State.init("test.read_value", loc);
124 read_state.addOperands(&.{producer.getResult(0).?});
125 read_state.addTypes(&.{i64_type});
126 const read = try ctx.createOperation(read_state);
127
128 try testing.expect(!try passes.saturation.defaultCandidate(null, effect_free));
129 try testing.expect(!try passes.saturation.defaultCandidate(null, read));
130 }
131
132 test "egraph pass replaces duplicate block-local expressions" {
133 const test_dialect = test_support.TestDialect;
134 const testing = std.testing;
135
136 var arena = alloc_arena.Arena.init(testing.allocator);
137 defer arena.deinit();
138 const allocator = arena.allocator();
139
140 var ctx = try buildTestContext(allocator);
141 defer ctx.deinit(allocator);
142
143 const loc = ir.Location.getUnknown();
144 const i64_type = try arith.getScalarType(&ctx, .i64);
145 const module = try test_dialect.ModuleOp.create(&ctx, loc);
146 const block = module.getBodyBlock();
147
148 var lhs = try arith.ConstantOp.createInt(&ctx, loc, i64_type, 7);
149 var rhs = try arith.ConstantOp.createInt(&ctx, loc, i64_type, 11);
150 var add1 = try arith.AddOp.create(&ctx, loc, lhs.getResult(), rhs.getResult());
151 var add2 = try arith.AddOp.create(&ctx, loc, rhs.getResult(), lhs.getResult());
152 const ret = try test_dialect.ReturnOp.create(&ctx, loc, &.{add2.getResult()});
153
154 try block.addOperation(lhs.op);
155 try block.addOperation(rhs.op);
156 try block.addOperation(add1.op);
157 try block.addOperation(add2.op);
158 try block.addOperation(ret.op);
159
160 var pm = pass_mod.PassManager.init(allocator);
161 defer pm.deinit();
162 try pm.addPass(EmptyTestPass.create());
163 try std.testing.expectEqual(pass_mod.PassResult.success, pm.run(module.op, &ctx));
164
165 try testing.expect(ret.op.getOperand(0).? == add1.getResult());
166 try testing.expect(ret.op.prev_op == add1.op);
167 }
168
169 test "egraph rules rewrite expressions to cheaper equivalent values" {
170 const test_dialect = test_support.TestDialect;
171 const testing = std.testing;
172
173 var arena = alloc_arena.Arena.init(testing.allocator);
174 defer arena.deinit();
175 const allocator = arena.allocator();
176
177 var ctx = try buildTestContext(allocator);
178 defer ctx.deinit(allocator);
179
180 const loc = ir.Location.getUnknown();
181 const i64_type = try arith.getScalarType(&ctx, .i64);
182 const module = try test_dialect.ModuleOp.create(&ctx, loc);
183 const block = module.getBodyBlock();
184
185 var value = try arith.ConstantOp.createInt(&ctx, loc, i64_type, 7);
186 var zero = try arith.ConstantOp.createInt(&ctx, loc, i64_type, 0);
187 var add = try arith.AddOp.create(&ctx, loc, value.getResult(), zero.getResult());
188 const ret = try test_dialect.ReturnOp.create(&ctx, loc, &.{add.getResult()});
189
190 try block.addOperation(value.op);
191 try block.addOperation(zero.op);
192 try block.addOperation(add.op);
193 try block.addOperation(ret.op);
194
195 var pm = pass_mod.PassManager.init(allocator);
196 defer pm.deinit();
197 try pm.addPass(ZeroTestPass.create());
198 try std.testing.expectEqual(pass_mod.PassResult.success, pm.run(module.op, &ctx));
199
200 try testing.expect(ret.op.getOperand(0).? == value.getResult());
201 try testing.expect(ret.op.prev_op == zero.op);
202 }
203
204 fn testClassifyType(ty: ir.Type) TypeClass {
205 const name = ty.getDialectTypeName() orelse return .other;
206 if (std.mem.eql(u8, name, "arith.i64")) return .integer;
207 return .other;
208 }
209
210 fn testReadScalar(attr: ir.Attribute) ?Scalar {
211 const value = arith.getIntValue(attr) orelse return null;
212 return .{ .int = value };
213 }
214
215 fn testBuildScalar(ir_ctx: *ir.Context, value: Scalar) anyerror!ir.Attribute {
216 return switch (value) {
217 .int => |int_value| arith.getIntAttr(ir_ctx, int_value),
218 else => error.UnsupportedScalar,
219 };
220 }
221
222 fn patternZeroRules(rule_set: *RewriteSet) anyerror!void {
223 rule_set.setConstantModel(.{
224 .op_name = "arith.constant",
225 .classify = testClassifyType,
226 .read = testReadScalar,
227 .build = testBuildScalar,
228 });
229 try rule_set.addPattern(.{
230 .name = "test-binary-zero-pattern",
231 .benefit = 10,
232 .classes = &.{.integer},
233 .lhs = .{
234 .name = "arith.add",
235 .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } },
236 },
237 .rhs = .{ .variable = 0 },
238 });
239 }
240
241 const PatternTestPass = EGraphPass(.{
242 .name = "choir-egraph-pattern-test",
243 .description = "test egraph pattern rule pass",
244 .populate_rules = patternZeroRules,
245 .options = .{ .candidate = testCandidate },
246 });
247
248 test "egraph pattern rules rewrite expressions as data" {
249 const test_dialect = test_support.TestDialect;
250 const testing = std.testing;
251
252 var arena = alloc_arena.Arena.init(testing.allocator);
253 defer arena.deinit();
254 const allocator = arena.allocator();
255
256 var ctx = try buildTestContext(allocator);
257 defer ctx.deinit(allocator);
258
259 const loc = ir.Location.getUnknown();
260 const i64_type = try arith.getScalarType(&ctx, .i64);
261 const module = try test_dialect.ModuleOp.create(&ctx, loc);
262 const block = module.getBodyBlock();
263
264 var value = try arith.ConstantOp.createInt(&ctx, loc, i64_type, 7);
265 var zero = try arith.ConstantOp.createInt(&ctx, loc, i64_type, 0);
266 var add = try arith.AddOp.create(&ctx, loc, zero.getResult(), value.getResult());
267 const ret = try test_dialect.ReturnOp.create(&ctx, loc, &.{add.getResult()});
268
269 try block.addOperation(value.op);
270 try block.addOperation(zero.op);
271 try block.addOperation(add.op);
272 try block.addOperation(ret.op);
273
274 var pm = pass_mod.PassManager.init(allocator);
275 defer pm.deinit();
276 try pm.addPass(PatternTestPass.create());
277 try testing.expectEqual(pass_mod.PassResult.success, pm.run(module.op, &ctx));
278
279 try testing.expect(ret.op.getOperand(0).? == value.getResult());
280 }
281
282 fn constantAddRules(rule_set: *RewriteSet) anyerror!void {
283 rule_set.setConstantModel(.{
284 .op_name = "arith.constant",
285 .classify = testClassifyType,
286 .read = testReadScalar,
287 .build = testBuildScalar,
288 });
289 try rule_set.addPattern(.{
290 .name = "test-constant-add",
291 .benefit = 10,
292 .classes = &.{.integer},
293 .lhs = .{
294 .name = "arith.add",
295 .operands = &.{ .{ .constant = .{ .int = 7 } }, .{ .constant = .{ .int = 7 } } },
296 },
297 .rhs = .{ .constant = .{ .int = 14 } },
298 });
299 }
300
301 const ConstantAddTestPass = EGraphPass(.{
302 .name = "choir-egraph-self-cancel-test",
303 .description = "test egraph constant materialization pass",
304 .populate_rules = constantAddRules,
305 .options = .{ .candidate = testCandidate },
306 });
307
308 test "egraph extraction materializes constants discovered by rules" {
309 const test_dialect = test_support.TestDialect;
310 const testing = std.testing;
311
312 var arena = alloc_arena.Arena.init(testing.allocator);
313 defer arena.deinit();
314 const allocator = arena.allocator();
315
316 var ctx = try buildTestContext(allocator);
317 defer ctx.deinit(allocator);
318
319 const loc = ir.Location.getUnknown();
320 const i64_type = try arith.getScalarType(&ctx, .i64);
321 const module = try test_dialect.ModuleOp.create(&ctx, loc);
322 const block = module.getBodyBlock();
323
324 var value = try arith.ConstantOp.createInt(&ctx, loc, i64_type, 7);
325 var diff = try arith.AddOp.create(&ctx, loc, value.getResult(), value.getResult());
326 const ret = try test_dialect.ReturnOp.create(&ctx, loc, &.{diff.getResult()});
327
328 try block.addOperation(value.op);
329 try block.addOperation(diff.op);
330 try block.addOperation(ret.op);
331
332 var pm = pass_mod.PassManager.init(allocator);
333 defer pm.deinit();
334 try pm.addPass(ConstantAddTestPass.create());
335 try testing.expectEqual(pass_mod.PassResult.success, pm.run(module.op, &ctx));
336
337 const replacement = ret.op.getOperand(0).?;
338 try testing.expect(replacement != diff.getResult());
339 const def_any = replacement.getDefiningOp() orelse return error.TestExpectedResult;
340 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
341 try testing.expectEqualStrings("arith.constant", def_op.name.name);
342 const constant = arith.ConstantOp{ .op = def_op };
343 try testing.expectEqual(@as(?i64, 14), constant.getIntValue());
344 try ir.verifyOperation(module.op, ir.verify.default_options);
345 }
346
347 fn reductionRules(rule_set: *RewriteSet) anyerror!void {
348 rule_set.setConstantModel(.{
349 .op_name = "arith.constant",
350 .classify = testClassifyType,
351 .read = testReadScalar,
352 .build = testBuildScalar,
353 });
354 rule_set.setCostModel(.{
355 .overrides = &.{.{ .name = "test.identity", .cost = 1 }},
356 });
357 try rule_set.addPattern(.{
358 .name = "test-add-zero-to-identity",
359 .benefit = 10,
360 .classes = &.{.integer},
361 .lhs = .{
362 .name = "arith.add",
363 .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } },
364 },
365 .rhs = .{ .operation = .{ .name = "test.identity", .operands = &.{.{ .variable = 0 }} } },
366 });
367 }
368
369 const ReductionTestPass = EGraphPass(.{
370 .name = "choir-egraph-reduction-test",
371 .description = "test egraph operation materialization pass",
372 .populate_rules = reductionRules,
373 .options = .{ .candidate = testCandidate },
374 });
375
376 test "egraph extraction rejects unqualified cheaper rule products" {
377 const test_dialect = test_support.TestDialect;
378 const testing = std.testing;
379
380 var arena = alloc_arena.Arena.init(testing.allocator);
381 defer arena.deinit();
382 const allocator = arena.allocator();
383
384 var ctx = try buildTestContext(allocator);
385 defer ctx.deinit(allocator);
386
387 const loc = ir.Location.getUnknown();
388 const i64_type = try arith.getScalarType(&ctx, .i64);
389 const module = try test_dialect.ModuleOp.create(&ctx, loc);
390 const block = module.getBodyBlock();
391
392 var value = try arith.ConstantOp.createInt(&ctx, loc, i64_type, 7);
393 var two = try arith.ConstantOp.createInt(&ctx, loc, i64_type, 0);
394 var double = try arith.AddOp.create(&ctx, loc, value.getResult(), two.getResult());
395 const ret = try test_dialect.ReturnOp.create(&ctx, loc, &.{double.getResult()});
396
397 try block.addOperation(value.op);
398 try block.addOperation(two.op);
399 try block.addOperation(double.op);
400 try block.addOperation(ret.op);
401
402 var pm = pass_mod.PassManager.init(allocator);
403 defer pm.deinit();
404 try pm.addPass(ReductionTestPass.create());
405 try testing.expectEqual(pass_mod.PassResult.success, pm.run(module.op, &ctx));
406
407 const replacement = ret.op.getOperand(0).?;
408 try testing.expect(replacement == double.getResult());
409 const def_any = replacement.getDefiningOp() orelse return error.TestExpectedResult;
410 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
411 try testing.expectEqualStrings("arith.add", def_op.name.name);
412 try testing.expect(def_op.getOperand(0).? == value.getResult());
413 try ir.verifyOperation(module.op, ir.verify.default_options);
414 }
415
416 test "egraph rewrite budget stops saturation without failing the pass" {
417 const test_dialect = test_support.TestDialect;
418 const testing = std.testing;
419
420 var arena = alloc_arena.Arena.init(testing.allocator);
421 defer arena.deinit();
422 const allocator = arena.allocator();
423
424 var ctx = try buildTestContext(allocator);
425 defer ctx.deinit(allocator);
426
427 const loc = ir.Location.getUnknown();
428 const i64_type = try arith.getScalarType(&ctx, .i64);
429 const module = try test_dialect.ModuleOp.create(&ctx, loc);
430 const block = module.getBodyBlock();
431
432 var seven = try arith.ConstantOp.createInt(&ctx, loc, i64_type, 7);
433 var eleven = try arith.ConstantOp.createInt(&ctx, loc, i64_type, 11);
434 var zero = try arith.ConstantOp.createInt(&ctx, loc, i64_type, 0);
435 var add1 = try arith.AddOp.create(&ctx, loc, seven.getResult(), zero.getResult());
436 var add2 = try arith.AddOp.create(&ctx, loc, eleven.getResult(), zero.getResult());
437 const ret = try test_dialect.ReturnOp.create(&ctx, loc, &.{ add1.getResult(), add2.getResult() });
438
439 try block.addOperation(seven.op);
440 try block.addOperation(eleven.op);
441 try block.addOperation(zero.op);
442 try block.addOperation(add1.op);
443 try block.addOperation(add2.op);
444 try block.addOperation(ret.op);
445
446 var pm = pass_mod.PassManager.init(allocator);
447 defer pm.deinit();
448 try pm.addPass(BudgetTestPass.create());
449 try testing.expectEqual(pass_mod.PassResult.success, pm.run(module.op, &ctx));
450
451 try testing.expect(ret.op.getOperand(0).? == seven.getResult());
452 try testing.expect(ret.op.getOperand(1).? == add2.getResult());
453 }