tiny.choir.dialects.tile.TileDialect

1pub const TileDialect = struct {2    pub const name = type_names.tile;3    const op_specs = ir.dialects.opSpec.dialect(@This());45    pub const spec = ir.dialects.dialectSpec(@This(), .{6        .types = ir.dialects.typeNames(type_specs),7    });89    const type_specs = struct {10        pub const tile = ir.dialects.TypeSpec{11            .name = type_names.tile,12            .interfaces = &.{13                interfaces.TypeParamInterface.entry(&tile_type_param_vtable),14            },15        };16        pub const barrier = type_names.barrier;17    };1819    pub const TileTypePayload = struct {20        element_type_name: []const u8,21        element_type: ?ir.Type,22        mem_level: TileMemLevel,23        shape: []u32,24    };2526    const tile_type_param_vtable = interfaces.TypeParamInterface.VTable{27        .parse = parseTileTypeParams,28    };2930    pub const LoadOp = struct {31        op: *ir.Operation,3233        pub const operation_spec = op_specs.leaf(.{34            .mnemonic = "load",35            .operands = 2,36            .results = 1,37        });38        pub const operation_name = operation_spec.name;3940        pub fn create(41            ctx: *ir.Context,42            loc: ir.Location,43            buffer: *ir.Value,44            coord: *ir.Value,45            result_type: ir.Type,46        ) !LoadOp {47            var builder = ir.OperationBuilder.init(ctx);48            var state = op_specs.state(@This(), loc);49            state.addOperands(&.{ buffer, coord });50            state.addTypes(&.{result_type});51            const op = try builder.create(state);52            return .{ .op = op };53        }5455        pub fn getResult(self: *const LoadOp) *ir.Value {56            return self.op.getResult(0).?;57        }5859        pub fn getBuffer(self: LoadOp) *ir.Value {60            return self.op.operands.items[0].value;61        }6263        pub fn getCoord(self: LoadOp) *ir.Value {64            return self.op.operands.items[1].value;65        }66    };6768    pub const StoreOp = struct {69        op: *ir.Operation,7071        pub const operation_spec = op_specs.leaf(.{72            .mnemonic = "store",73            .operands = 3,74            .results = 0,75        });76        pub const operation_name = operation_spec.name;7778        pub fn create(79            ctx: *ir.Context,80            loc: ir.Location,81            tile: *ir.Value,82            buffer: *ir.Value,83            coord: *ir.Value,84        ) !StoreOp {85            var builder = ir.OperationBuilder.init(ctx);86            var state = op_specs.state(@This(), loc);87            state.addOperands(&.{ tile, buffer, coord });88            const op = try builder.create(state);89            return .{ .op = op };90        }9192        pub fn getTile(self: StoreOp) *ir.Value {93            return self.op.operands.items[0].value;94        }9596        pub fn getBuffer(self: StoreOp) *ir.Value {97            return self.op.operands.items[1].value;98        }99100        pub fn getCoord(self: StoreOp) *ir.Value {101            return self.op.operands.items[2].value;102        }103    };104105    pub const MmaOp = struct {106        op: *ir.Operation,107108        pub const operation_spec = op_specs.leaf(.{109            .mnemonic = "mma",110            .operands = 3,111            .results = 1,112        });113        pub const operation_name = operation_spec.name;114115        pub fn create(116            ctx: *ir.Context,117            loc: ir.Location,118            a: *ir.Value,119            b: *ir.Value,120            c: *ir.Value,121            result_type: ir.Type,122        ) !MmaOp {123            var builder = ir.OperationBuilder.init(ctx);124            var state = op_specs.state(@This(), loc);125            state.addOperands(&.{ a, b, c });126            state.addTypes(&.{result_type});127            const op = try builder.create(state);128            return .{ .op = op };129        }130131        pub fn getResult(self: *const MmaOp) *ir.Value {132            return self.op.getResult(0).?;133        }134135        pub fn getA(self: MmaOp) *ir.Value {136            return self.op.operands.items[0].value;137        }138139        pub fn getB(self: MmaOp) *ir.Value {140            return self.op.operands.items[1].value;141        }142143        pub fn getC(self: MmaOp) *ir.Value {144            return self.op.operands.items[2].value;145        }146    };147148    pub const CopyOp = struct {149        op: *ir.Operation,150151        pub const operation_spec = op_specs.leaf(.{152            .mnemonic = "copy",153            .operands = 2,154            .results = 0,155        });156        pub const operation_name = operation_spec.name;157158        pub fn create(159            ctx: *ir.Context,160            loc: ir.Location,161            src: *ir.Value,162            dst: *ir.Value,163        ) !CopyOp {164            var builder = ir.OperationBuilder.init(ctx);165            var state = op_specs.state(@This(), loc);166            state.addOperands(&.{ src, dst });167            const op = try builder.create(state);168            return .{ .op = op };169        }170171        pub fn getSrc(self: CopyOp) *ir.Value {172            return self.op.operands.items[0].value;173        }174175        pub fn getDst(self: CopyOp) *ir.Value {176            return self.op.operands.items[1].value;177        }178    };179180    pub const BarrierOp = struct {181        op: *ir.Operation,182183        pub const operation_spec = op_specs.leaf(.{184            .mnemonic = "barrier",185            .operands = 0,186            .results = 0,187        });188        pub const operation_name = operation_spec.name;189190        pub fn create(ctx: *ir.Context, loc: ir.Location) !BarrierOp {191            var builder = ir.OperationBuilder.init(ctx);192            const state = op_specs.state(@This(), loc);193            const op = try builder.create(state);194            return .{ .op = op };195        }196    };197198    pub const ArriveOp = struct {199        op: *ir.Operation,200201        pub const operation_spec = op_specs.leaf(.{202            .mnemonic = "arrive",203            .operands = 1,204            .results = 0,205        });206        pub const operation_name = operation_spec.name;207208        pub fn create(ctx: *ir.Context, loc: ir.Location, barrier: *ir.Value) !ArriveOp {209            var builder = ir.OperationBuilder.init(ctx);210            var state = op_specs.state(@This(), loc);211            state.addOperands(&.{barrier});212            const op = try builder.create(state);213            return .{ .op = op };214        }215216        pub fn getBarrier(self: ArriveOp) *ir.Value {217            return self.op.operands.items[0].value;218        }219    };220221    pub const WaitOp = struct {222        op: *ir.Operation,223224        pub const operation_spec = op_specs.leaf(.{225            .mnemonic = "wait",226            .operands = 2,227            .results = 0,228        });229        pub const operation_name = operation_spec.name;230231        pub fn create(ctx: *ir.Context, loc: ir.Location, barrier: *ir.Value, phase: *ir.Value) !WaitOp {232            var builder = ir.OperationBuilder.init(ctx);233            var state = op_specs.state(@This(), loc);234            state.addOperands(&.{ barrier, phase });235            const op = try builder.create(state);236            return .{ .op = op };237        }238239        pub fn getBarrier(self: WaitOp) *ir.Value {240            return self.op.operands.items[0].value;241        }242243        pub fn getPhase(self: WaitOp) *ir.Value {244            return self.op.operands.items[1].value;245        }246    };247248    fn deinitTilePayload(allocator: std.mem.Allocator, ptr: *anyopaque) void {249        const payload: *TileTypePayload = @ptrCast(@alignCast(ptr));250        if (payload.shape.len > 0) {251            allocator.free(payload.shape);252        }253        allocator.destroy(payload);254    }255256    fn parseShape(allocator: std.mem.Allocator, shape_str: []const u8) ![]u32 {257        if (shape_str.len == 0) return error.InvalidTileShape;258259        var dims = std.ArrayListUnmanaged(u32).empty;260        errdefer dims.deinit(allocator);261262        var iter = std.mem.splitScalar(u8, shape_str, 'x');263        while (iter.next()) |part| {264            if (part.len == 0) return error.InvalidTileShape;265            const value = std.fmt.parseInt(u32, part, 10) catch return error.InvalidTileShape;266            try dims.append(allocator, value);267        }268269        if (dims.items.len == 0) return error.InvalidTileShape;270        return dims.toOwnedSlice(allocator);271    }272273    fn parseTileTypeParams(274        type_ptr: *const anyopaque,275        ctx_opaque: *const interfaces.ContextOpaque,276    ) anyerror!?interfaces.TypeParamPayload {277        const ctx = interfaces.castContext(ir.Context, ctx_opaque);278        const storage: *const ir.Type.DialectTypeStorage = @ptrCast(@alignCast(type_ptr));279        if (storage.param_key.len == 0) return null;280281        var iter = std.mem.splitScalar(u8, storage.param_key, ',');282        const elem_name = iter.next() orelse return null;283        const mem_level_str = iter.next() orelse return null;284        const shape_str = iter.next() orelse return null;285        if (iter.next() != null) return null;286287        const mem_level = TileMemLevel.fromString(mem_level_str) orelse return null;288        const allocator = ir.context.typePayloadAllocator(ctx);289        const shape = parseShape(allocator, shape_str) catch return null;290        errdefer allocator.free(shape);291292        const payload = try allocator.create(TileTypePayload);293        payload.* = .{294            .element_type_name = elem_name,295            .element_type = ctx.getDialectTypeFromName(elem_name) catch null,296            .mem_level = mem_level,297            .shape = shape,298        };299        return .{ .ptr = payload, .deinit = deinitTilePayload };300    }301302    fn loadSpec(ctx: *ir.Context) !void {303        ir.dialects.loadDialectSpec(ctx, spec) catch |err| switch (err) {304            error.ContextFrozen => {},305            else => return err,306        };307    }308309    fn payloadFromType(ctx: *ir.Context, typ: ir.Type) ?*const TileTypePayload {310        loadSpec(ctx) catch return null;311        return ctx.getTypeParamPayload(typ, TileTypePayload) catch null;312    }313314    pub fn getTileType(315        ctx: *ir.Context,316        element_type: ir.Type,317        mem_level: TileMemLevel,318        shape: []const u32,319    ) !ir.Type {320        try loadSpec(ctx);321        if (shape.len == 0) return error.InvalidTileShape;322323        var buf: [128]u8 = undefined;324        const elem_name = element_type.getDialectTypeName() orelse "unknown";325        var pos: usize = 0;326        pos = try ir.format.appendFmt(buf[0..], pos, "{s},{s},", .{ elem_name, mem_level.toString() });327        for (shape, 0..) |dim, idx| {328            if (idx > 0) {329                pos = try ir.format.appendFmt(buf[0..], pos, "x", .{});330            }331            pos = try ir.format.appendFmt(buf[0..], pos, "{d}", .{dim});332        }333        return ctx.getDialectTypeFromNameWithKey(type_names.tile, buf[0..pos]);334    }335336    pub fn getTileElementType(ctx: *ir.Context, typ: ir.Type) ?ir.Type {337        const payload = payloadFromType(ctx, typ) orelse return null;338        return payload.element_type;339    }340341    pub fn getTileMemLevel(ctx: *ir.Context, typ: ir.Type) ?TileMemLevel {342        const payload = payloadFromType(ctx, typ) orelse return null;343        return payload.mem_level;344    }345346    pub fn getTileShape(ctx: *ir.Context, typ: ir.Type) ?[]const u32 {347        const payload = payloadFromType(ctx, typ) orelse return null;348        return payload.shape;349    }350351    pub fn getBarrierType(ctx: *ir.Context) !ir.Type {352        try loadSpec(ctx);353        return ctx.getDialectTypeFromName(type_names.barrier);354    }355};