tiny.choir.dialects.tile

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

1const std = @import("std");2const alloc_arena = @import("alloc_arena");3const ir = @import("../core/root.zig");4const interfaces = @import("../core/root.zig").interfaces;56pub const type_names = struct {7    pub const tile = "tile";8    pub const barrier = "tile.barrier";9};1011pub const TileMemLevel = enum {12    register,13    shared,14    global,1516    pub fn toString(self: TileMemLevel) []const u8 {17        return @tagName(self);18    }1920    pub fn fromString(s: []const u8) ?TileMemLevel {21        inline for (@typeInfo(TileMemLevel).@"enum".field_names, std.meta.tags(TileMemLevel)) |field_name, value| {22            if (std.mem.eql(u8, s, field_name)) return value;23        }24        return null;25    }26};2728pub const TileDialect = struct {29    pub const name = type_names.tile;30    const op_specs = ir.dialects.opSpec.dialect(@This());3132    pub const spec = ir.dialects.dialectSpec(@This(), .{33        .types = ir.dialects.typeNames(type_specs),34    });3536    const type_specs = struct {37        pub const tile = ir.dialects.TypeSpec{38            .name = type_names.tile,39            .interfaces = &.{40                interfaces.TypeParamInterface.entry(&tile_type_param_vtable),41            },42        };43        pub const barrier = type_names.barrier;44    };4546    pub const TileTypePayload = struct {47        element_type_name: []const u8,48        element_type: ?ir.Type,49        mem_level: TileMemLevel,50        shape: []u32,51    };5253    const tile_type_param_vtable = interfaces.TypeParamInterface.VTable{54        .parse = parseTileTypeParams,55    };5657    pub const LoadOp = struct {58        op: *ir.Operation,5960        pub const operation_spec = op_specs.leaf(.{61            .mnemonic = "load",62            .operands = 2,63            .results = 1,64        });65        pub const operation_name = operation_spec.name;6667        pub fn create(68            ctx: *ir.Context,69            loc: ir.Location,70            buffer: *ir.Value,71            coord: *ir.Value,72            result_type: ir.Type,73        ) !LoadOp {74            var builder = ir.OperationBuilder.init(ctx);75            var state = op_specs.state(@This(), loc);76            state.addOperands(&.{ buffer, coord });77            state.addTypes(&.{result_type});78            const op = try builder.create(state);79            return .{ .op = op };80        }8182        pub fn getResult(self: *const LoadOp) *ir.Value {83            return self.op.getResult(0).?;84        }8586        pub fn getBuffer(self: LoadOp) *ir.Value {87            return self.op.operands.items[0].value;88        }8990        pub fn getCoord(self: LoadOp) *ir.Value {91            return self.op.operands.items[1].value;92        }93    };9495    pub const StoreOp = struct {96        op: *ir.Operation,9798        pub const operation_spec = op_specs.leaf(.{99            .mnemonic = "store",100            .operands = 3,101            .results = 0,102        });103        pub const operation_name = operation_spec.name;104105        pub fn create(106            ctx: *ir.Context,107            loc: ir.Location,108            tile: *ir.Value,109            buffer: *ir.Value,110            coord: *ir.Value,111        ) !StoreOp {112            var builder = ir.OperationBuilder.init(ctx);113            var state = op_specs.state(@This(), loc);114            state.addOperands(&.{ tile, buffer, coord });115            const op = try builder.create(state);116            return .{ .op = op };117        }118119        pub fn getTile(self: StoreOp) *ir.Value {120            return self.op.operands.items[0].value;121        }122123        pub fn getBuffer(self: StoreOp) *ir.Value {124            return self.op.operands.items[1].value;125        }126127        pub fn getCoord(self: StoreOp) *ir.Value {128            return self.op.operands.items[2].value;129        }130    };131132    pub const MmaOp = struct {133        op: *ir.Operation,134135        pub const operation_spec = op_specs.leaf(.{136            .mnemonic = "mma",137            .operands = 3,138            .results = 1,139        });140        pub const operation_name = operation_spec.name;141142        pub fn create(143            ctx: *ir.Context,144            loc: ir.Location,145            a: *ir.Value,146            b: *ir.Value,147            c: *ir.Value,148            result_type: ir.Type,149        ) !MmaOp {150            var builder = ir.OperationBuilder.init(ctx);151            var state = op_specs.state(@This(), loc);152            state.addOperands(&.{ a, b, c });153            state.addTypes(&.{result_type});154            const op = try builder.create(state);155            return .{ .op = op };156        }157158        pub fn getResult(self: *const MmaOp) *ir.Value {159            return self.op.getResult(0).?;160        }161162        pub fn getA(self: MmaOp) *ir.Value {163            return self.op.operands.items[0].value;164        }165166        pub fn getB(self: MmaOp) *ir.Value {167            return self.op.operands.items[1].value;168        }169170        pub fn getC(self: MmaOp) *ir.Value {171            return self.op.operands.items[2].value;172        }173    };174175    pub const CopyOp = struct {176        op: *ir.Operation,177178        pub const operation_spec = op_specs.leaf(.{179            .mnemonic = "copy",180            .operands = 2,181            .results = 0,182        });183        pub const operation_name = operation_spec.name;184185        pub fn create(186            ctx: *ir.Context,187            loc: ir.Location,188            src: *ir.Value,189            dst: *ir.Value,190        ) !CopyOp {191            var builder = ir.OperationBuilder.init(ctx);192            var state = op_specs.state(@This(), loc);193            state.addOperands(&.{ src, dst });194            const op = try builder.create(state);195            return .{ .op = op };196        }197198        pub fn getSrc(self: CopyOp) *ir.Value {199            return self.op.operands.items[0].value;200        }201202        pub fn getDst(self: CopyOp) *ir.Value {203            return self.op.operands.items[1].value;204        }205    };206207    pub const BarrierOp = struct {208        op: *ir.Operation,209210        pub const operation_spec = op_specs.leaf(.{211            .mnemonic = "barrier",212            .operands = 0,213            .results = 0,214        });215        pub const operation_name = operation_spec.name;216217        pub fn create(ctx: *ir.Context, loc: ir.Location) !BarrierOp {218            var builder = ir.OperationBuilder.init(ctx);219            const state = op_specs.state(@This(), loc);220            const op = try builder.create(state);221            return .{ .op = op };222        }223    };224225    pub const ArriveOp = struct {226        op: *ir.Operation,227228        pub const operation_spec = op_specs.leaf(.{229            .mnemonic = "arrive",230            .operands = 1,231            .results = 0,232        });233        pub const operation_name = operation_spec.name;234235        pub fn create(ctx: *ir.Context, loc: ir.Location, barrier: *ir.Value) !ArriveOp {236            var builder = ir.OperationBuilder.init(ctx);237            var state = op_specs.state(@This(), loc);238            state.addOperands(&.{barrier});239            const op = try builder.create(state);240            return .{ .op = op };241        }242243        pub fn getBarrier(self: ArriveOp) *ir.Value {244            return self.op.operands.items[0].value;245        }246    };247248    pub const WaitOp = struct {249        op: *ir.Operation,250251        pub const operation_spec = op_specs.leaf(.{252            .mnemonic = "wait",253            .operands = 2,254            .results = 0,255        });256        pub const operation_name = operation_spec.name;257258        pub fn create(ctx: *ir.Context, loc: ir.Location, barrier: *ir.Value, phase: *ir.Value) !WaitOp {259            var builder = ir.OperationBuilder.init(ctx);260            var state = op_specs.state(@This(), loc);261            state.addOperands(&.{ barrier, phase });262            const op = try builder.create(state);263            return .{ .op = op };264        }265266        pub fn getBarrier(self: WaitOp) *ir.Value {267            return self.op.operands.items[0].value;268        }269270        pub fn getPhase(self: WaitOp) *ir.Value {271            return self.op.operands.items[1].value;272        }273    };274275    fn deinitTilePayload(allocator: std.mem.Allocator, ptr: *anyopaque) void {276        const payload: *TileTypePayload = @ptrCast(@alignCast(ptr));277        if (payload.shape.len > 0) {278            allocator.free(payload.shape);279        }280        allocator.destroy(payload);281    }282283    fn parseShape(allocator: std.mem.Allocator, shape_str: []const u8) ![]u32 {284        if (shape_str.len == 0) return error.InvalidTileShape;285286        var dims = std.ArrayListUnmanaged(u32).empty;287        errdefer dims.deinit(allocator);288289        var iter = std.mem.splitScalar(u8, shape_str, 'x');290        while (iter.next()) |part| {291            if (part.len == 0) return error.InvalidTileShape;292            const value = std.fmt.parseInt(u32, part, 10) catch return error.InvalidTileShape;293            try dims.append(allocator, value);294        }295296        if (dims.items.len == 0) return error.InvalidTileShape;297        return dims.toOwnedSlice(allocator);298    }299300    fn parseTileTypeParams(301        type_ptr: *const anyopaque,302        ctx_opaque: *const interfaces.ContextOpaque,303    ) anyerror!?interfaces.TypeParamPayload {304        const ctx = interfaces.castContext(ir.Context, ctx_opaque);305        const storage: *const ir.Type.DialectTypeStorage = @ptrCast(@alignCast(type_ptr));306        if (storage.param_key.len == 0) return null;307308        var iter = std.mem.splitScalar(u8, storage.param_key, ',');309        const elem_name = iter.next() orelse return null;310        const mem_level_str = iter.next() orelse return null;311        const shape_str = iter.next() orelse return null;312        if (iter.next() != null) return null;313314        const mem_level = TileMemLevel.fromString(mem_level_str) orelse return null;315        const allocator = ir.context.typePayloadAllocator(ctx);316        const shape = parseShape(allocator, shape_str) catch return null;317        errdefer allocator.free(shape);318319        const payload = try allocator.create(TileTypePayload);320        payload.* = .{321            .element_type_name = elem_name,322            .element_type = ctx.getDialectTypeFromName(elem_name) catch null,323            .mem_level = mem_level,324            .shape = shape,325        };326        return .{ .ptr = payload, .deinit = deinitTilePayload };327    }328329    fn loadSpec(ctx: *ir.Context) !void {330        ir.dialects.loadDialectSpec(ctx, spec) catch |err| switch (err) {331            error.ContextFrozen => {},332            else => return err,333        };334    }335336    fn payloadFromType(ctx: *ir.Context, typ: ir.Type) ?*const TileTypePayload {337        loadSpec(ctx) catch return null;338        return ctx.getTypeParamPayload(typ, TileTypePayload) catch null;339    }340341    pub fn getTileType(342        ctx: *ir.Context,343        element_type: ir.Type,344        mem_level: TileMemLevel,345        shape: []const u32,346    ) !ir.Type {347        try loadSpec(ctx);348        if (shape.len == 0) return error.InvalidTileShape;349350        var buf: [128]u8 = undefined;351        const elem_name = element_type.getDialectTypeName() orelse "unknown";352        var pos: usize = 0;353        pos = try ir.format.appendFmt(buf[0..], pos, "{s},{s},", .{ elem_name, mem_level.toString() });354        for (shape, 0..) |dim, idx| {355            if (idx > 0) {356                pos = try ir.format.appendFmt(buf[0..], pos, "x", .{});357            }358            pos = try ir.format.appendFmt(buf[0..], pos, "{d}", .{dim});359        }360        return ctx.getDialectTypeFromNameWithKey(type_names.tile, buf[0..pos]);361    }362363    pub fn getTileElementType(ctx: *ir.Context, typ: ir.Type) ?ir.Type {364        const payload = payloadFromType(ctx, typ) orelse return null;365        return payload.element_type;366    }367368    pub fn getTileMemLevel(ctx: *ir.Context, typ: ir.Type) ?TileMemLevel {369        const payload = payloadFromType(ctx, typ) orelse return null;370        return payload.mem_level;371    }372373    pub fn getTileShape(ctx: *ir.Context, typ: ir.Type) ?[]const u32 {374        const payload = payloadFromType(ctx, typ) orelse return null;375        return payload.shape;376    }377378    pub fn getBarrierType(ctx: *ir.Context) !ir.Type {379        try loadSpec(ctx);380        return ctx.getDialectTypeFromName(type_names.barrier);381    }382};383384test "TileDialect.TileType roundtrip" {385    const testing = std.testing;386387    var arena = alloc_arena.Arena.init(std.testing.allocator);388    defer arena.deinit();389    const allocator = arena.allocator();390391    var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);392    defer ctx.deinit(allocator);393394    const arith = @import("arith/root.zig");395    const elem_type = try arith.ArithDialect.getScalarType(&ctx, .f16);396    const shape = [_]u32{ 16, 8 };397    const tile_type = try TileDialect.getTileType(&ctx, elem_type, .shared, &shape);398399    try testing.expectEqual(TileMemLevel.shared, TileDialect.getTileMemLevel(&ctx, tile_type).?);400    const parsed_shape = TileDialect.getTileShape(&ctx, tile_type).?;401    try testing.expectEqual(@as(usize, shape.len), parsed_shape.len);402    try testing.expectEqual(shape[0], parsed_shape[0]);403    try testing.expectEqual(shape[1], parsed_shape[1]);404    try testing.expect(TileDialect.getTileElementType(&ctx, tile_type).?.eql(elem_type));405}406407test "TileDialect spec owns tile type params" {408    const testing = std.testing;409410    var arena = alloc_arena.Arena.init(std.testing.allocator);411    defer arena.deinit();412    const allocator = arena.allocator();413414    var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);415    defer ctx.deinit(allocator);416417    try ir.dialects.loadDialectSpec(&ctx, TileDialect.spec);418419    const type_info = ctx.lookupType(type_names.tile) orelse return error.TestExpectedType;420    try testing.expect(type_info.hasInterface(interfaces.TypeParamInterface.id));421    try testing.expect(ctx.lookupType(type_names.barrier) != null);422}423424test "TileDialect specs own operation shapes and traits" {425    const testing = std.testing;426427    var arena = alloc_arena.Arena.init(std.testing.allocator);428    defer arena.deinit();429    const allocator = arena.allocator();430431    var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);432    defer ctx.deinit(allocator);433434    try ir.dialects.loadDialectSpec(&ctx, TileDialect.spec);435436    const Helpers = struct {437        fn expectLeaf(op_info: *const interfaces.OperationInfo, operands: usize, results: usize) !void {438            try testing.expect(op_info.shape.operands.allows(operands));439            try testing.expect(!op_info.shape.operands.allows(operands + 1));440            if (operands > 0) try testing.expect(!op_info.shape.operands.allows(operands - 1));441            try testing.expect(op_info.shape.results.allows(results));442            try testing.expect(!op_info.shape.results.allows(results + 1));443            if (results > 0) try testing.expect(!op_info.shape.results.allows(results - 1));444            try testing.expect(op_info.shape.regions.allows(0));445            try testing.expect(!op_info.shape.regions.allows(1));446            try testing.expect(op_info.shape.successors.allows(0));447            try testing.expect(!op_info.shape.successors.allows(1));448        }449450        fn lookup(context: *ir.Context, name: []const u8) !*const interfaces.OperationInfo {451            return context.lookupOperation(name) orelse error.TestExpectedOperation;452        }453    };454455    const load_info = try Helpers.lookup(&ctx, TileDialect.LoadOp.operation_name);456    try Helpers.expectLeaf(load_info, 2, 1);457    try testing.expect(load_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);458459    const store_info = try Helpers.lookup(&ctx, TileDialect.StoreOp.operation_name);460    try Helpers.expectLeaf(store_info, 3, 0);461    try testing.expect(store_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);462463    const mma_info = try Helpers.lookup(&ctx, TileDialect.MmaOp.operation_name);464    try Helpers.expectLeaf(mma_info, 3, 1);465    try testing.expect(mma_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);466467    const copy_info = try Helpers.lookup(&ctx, TileDialect.CopyOp.operation_name);468    try Helpers.expectLeaf(copy_info, 2, 0);469    try testing.expect(copy_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);470471    const barrier_info = try Helpers.lookup(&ctx, TileDialect.BarrierOp.operation_name);472    try Helpers.expectLeaf(barrier_info, 0, 0);473    try testing.expect(barrier_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);474475    const arrive_info = try Helpers.lookup(&ctx, TileDialect.ArriveOp.operation_name);476    try Helpers.expectLeaf(arrive_info, 1, 0);477    try testing.expect(arrive_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);478479    const wait_info = try Helpers.lookup(&ctx, TileDialect.WaitOp.operation_name);480    try Helpers.expectLeaf(wait_info, 2, 0);481    try testing.expect(wait_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);482}483484test "TileDialect verifier rejects malformed operation shapes" {485    const testing = std.testing;486487    var arena = alloc_arena.Arena.init(std.testing.allocator);488    defer arena.deinit();489    const allocator = arena.allocator();490491    var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);492    defer ctx.deinit(allocator);493494    const arith = @import("arith/root.zig");495    const loc = ir.Location.getUnknown();496    const elem_type = try arith.ArithDialect.getScalarType(&ctx, .f32);497    const shape = [_]u32{ 16, 8 };498    const tile_type = try TileDialect.getTileType(&ctx, elem_type, .shared, &shape);499500    var state = ir.Operation.State.init(TileDialect.LoadOp.operation_name, loc);501    state.addTypes(&.{tile_type});502    const malformed_load = try ctx.createOperation(state);503504    try testing.expectError(error.OperandCountMismatch, ir.verifyOperation(malformed_load, .{ .recursive = false }));505}506507test "TileDialect ops capture operands and results" {508    const testing = std.testing;509510    var arena = alloc_arena.Arena.init(std.testing.allocator);511    defer arena.deinit();512    const allocator = arena.allocator();513514    var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);515    defer ctx.deinit(allocator);516517    const arith = @import("arith/root.zig");518    const memref = @import("memref.zig");519520    const loc = ir.Location.getUnknown();521    const elem_type = try arith.ArithDialect.getScalarType(&ctx, .f32);522    const index_type = try arith.ArithDialect.getIndexType(&ctx);523    const memref_type = try memref.MemrefDialect.getMemrefType1D(&ctx, 128, elem_type, .shared);524    var alloc = try memref.MemrefDialect.AllocOp.createStatic(&ctx, loc, memref_type);525526    var coord = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 0);527    const shape = [_]u32{ 16, 8 };528    const tile_type = try TileDialect.getTileType(&ctx, elem_type, .shared, &shape);529530    const load = try TileDialect.LoadOp.create(&ctx, loc, alloc.getResult(), coord.getResult(), tile_type);531    try testing.expect(load.getResult().type.eql(tile_type));532    try testing.expect(load.getBuffer() == alloc.getResult());533    try testing.expect(load.getCoord() == coord.getResult());534535    const store = try TileDialect.StoreOp.create(&ctx, loc, load.getResult(), alloc.getResult(), coord.getResult());536    try testing.expect(store.getTile() == load.getResult());537    try testing.expect(store.getBuffer() == alloc.getResult());538    try testing.expect(store.getCoord() == coord.getResult());539540    const mma = try TileDialect.MmaOp.create(&ctx, loc, load.getResult(), load.getResult(), load.getResult(), tile_type);541    try testing.expect(mma.getResult().type.eql(tile_type));542    try testing.expect(mma.getA() == load.getResult());543544    const copy = try TileDialect.CopyOp.create(&ctx, loc, load.getResult(), load.getResult());545    try testing.expect(copy.getSrc() == load.getResult());546    try testing.expect(copy.getDst() == load.getResult());547548    const barrier_type = try TileDialect.getBarrierType(&ctx);549    try @import("fixture/root.zig").registerTestDialect(&ctx);550    _ = try ctx.registerOperation("test.tile.barrier", .{});551    var builder = ir.OperationBuilder.init(&ctx);552    var barrier_state = ir.Operation.State.init("test.tile.barrier", loc);553    barrier_state.addTypes(&.{barrier_type});554    const barrier_op = try builder.create(barrier_state);555    const barrier_val = barrier_op.getResult(0).?;556557    _ = try TileDialect.BarrierOp.create(&ctx, loc);558    const arrive = try TileDialect.ArriveOp.create(&ctx, loc, barrier_val);559    try testing.expect(arrive.getBarrier() == barrier_val);560561    var phase = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 1);562    const wait = try TileDialect.WaitOp.create(&ctx, loc, barrier_val, phase.getResult());563    try testing.expect(wait.getBarrier() == barrier_val);564    try testing.expect(wait.getPhase() == phase.getResult());565}