tiny.choir.dialects.arith.eval

1const std = @import("std");2const ir = @import("../../core/root.zig");3const predicate_mod = @import("predicate.zig");4const scalar = @import("scalar.zig");5const types = @import("types.zig");67const Operation = ir.Operation;8const Attribute = ir.Attribute;9const interfaces = ir.interfaces;10const EvalError = interfaces.EvalError;11const EvalContext = interfaces.EvalContext;12const CmpPredicate = predicate_mod.CmpPredicate;1314const arith_prefix = "arith.";1516const 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};4445fn 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}7778fn 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}8586fn 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}9394fn 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}100101fn intOperand(attr: Attribute) ?i64 {102    const int_attr = attr.cast(Attribute.IntegerAttr) orelse return null;103    return int_attr.getValue();104}105106fn 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}111112fn makeInt(ctx: *ir.Context, value: i64) EvalError!Attribute {113    return ctx.getI64Attr(value) catch error.OutOfMemory;114}115116fn makeBool(ctx: *ir.Context, value: bool) EvalError!Attribute {117    return ctx.getBoolAttr(value) catch error.OutOfMemory;118}119120fn 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}130131fn usesUnsignedOrder(kind: types.ScalarKind) bool {132    return switch (kind) {133        .u8, .u16, .u32, .u64, .index => true,134        else => false,135    };136}137138fn 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    }150151    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}159160fn 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}178179fn 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}188189fn 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();198199    if (kind == .constant) {200        return op.getAttr("value") orelse error.InvalidConstant;201    }202203    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    }221222    if (kind == .cmp) {223        if (operands.len != 2) return error.InvalidOperand;224        const operand_type = types.scalarKindFromType(op.operands.items[0].value.type) orelse225            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    }236237    if (kind == .cast) return evaluateCast(op, operands);238    const result_class = types.scalarKindFromType(op.results.items[0].type) orelse239        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}247248fn 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}257258fn 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);265266    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]) orelse270            return error.InvalidOperand;271        return makeInt(ctx, scalar.truncate(selected, bits));272    }273274    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    }289290    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.InvalidOperand303        else304            scalar.divTruncChecked(lhs, rhs, bits) orelse return error.InvalidOperand,305        .rem => if (usesUnsignedOrder(result_kind))306            scalar.remTruncUnsignedChecked(lhs, rhs, bits) orelse return error.InvalidOperand307        else308            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}326327fn floatOperand(attr: Attribute) EvalError!f64 {328    return (attr.cast(Attribute.FloatAttr) orelse return error.InvalidOperand).getValue();329}330331fn makeFloat(ctx: *ir.Context, value: f64) EvalError!Attribute {332    return ctx.getF64Attr(value) catch error.OutOfMemory;333}334335fn 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}353354fn 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}364365fn 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}386387fn 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}408409fn 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}428429pub 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}438439fn integerValue(value: i64, kind: types.ScalarKind) i128 {440    const bits = types.scalarBitWidth(kind);441    return if (usesUnsignedOrder(kind))442        scalar.maskToBits(value, bits)443    else444        scalar.truncate(value, bits);445}446447fn 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}456457/// Round an exact integer directly to eight significant bits, avoiding double rounding.458fn 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}470471fn 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) orelse474        return error.UnsupportedOperation;475    const result = types.scalarKindFromType(op.results.items[0].type) orelse476        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            narrowed496        else497            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}515516const evaluatable_vtable = interfaces.Evaluatable.VTable{517    .canEval = canEval,518    .evaluate = evaluate,519};520521pub fn fallback(op: *const Operation) ?*const anyopaque {522    _ = op;523    return &evaluatable_vtable;524}525526test "arith evaluator minmax respects signed and index ordering" {527    const testing = std.testing;528529    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}

1pub const eval = @import("eval.zig");