tiny.choir.dialects.arith.types

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

1const std = @import("std");2const ir = @import("../../core/root.zig");34pub const ScalarKind = enum {5    i8,6    i16,7    i32,8    i64,9    u8,10    u16,11    u32,12    u64,13    f16,14    bf16,15    f32,16    f64,17    index,18    bool,19};2021pub const ScalarClass = enum {22    boolean,23    signed_integer,24    unsigned_integer,25    float,26    index,27};2829pub const ScalarDescriptor = struct {30    kind: ScalarKind,31    suffix: []const u8,32    name: []const u8,33    class: ScalarClass,34    bit_width: u8,35    vector_lanes: []const u32,36};3738pub const ScalarSet = struct {39    bits: u16 = 0,4041    pub fn init(comptime kinds: []const ScalarKind) ScalarSet {42        comptime {43            for (kinds, 0..) |kind, index| {44                for (kinds[0..index]) |previous| {45                    if (kind == previous) @compileError("duplicate scalar kind in set");46                }47            }48        }49        var set: ScalarSet = .{};50        inline for (kinds) |kind| {51            set.insert(kind);52        }53        return set;54    }5556    pub fn insert(self: *ScalarSet, kind: ScalarKind) void {57        self.bits |= scalarBit(kind);58    }5960    pub fn contains(self: ScalarSet, kind: ScalarKind) bool {61        return self.bits & scalarBit(kind) != 0;62    }6364    pub fn containsAll(self: ScalarSet, required: ScalarSet) bool {65        return self.bits & required.bits == required.bits;66    }6768    pub fn kindFromTypeName(self: ScalarSet, name: []const u8) ?ScalarKind {69        const kind = scalarKindFromTypeName(name) orelse return null;70        return if (self.contains(kind)) kind else null;71    }7273    pub fn kindFromType(self: ScalarSet, ty: ir.Type) ?ScalarKind {74        return self.kindFromTypeName(ty.getDialectTypeName() orelse return null);75    }76};7778const VectorTypeSpec = struct {79    kind: ScalarKind,80    width: u32,81    name: []const u8,82};8384fn arithTypeName(comptime suffix: []const u8) []const u8 {85    return "arith." ++ suffix;86}8788fn scalarDescriptorValue(89    comptime kind: ScalarKind,90    class: ScalarClass,91    bit_width: u8,92    vector_lanes: []const u32,93) ScalarDescriptor {94    const suffix = @tagName(kind);95    return .{96        .kind = kind,97        .suffix = suffix,98        .name = arithTypeName(suffix),99        .class = class,100        .bit_width = bit_width,101        .vector_lanes = vector_lanes,102    };103}104105const scalar_descriptors = [_]ScalarDescriptor{106    scalarDescriptorValue(.i8, .signed_integer, 8, &.{ 2, 4, 8, 16 }),107    scalarDescriptorValue(.i16, .signed_integer, 16, &.{ 2, 4, 8, 16 }),108    scalarDescriptorValue(.i32, .signed_integer, 32, &.{ 2, 4, 8, 16 }),109    scalarDescriptorValue(.i64, .signed_integer, 64, &.{ 2, 4, 8 }),110    scalarDescriptorValue(.u8, .unsigned_integer, 8, &.{ 2, 4, 8, 16 }),111    scalarDescriptorValue(.u16, .unsigned_integer, 16, &.{ 2, 4, 8, 16 }),112    scalarDescriptorValue(.u32, .unsigned_integer, 32, &.{ 2, 4, 8, 16 }),113    scalarDescriptorValue(.u64, .unsigned_integer, 64, &.{ 2, 4, 8 }),114    scalarDescriptorValue(.f16, .float, 16, &.{ 2, 4, 8, 16 }),115    scalarDescriptorValue(.bf16, .float, 16, &.{ 2, 4, 8, 16 }),116    scalarDescriptorValue(.f32, .float, 32, &.{ 2, 4, 8, 16 }),117    scalarDescriptorValue(.f64, .float, 64, &.{ 2, 4, 8 }),118    scalarDescriptorValue(.index, .index, 64, &.{ 2, 4, 8, 16 }),119    scalarDescriptorValue(.bool, .boolean, 1, &.{}),120};121122comptime {123    const field_names = @typeInfo(ScalarKind).@"enum".field_names;124    if (field_names.len > @bitSizeOf(u16)) @compileError("ScalarSet storage is too small");125    if (field_names.len != scalar_descriptors.len) @compileError("scalar descriptor catalog is incomplete");126    for (scalar_descriptors, 0..) |descriptor, index| {127        if (@backingInt(descriptor.kind) != index) @compileError("scalar descriptor catalog order drifted");128    }129}130131fn scalarBit(kind: ScalarKind) u16 {132    return @as(u16, 1) << @intCast(@backingInt(kind));133}134135pub fn scalarDescriptor(kind: ScalarKind) *const ScalarDescriptor {136    return &scalar_descriptors[@backingInt(kind)];137}138139pub fn scalarTypeName(kind: ScalarKind) []const u8 {140    return scalarDescriptor(kind).name;141}142143pub fn scalarTypeSuffix(kind: ScalarKind) []const u8 {144    return scalarDescriptor(kind).suffix;145}146147pub fn scalarKindFromTypeName(name: []const u8) ?ScalarKind {148    const prefix = "arith.";149    if (name.len <= prefix.len or !std.mem.eql(u8, name[0..prefix.len], prefix)) return null;150    return scalarKindFromSuffix(name[prefix.len..]);151}152153pub fn scalarKindFromType(ty: ir.Type) ?ScalarKind {154    return scalarKindFromTypeName(ty.getDialectTypeName() orelse return null);155}156157pub fn scalarKindFromSuffix(suffix: []const u8) ?ScalarKind {158    return std.meta.stringToEnum(ScalarKind, suffix);159}160161pub fn scalarKindIsInteger(kind: ScalarKind) bool {162    return scalarKindIsSignedInteger(kind) or scalarKindIsUnsignedInteger(kind);163}164165pub fn scalarKindIsSignedInteger(kind: ScalarKind) bool {166    const class = scalarDescriptor(kind).class;167    return class == .signed_integer or class == .index;168}169170pub fn scalarKindIsUnsignedInteger(kind: ScalarKind) bool {171    return scalarDescriptor(kind).class == .unsigned_integer;172}173174pub fn scalarKindIsFloat(kind: ScalarKind) bool {175    return scalarDescriptor(kind).class == .float;176}177178pub fn scalarBitWidth(kind: ScalarKind) u8 {179    return scalarDescriptor(kind).bit_width;180}181182fn vectorTypeNameStatic(comptime width: u32, comptime kind: ScalarKind) []const u8 {183    @setEvalBranchQuota(32_000);184    return std.fmt.comptimePrint("arith.vec{d}x{s}", .{ width, scalarTypeSuffix(kind) });185}186187fn vectorTypeSpecCount() usize {188    comptime var count: usize = 0;189    inline for (scalar_descriptors) |spec| {190        count += spec.vector_lanes.len;191    }192    return count;193}194195const vector_type_specs = blk: {196    const count = vectorTypeSpecCount();197    var out: [count]VectorTypeSpec = undefined;198    var index: usize = 0;199    for (scalar_descriptors) |spec| {200        for (spec.vector_lanes) |width| {201            out[index] = .{202                .kind = spec.kind,203                .width = width,204                .name = vectorTypeNameStatic(width, spec.kind),205            };206            index += 1;207        }208    }209    break :blk out;210};211212const registered_type_specs = blk: {213    const count = scalar_descriptors.len + vector_type_specs.len;214    var out: [count]ir.dialects.TypeSpec = undefined;215    var index: usize = 0;216    for (scalar_descriptors) |spec| {217        out[index] = ir.dialects.typeName(spec.name);218        index += 1;219    }220    for (vector_type_specs) |spec| {221        out[index] = ir.dialects.typeName(spec.name);222        index += 1;223    }224    break :blk out;225};226227pub fn vectorTypeName(width: u32, kind: ScalarKind) ?[]const u8 {228    inline for (vector_type_specs) |spec| {229        if (spec.width == width and spec.kind == kind) return spec.name;230    }231    return null;232}233234pub fn vectorTypeNameForElement(width: u32, element_type_name: []const u8) ?[]const u8 {235    const kind = scalarKindFromTypeName(element_type_name) orelse return null;236    return vectorTypeName(width, kind);237}238239pub fn registeredTypeSpecs() []const ir.dialects.TypeSpec {240    return &registered_type_specs;241}242243pub const type_names = struct {244    pub const int8 = scalarTypeName(.i8);245    pub const int16 = scalarTypeName(.i16);246    pub const int32 = scalarTypeName(.i32);247    pub const int64 = scalarTypeName(.i64);248    pub const uint8 = scalarTypeName(.u8);249    pub const uint16 = scalarTypeName(.u16);250    pub const uint32 = scalarTypeName(.u32);251    pub const uint64 = scalarTypeName(.u64);252    pub const float16 = scalarTypeName(.f16);253    pub const bfloat16 = scalarTypeName(.bf16);254    pub const float32 = scalarTypeName(.f32);255    pub const float64 = scalarTypeName(.f64);256    pub const index = scalarTypeName(.index);257    pub const boolean = scalarTypeName(.bool);258259    pub const vec2xi8 = vectorTypeName(2, .i8).?;260    pub const vec4xi8 = vectorTypeName(4, .i8).?;261    pub const vec8xi8 = vectorTypeName(8, .i8).?;262    pub const vec16xi8 = vectorTypeName(16, .i8).?;263    pub const vec2xi16 = vectorTypeName(2, .i16).?;264    pub const vec4xi16 = vectorTypeName(4, .i16).?;265    pub const vec8xi16 = vectorTypeName(8, .i16).?;266    pub const vec16xi16 = vectorTypeName(16, .i16).?;267    pub const vec2xi32 = vectorTypeName(2, .i32).?;268    pub const vec4xi32 = vectorTypeName(4, .i32).?;269    pub const vec8xi32 = vectorTypeName(8, .i32).?;270    pub const vec16xi32 = vectorTypeName(16, .i32).?;271    pub const vec2xi64 = vectorTypeName(2, .i64).?;272    pub const vec4xi64 = vectorTypeName(4, .i64).?;273    pub const vec8xi64 = vectorTypeName(8, .i64).?;274    pub const vec2xu8 = vectorTypeName(2, .u8).?;275    pub const vec4xu8 = vectorTypeName(4, .u8).?;276    pub const vec8xu8 = vectorTypeName(8, .u8).?;277    pub const vec16xu8 = vectorTypeName(16, .u8).?;278    pub const vec2xu16 = vectorTypeName(2, .u16).?;279    pub const vec4xu16 = vectorTypeName(4, .u16).?;280    pub const vec8xu16 = vectorTypeName(8, .u16).?;281    pub const vec16xu16 = vectorTypeName(16, .u16).?;282    pub const vec2xu32 = vectorTypeName(2, .u32).?;283    pub const vec4xu32 = vectorTypeName(4, .u32).?;284    pub const vec8xu32 = vectorTypeName(8, .u32).?;285    pub const vec16xu32 = vectorTypeName(16, .u32).?;286    pub const vec2xu64 = vectorTypeName(2, .u64).?;287    pub const vec4xu64 = vectorTypeName(4, .u64).?;288    pub const vec8xu64 = vectorTypeName(8, .u64).?;289    pub const vec2xf16 = vectorTypeName(2, .f16).?;290    pub const vec4xf16 = vectorTypeName(4, .f16).?;291    pub const vec8xf16 = vectorTypeName(8, .f16).?;292    pub const vec16xf16 = vectorTypeName(16, .f16).?;293    pub const vec2xbf16 = vectorTypeName(2, .bf16).?;294    pub const vec4xbf16 = vectorTypeName(4, .bf16).?;295    pub const vec8xbf16 = vectorTypeName(8, .bf16).?;296    pub const vec16xbf16 = vectorTypeName(16, .bf16).?;297    pub const vec2xf32 = vectorTypeName(2, .f32).?;298    pub const vec4xf32 = vectorTypeName(4, .f32).?;299    pub const vec8xf32 = vectorTypeName(8, .f32).?;300    pub const vec16xf32 = vectorTypeName(16, .f32).?;301    pub const vec2xf64 = vectorTypeName(2, .f64).?;302    pub const vec4xf64 = vectorTypeName(4, .f64).?;303    pub const vec8xf64 = vectorTypeName(8, .f64).?;304    pub const vec2xindex = vectorTypeName(2, .index).?;305    pub const vec4xindex = vectorTypeName(4, .index).?;306    pub const vec8xindex = vectorTypeName(8, .index).?;307    pub const vec16xindex = vectorTypeName(16, .index).?;308};309310pub const VectorTypeInfo = struct {311    width: u32,312    elem_type_name: []const u8,313};314315pub fn parseVectorTypeName(type_name: []const u8) ?VectorTypeInfo {316    const prefix = "arith.vec";317    if (!std.mem.startsWith(u8, type_name, prefix)) return null;318    const rest = type_name[prefix.len..];319    const x_index = std.mem.indexOfScalar(u8, rest, 'x') orelse return null;320    const width_str = rest[0..x_index];321    if (width_str.len == 0) return null;322    const width = std.fmt.parseInt(u32, width_str, 10) catch return null;323    if (width == 0) return null;324325    const elem_suffix = rest[x_index + 1 ..];326    const kind = scalarKindFromSuffix(elem_suffix) orelse return null;327    if (vectorTypeName(width, kind)) |canonical| {328        if (!std.mem.eql(u8, type_name, canonical)) return null;329    } else {330        return null;331    }332333    return .{ .width = width, .elem_type_name = scalarTypeName(kind) };334}