lib/choir/src/dialects/arith/rules.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const ir = @import("../../core/root.zig");
3 const egraph = @import("../../egraph/root.zig");
4 const types = @import("types.zig");
5 const fold_mod = @import("folds.zig");
6
7 pub fn Rules(comptime Dialect: type) type {
8 return struct {
9 const folds = fold_mod.Folds(Dialect);
10 fn classifyEGraphType(ty: ir.Type) egraph.TypeClass {
11 if (folds.isBoolType(ty)) return .boolean;
12 if (folds.isIntegerLikeType(ty)) return .integer;
13 const type_name = ty.getDialectTypeName() orelse return .other;
14 if (std.mem.eql(u8, type_name, types.type_names.float16) or
15 std.mem.eql(u8, type_name, types.type_names.float32) or
16 std.mem.eql(u8, type_name, types.type_names.float64))
17 {
18 return .float;
19 }
20 return .other;
21 }
22
23 pub fn egraphConstantModel() egraph.ConstantModel {
24 return .{
25 .op_name = Dialect.ConstantOp.operation_name,
26 .classify = classifyEGraphType,
27 .storage = .properties,
28 };
29 }
30
31 pub fn egraphCostModel() egraph.CostModel {
32 return .{
33 .operation = 4,
34 .constant = 1,
35 .constant_op_name = Dialect.ConstantOp.operation_name,
36 .overrides = &.{
37 .{ .name = Dialect.ShlOp.operation_name, .cost = 2 },
38 .{ .name = Dialect.ShrOp.operation_name, .cost = 2 },
39 .{ .name = Dialect.UshrOp.operation_name, .cost = 2 },
40 .{ .name = Dialect.AndOp.operation_name, .cost = 2 },
41 .{ .name = Dialect.OrOp.operation_name, .cost = 2 },
42 .{ .name = Dialect.XorOp.operation_name, .cost = 2 },
43 .{ .name = Dialect.NotOp.operation_name, .cost = 2 },
44 .{ .name = Dialect.NegOp.operation_name, .cost = 2 },
45 .{ .name = Dialect.MulOp.operation_name, .cost = 5 },
46 .{ .name = Dialect.DivOp.operation_name, .cost = 16 },
47 .{ .name = Dialect.RemOp.operation_name, .cost = 16 },
48 },
49 };
50 }
51
52 const integer_only: []const egraph.TypeClass = &.{.integer};
53 const boolean_only: []const egraph.TypeClass = &.{.boolean};
54 const integer_or_boolean: []const egraph.TypeClass = &.{ .integer, .boolean };
55 const integer_or_float: []const egraph.TypeClass = &.{ .integer, .float };
56 const float_only: []const egraph.TypeClass = &.{.float};
57
58 pub const egraph_rules = [_]egraph.PatternRule{
59 .{
60 .name = "arith-add-zero",
61 .classes = integer_only,
62 .lhs = .{ .name = Dialect.AddOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } },
63 .rhs = .{ .variable = 0 },
64 },
65 .{
66 .name = "arith-sub-zero",
67 .classes = integer_only,
68 .lhs = .{ .name = Dialect.SubOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } },
69 .rhs = .{ .variable = 0 },
70 },
71 .{
72 .name = "arith-sub-self",
73 .classes = integer_only,
74 .lhs = .{ .name = Dialect.SubOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .variable = 0 } } },
75 .rhs = .{ .constant = .{ .int = 0 } },
76 },
77 .{
78 .name = "arith-mul-one",
79 .classes = integer_only,
80 .lhs = .{ .name = Dialect.MulOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 1 } } } },
81 .rhs = .{ .variable = 0 },
82 },
83 .{
84 .name = "arith-mul-zero",
85 .classes = integer_only,
86 .lhs = .{ .name = Dialect.MulOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } },
87 .rhs = .{ .constant = .{ .int = 0 } },
88 },
89 .{
90 .name = "arith-div-one",
91 .classes = integer_only,
92 .lhs = .{ .name = Dialect.DivOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 1 } } } },
93 .rhs = .{ .variable = 0 },
94 },
95 .{
96 .name = "arith-rem-one",
97 .classes = integer_only,
98 .lhs = .{ .name = Dialect.RemOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 1 } } } },
99 .rhs = .{ .constant = .{ .int = 0 } },
100 },
101 .{
102 .name = "arith-shl-zero",
103 .classes = integer_only,
104 .lhs = .{ .name = Dialect.ShlOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } },
105 .rhs = .{ .variable = 0 },
106 },
107 .{
108 .name = "arith-shr-zero",
109 .classes = integer_only,
110 .lhs = .{ .name = Dialect.ShrOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } },
111 .rhs = .{ .variable = 0 },
112 },
113 .{
114 .name = "arith-ushr-zero",
115 .classes = integer_only,
116 .lhs = .{ .name = Dialect.UshrOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } },
117 .rhs = .{ .variable = 0 },
118 },
119 .{
120 .name = "arith-and-self",
121 .classes = integer_or_boolean,
122 .lhs = .{ .name = Dialect.AndOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .variable = 0 } } },
123 .rhs = .{ .variable = 0 },
124 },
125 .{
126 .name = "arith-or-self",
127 .classes = integer_or_boolean,
128 .lhs = .{ .name = Dialect.OrOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .variable = 0 } } },
129 .rhs = .{ .variable = 0 },
130 },
131 .{
132 .name = "arith-xor-self",
133 .classes = integer_only,
134 .lhs = .{ .name = Dialect.XorOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .variable = 0 } } },
135 .rhs = .{ .constant = .{ .int = 0 } },
136 },
137 .{
138 .name = "arith-xor-self-bool",
139 .classes = boolean_only,
140 .lhs = .{ .name = Dialect.XorOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .variable = 0 } } },
141 .rhs = .{ .constant = .{ .boolean = false } },
142 },
143 .{
144 .name = "arith-and-zero",
145 .classes = integer_only,
146 .lhs = .{ .name = Dialect.AndOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } },
147 .rhs = .{ .constant = .{ .int = 0 } },
148 },
149 .{
150 .name = "arith-or-zero",
151 .classes = integer_only,
152 .lhs = .{ .name = Dialect.OrOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } },
153 .rhs = .{ .variable = 0 },
154 },
155 .{
156 .name = "arith-xor-zero",
157 .classes = integer_only,
158 .lhs = .{ .name = Dialect.XorOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } },
159 .rhs = .{ .variable = 0 },
160 },
161 .{
162 .name = "arith-and-true",
163 .classes = boolean_only,
164 .lhs = .{ .name = Dialect.AndOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .boolean = true } } } },
165 .rhs = .{ .variable = 0 },
166 },
167 .{
168 .name = "arith-and-false",
169 .classes = boolean_only,
170 .lhs = .{ .name = Dialect.AndOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .boolean = false } } } },
171 .rhs = .{ .constant = .{ .boolean = false } },
172 },
173 .{
174 .name = "arith-or-true",
175 .classes = boolean_only,
176 .lhs = .{ .name = Dialect.OrOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .boolean = true } } } },
177 .rhs = .{ .constant = .{ .boolean = true } },
178 },
179 .{
180 .name = "arith-or-false",
181 .classes = boolean_only,
182 .lhs = .{ .name = Dialect.OrOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .boolean = false } } } },
183 .rhs = .{ .variable = 0 },
184 },
185 .{
186 .name = "arith-xor-false",
187 .classes = boolean_only,
188 .lhs = .{ .name = Dialect.XorOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .boolean = false } } } },
189 .rhs = .{ .variable = 0 },
190 },
191 .{
192 .name = "arith-not-not",
193 .classes = integer_or_boolean,
194 .lhs = .{ .name = Dialect.NotOp.operation_name, .operands = &.{
195 .{ .operation = .{ .name = Dialect.NotOp.operation_name, .operands = &.{.{ .variable = 0 }} } },
196 } },
197 .rhs = .{ .variable = 0 },
198 },
199 .{
200 .name = "arith-neg-neg",
201 .classes = integer_or_float,
202 .lhs = .{ .name = Dialect.NegOp.operation_name, .operands = &.{
203 .{ .operation = .{ .name = Dialect.NegOp.operation_name, .operands = &.{.{ .variable = 0 }} } },
204 } },
205 .rhs = .{ .variable = 0 },
206 },
207 .{
208 .name = "arith-mul-two-shl",
209 .classes = integer_only,
210 .lhs = .{ .name = Dialect.MulOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 2 } } } },
211 .rhs = .{ .operation = .{ .name = Dialect.ShlOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 1 } } } } },
212 },
213 .{
214 .name = "arith-mul-four-shl",
215 .classes = integer_only,
216 .lhs = .{ .name = Dialect.MulOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 4 } } } },
217 .rhs = .{ .operation = .{ .name = Dialect.ShlOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 2 } } } } },
218 },
219 .{
220 .name = "arith-mul-eight-shl",
221 .classes = integer_only,
222 .lhs = .{ .name = Dialect.MulOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 8 } } } },
223 .rhs = .{ .operation = .{ .name = Dialect.ShlOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 3 } } } } },
224 },
225 .{
226 .name = "arith-fmul-one",
227 .classes = float_only,
228 .lhs = .{ .name = Dialect.MulOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .float = 1.0 } } } },
229 .rhs = .{ .variable = 0 },
230 },
231 .{
232 .name = "arith-fdiv-one",
233 .classes = float_only,
234 .lhs = .{ .name = Dialect.DivOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .float = 1.0 } } } },
235 .rhs = .{ .variable = 0 },
236 },
237 .{
238 .name = "arith-fsub-zero",
239 .classes = float_only,
240 .lhs = .{ .name = Dialect.SubOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .float = 0.0 } } } },
241 .rhs = .{ .variable = 0 },
242 },
243 .{
244 .name = "arith-fadd-negzero",
245 .classes = float_only,
246 .lhs = .{ .name = Dialect.AddOp.operation_name, .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .float = -0.0 } } } },
247 .rhs = .{ .variable = 0 },
248 },
249 };
250
251 pub fn populateEGraphRules(rule_set: *egraph.RewriteSet) anyerror!void {
252 rule_set.setConstantModel(egraphConstantModel());
253 rule_set.setCostModel(egraphCostModel());
254 try rule_set.addPatterns(&egraph_rules);
255 }
256 };
257 }