lib/choir/src/dialects/arith/folds.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const ir = @import("../../core/root.zig");
3 const predicate_mod = @import("predicate.zig");
4 const types = @import("types.zig");
5
6 const CmpPredicate = predicate_mod.CmpPredicate;
7
8 const FoldBinaryOperands = struct {
9 lhs: *ir.Value,
10 rhs: *ir.Value,
11 };
12
13 pub fn Folds(comptime Dialect: type) type {
14 return struct {
15 fn appendFoldValue(
16 results: *ir.interfaces.FoldResults,
17 value: *ir.Value,
18 ) !void {
19 try results.append(.{ .value = value });
20 }
21
22 fn appendFoldAttribute(
23 results: *ir.interfaces.FoldResults,
24 attr: ir.Attribute,
25 ) !void {
26 try results.append(.{ .attribute = attr });
27 }
28
29 fn foldBinaryOperands(op: *const ir.Operation) ?FoldBinaryOperands {
30 if (op.operands.items.len != 2) return null;
31 if (op.results.items.len != 1) return null;
32 return .{
33 .lhs = op.operands.items[0].value,
34 .rhs = op.operands.items[1].value,
35 };
36 }
37
38 fn appendFoldValueForResult(
39 op: *const ir.Operation,
40 results: *ir.interfaces.FoldResults,
41 value: *ir.Value,
42 ) !void {
43 if (op.results.items.len != 1) return;
44 if (!op.results.items[0].type.eql(value.type)) return;
45 try appendFoldValue(results, value);
46 }
47
48 fn appendIntFoldAttributeForResult(
49 op: *const ir.Operation,
50 results: *ir.interfaces.FoldResults,
51 value: i64,
52 ) !void {
53 if (op.results.items.len != 1) return;
54 if (!isIntegerLikeType(op.results.items[0].type)) return;
55 const attr = try Dialect.getIntAttr(op.getContext(), value);
56 try appendFoldAttribute(results, attr);
57 }
58
59 fn appendBoolFoldAttributeForResult(
60 op: *const ir.Operation,
61 results: *ir.interfaces.FoldResults,
62 value: bool,
63 ) !void {
64 if (op.results.items.len != 1) return;
65 if (!isBoolType(op.results.items[0].type)) return;
66 const attr = try Dialect.getBoolAttr(op.getContext(), value);
67 try appendFoldAttribute(results, attr);
68 }
69
70 fn constantDefiningOp(value: *ir.Value) ?*ir.Operation {
71 const def_any = value.getDefiningOp() orelse return null;
72 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
73 if (!std.mem.eql(u8, def_op.name.name, Dialect.ConstantOp.operation_name)) return null;
74 return def_op;
75 }
76
77 pub fn constantIntFromValue(value: *ir.Value) ?i64 {
78 if (!isIntegerLikeType(value.type)) return null;
79 const def_op = constantDefiningOp(value) orelse return null;
80 const int_attr = def_op.getAttrAs(ir.Attribute.IntegerAttr, "value") orelse return null;
81 return int_attr.getValue();
82 }
83
84 pub fn constantBoolFromValue(value: *ir.Value) ?bool {
85 const def_op = constantDefiningOp(value) orelse return null;
86 if (def_op.getAttrAs(ir.Attribute.BoolAttr, "value")) |bool_attr| return bool_attr.getValue();
87 if (def_op.getAttrAs(ir.Attribute.IntegerAttr, "value")) |int_attr| {
88 const int_value = int_attr.getValue();
89 if (!isBoolType(value.type)) return null;
90 return int_value != 0;
91 }
92 return null;
93 }
94
95 pub fn constantIntEquals(value: *ir.Value, expected: i64) bool {
96 return (constantIntFromValue(value) orelse return false) == expected;
97 }
98
99 pub fn constantBoolEquals(value: *ir.Value, expected: bool) bool {
100 return (constantBoolFromValue(value) orelse return false) == expected;
101 }
102
103 pub fn isBoolType(ty: ir.Type) bool {
104 return types.scalarKindFromType(ty) == .bool;
105 }
106
107 pub fn isIntegerLikeType(ty: ir.Type) bool {
108 const kind = types.scalarKindFromType(ty) orelse return false;
109 return types.scalarKindIsInteger(kind);
110 }
111
112 pub fn foldSameTypeUnary(
113 op_ptr: *const anyopaque,
114 results: *ir.interfaces.FoldResults,
115 ) anyerror!void {
116 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
117 var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));
118 defer declaration.deinit(op.allocator);
119 if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;
120 if (op.operands.items.len != 1) return;
121 if (op.results.items.len != 1) return;
122 const input = op.operands.items[0].value;
123 if (!op.results.items[0].type.eql(input.type)) return;
124 try appendFoldValue(results, input);
125 }
126
127 pub fn foldSelect(
128 op_ptr: *const anyopaque,
129 results: *ir.interfaces.FoldResults,
130 ) anyerror!void {
131 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
132 var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));
133 defer declaration.deinit(op.allocator);
134 if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;
135 if (op.operands.items.len != 3) return;
136 if (op.results.items.len != 1) return;
137 const true_value = op.operands.items[1].value;
138 const false_value = op.operands.items[2].value;
139 if (true_value != false_value) return;
140 if (!op.results.items[0].type.eql(true_value.type)) return;
141 try appendFoldValue(results, true_value);
142 }
143
144 pub fn foldAdd(
145 op_ptr: *const anyopaque,
146 results: *ir.interfaces.FoldResults,
147 ) anyerror!void {
148 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
149 var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));
150 defer declaration.deinit(op.allocator);
151 if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;
152 const operands = foldBinaryOperands(op) orelse return;
153 if (constantIntEquals(operands.rhs, 0)) return appendFoldValueForResult(op, results, operands.lhs);
154 if (constantIntEquals(operands.lhs, 0)) return appendFoldValueForResult(op, results, operands.rhs);
155 }
156
157 pub fn foldSub(
158 op_ptr: *const anyopaque,
159 results: *ir.interfaces.FoldResults,
160 ) anyerror!void {
161 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
162 var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));
163 defer declaration.deinit(op.allocator);
164 if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;
165 const operands = foldBinaryOperands(op) orelse return;
166 if (constantIntEquals(operands.rhs, 0)) return appendFoldValueForResult(op, results, operands.lhs);
167 if (operands.lhs == operands.rhs) return appendIntFoldAttributeForResult(op, results, 0);
168 }
169
170 pub fn foldMul(
171 op_ptr: *const anyopaque,
172 results: *ir.interfaces.FoldResults,
173 ) anyerror!void {
174 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
175 var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));
176 defer declaration.deinit(op.allocator);
177 if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;
178 const operands = foldBinaryOperands(op) orelse return;
179 if (constantIntEquals(operands.rhs, 0)) return appendFoldValueForResult(op, results, operands.rhs);
180 if (constantIntEquals(operands.lhs, 0)) return appendFoldValueForResult(op, results, operands.lhs);
181 if (constantIntEquals(operands.rhs, 1)) return appendFoldValueForResult(op, results, operands.lhs);
182 if (constantIntEquals(operands.lhs, 1)) return appendFoldValueForResult(op, results, operands.rhs);
183 }
184
185 pub fn foldDiv(
186 op_ptr: *const anyopaque,
187 results: *ir.interfaces.FoldResults,
188 ) anyerror!void {
189 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
190 var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));
191 defer declaration.deinit(op.allocator);
192 if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;
193 const operands = foldBinaryOperands(op) orelse return;
194 if (constantIntEquals(operands.rhs, 1)) return appendFoldValueForResult(op, results, operands.lhs);
195 }
196
197 pub fn foldAnd(
198 op_ptr: *const anyopaque,
199 results: *ir.interfaces.FoldResults,
200 ) anyerror!void {
201 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
202 var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));
203 defer declaration.deinit(op.allocator);
204 if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;
205 const operands = foldBinaryOperands(op) orelse return;
206 if (operands.lhs == operands.rhs) return appendFoldValueForResult(op, results, operands.lhs);
207 if (constantBoolEquals(operands.rhs, false) or constantIntEquals(operands.rhs, 0)) {
208 return appendFoldValueForResult(op, results, operands.rhs);
209 }
210 if (constantBoolEquals(operands.lhs, false) or constantIntEquals(operands.lhs, 0)) {
211 return appendFoldValueForResult(op, results, operands.lhs);
212 }
213 if (constantBoolEquals(operands.rhs, true) or constantIntEquals(operands.rhs, -1)) {
214 return appendFoldValueForResult(op, results, operands.lhs);
215 }
216 if (constantBoolEquals(operands.lhs, true) or constantIntEquals(operands.lhs, -1)) {
217 return appendFoldValueForResult(op, results, operands.rhs);
218 }
219 }
220
221 pub fn foldOr(
222 op_ptr: *const anyopaque,
223 results: *ir.interfaces.FoldResults,
224 ) anyerror!void {
225 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
226 var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));
227 defer declaration.deinit(op.allocator);
228 if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;
229 const operands = foldBinaryOperands(op) orelse return;
230 if (operands.lhs == operands.rhs) return appendFoldValueForResult(op, results, operands.lhs);
231 if (constantBoolEquals(operands.rhs, true) or constantIntEquals(operands.rhs, -1)) {
232 return appendFoldValueForResult(op, results, operands.rhs);
233 }
234 if (constantBoolEquals(operands.lhs, true) or constantIntEquals(operands.lhs, -1)) {
235 return appendFoldValueForResult(op, results, operands.lhs);
236 }
237 if (constantBoolEquals(operands.rhs, false) or constantIntEquals(operands.rhs, 0)) {
238 return appendFoldValueForResult(op, results, operands.lhs);
239 }
240 if (constantBoolEquals(operands.lhs, false) or constantIntEquals(operands.lhs, 0)) {
241 return appendFoldValueForResult(op, results, operands.rhs);
242 }
243 }
244
245 pub fn foldXor(
246 op_ptr: *const anyopaque,
247 results: *ir.interfaces.FoldResults,
248 ) anyerror!void {
249 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
250 var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));
251 defer declaration.deinit(op.allocator);
252 if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;
253 const operands = foldBinaryOperands(op) orelse return;
254 if (operands.lhs == operands.rhs) {
255 if (isBoolType(op.results.items[0].type)) return appendBoolFoldAttributeForResult(op, results, false);
256 return appendIntFoldAttributeForResult(op, results, 0);
257 }
258 if (constantBoolEquals(operands.rhs, false) or constantIntEquals(operands.rhs, 0)) {
259 return appendFoldValueForResult(op, results, operands.lhs);
260 }
261 if (constantBoolEquals(operands.lhs, false) or constantIntEquals(operands.lhs, 0)) {
262 return appendFoldValueForResult(op, results, operands.rhs);
263 }
264 }
265
266 pub fn foldShift(
267 op_ptr: *const anyopaque,
268 results: *ir.interfaces.FoldResults,
269 ) anyerror!void {
270 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
271 var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));
272 defer declaration.deinit(op.allocator);
273 if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;
274 const operands = foldBinaryOperands(op) orelse return;
275 if (constantIntEquals(operands.rhs, 0)) return appendFoldValueForResult(op, results, operands.lhs);
276 }
277
278 pub fn foldNot(
279 op_ptr: *const anyopaque,
280 results: *ir.interfaces.FoldResults,
281 ) anyerror!void {
282 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
283 var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));
284 defer declaration.deinit(op.allocator);
285 if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;
286 if (op.operands.items.len != 1) return;
287 if (op.results.items.len != 1) return;
288 const value = constantBoolFromValue(op.operands.items[0].value) orelse return;
289 try appendBoolFoldAttributeForResult(op, results, !value);
290 }
291
292 fn cmpSelfResult(predicate_value: CmpPredicate, ty: ir.Type) ?bool {
293 if (isBoolType(ty)) {
294 return switch (predicate_value) {
295 .eq => true,
296 .ne => false,
297 else => null,
298 };
299 }
300 if (!isIntegerLikeType(ty)) return null;
301 return switch (predicate_value) {
302 .eq, .le, .ge, .sle, .sge, .ule, .uge => true,
303 .ne, .lt, .gt, .slt, .sgt, .ult, .ugt => false,
304 };
305 }
306
307 pub fn foldCmp(
308 op_ptr: *const anyopaque,
309 results: *ir.interfaces.FoldResults,
310 ) anyerror!void {
311 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
312 var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));
313 defer declaration.deinit(op.allocator);
314 if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;
315 const operands = foldBinaryOperands(op) orelse return;
316 if (operands.lhs != operands.rhs) return;
317 const cmp = Dialect.CmpOp{ .op = @constCast(op) };
318 const predicate = cmp.getPredicate() orelse return;
319 const folded = cmpSelfResult(predicate, operands.lhs.type) orelse return;
320 try appendBoolFoldAttributeForResult(op, results, folded);
321 }
322 };
323 }