lib/choir/src/dialects/arith/eval.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const ir = @import("../../core/root.zig");
3 const predicate_mod = @import("predicate.zig");
4 const scalar = @import("scalar.zig");
5 const types = @import("types.zig");
6
7 const Operation = ir.Operation;
8 const Attribute = ir.Attribute;
9 const interfaces = ir.interfaces;
10 const EvalError = interfaces.EvalError;
11 const EvalContext = interfaces.EvalContext;
12 const CmpPredicate = predicate_mod.CmpPredicate;
13
14 const arith_prefix = "arith.";
15
16 const ArithOp = enum {
17 constant,
18 add,
19 sub,
20 mul,
21 div,
22 rem,
23 max,
24 min,
25 cmp,
26 select,
27 neg,
28 abs,
29 not,
30 band,
31 bor,
32 bxor,
33 shl,
34 shr,
35 ushr,
36 cast,
37 sqrt,
38 popcount,
39 umulhi,
40 addo,
41 subo,
42 mulo,
43 };
44
45 fn classify(op: *const Operation) ?ArithOp {
46 const name = op.name.name;
47 if (!std.mem.startsWith(u8, name, arith_prefix)) return null;
48 const suffix = name[arith_prefix.len..];
49 if (std.mem.eql(u8, suffix, "constant")) return .constant;
50 if (std.mem.eql(u8, suffix, "add")) return .add;
51 if (std.mem.eql(u8, suffix, "sub")) return .sub;
52 if (std.mem.eql(u8, suffix, "mul")) return .mul;
53 if (std.mem.eql(u8, suffix, "div")) return .div;
54 if (std.mem.eql(u8, suffix, "rem")) return .rem;
55 if (std.mem.eql(u8, suffix, "max")) return .max;
56 if (std.mem.eql(u8, suffix, "min")) return .min;
57 if (std.mem.eql(u8, suffix, "cmp")) return .cmp;
58 if (std.mem.eql(u8, suffix, "select")) return .select;
59 if (std.mem.eql(u8, suffix, "neg")) return .neg;
60 if (std.mem.eql(u8, suffix, "abs")) return .abs;
61 if (std.mem.eql(u8, suffix, "not")) return .not;
62 if (std.mem.eql(u8, suffix, "and")) return .band;
63 if (std.mem.eql(u8, suffix, "or")) return .bor;
64 if (std.mem.eql(u8, suffix, "xor")) return .bxor;
65 if (std.mem.eql(u8, suffix, "shl")) return .shl;
66 if (std.mem.eql(u8, suffix, "shr")) return .shr;
67 if (std.mem.eql(u8, suffix, "ushr")) return .ushr;
68 if (std.mem.eql(u8, suffix, "cast")) return .cast;
69 if (std.mem.eql(u8, suffix, "sqrt")) return .sqrt;
70 if (std.mem.eql(u8, suffix, "popcount")) return .popcount;
71 if (std.mem.eql(u8, suffix, "umulhi")) return .umulhi;
72 if (std.mem.eql(u8, suffix, "addo")) return .addo;
73 if (std.mem.eql(u8, suffix, "subo")) return .subo;
74 if (std.mem.eql(u8, suffix, "mulo")) return .mulo;
75 return null;
76 }
77
78 fn intResultKind(op: *const Operation) ?types.ScalarKind {
79 if (op.results.items.len != 1) return null;
80 const type_name = op.results.items[0].type.getDialectTypeName() orelse return null;
81 const kind = types.scalarKindFromTypeName(type_name) orelse return null;
82 if (!types.scalarKindIsInteger(kind)) return null;
83 return kind;
84 }
85
86 fn intOperandKind(op: *const Operation, index: usize) ?types.ScalarKind {
87 if (op.operands.items.len <= index) return null;
88 const type_name = op.operands.items[index].value.type.getDialectTypeName() orelse return null;
89 const kind = types.scalarKindFromTypeName(type_name) orelse return null;
90 if (!types.scalarKindIsInteger(kind)) return null;
91 return kind;
92 }
93
94 fn cmpOperandKind(op: *const Operation) ?types.ScalarKind {
95 const lhs_kind = intOperandKind(op, 0) orelse return null;
96 const rhs_kind = intOperandKind(op, 1) orelse return null;
97 if (lhs_kind != rhs_kind) return null;
98 return lhs_kind;
99 }
100
101 fn intOperand(attr: Attribute) ?i64 {
102 const int_attr = attr.cast(Attribute.IntegerAttr) orelse return null;
103 return int_attr.getValue();
104 }
105
106 fn boolOperand(attr: Attribute) ?bool {
107 if (attr.cast(Attribute.BoolAttr)) |value| return value.getValue();
108 if (attr.cast(Attribute.IntegerAttr)) |value| return value.getValue() != 0;
109 return null;
110 }
111
112 fn makeInt(ctx: *ir.Context, value: i64) EvalError!Attribute {
113 return ctx.getI64Attr(value) catch error.OutOfMemory;
114 }
115
116 fn makeBool(ctx: *ir.Context, value: bool) EvalError!Attribute {
117 return ctx.getBoolAttr(value) catch error.OutOfMemory;
118 }
119
120 fn cmpPredicate(op: *const Operation) ?CmpPredicate {
121 const attr = op.getAttrAs(Attribute.DialectAttr, "predicate") orelse return null;
122 inline for (
123 @typeInfo(CmpPredicate).@"enum".field_names,
124 std.meta.tags(CmpPredicate),
125 ) |field_name, value| {
126 if (std.mem.eql(u8, attr.payload, field_name)) return value;
127 }
128 return null;
129 }
130
131 fn usesUnsignedOrder(kind: types.ScalarKind) bool {
132 return switch (kind) {
133 .u8, .u16, .u32, .u64, .index => true,
134 else => false,
135 };
136 }
137
138 fn minMaxInt(kind: ArithOp, lhs: i64, rhs: i64, result_kind: types.ScalarKind) i64 {
139 const bits = types.scalarBitWidth(result_kind);
140 if (usesUnsignedOrder(result_kind)) {
141 const left = scalar.maskToBits(lhs, bits);
142 const right = scalar.maskToBits(rhs, bits);
143 const selected = switch (kind) {
144 .max => if (left >= right) left else right,
145 .min => if (left <= right) left else right,
146 else => unreachable,
147 };
148 return scalar.unsignedResult(selected, bits);
149 }
150
151 const left = scalar.truncate(lhs, bits);
152 const right = scalar.truncate(rhs, bits);
153 return switch (kind) {
154 .max => if (left >= right) left else right,
155 .min => if (left <= right) left else right,
156 else => unreachable,
157 };
158 }
159
160 fn cmpInt(predicate: CmpPredicate, lhs: i64, rhs: i64, bits: u8) bool {
161 const lhs_u = scalar.maskToBits(lhs, bits);
162 const rhs_u = scalar.maskToBits(rhs, bits);
163 const lhs_s = scalar.truncate(lhs, bits);
164 const rhs_s = scalar.truncate(rhs, bits);
165 return switch (predicate) {
166 .eq => lhs_u == rhs_u,
167 .ne => lhs_u != rhs_u,
168 .lt, .slt => lhs_s < rhs_s,
169 .le, .sle => lhs_s <= rhs_s,
170 .gt, .sgt => lhs_s > rhs_s,
171 .ge, .sge => lhs_s >= rhs_s,
172 .ult => lhs_u < rhs_u,
173 .ule => lhs_u <= rhs_u,
174 .ugt => lhs_u > rhs_u,
175 .uge => lhs_u >= rhs_u,
176 };
177 }
178
179 fn canEval(op_ptr: *const anyopaque) bool {
180 const op: *const Operation = @ptrCast(@alignCast(op_ptr));
181 const kind = classify(op) orelse return false;
182 if (kind == .constant) return true;
183 if (kind == .addo or kind == .subo or kind == .mulo) return overflowShape(op);
184 if (op.getNumResults() != 1) return false;
185 if (kind == .cmp) return op.getNumOperands() == 2;
186 return types.scalarKindFromType(op.results.items[0].type) != null;
187 }
188
189 fn evaluate(
190 op_ptr: *const anyopaque,
191 operands: []const Attribute,
192 eval_ctx: *const EvalContext,
193 ) EvalError!Attribute {
194 _ = eval_ctx;
195 const op: *const Operation = @ptrCast(@alignCast(op_ptr));
196 const kind = classify(op) orelse return error.UnsupportedOperation;
197 const ctx = op.getContext();
198
199 if (kind == .constant) {
200 return op.getAttr("value") orelse error.InvalidConstant;
201 }
202
203 if (kind == .addo or kind == .subo or kind == .mulo) {
204 if (!overflowShape(op)) return error.UnsupportedOperation;
205 if (operands.len != 2) return error.InvalidOperand;
206 const lhs: i128 = intOperand(operands[0]) orelse return error.InvalidOperand;
207 const rhs: i128 = intOperand(operands[1]) orelse return error.InvalidOperand;
208 const exact = switch (kind) {
209 .addo => lhs + rhs,
210 .subo => lhs - rhs,
211 .mulo => lhs * rhs,
212 else => unreachable,
213 };
214 const wrapped: i64 = @truncate(exact);
215 const overflow = exact < std.math.minInt(i64) or exact > std.math.maxInt(i64);
216 return ctx.getArrayAttr(&.{
217 try makeInt(ctx, wrapped),
218 try makeBool(ctx, overflow),
219 }) catch error.OutOfMemory;
220 }
221
222 if (kind == .cmp) {
223 if (operands.len != 2) return error.InvalidOperand;
224 const operand_type = types.scalarKindFromType(op.operands.items[0].value.type) orelse
225 return error.UnsupportedOperation;
226 if (types.scalarKindIsFloat(operand_type)) {
227 return evaluateFloatCmp(op, operands, operand_type);
228 }
229 if (operand_type == .bool) return evaluateBoolCmp(op, operands);
230 const operand_kind = cmpOperandKind(op) orelse return error.UnsupportedOperation;
231 const lhs = intOperand(operands[0]) orelse return error.InvalidOperand;
232 const rhs = intOperand(operands[1]) orelse return error.InvalidOperand;
233 const predicate = cmpPredicate(op) orelse return error.InvalidOperand;
234 return makeBool(ctx, cmpInt(predicate, lhs, rhs, types.scalarBitWidth(operand_kind)));
235 }
236
237 if (kind == .cast) return evaluateCast(op, operands);
238 const result_class = types.scalarKindFromType(op.results.items[0].type) orelse
239 return error.UnsupportedOperation;
240 if (types.scalarKindIsFloat(result_class)) {
241 return evaluateFloat(op, kind, operands, result_class);
242 }
243 if (result_class == .bool) return evaluateBool(ctx, kind, operands);
244 const result_kind = intResultKind(op) orelse return error.UnsupportedOperation;
245 return evaluateInteger(ctx, kind, operands, result_kind);
246 }
247
248 fn overflowShape(op: *const Operation) bool {
249 if (op.getNumOperands() != 2 or op.getNumResults() != 2) return false;
250 if (types.scalarKindFromType(op.results.items[0].type) != .i64) return false;
251 if (types.scalarKindFromType(op.results.items[1].type) != .bool) return false;
252 for (op.getOperandValues()) |operand| {
253 if (types.scalarKindFromType(operand.type) != .i64) return false;
254 }
255 return op.regions.items.len == 0;
256 }
257
258 fn evaluateInteger(
259 ctx: *ir.Context,
260 kind: ArithOp,
261 operands: []const Attribute,
262 result_kind: types.ScalarKind,
263 ) EvalError!Attribute {
264 const bits = types.scalarBitWidth(result_kind);
265
266 if (kind == .select) {
267 if (operands.len != 3) return error.InvalidOperand;
268 const cond = boolOperand(operands[0]) orelse return error.InvalidOperand;
269 const selected = intOperand(if (cond) operands[1] else operands[2]) orelse
270 return error.InvalidOperand;
271 return makeInt(ctx, scalar.truncate(selected, bits));
272 }
273
274 switch (kind) {
275 .neg, .abs, .not, .popcount => {
276 if (operands.len != 1) return error.InvalidOperand;
277 const value = intOperand(operands[0]) orelse return error.InvalidOperand;
278 const folded: i64 = switch (kind) {
279 .neg => scalar.negWrap(value, bits),
280 .abs => scalar.absWrap(value, bits),
281 .not => scalar.bitNot(value, bits),
282 .popcount => @intCast(@popCount(scalar.maskToBits(value, bits))),
283 else => unreachable,
284 };
285 return makeInt(ctx, folded);
286 },
287 else => {},
288 }
289
290 if (operands.len != 2) return error.InvalidOperand;
291 const lhs = intOperand(operands[0]) orelse return error.InvalidOperand;
292 const rhs = intOperand(operands[1]) orelse return error.InvalidOperand;
293 const folded: i64 = switch (kind) {
294 .add => scalar.addWrap(lhs, rhs, bits),
295 .sub => scalar.subWrap(lhs, rhs, bits),
296 .mul => scalar.mulWrap(lhs, rhs, bits),
297 .umulhi => scalar.unsignedResult(@intCast(
298 (@as(u128, scalar.maskToBits(lhs, bits)) * scalar.maskToBits(rhs, bits)) >>
299 @intCast(bits),
300 ), bits),
301 .div => if (usesUnsignedOrder(result_kind))
302 scalar.divTruncUnsignedChecked(lhs, rhs, bits) orelse return error.InvalidOperand
303 else
304 scalar.divTruncChecked(lhs, rhs, bits) orelse return error.InvalidOperand,
305 .rem => if (usesUnsignedOrder(result_kind))
306 scalar.remTruncUnsignedChecked(lhs, rhs, bits) orelse return error.InvalidOperand
307 else
308 scalar.remTruncChecked(lhs, rhs, bits) orelse return error.InvalidOperand,
309 .max, .min => minMaxInt(kind, lhs, rhs, result_kind),
310 .band => scalar.bitAnd(lhs, rhs, bits),
311 .bor => scalar.bitOr(lhs, rhs, bits),
312 .bxor => scalar.bitXor(lhs, rhs, bits),
313 .shl, .shr, .ushr => blk: {
314 const count = scalar.shiftCount(rhs, bits) orelse return error.InvalidOperand;
315 break :blk switch (kind) {
316 .shl => scalar.shiftLeftWrap(lhs, count, bits),
317 .shr => scalar.shiftRightArithmetic(lhs, count, bits),
318 .ushr => scalar.shiftRightLogical(lhs, count, bits),
319 else => unreachable,
320 };
321 },
322 else => return error.UnsupportedOperation,
323 };
324 return makeInt(ctx, folded);
325 }
326
327 fn floatOperand(attr: Attribute) EvalError!f64 {
328 return (attr.cast(Attribute.FloatAttr) orelse return error.InvalidOperand).getValue();
329 }
330
331 fn makeFloat(ctx: *ir.Context, value: f64) EvalError!Attribute {
332 return ctx.getF64Attr(value) catch error.OutOfMemory;
333 }
334
335 fn evaluateBool(ctx: *ir.Context, kind: ArithOp, operands: []const Attribute) EvalError!Attribute {
336 if (kind == .select) {
337 if (operands.len != 3) return error.InvalidOperand;
338 const condition = boolOperand(operands[0]) orelse return error.InvalidOperand;
339 return operands[if (condition) @as(usize, 1) else 2];
340 }
341 if (operands.len == 0) return error.InvalidOperand;
342 const lhs = boolOperand(operands[0]) orelse return error.InvalidOperand;
343 if (kind == .not and operands.len == 1) return makeBool(ctx, !lhs);
344 if (operands.len != 2) return error.InvalidOperand;
345 const rhs = boolOperand(operands[1]) orelse return error.InvalidOperand;
346 return makeBool(ctx, switch (kind) {
347 .band => lhs and rhs,
348 .bor => lhs or rhs,
349 .bxor => lhs != rhs,
350 else => return error.UnsupportedOperation,
351 });
352 }
353
354 fn evaluateBoolCmp(op: *const Operation, operands: []const Attribute) EvalError!Attribute {
355 const lhs = boolOperand(operands[0]) orelse return error.InvalidOperand;
356 const rhs = boolOperand(operands[1]) orelse return error.InvalidOperand;
357 const predicate = cmpPredicate(op) orelse return error.InvalidPredicate;
358 return makeBool(op.getContext(), switch (predicate) {
359 .eq => lhs == rhs,
360 .ne => lhs != rhs,
361 else => return error.InvalidPredicate,
362 });
363 }
364
365 fn evaluateFloatCmp(
366 op: *const Operation,
367 operands: []const Attribute,
368 kind: types.ScalarKind,
369 ) EvalError!Attribute {
370 if (!op.getContext().arithmetic_policy.permitsFloatingValues()) {
371 return error.RequiresDynamicInfo;
372 }
373 const lhs = try roundedFloat(try floatOperand(operands[0]), kind);
374 const rhs = try roundedFloat(try floatOperand(operands[1]), kind);
375 const predicate = cmpPredicate(op) orelse return error.InvalidPredicate;
376 return makeBool(op.getContext(), switch (predicate) {
377 .eq => lhs == rhs,
378 .ne => lhs != rhs,
379 .lt => lhs < rhs,
380 .le => lhs <= rhs,
381 .gt => lhs > rhs,
382 .ge => lhs >= rhs,
383 else => return error.InvalidPredicate,
384 });
385 }
386
387 fn evaluateFloat(
388 op: *const Operation,
389 kind: ArithOp,
390 operands: []const Attribute,
391 typ: types.ScalarKind,
392 ) EvalError!Attribute {
393 if (!op.getContext().arithmetic_policy.permitsFloatingValues()) {
394 return error.RequiresDynamicInfo;
395 }
396 if (kind == .select) {
397 if (operands.len != 3) return error.InvalidOperand;
398 const condition = boolOperand(operands[0]) orelse return error.InvalidOperand;
399 return operands[if (condition) @as(usize, 1) else 2];
400 }
401 const value = switch (typ) {
402 .f32 => try floatArithmetic(f32, kind, operands),
403 .f64 => try floatArithmetic(f64, kind, operands),
404 else => return error.UnsupportedOperation,
405 };
406 return makeFloat(op.getContext(), value);
407 }
408
409 fn floatArithmetic(comptime T: type, kind: ArithOp, operands: []const Attribute) EvalError!f64 {
410 if (operands.len == 0) return error.InvalidOperand;
411 const lhs: T = @floatCast(try floatOperand(operands[0]));
412 if (operands.len == 1) return switch (kind) {
413 .neg => -lhs,
414 .abs => @abs(lhs),
415 .sqrt => if (lhs < 0) error.UnsupportedOperation else @sqrt(lhs),
416 else => error.UnsupportedOperation,
417 };
418 if (operands.len != 2) return error.InvalidOperand;
419 const rhs: T = @floatCast(try floatOperand(operands[1]));
420 return switch (kind) {
421 .add => lhs + rhs,
422 .sub => lhs - rhs,
423 .mul => lhs * rhs,
424 .div => lhs / rhs,
425 else => error.UnsupportedOperation,
426 };
427 }
428
429 pub fn roundedFloat(value: f64, kind: types.ScalarKind) EvalError!f64 {
430 return switch (kind) {
431 .f16 => @as(f16, @floatCast(value)),
432 .bf16 => scalar.roundBfloat(value),
433 .f32 => @as(f32, @floatCast(value)),
434 .f64 => value,
435 else => error.UnsupportedOperation,
436 };
437 }
438
439 fn integerValue(value: i64, kind: types.ScalarKind) i128 {
440 const bits = types.scalarBitWidth(kind);
441 return if (usesUnsignedOrder(kind))
442 scalar.maskToBits(value, bits)
443 else
444 scalar.truncate(value, bits);
445 }
446
447 fn integerFloat(value: i128, kind: types.ScalarKind) EvalError!f64 {
448 return switch (kind) {
449 .f16 => @as(f16, @floatFromInt(value)),
450 .f32 => @as(f32, @floatFromInt(value)),
451 .f64 => @floatFromInt(value),
452 .bf16 => integerBfloat(value),
453 else => error.UnsupportedOperation,
454 };
455 }
456
457 /// Round an exact integer directly to eight significant bits, avoiding double rounding.
458 fn integerBfloat(value: i128) f64 {
459 const magnitude: u128 = @intCast(if (value < 0) -value else value);
460 if (magnitude < 256) return @floatFromInt(value);
461 const exponent = std.math.log2_int(u128, magnitude);
462 const shift: u7 = @intCast(exponent - 7);
463 const half = @as(u128, 1) << (shift - 1);
464 const remainder = magnitude & ((half << 1) - 1);
465 var significant = magnitude >> shift;
466 if (remainder > half or (remainder == half and significant & 1 != 0)) significant += 1;
467 const rounded: f64 = @floatFromInt(significant << shift);
468 return if (value < 0) -rounded else rounded;
469 }
470
471 fn evaluateCast(op: *const Operation, operands: []const Attribute) EvalError!Attribute {
472 if (operands.len != 1) return error.InvalidOperand;
473 const source = types.scalarKindFromType(op.operands.items[0].value.type) orelse
474 return error.UnsupportedOperation;
475 const result = types.scalarKindFromType(op.results.items[0].type) orelse
476 return error.UnsupportedOperation;
477 if (source == .bool or result == .bool) {
478 if (source == result) return operands[0];
479 return error.UnsupportedOperation;
480 }
481 if (types.scalarKindIsFloat(source) or types.scalarKindIsFloat(result)) {
482 if (!op.getContext().arithmetic_policy.permitsFloatingValues()) {
483 return error.RequiresDynamicInfo;
484 }
485 }
486 if (types.scalarKindIsInteger(source)) {
487 const raw = intOperand(operands[0]) orelse return error.InvalidOperand;
488 const value = integerValue(raw, source);
489 if (types.scalarKindIsFloat(result)) {
490 return makeFloat(op.getContext(), try integerFloat(value, result));
491 }
492 const bits: u64 = @truncate(@as(u128, @bitCast(value)));
493 const narrowed = scalar.unsignedResult(bits, types.scalarBitWidth(result));
494 const result_value = if (usesUnsignedOrder(result))
495 narrowed
496 else
497 scalar.truncate(narrowed, types.scalarBitWidth(result));
498 return makeInt(op.getContext(), result_value);
499 }
500 if (types.scalarKindIsFloat(source) and types.scalarKindIsInteger(result)) {
501 const value = try roundedFloat(try floatOperand(operands[0]), source);
502 const integer = scalar.floatToInt(
503 value,
504 types.scalarBitWidth(result),
505 types.scalarKindIsSignedInteger(result) and result != .index,
506 ) orelse return error.UnsupportedOperation;
507 return makeInt(op.getContext(), integer);
508 }
509 if (types.scalarKindIsFloat(source) and types.scalarKindIsFloat(result)) {
510 const value = try roundedFloat(try floatOperand(operands[0]), source);
511 return makeFloat(op.getContext(), try roundedFloat(value, result));
512 }
513 return error.UnsupportedOperation;
514 }
515
516 const evaluatable_vtable = interfaces.Evaluatable.VTable{
517 .canEval = canEval,
518 .evaluate = evaluate,
519 };
520
521 pub fn fallback(op: *const Operation) ?*const anyopaque {
522 _ = op;
523 return &evaluatable_vtable;
524 }
525
526 test "arith evaluator minmax respects signed and index ordering" {
527 const testing = std.testing;
528
529 try testing.expectEqual(@as(i64, 7), minMaxInt(.max, -5, 7, .i64));
530 try testing.expectEqual(@as(i64, -5), minMaxInt(.min, -5, 7, .i64));
531 try testing.expectEqual(@as(i64, -1), minMaxInt(.max, -1, 1, .index));
532 try testing.expectEqual(@as(i64, 1), minMaxInt(.min, -1, 1, .index));
533 try testing.expectEqual(@as(i64, 255), minMaxInt(.max, -1, 1, .u8));
534 try testing.expectEqual(@as(i64, 1), minMaxInt(.min, -1, 1, .u8));
535 }