tiny.choir.backends.gpu.nvptx.dialect

1const std = @import("std");2const choir = @import("../../../root.zig");34const ir = choir.ir;5const arith = choir.dialects.arith;6const gpu = @import("../../../dialects/gpu/root.zig");78const Dimension = gpu.Dimension;9const MmaShape = gpu.MmaShape;10const ShuffleMode = gpu.ShuffleMode;11const WarpOpKind = gpu.WarpOpKind;1213pub const NvptxDialect = struct {14    pub const name = "nvptx";15    const op_specs = ir.dialects.opSpec.dialect(@This());16    pub const spec = ir.dialects.dialectSpec(@This(), .{17        .dialect_attributes = &.{18            "nvptx.dim",       "nvptx.shuffle_mode", "nvptx.warp_op",19            "nvptx.mma_shape", "nvptx.atomic_kind",20        },21    });2223    pub const ThreadIdxOp = struct {24        op: *ir.Operation,25        pub const operation_spec = op_specs.leaf(.{26            .mnemonic = "tid",27            .operands = 0,28            .results = 1,29            .required_attrs = &.{"dim"},30        });31        pub const operation_name = operation_spec.name;3233        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !ThreadIdxOp {34            try loadSpec(ctx);35            var builder = ir.OperationBuilder.init(ctx);36            const index_type = try arith.ArithDialect.getIndexType(ctx);37            var state = op_specs.state(@This(), loc);38            state.addTypes(&.{index_type});3940            const op = try builder.create(state);41            try setDimensionAttr(op, ctx, dim);42            return .{ .op = op };43        }4445        pub fn getResult(self: *const ThreadIdxOp) *ir.Value {46            return self.op.getResult(0).?;47        }4849        pub fn getDimension(self: ThreadIdxOp) ?Dimension {50            return getDimensionAttr(self.op);51        }52    };5354    pub const BlockIdxOp = struct {55        op: *ir.Operation,56        pub const operation_spec = op_specs.leaf(.{57            .mnemonic = "ctaid",58            .operands = 0,59            .results = 1,60            .required_attrs = &.{"dim"},61        });62        pub const operation_name = operation_spec.name;6364        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !BlockIdxOp {65            try loadSpec(ctx);66            var builder = ir.OperationBuilder.init(ctx);67            const index_type = try arith.ArithDialect.getIndexType(ctx);68            var state = op_specs.state(@This(), loc);69            state.addTypes(&.{index_type});7071            const op = try builder.create(state);72            try setDimensionAttr(op, ctx, dim);73            return .{ .op = op };74        }7576        pub fn getResult(self: *const BlockIdxOp) *ir.Value {77            return self.op.getResult(0).?;78        }7980        pub fn getDimension(self: BlockIdxOp) ?Dimension {81            return getDimensionAttr(self.op);82        }83    };8485    pub const BlockDimOp = struct {86        op: *ir.Operation,87        pub const operation_spec = op_specs.leaf(.{88            .mnemonic = "ntid",89            .operands = 0,90            .results = 1,91            .required_attrs = &.{"dim"},92        });93        pub const operation_name = operation_spec.name;9495        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !BlockDimOp {96            try loadSpec(ctx);97            var builder = ir.OperationBuilder.init(ctx);98            const index_type = try arith.ArithDialect.getIndexType(ctx);99            var state = op_specs.state(@This(), loc);100            state.addTypes(&.{index_type});101102            const op = try builder.create(state);103            try setDimensionAttr(op, ctx, dim);104            return .{ .op = op };105        }106107        pub fn getResult(self: *const BlockDimOp) *ir.Value {108            return self.op.getResult(0).?;109        }110111        pub fn getDimension(self: BlockDimOp) ?Dimension {112            return getDimensionAttr(self.op);113        }114    };115116    pub const GridDimOp = struct {117        op: *ir.Operation,118        pub const operation_spec = op_specs.leaf(.{119            .mnemonic = "nctaid",120            .operands = 0,121            .results = 1,122            .required_attrs = &.{"dim"},123        });124        pub const operation_name = operation_spec.name;125126        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !GridDimOp {127            try loadSpec(ctx);128            var builder = ir.OperationBuilder.init(ctx);129            const index_type = try arith.ArithDialect.getIndexType(ctx);130            var state = op_specs.state(@This(), loc);131            state.addTypes(&.{index_type});132133            const op = try builder.create(state);134            try setDimensionAttr(op, ctx, dim);135            return .{ .op = op };136        }137138        pub fn getResult(self: *const GridDimOp) *ir.Value {139            return self.op.getResult(0).?;140        }141142        pub fn getDimension(self: GridDimOp) ?Dimension {143            return getDimensionAttr(self.op);144        }145    };146147    pub const Barrier0Op = struct {148        op: *ir.Operation,149        pub const operation_spec = op_specs.leaf(.{150            .mnemonic = "barrier0",151            .operands = 0,152            .results = 0,153        });154        pub const operation_name = operation_spec.name;155156        pub fn create(ctx: *ir.Context, loc: ir.Location) !Barrier0Op {157            try loadSpec(ctx);158            var builder = ir.OperationBuilder.init(ctx);159            const state = op_specs.state(@This(), loc);160            const op = try builder.create(state);161            return .{ .op = op };162        }163    };164165    pub const WarpBarrierAllOp = struct {166        op: *ir.Operation,167        pub const operation_spec = op_specs.leaf(.{168            .mnemonic = "bar.warp.sync.all",169            .operands = 0,170            .results = 0,171        });172        pub const operation_name = operation_spec.name;173174        pub fn create(ctx: *ir.Context, loc: ir.Location) !WarpBarrierAllOp {175            try loadSpec(ctx);176            var builder = ir.OperationBuilder.init(ctx);177            const state = op_specs.state(@This(), loc);178            const op = try builder.create(state);179            return .{ .op = op };180        }181    };182183    pub const LaneIdOp = struct {184        op: *ir.Operation,185        pub const operation_spec = op_specs.leaf(.{186            .mnemonic = "laneid",187            .operands = 0,188            .results = 1,189        });190        pub const operation_name = operation_spec.name;191192        pub fn create(ctx: *ir.Context, loc: ir.Location) !LaneIdOp {193            try loadSpec(ctx);194            var builder = ir.OperationBuilder.init(ctx);195            const index_type = try arith.ArithDialect.getIndexType(ctx);196            var state = op_specs.state(@This(), loc);197            state.addTypes(&.{index_type});198            const op = try builder.create(state);199            return .{ .op = op };200        }201202        pub fn getResult(self: *const LaneIdOp) *ir.Value {203            return self.op.getResult(0).?;204        }205    };206207    pub const WarpIdOp = struct {208        op: *ir.Operation,209        pub const operation_spec = op_specs.leaf(.{210            .mnemonic = "warpid",211            .operands = 0,212            .results = 1,213        });214        pub const operation_name = operation_spec.name;215216        pub fn create(ctx: *ir.Context, loc: ir.Location) !WarpIdOp {217            try loadSpec(ctx);218            var builder = ir.OperationBuilder.init(ctx);219            const index_type = try arith.ArithDialect.getIndexType(ctx);220            var state = op_specs.state(@This(), loc);221            state.addTypes(&.{index_type});222            const op = try builder.create(state);223            return .{ .op = op };224        }225226        pub fn getResult(self: *const WarpIdOp) *ir.Value {227            return self.op.getResult(0).?;228        }229    };230231    pub const BarrierSyncOp = struct {232        op: *ir.Operation,233        pub const operation_spec = op_specs.leaf(.{234            .mnemonic = "bar.sync",235            .operands = 1,236            .results = 0,237        });238        pub const operation_name = operation_spec.name;239240        pub fn create(ctx: *ir.Context, loc: ir.Location, barrier_id: *ir.Value) !BarrierSyncOp {241            try loadSpec(ctx);242            var builder = ir.OperationBuilder.init(ctx);243            var state = op_specs.state(@This(), loc);244            state.addOperands(&.{barrier_id});245            const op = try builder.create(state);246            return .{ .op = op };247        }248249        pub fn getBarrierId(self: BarrierSyncOp) *ir.Value {250            return self.op.operands.items[0].value;251        }252    };253254    pub const SyncWarpOp = struct {255        op: *ir.Operation,256        pub const operation_spec = op_specs.leaf(.{257            .mnemonic = "bar.warp.sync",258            .operands = 1,259            .results = 0,260        });261        pub const operation_name = operation_spec.name;262263        pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value) !SyncWarpOp {264            try loadSpec(ctx);265            var builder = ir.OperationBuilder.init(ctx);266            var state = op_specs.state(@This(), loc);267            state.addOperands(&.{mask});268            const op = try builder.create(state);269            return .{ .op = op };270        }271272        pub fn getMask(self: SyncWarpOp) *ir.Value {273            return self.op.operands.items[0].value;274        }275    };276277    pub const ActiveMaskOp = struct {278        op: *ir.Operation,279        pub const operation_spec = op_specs.leaf(.{280            .mnemonic = "activemask",281            .operands = 0,282            .results = 1,283        });284        pub const operation_name = operation_spec.name;285286        pub fn create(ctx: *ir.Context, loc: ir.Location) !ActiveMaskOp {287            try loadSpec(ctx);288            var builder = ir.OperationBuilder.init(ctx);289            const i32_type = try arith.ArithDialect.getI32Type(ctx);290            var state = op_specs.state(@This(), loc);291            state.addTypes(&.{i32_type});292            const op = try builder.create(state);293            return .{ .op = op };294        }295296        pub fn getResult(self: *const ActiveMaskOp) *ir.Value {297            return self.op.getResult(0).?;298        }299    };300301    pub const AllSyncOp = struct {302        op: *ir.Operation,303        pub const operation_spec = op_specs.leaf(.{304            .mnemonic = "vote.sync.all",305            .operands = .{ "mask", "predicate" },306            .results = .{"result"},307        });308        pub const operation_name = operation_spec.name;309310        pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !AllSyncOp {311            try loadSpec(ctx);312            var builder = ir.OperationBuilder.init(ctx);313            const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);314            var state = op_specs.state(@This(), loc);315            state.addOperands(&.{ mask, pred });316            state.addTypes(&.{bool_type});317            const op = try builder.create(state);318            return .{ .op = op };319        }320321        pub fn getResult(self: *const AllSyncOp) *ir.Value {322            return self.op.getResult(0).?;323        }324325        pub fn getMask(self: AllSyncOp) *ir.Value {326            return self.op.operands.items[0].value;327        }328329        pub fn getPredicate(self: AllSyncOp) *ir.Value {330            return self.op.operands.items[1].value;331        }332    };333334    pub const AnySyncOp = struct {335        op: *ir.Operation,336        pub const operation_spec = op_specs.leaf(.{337            .mnemonic = "vote.sync.any",338            .operands = .{ "mask", "predicate" },339            .results = .{"result"},340        });341        pub const operation_name = operation_spec.name;342343        pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !AnySyncOp {344            try loadSpec(ctx);345            var builder = ir.OperationBuilder.init(ctx);346            const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);347            var state = op_specs.state(@This(), loc);348            state.addOperands(&.{ mask, pred });349            state.addTypes(&.{bool_type});350            const op = try builder.create(state);351            return .{ .op = op };352        }353354        pub fn getResult(self: *const AnySyncOp) *ir.Value {355            return self.op.getResult(0).?;356        }357358        pub fn getMask(self: AnySyncOp) *ir.Value {359            return self.op.operands.items[0].value;360        }361362        pub fn getPredicate(self: AnySyncOp) *ir.Value {363            return self.op.operands.items[1].value;364        }365    };366367    pub const BallotSyncOp = struct {368        op: *ir.Operation,369        pub const operation_spec = op_specs.leaf(.{370            .mnemonic = "vote.sync.ballot",371            .operands = .{ "mask", "predicate" },372            .results = .{"result"},373        });374        pub const operation_name = operation_spec.name;375376        pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !BallotSyncOp {377            try loadSpec(ctx);378            var builder = ir.OperationBuilder.init(ctx);379            const i32_type = try arith.ArithDialect.getI32Type(ctx);380            var state = op_specs.state(@This(), loc);381            state.addOperands(&.{ mask, pred });382            state.addTypes(&.{i32_type});383            const op = try builder.create(state);384            return .{ .op = op };385        }386387        pub fn getResult(self: *const BallotSyncOp) *ir.Value {388            return self.op.getResult(0).?;389        }390391        pub fn getMask(self: BallotSyncOp) *ir.Value {392            return self.op.operands.items[0].value;393        }394395        pub fn getPredicate(self: BallotSyncOp) *ir.Value {396            return self.op.operands.items[1].value;397        }398    };399400    pub const ShflSyncOp = struct {401        op: *ir.Operation,402        pub const operation_spec = op_specs.leaf(.{403            .mnemonic = "shfl.sync",404            .operands = .{ "mask", "src", "lane_or_delta" },405            .results = .{"result"},406            .required_attrs = &.{"mode"},407        });408        pub const operation_name = operation_spec.name;409410        pub fn create(411            ctx: *ir.Context,412            loc: ir.Location,413            mode: ShuffleMode,414            mask: *ir.Value,415            src: *ir.Value,416            lane_or_delta: *ir.Value,417        ) !ShflSyncOp {418            try loadSpec(ctx);419            var builder = ir.OperationBuilder.init(ctx);420            var state = op_specs.state(@This(), loc);421            state.addOperands(&.{ mask, src, lane_or_delta });422            state.addTypes(&.{src.type});423            const op = try builder.create(state);424            try setShuffleModeAttr(op, ctx, mode);425            return .{ .op = op };426        }427428        pub fn getResult(self: *const ShflSyncOp) *ir.Value {429            return self.op.getResult(0).?;430        }431432        pub fn getMask(self: ShflSyncOp) *ir.Value {433            return self.op.operands.items[0].value;434        }435436        pub fn getSrc(self: ShflSyncOp) *ir.Value {437            return self.op.operands.items[1].value;438        }439440        pub fn getLaneOrDelta(self: ShflSyncOp) *ir.Value {441            return self.op.operands.items[2].value;442        }443444        pub fn getMode(self: ShflSyncOp) ?ShuffleMode {445            return getShuffleModeAttr(self.op);446        }447    };448449    pub const WarpReduceOp = struct {450        op: *ir.Operation,451        pub const operation_spec = op_specs.leaf(.{452            .mnemonic = "warp.reduce",453            .operands = .{ "mask", "value" },454            .results = .{"result"},455            .required_attrs = &.{"op"},456        });457        pub const operation_name = operation_spec.name;458459        pub fn create(460            ctx: *ir.Context,461            loc: ir.Location,462            op_kind: WarpOpKind,463            mask: *ir.Value,464            value: *ir.Value,465        ) !WarpReduceOp {466            try loadSpec(ctx);467            var builder = ir.OperationBuilder.init(ctx);468            var state = op_specs.state(@This(), loc);469            state.addOperands(&.{ mask, value });470            state.addTypes(&.{value.type});471            const op = try builder.create(state);472            try setWarpOpAttr(op, ctx, op_kind);473            return .{ .op = op };474        }475476        pub fn getResult(self: *const WarpReduceOp) *ir.Value {477            return self.op.getResult(0).?;478        }479480        pub fn getMask(self: WarpReduceOp) *ir.Value {481            return self.op.operands.items[0].value;482        }483484        pub fn getValue(self: WarpReduceOp) *ir.Value {485            return self.op.operands.items[1].value;486        }487488        pub fn getOpKind(self: WarpReduceOp) ?WarpOpKind {489            return getWarpOpAttr(self.op);490        }491    };492493    pub const MmaSyncOp = struct {494        op: *ir.Operation,495        pub const operation_spec = op_specs.leaf(.{496            .mnemonic = "mma.sync",497            .operands = .{ "a0", "a1", "a2", "a3", "b0", "b1", "c0", "c1", "c2", "c3" },498            .results = .{ "d0", "d1", "d2", "d3" },499            .required_attrs = &.{"shape"},500        });501        pub const operation_name = operation_spec.name;502503        pub fn create(504            ctx: *ir.Context,505            loc: ir.Location,506            operands: [10]*ir.Value,507            shape: MmaShape,508        ) !MmaSyncOp {509            try loadSpec(ctx);510            var builder = ir.OperationBuilder.init(ctx);511            var state = op_specs.state(@This(), loc);512            state.addOperands(&operands);513            state.addTypes(&.{ operands[6].type, operands[7].type, operands[8].type, operands[9].type });514            const op = try builder.create(state);515            try setMmaShapeAttr(op, ctx, shape);516            return .{ .op = op };517        }518519        pub fn getOperandValue(self: MmaSyncOp, index: usize) *ir.Value {520            return self.op.operands.items[index].value;521        }522523        pub fn getD(self: *const MmaSyncOp, index: usize) *ir.Value {524            return self.op.getResult(index).?;525        }526527        pub fn getShape(self: MmaSyncOp) ?MmaShape {528            return getMmaShapeAttr(self.op);529        }530    };531532    pub const CpAsyncSharedOp = struct {533        op: *ir.Operation,534        pub const operation_spec = op_specs.leaf(.{535            .mnemonic = "cp.async.shared",536            .operands = .{ "dst", "dst_index", "src", "src_index" },537            .results = 0,538            .required_attrs = &.{"bytes"},539        });540        pub const operation_name = operation_spec.name;541542        pub fn create(543            ctx: *ir.Context,544            loc: ir.Location,545            dst: *ir.Value,546            dst_index: *ir.Value,547            src: *ir.Value,548            src_index: *ir.Value,549            bytes: u32,550        ) !CpAsyncSharedOp {551            try loadSpec(ctx);552            var builder = ir.OperationBuilder.init(ctx);553            var state = op_specs.state(@This(), loc);554            state.addOperands(&.{ dst, dst_index, src, src_index });555            const op = try builder.create(state);556            const bytes_attr = try ctx.getI64Attr(@intCast(bytes));557            try op.setAttr("bytes", bytes_attr);558            return .{ .op = op };559        }560561        pub fn getDst(self: CpAsyncSharedOp) *ir.Value {562            return ir.dialects.operand(operation_spec, self.op, "dst");563        }564565        pub fn getDstIndex(self: CpAsyncSharedOp) *ir.Value {566            return ir.dialects.operand(operation_spec, self.op, "dst_index");567        }568569        pub fn getSrc(self: CpAsyncSharedOp) *ir.Value {570            return ir.dialects.operand(operation_spec, self.op, "src");571        }572573        pub fn getSrcIndex(self: CpAsyncSharedOp) *ir.Value {574            return ir.dialects.operand(operation_spec, self.op, "src_index");575        }576577        pub fn getBytes(self: CpAsyncSharedOp) ?u32 {578            const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "bytes") orelse return null;579            const raw = int_attr.getUnsignedValue();580            if (raw > std.math.maxInt(u32)) return null;581            return @intCast(raw);582        }583    };584585    pub const FenceDeviceOp = struct {586        op: *ir.Operation,587        pub const operation_spec = op_specs.leaf(.{588            .mnemonic = "membar.gl",589            .operands = 0,590            .results = 0,591        });592        pub const operation_name = operation_spec.name;593594        pub fn create(ctx: *ir.Context, loc: ir.Location) !FenceDeviceOp {595            try loadSpec(ctx);596            var builder = ir.OperationBuilder.init(ctx);597            const state = op_specs.state(@This(), loc);598            const op = try builder.create(state);599            return .{ .op = op };600        }601    };602603    pub const CpAsyncCommitOp = struct {604        op: *ir.Operation,605        pub const operation_spec = op_specs.leaf(.{606            .mnemonic = "cp.async.commit",607            .operands = 0,608            .results = 0,609        });610        pub const operation_name = operation_spec.name;611612        pub fn create(ctx: *ir.Context, loc: ir.Location) !CpAsyncCommitOp {613            try loadSpec(ctx);614            var builder = ir.OperationBuilder.init(ctx);615            const state = op_specs.state(@This(), loc);616            const op = try builder.create(state);617            return .{ .op = op };618        }619    };620621    pub const CpAsyncWaitOp = struct {622        op: *ir.Operation,623        pub const operation_spec = op_specs.leaf(.{624            .mnemonic = "cp.async.wait",625            .operands = 0,626            .results = 0,627            .required_attrs = &.{"groups"},628        });629        pub const operation_name = operation_spec.name;630631        pub fn create(ctx: *ir.Context, loc: ir.Location, groups: u32) !CpAsyncWaitOp {632            try loadSpec(ctx);633            var builder = ir.OperationBuilder.init(ctx);634            const state = op_specs.state(@This(), loc);635            const op = try builder.create(state);636            const groups_attr = try ctx.getI64Attr(@intCast(groups));637            try op.setAttr("groups", groups_attr);638            return .{ .op = op };639        }640641        pub fn getGroups(self: CpAsyncWaitOp) ?u32 {642            const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "groups") orelse return null;643            const raw = int_attr.getUnsignedValue();644            if (raw > std.math.maxInt(u32)) return null;645            return @intCast(raw);646        }647    };648649    pub const WarpScanOp = struct {650        op: *ir.Operation,651        pub const operation_spec = op_specs.leaf(.{652            .mnemonic = "warp.scan",653            .operands = .{ "mask", "value" },654            .results = .{"result"},655            .required_attrs = &.{ "inclusive", "op" },656        });657        pub const operation_name = operation_spec.name;658659        pub fn create(660            ctx: *ir.Context,661            loc: ir.Location,662            op_kind: WarpOpKind,663            inclusive: bool,664            mask: *ir.Value,665            value: *ir.Value,666        ) !WarpScanOp {667            try loadSpec(ctx);668            var builder = ir.OperationBuilder.init(ctx);669            var state = op_specs.state(@This(), loc);670            state.addOperands(&.{ mask, value });671            state.addTypes(&.{value.type});672            const op = try builder.create(state);673            try setWarpOpAttr(op, ctx, op_kind);674            try setBoolAttr(op, ctx, "inclusive", inclusive);675            return .{ .op = op };676        }677678        pub fn getResult(self: *const WarpScanOp) *ir.Value {679            return self.op.getResult(0).?;680        }681682        pub fn getMask(self: WarpScanOp) *ir.Value {683            return self.op.operands.items[0].value;684        }685686        pub fn getValue(self: WarpScanOp) *ir.Value {687            return self.op.operands.items[1].value;688        }689690        pub fn getOpKind(self: WarpScanOp) ?WarpOpKind {691            return getWarpOpAttr(self.op);692        }693694        pub fn isInclusive(self: WarpScanOp) bool {695            return getBoolAttrValue(self.op, "inclusive");696        }697    };698699    pub const LoadLocalOp: type = loadOp("local");700    pub const StoreLocalOp: type = storeOp("local");701702    pub const LoadGlobalOp: type = loadOp("global");703704    pub const StoreGlobalOp: type = storeOp("global");705706    pub const AtomicGlobalOp = struct {707        op: *ir.Operation,708        pub const operation_spec = op_specs.leaf(.{709            .mnemonic = "atom.global",710            .operands = 3,711            .results = 1,712            .attrs = &.{"kind"},713        });714        pub const operation_name = operation_spec.name;715716        pub fn create(717            ctx: *ir.Context,718            loc: ir.Location,719            kind: choir.dialects.AtomicRmwKind,720            value: *ir.Value,721            memref: *ir.Value,722            index: *ir.Value,723            result_type: ir.Type,724        ) !AtomicGlobalOp {725            try loadSpec(ctx);726            var builder = ir.OperationBuilder.init(ctx);727            var state = op_specs.state(@This(), loc);728            state.addOperands(&.{ value, memref, index });729            state.addTypes(&.{result_type});730            const op = try builder.create(state);731            const kind_attr = try ctx.getDialectAttr("nvptx.atomic_kind", kind.toString());732            try op.setAttr("kind", kind_attr);733            return .{ .op = op };734        }735736        pub fn getKind(self: AtomicGlobalOp) ?choir.dialects.AtomicRmwKind {737            const dialect_attr = self.op.getAttrAs(ir.Attribute.DialectAttr, "kind") orelse return null;738            return choir.dialects.AtomicRmwKind.fromString(dialect_attr.payload);739        }740741        pub fn getValue(self: AtomicGlobalOp) *ir.Value {742            return self.op.operands.items[0].value;743        }744745        pub fn getMemref(self: AtomicGlobalOp) *ir.Value {746            return self.op.operands.items[1].value;747        }748749        pub fn getIndex(self: AtomicGlobalOp) *ir.Value {750            return self.op.operands.items[2].value;751        }752753        pub fn getResult(self: *const AtomicGlobalOp) *ir.Value {754            return self.op.getResult(0).?;755        }756    };757758    pub const AtomicSharedOp = struct {759        op: *ir.Operation,760        pub const operation_spec = op_specs.leaf(.{761            .mnemonic = "atom.shared",762            .operands = 3,763            .results = 1,764            .attrs = &.{"kind"},765        });766        pub const operation_name = operation_spec.name;767768        pub fn create(769            ctx: *ir.Context,770            loc: ir.Location,771            kind: choir.dialects.AtomicRmwKind,772            value: *ir.Value,773            memref: *ir.Value,774            index: *ir.Value,775            result_type: ir.Type,776        ) !AtomicSharedOp {777            try loadSpec(ctx);778            var builder = ir.OperationBuilder.init(ctx);779            var state = op_specs.state(@This(), loc);780            state.addOperands(&.{ value, memref, index });781            state.addTypes(&.{result_type});782            const op = try builder.create(state);783            const kind_attr = try ctx.getDialectAttr("nvptx.atomic_kind", kind.toString());784            try op.setAttr("kind", kind_attr);785            return .{ .op = op };786        }787788        pub fn getKind(self: AtomicSharedOp) ?choir.dialects.AtomicRmwKind {789            const dialect_attr = self.op.getAttrAs(ir.Attribute.DialectAttr, "kind") orelse return null;790            return choir.dialects.AtomicRmwKind.fromString(dialect_attr.payload);791        }792793        pub fn getValue(self: AtomicSharedOp) *ir.Value {794            return self.op.operands.items[0].value;795        }796797        pub fn getMemref(self: AtomicSharedOp) *ir.Value {798            return self.op.operands.items[1].value;799        }800801        pub fn getIndex(self: AtomicSharedOp) *ir.Value {802            return self.op.operands.items[2].value;803        }804805        pub fn getResult(self: *const AtomicSharedOp) *ir.Value {806            return self.op.getResult(0).?;807        }808    };809810    pub const AtomicCasGlobalOp = struct {811        op: *ir.Operation,812        pub const operation_spec = op_specs.leaf(.{813            .mnemonic = "atom.global.cas",814            .operands = 4,815            .results = 1,816        });817        pub const operation_name = operation_spec.name;818819        pub fn create(820            ctx: *ir.Context,821            loc: ir.Location,822            expected: *ir.Value,823            desired: *ir.Value,824            memref: *ir.Value,825            index: *ir.Value,826            result_type: ir.Type,827        ) !AtomicCasGlobalOp {828            try loadSpec(ctx);829            var builder = ir.OperationBuilder.init(ctx);830            var state = op_specs.state(@This(), loc);831            state.addOperands(&.{ expected, desired, memref, index });832            state.addTypes(&.{result_type});833            const op = try builder.create(state);834            return .{ .op = op };835        }836837        pub fn getExpected(self: AtomicCasGlobalOp) *ir.Value {838            return self.op.operands.items[0].value;839        }840841        pub fn getDesired(self: AtomicCasGlobalOp) *ir.Value {842            return self.op.operands.items[1].value;843        }844845        pub fn getMemref(self: AtomicCasGlobalOp) *ir.Value {846            return self.op.operands.items[2].value;847        }848849        pub fn getIndex(self: AtomicCasGlobalOp) *ir.Value {850            return self.op.operands.items[3].value;851        }852853        pub fn getResult(self: *const AtomicCasGlobalOp) *ir.Value {854            return self.op.getResult(0).?;855        }856    };857858    pub const AtomicCasSharedOp = struct {859        op: *ir.Operation,860        pub const operation_spec = op_specs.leaf(.{861            .mnemonic = "atom.shared.cas",862            .operands = 4,863            .results = 1,864        });865        pub const operation_name = operation_spec.name;866867        pub fn create(868            ctx: *ir.Context,869            loc: ir.Location,870            expected: *ir.Value,871            desired: *ir.Value,872            memref: *ir.Value,873            index: *ir.Value,874            result_type: ir.Type,875        ) !AtomicCasSharedOp {876            try loadSpec(ctx);877            var builder = ir.OperationBuilder.init(ctx);878            var state = op_specs.state(@This(), loc);879            state.addOperands(&.{ expected, desired, memref, index });880            state.addTypes(&.{result_type});881            const op = try builder.create(state);882            return .{ .op = op };883        }884885        pub fn getExpected(self: AtomicCasSharedOp) *ir.Value {886            return self.op.operands.items[0].value;887        }888889        pub fn getDesired(self: AtomicCasSharedOp) *ir.Value {890            return self.op.operands.items[1].value;891        }892893        pub fn getMemref(self: AtomicCasSharedOp) *ir.Value {894            return self.op.operands.items[2].value;895        }896897        pub fn getIndex(self: AtomicCasSharedOp) *ir.Value {898            return self.op.operands.items[3].value;899        }900901        pub fn getResult(self: *const AtomicCasSharedOp) *ir.Value {902            return self.op.getResult(0).?;903        }904    };905906    pub const LoadSharedOp: type = loadOp("shared");907908    pub const StoreSharedOp: type = storeOp("shared");909910    pub const SinApproxF32Op = struct {911        op: *ir.Operation,912        pub const operation_spec = op_specs.leaf(.{913            .mnemonic = "sin.approx.f32",914            .operands = 1,915            .results = 1,916        });917        pub const operation_name = operation_spec.name;918919        pub fn create(ctx: *ir.Context, loc: ir.Location, value: *ir.Value) !SinApproxF32Op {920            try loadSpec(ctx);921            var builder = ir.OperationBuilder.init(ctx);922            var state = op_specs.state(@This(), loc);923            state.addOperands(&.{value});924            state.addTypes(&.{value.type});925            const op = try builder.create(state);926            return .{ .op = op };927        }928929        pub fn getValue(self: SinApproxF32Op) *ir.Value {930            return self.op.operands.items[0].value;931        }932933        pub fn getResult(self: *const SinApproxF32Op) *ir.Value {934            return self.op.getResult(0).?;935        }936    };937938    pub const CosApproxF32Op = struct {939        op: *ir.Operation,940        pub const operation_spec = op_specs.leaf(.{941            .mnemonic = "cos.approx.f32",942            .operands = 1,943            .results = 1,944        });945        pub const operation_name = operation_spec.name;946947        pub fn create(ctx: *ir.Context, loc: ir.Location, value: *ir.Value) !CosApproxF32Op {948            try loadSpec(ctx);949            var builder = ir.OperationBuilder.init(ctx);950            var state = op_specs.state(@This(), loc);951            state.addOperands(&.{value});952            state.addTypes(&.{value.type});953            const op = try builder.create(state);954            return .{ .op = op };955        }956957        pub fn getValue(self: CosApproxF32Op) *ir.Value {958            return self.op.operands.items[0].value;959        }960961        pub fn getResult(self: *const CosApproxF32Op) *ir.Value {962            return self.op.getResult(0).?;963        }964    };965966    pub const FmaRnF32Op = struct {967        op: *ir.Operation,968        pub const operation_spec = op_specs.leaf(.{969            .mnemonic = "fma.rn.f32",970            .operands = 3,971            .results = 1,972        });973        pub const operation_name = operation_spec.name;974975        pub fn create(ctx: *ir.Context, loc: ir.Location, a: *ir.Value, b: *ir.Value, c: *ir.Value) !FmaRnF32Op {976            try loadSpec(ctx);977            var builder = ir.OperationBuilder.init(ctx);978            var state = op_specs.state(@This(), loc);979            state.addOperands(&.{ a, b, c });980            state.addTypes(&.{a.type});981            const op = try builder.create(state);982            return .{ .op = op };983        }984985        pub fn getA(self: FmaRnF32Op) *ir.Value {986            return self.op.operands.items[0].value;987        }988989        pub fn getB(self: FmaRnF32Op) *ir.Value {990            return self.op.operands.items[1].value;991        }992993        pub fn getC(self: FmaRnF32Op) *ir.Value {994            return self.op.operands.items[2].value;995        }996997        pub fn getResult(self: *const FmaRnF32Op) *ir.Value {998            return self.op.getResult(0).?;999        }1000    };10011002    fn loadOp(comptime space: []const u8) type {1003        return struct {1004            op: *ir.Operation,1005            pub const operation_spec = op_specs.leaf(.{1006                .mnemonic = "ld." ++ space,1007                .operands = 2,1008                .results = 1,1009            });1010            pub const operation_name = operation_spec.name;10111012            pub fn create(1013                ctx: *ir.Context,1014                loc: ir.Location,1015                memref: *ir.Value,1016                index: *ir.Value,1017                result_type: ir.Type,1018            ) !@This() {1019                try loadSpec(ctx);1020                var builder = ir.OperationBuilder.init(ctx);1021                var state = op_specs.state(@This(), loc);1022                state.addOperands(&.{ memref, index });1023                state.addTypes(&.{result_type});1024                const op = try builder.create(state);1025                return .{ .op = op };1026            }10271028            pub fn getResult(self: *const @This()) *ir.Value {1029                return self.op.getResult(0).?;1030            }10311032            pub fn getMemref(self: @This()) *ir.Value {1033                return self.op.operands.items[0].value;1034            }10351036            pub fn getIndex(self: @This()) *ir.Value {1037                return self.op.operands.items[1].value;1038            }1039        };1040    }10411042    fn storeOp(comptime space: []const u8) type {1043        return struct {1044            op: *ir.Operation,1045            pub const operation_spec = op_specs.leaf(.{1046                .mnemonic = "st." ++ space,1047                .operands = 3,1048                .results = 0,1049            });1050            pub const operation_name = operation_spec.name;10511052            pub fn create(1053                ctx: *ir.Context,1054                loc: ir.Location,1055                value: *ir.Value,1056                memref: *ir.Value,1057                index: *ir.Value,1058            ) !@This() {1059                try loadSpec(ctx);1060                var builder = ir.OperationBuilder.init(ctx);1061                var state = op_specs.state(@This(), loc);1062                state.addOperands(&.{ value, memref, index });1063                const op = try builder.create(state);1064                return .{ .op = op };1065            }10661067            pub fn getValue(self: @This()) *ir.Value {1068                return self.op.operands.items[0].value;1069            }10701071            pub fn getMemref(self: @This()) *ir.Value {1072                return self.op.operands.items[1].value;1073            }10741075            pub fn getIndex(self: @This()) *ir.Value {1076                return self.op.operands.items[2].value;1077            }1078        };1079    }10801081    fn loadSpec(ctx: *ir.Context) !void {1082        try ir.dialects.loadDialectSpec(ctx, spec);1083    }10841085    fn setDimensionAttr(op: *ir.Operation, ctx: *ir.Context, dim: Dimension) !void {1086        const attr = try ctx.getDialectAttr("nvptx.dim", dim.toString());1087        try op.setAttr("dim", attr);1088    }10891090    fn getDimensionAttr(op: *const ir.Operation) ?Dimension {1091        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "dim") orelse return null;1092        return Dimension.fromString(dialect_attr.payload);1093    }10941095    fn setShuffleModeAttr(op: *ir.Operation, ctx: *ir.Context, mode: ShuffleMode) !void {1096        const attr = try ctx.getDialectAttr("nvptx.shuffle_mode", mode.toString());1097        try op.setAttr("mode", attr);1098    }10991100    fn getShuffleModeAttr(op: *const ir.Operation) ?ShuffleMode {1101        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "mode") orelse return null;1102        return ShuffleMode.fromString(dialect_attr.payload);1103    }11041105    fn setWarpOpAttr(op: *ir.Operation, ctx: *ir.Context, op_kind: WarpOpKind) !void {1106        const attr = try ctx.getDialectAttr("nvptx.warp_op", op_kind.toString());1107        try op.setAttr("op", attr);1108    }11091110    fn getWarpOpAttr(op: *const ir.Operation) ?WarpOpKind {1111        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "op") orelse return null;1112        return WarpOpKind.fromString(dialect_attr.payload);1113    }11141115    fn setMmaShapeAttr(op: *ir.Operation, ctx: *ir.Context, shape: MmaShape) !void {1116        var buf: [32]u8 = undefined;1117        const shape_str = try shape.toString(buf[0..]);1118        const shape_attr = try ctx.getDialectAttr("nvptx.mma_shape", shape_str);1119        try op.setAttr("shape", shape_attr);1120    }11211122    fn getMmaShapeAttr(op: *const ir.Operation) ?MmaShape {1123        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "shape") orelse return null;1124        return MmaShape.parse(dialect_attr.payload);1125    }11261127    fn setBoolAttr(op: *ir.Operation, ctx: *ir.Context, attr_name: []const u8, value: bool) !void {1128        const bool_attr = try ctx.getBoolAttr(value);1129        try op.setAttr(attr_name, bool_attr);1130    }11311132    fn getBoolAttrValue(op: *const ir.Operation, attr_name: []const u8) bool {1133        const bool_attr = op.getAttrAs(ir.Attribute.BoolAttr, attr_name) orelse return false;1134        return bool_attr.getValue();1135    }1136};11371138test {1139    std.testing.refAllDecls(@This());1140}11411142test "NvptxDialect operation specs register shapes and attributes" {1143    const testing = std.testing;11441145    var ctx = try ir.Context.init(testing.allocator, ir.Context.Limits.testing);1146    defer ctx.deinit(testing.allocator);11471148    try ir.dialects.loadDialectSpec(&ctx, NvptxDialect.spec);11491150    const tid_info = ctx.lookupOperation(NvptxDialect.ThreadIdxOp.operation_name) orelse return error.TestExpectedOperation;1151    try testing.expect(tid_info.shape.operands.allows(0));1152    try testing.expect(!tid_info.shape.operands.allows(1));1153    try testing.expect(tid_info.shape.results.allows(1));1154    try testing.expect(!tid_info.shape.results.allows(0));1155    try testing.expect(tid_info.shape.regions.allows(0));1156    try testing.expect(!tid_info.shape.regions.allows(1));1157    try testing.expect(tid_info.shape.successors.allows(0));1158    try testing.expect(!tid_info.shape.successors.allows(1));1159    try testing.expect(tid_info.hasInherentAttributeName("dim"));11601161    const load_info = ctx.lookupOperation(NvptxDialect.LoadGlobalOp.operation_name) orelse return error.TestExpectedOperation;1162    try testing.expect(load_info.shape.operands.allows(2));1163    try testing.expect(!load_info.shape.operands.allows(1));1164    try testing.expect(load_info.shape.results.allows(1));1165    try testing.expect(!load_info.shape.results.allows(0));11661167    const store_info = ctx.lookupOperation(NvptxDialect.StoreGlobalOp.operation_name) orelse return error.TestExpectedOperation;1168    try testing.expect(store_info.shape.operands.allows(3));1169    try testing.expect(!store_info.shape.operands.allows(2));1170    try testing.expect(store_info.shape.results.allows(0));1171    try testing.expect(!store_info.shape.results.allows(1));11721173    const cas_info = ctx.lookupOperation(NvptxDialect.AtomicCasGlobalOp.operation_name) orelse return error.TestExpectedOperation;1174    try testing.expect(cas_info.shape.operands.allows(4));1175    try testing.expect(!cas_info.shape.operands.allows(3));1176    try testing.expect(cas_info.shape.results.allows(1));1177    try testing.expect(!cas_info.shape.results.allows(0));1178}11791180test "NvptxDialect op-spec state supports strict contexts and verifier shape checks" {1181    const testing = std.testing;11821183    var ctx = try ir.Context.init(testing.allocator, ir.Context.Limits.testing);1184    defer ctx.deinit(testing.allocator);11851186    try ctx.requireRegistered();1187    try ir.dialects.loadDialectSpec(&ctx, choir.dialects.arith.spec);1188    try ir.dialects.loadDialectSpec(&ctx, NvptxDialect.spec);11891190    const loc = ir.Location.getUnknown();1191    const tid = try NvptxDialect.ThreadIdxOp.create(&ctx, loc, .x);1192    try testing.expectEqual(Dimension.x, tid.getDimension().?);1193    try ir.verifyOperation(tid.op, .{ .recursive = false });11941195    const index_type = try choir.dialects.arith.ArithDialect.getIndexType(&ctx);1196    var bad_load_state = ir.Operation.State.init(NvptxDialect.LoadGlobalOp.operation_name, loc);1197    bad_load_state.addTypes(&.{index_type});1198    const bad_load = try ctx.createOperation(bad_load_state);1199    try testing.expectError(error.OperandCountMismatch, ir.verifyOperation(bad_load, .{ .recursive = false }));1200}

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