tiny.choir.backends.gpu.nvptx.dialect.NvptxDialect

1pub const NvptxDialect = struct {2    pub const name = "nvptx";3    const op_specs = ir.dialects.opSpec.dialect(@This());4    pub const spec = ir.dialects.dialectSpec(@This(), .{5        .dialect_attributes = &.{6            "nvptx.dim",       "nvptx.shuffle_mode", "nvptx.warp_op",7            "nvptx.mma_shape", "nvptx.atomic_kind",8        },9    });1011    pub const ThreadIdxOp = struct {12        op: *ir.Operation,13        pub const operation_spec = op_specs.leaf(.{14            .mnemonic = "tid",15            .operands = 0,16            .results = 1,17            .required_attrs = &.{"dim"},18        });19        pub const operation_name = operation_spec.name;2021        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !ThreadIdxOp {22            try loadSpec(ctx);23            var builder = ir.OperationBuilder.init(ctx);24            const index_type = try arith.ArithDialect.getIndexType(ctx);25            var state = op_specs.state(@This(), loc);26            state.addTypes(&.{index_type});2728            const op = try builder.create(state);29            try setDimensionAttr(op, ctx, dim);30            return .{ .op = op };31        }3233        pub fn getResult(self: *const ThreadIdxOp) *ir.Value {34            return self.op.getResult(0).?;35        }3637        pub fn getDimension(self: ThreadIdxOp) ?Dimension {38            return getDimensionAttr(self.op);39        }40    };4142    pub const BlockIdxOp = struct {43        op: *ir.Operation,44        pub const operation_spec = op_specs.leaf(.{45            .mnemonic = "ctaid",46            .operands = 0,47            .results = 1,48            .required_attrs = &.{"dim"},49        });50        pub const operation_name = operation_spec.name;5152        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !BlockIdxOp {53            try loadSpec(ctx);54            var builder = ir.OperationBuilder.init(ctx);55            const index_type = try arith.ArithDialect.getIndexType(ctx);56            var state = op_specs.state(@This(), loc);57            state.addTypes(&.{index_type});5859            const op = try builder.create(state);60            try setDimensionAttr(op, ctx, dim);61            return .{ .op = op };62        }6364        pub fn getResult(self: *const BlockIdxOp) *ir.Value {65            return self.op.getResult(0).?;66        }6768        pub fn getDimension(self: BlockIdxOp) ?Dimension {69            return getDimensionAttr(self.op);70        }71    };7273    pub const BlockDimOp = struct {74        op: *ir.Operation,75        pub const operation_spec = op_specs.leaf(.{76            .mnemonic = "ntid",77            .operands = 0,78            .results = 1,79            .required_attrs = &.{"dim"},80        });81        pub const operation_name = operation_spec.name;8283        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !BlockDimOp {84            try loadSpec(ctx);85            var builder = ir.OperationBuilder.init(ctx);86            const index_type = try arith.ArithDialect.getIndexType(ctx);87            var state = op_specs.state(@This(), loc);88            state.addTypes(&.{index_type});8990            const op = try builder.create(state);91            try setDimensionAttr(op, ctx, dim);92            return .{ .op = op };93        }9495        pub fn getResult(self: *const BlockDimOp) *ir.Value {96            return self.op.getResult(0).?;97        }9899        pub fn getDimension(self: BlockDimOp) ?Dimension {100            return getDimensionAttr(self.op);101        }102    };103104    pub const GridDimOp = struct {105        op: *ir.Operation,106        pub const operation_spec = op_specs.leaf(.{107            .mnemonic = "nctaid",108            .operands = 0,109            .results = 1,110            .required_attrs = &.{"dim"},111        });112        pub const operation_name = operation_spec.name;113114        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !GridDimOp {115            try loadSpec(ctx);116            var builder = ir.OperationBuilder.init(ctx);117            const index_type = try arith.ArithDialect.getIndexType(ctx);118            var state = op_specs.state(@This(), loc);119            state.addTypes(&.{index_type});120121            const op = try builder.create(state);122            try setDimensionAttr(op, ctx, dim);123            return .{ .op = op };124        }125126        pub fn getResult(self: *const GridDimOp) *ir.Value {127            return self.op.getResult(0).?;128        }129130        pub fn getDimension(self: GridDimOp) ?Dimension {131            return getDimensionAttr(self.op);132        }133    };134135    pub const Barrier0Op = struct {136        op: *ir.Operation,137        pub const operation_spec = op_specs.leaf(.{138            .mnemonic = "barrier0",139            .operands = 0,140            .results = 0,141        });142        pub const operation_name = operation_spec.name;143144        pub fn create(ctx: *ir.Context, loc: ir.Location) !Barrier0Op {145            try loadSpec(ctx);146            var builder = ir.OperationBuilder.init(ctx);147            const state = op_specs.state(@This(), loc);148            const op = try builder.create(state);149            return .{ .op = op };150        }151    };152153    pub const WarpBarrierAllOp = struct {154        op: *ir.Operation,155        pub const operation_spec = op_specs.leaf(.{156            .mnemonic = "bar.warp.sync.all",157            .operands = 0,158            .results = 0,159        });160        pub const operation_name = operation_spec.name;161162        pub fn create(ctx: *ir.Context, loc: ir.Location) !WarpBarrierAllOp {163            try loadSpec(ctx);164            var builder = ir.OperationBuilder.init(ctx);165            const state = op_specs.state(@This(), loc);166            const op = try builder.create(state);167            return .{ .op = op };168        }169    };170171    pub const LaneIdOp = struct {172        op: *ir.Operation,173        pub const operation_spec = op_specs.leaf(.{174            .mnemonic = "laneid",175            .operands = 0,176            .results = 1,177        });178        pub const operation_name = operation_spec.name;179180        pub fn create(ctx: *ir.Context, loc: ir.Location) !LaneIdOp {181            try loadSpec(ctx);182            var builder = ir.OperationBuilder.init(ctx);183            const index_type = try arith.ArithDialect.getIndexType(ctx);184            var state = op_specs.state(@This(), loc);185            state.addTypes(&.{index_type});186            const op = try builder.create(state);187            return .{ .op = op };188        }189190        pub fn getResult(self: *const LaneIdOp) *ir.Value {191            return self.op.getResult(0).?;192        }193    };194195    pub const WarpIdOp = struct {196        op: *ir.Operation,197        pub const operation_spec = op_specs.leaf(.{198            .mnemonic = "warpid",199            .operands = 0,200            .results = 1,201        });202        pub const operation_name = operation_spec.name;203204        pub fn create(ctx: *ir.Context, loc: ir.Location) !WarpIdOp {205            try loadSpec(ctx);206            var builder = ir.OperationBuilder.init(ctx);207            const index_type = try arith.ArithDialect.getIndexType(ctx);208            var state = op_specs.state(@This(), loc);209            state.addTypes(&.{index_type});210            const op = try builder.create(state);211            return .{ .op = op };212        }213214        pub fn getResult(self: *const WarpIdOp) *ir.Value {215            return self.op.getResult(0).?;216        }217    };218219    pub const BarrierSyncOp = struct {220        op: *ir.Operation,221        pub const operation_spec = op_specs.leaf(.{222            .mnemonic = "bar.sync",223            .operands = 1,224            .results = 0,225        });226        pub const operation_name = operation_spec.name;227228        pub fn create(ctx: *ir.Context, loc: ir.Location, barrier_id: *ir.Value) !BarrierSyncOp {229            try loadSpec(ctx);230            var builder = ir.OperationBuilder.init(ctx);231            var state = op_specs.state(@This(), loc);232            state.addOperands(&.{barrier_id});233            const op = try builder.create(state);234            return .{ .op = op };235        }236237        pub fn getBarrierId(self: BarrierSyncOp) *ir.Value {238            return self.op.operands.items[0].value;239        }240    };241242    pub const SyncWarpOp = struct {243        op: *ir.Operation,244        pub const operation_spec = op_specs.leaf(.{245            .mnemonic = "bar.warp.sync",246            .operands = 1,247            .results = 0,248        });249        pub const operation_name = operation_spec.name;250251        pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value) !SyncWarpOp {252            try loadSpec(ctx);253            var builder = ir.OperationBuilder.init(ctx);254            var state = op_specs.state(@This(), loc);255            state.addOperands(&.{mask});256            const op = try builder.create(state);257            return .{ .op = op };258        }259260        pub fn getMask(self: SyncWarpOp) *ir.Value {261            return self.op.operands.items[0].value;262        }263    };264265    pub const ActiveMaskOp = struct {266        op: *ir.Operation,267        pub const operation_spec = op_specs.leaf(.{268            .mnemonic = "activemask",269            .operands = 0,270            .results = 1,271        });272        pub const operation_name = operation_spec.name;273274        pub fn create(ctx: *ir.Context, loc: ir.Location) !ActiveMaskOp {275            try loadSpec(ctx);276            var builder = ir.OperationBuilder.init(ctx);277            const i32_type = try arith.ArithDialect.getI32Type(ctx);278            var state = op_specs.state(@This(), loc);279            state.addTypes(&.{i32_type});280            const op = try builder.create(state);281            return .{ .op = op };282        }283284        pub fn getResult(self: *const ActiveMaskOp) *ir.Value {285            return self.op.getResult(0).?;286        }287    };288289    pub const AllSyncOp = struct {290        op: *ir.Operation,291        pub const operation_spec = op_specs.leaf(.{292            .mnemonic = "vote.sync.all",293            .operands = .{ "mask", "predicate" },294            .results = .{"result"},295        });296        pub const operation_name = operation_spec.name;297298        pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !AllSyncOp {299            try loadSpec(ctx);300            var builder = ir.OperationBuilder.init(ctx);301            const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);302            var state = op_specs.state(@This(), loc);303            state.addOperands(&.{ mask, pred });304            state.addTypes(&.{bool_type});305            const op = try builder.create(state);306            return .{ .op = op };307        }308309        pub fn getResult(self: *const AllSyncOp) *ir.Value {310            return self.op.getResult(0).?;311        }312313        pub fn getMask(self: AllSyncOp) *ir.Value {314            return self.op.operands.items[0].value;315        }316317        pub fn getPredicate(self: AllSyncOp) *ir.Value {318            return self.op.operands.items[1].value;319        }320    };321322    pub const AnySyncOp = struct {323        op: *ir.Operation,324        pub const operation_spec = op_specs.leaf(.{325            .mnemonic = "vote.sync.any",326            .operands = .{ "mask", "predicate" },327            .results = .{"result"},328        });329        pub const operation_name = operation_spec.name;330331        pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !AnySyncOp {332            try loadSpec(ctx);333            var builder = ir.OperationBuilder.init(ctx);334            const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);335            var state = op_specs.state(@This(), loc);336            state.addOperands(&.{ mask, pred });337            state.addTypes(&.{bool_type});338            const op = try builder.create(state);339            return .{ .op = op };340        }341342        pub fn getResult(self: *const AnySyncOp) *ir.Value {343            return self.op.getResult(0).?;344        }345346        pub fn getMask(self: AnySyncOp) *ir.Value {347            return self.op.operands.items[0].value;348        }349350        pub fn getPredicate(self: AnySyncOp) *ir.Value {351            return self.op.operands.items[1].value;352        }353    };354355    pub const BallotSyncOp = struct {356        op: *ir.Operation,357        pub const operation_spec = op_specs.leaf(.{358            .mnemonic = "vote.sync.ballot",359            .operands = .{ "mask", "predicate" },360            .results = .{"result"},361        });362        pub const operation_name = operation_spec.name;363364        pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !BallotSyncOp {365            try loadSpec(ctx);366            var builder = ir.OperationBuilder.init(ctx);367            const i32_type = try arith.ArithDialect.getI32Type(ctx);368            var state = op_specs.state(@This(), loc);369            state.addOperands(&.{ mask, pred });370            state.addTypes(&.{i32_type});371            const op = try builder.create(state);372            return .{ .op = op };373        }374375        pub fn getResult(self: *const BallotSyncOp) *ir.Value {376            return self.op.getResult(0).?;377        }378379        pub fn getMask(self: BallotSyncOp) *ir.Value {380            return self.op.operands.items[0].value;381        }382383        pub fn getPredicate(self: BallotSyncOp) *ir.Value {384            return self.op.operands.items[1].value;385        }386    };387388    pub const ShflSyncOp = struct {389        op: *ir.Operation,390        pub const operation_spec = op_specs.leaf(.{391            .mnemonic = "shfl.sync",392            .operands = .{ "mask", "src", "lane_or_delta" },393            .results = .{"result"},394            .required_attrs = &.{"mode"},395        });396        pub const operation_name = operation_spec.name;397398        pub fn create(399            ctx: *ir.Context,400            loc: ir.Location,401            mode: ShuffleMode,402            mask: *ir.Value,403            src: *ir.Value,404            lane_or_delta: *ir.Value,405        ) !ShflSyncOp {406            try loadSpec(ctx);407            var builder = ir.OperationBuilder.init(ctx);408            var state = op_specs.state(@This(), loc);409            state.addOperands(&.{ mask, src, lane_or_delta });410            state.addTypes(&.{src.type});411            const op = try builder.create(state);412            try setShuffleModeAttr(op, ctx, mode);413            return .{ .op = op };414        }415416        pub fn getResult(self: *const ShflSyncOp) *ir.Value {417            return self.op.getResult(0).?;418        }419420        pub fn getMask(self: ShflSyncOp) *ir.Value {421            return self.op.operands.items[0].value;422        }423424        pub fn getSrc(self: ShflSyncOp) *ir.Value {425            return self.op.operands.items[1].value;426        }427428        pub fn getLaneOrDelta(self: ShflSyncOp) *ir.Value {429            return self.op.operands.items[2].value;430        }431432        pub fn getMode(self: ShflSyncOp) ?ShuffleMode {433            return getShuffleModeAttr(self.op);434        }435    };436437    pub const WarpReduceOp = struct {438        op: *ir.Operation,439        pub const operation_spec = op_specs.leaf(.{440            .mnemonic = "warp.reduce",441            .operands = .{ "mask", "value" },442            .results = .{"result"},443            .required_attrs = &.{"op"},444        });445        pub const operation_name = operation_spec.name;446447        pub fn create(448            ctx: *ir.Context,449            loc: ir.Location,450            op_kind: WarpOpKind,451            mask: *ir.Value,452            value: *ir.Value,453        ) !WarpReduceOp {454            try loadSpec(ctx);455            var builder = ir.OperationBuilder.init(ctx);456            var state = op_specs.state(@This(), loc);457            state.addOperands(&.{ mask, value });458            state.addTypes(&.{value.type});459            const op = try builder.create(state);460            try setWarpOpAttr(op, ctx, op_kind);461            return .{ .op = op };462        }463464        pub fn getResult(self: *const WarpReduceOp) *ir.Value {465            return self.op.getResult(0).?;466        }467468        pub fn getMask(self: WarpReduceOp) *ir.Value {469            return self.op.operands.items[0].value;470        }471472        pub fn getValue(self: WarpReduceOp) *ir.Value {473            return self.op.operands.items[1].value;474        }475476        pub fn getOpKind(self: WarpReduceOp) ?WarpOpKind {477            return getWarpOpAttr(self.op);478        }479    };480481    pub const MmaSyncOp = struct {482        op: *ir.Operation,483        pub const operation_spec = op_specs.leaf(.{484            .mnemonic = "mma.sync",485            .operands = .{ "a0", "a1", "a2", "a3", "b0", "b1", "c0", "c1", "c2", "c3" },486            .results = .{ "d0", "d1", "d2", "d3" },487            .required_attrs = &.{"shape"},488        });489        pub const operation_name = operation_spec.name;490491        pub fn create(492            ctx: *ir.Context,493            loc: ir.Location,494            operands: [10]*ir.Value,495            shape: MmaShape,496        ) !MmaSyncOp {497            try loadSpec(ctx);498            var builder = ir.OperationBuilder.init(ctx);499            var state = op_specs.state(@This(), loc);500            state.addOperands(&operands);501            state.addTypes(&.{ operands[6].type, operands[7].type, operands[8].type, operands[9].type });502            const op = try builder.create(state);503            try setMmaShapeAttr(op, ctx, shape);504            return .{ .op = op };505        }506507        pub fn getOperandValue(self: MmaSyncOp, index: usize) *ir.Value {508            return self.op.operands.items[index].value;509        }510511        pub fn getD(self: *const MmaSyncOp, index: usize) *ir.Value {512            return self.op.getResult(index).?;513        }514515        pub fn getShape(self: MmaSyncOp) ?MmaShape {516            return getMmaShapeAttr(self.op);517        }518    };519520    pub const CpAsyncSharedOp = struct {521        op: *ir.Operation,522        pub const operation_spec = op_specs.leaf(.{523            .mnemonic = "cp.async.shared",524            .operands = .{ "dst", "dst_index", "src", "src_index" },525            .results = 0,526            .required_attrs = &.{"bytes"},527        });528        pub const operation_name = operation_spec.name;529530        pub fn create(531            ctx: *ir.Context,532            loc: ir.Location,533            dst: *ir.Value,534            dst_index: *ir.Value,535            src: *ir.Value,536            src_index: *ir.Value,537            bytes: u32,538        ) !CpAsyncSharedOp {539            try loadSpec(ctx);540            var builder = ir.OperationBuilder.init(ctx);541            var state = op_specs.state(@This(), loc);542            state.addOperands(&.{ dst, dst_index, src, src_index });543            const op = try builder.create(state);544            const bytes_attr = try ctx.getI64Attr(@intCast(bytes));545            try op.setAttr("bytes", bytes_attr);546            return .{ .op = op };547        }548549        pub fn getDst(self: CpAsyncSharedOp) *ir.Value {550            return ir.dialects.operand(operation_spec, self.op, "dst");551        }552553        pub fn getDstIndex(self: CpAsyncSharedOp) *ir.Value {554            return ir.dialects.operand(operation_spec, self.op, "dst_index");555        }556557        pub fn getSrc(self: CpAsyncSharedOp) *ir.Value {558            return ir.dialects.operand(operation_spec, self.op, "src");559        }560561        pub fn getSrcIndex(self: CpAsyncSharedOp) *ir.Value {562            return ir.dialects.operand(operation_spec, self.op, "src_index");563        }564565        pub fn getBytes(self: CpAsyncSharedOp) ?u32 {566            const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "bytes") orelse return null;567            const raw = int_attr.getUnsignedValue();568            if (raw > std.math.maxInt(u32)) return null;569            return @intCast(raw);570        }571    };572573    pub const FenceDeviceOp = struct {574        op: *ir.Operation,575        pub const operation_spec = op_specs.leaf(.{576            .mnemonic = "membar.gl",577            .operands = 0,578            .results = 0,579        });580        pub const operation_name = operation_spec.name;581582        pub fn create(ctx: *ir.Context, loc: ir.Location) !FenceDeviceOp {583            try loadSpec(ctx);584            var builder = ir.OperationBuilder.init(ctx);585            const state = op_specs.state(@This(), loc);586            const op = try builder.create(state);587            return .{ .op = op };588        }589    };590591    pub const CpAsyncCommitOp = struct {592        op: *ir.Operation,593        pub const operation_spec = op_specs.leaf(.{594            .mnemonic = "cp.async.commit",595            .operands = 0,596            .results = 0,597        });598        pub const operation_name = operation_spec.name;599600        pub fn create(ctx: *ir.Context, loc: ir.Location) !CpAsyncCommitOp {601            try loadSpec(ctx);602            var builder = ir.OperationBuilder.init(ctx);603            const state = op_specs.state(@This(), loc);604            const op = try builder.create(state);605            return .{ .op = op };606        }607    };608609    pub const CpAsyncWaitOp = struct {610        op: *ir.Operation,611        pub const operation_spec = op_specs.leaf(.{612            .mnemonic = "cp.async.wait",613            .operands = 0,614            .results = 0,615            .required_attrs = &.{"groups"},616        });617        pub const operation_name = operation_spec.name;618619        pub fn create(ctx: *ir.Context, loc: ir.Location, groups: u32) !CpAsyncWaitOp {620            try loadSpec(ctx);621            var builder = ir.OperationBuilder.init(ctx);622            const state = op_specs.state(@This(), loc);623            const op = try builder.create(state);624            const groups_attr = try ctx.getI64Attr(@intCast(groups));625            try op.setAttr("groups", groups_attr);626            return .{ .op = op };627        }628629        pub fn getGroups(self: CpAsyncWaitOp) ?u32 {630            const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "groups") orelse return null;631            const raw = int_attr.getUnsignedValue();632            if (raw > std.math.maxInt(u32)) return null;633            return @intCast(raw);634        }635    };636637    pub const WarpScanOp = struct {638        op: *ir.Operation,639        pub const operation_spec = op_specs.leaf(.{640            .mnemonic = "warp.scan",641            .operands = .{ "mask", "value" },642            .results = .{"result"},643            .required_attrs = &.{ "inclusive", "op" },644        });645        pub const operation_name = operation_spec.name;646647        pub fn create(648            ctx: *ir.Context,649            loc: ir.Location,650            op_kind: WarpOpKind,651            inclusive: bool,652            mask: *ir.Value,653            value: *ir.Value,654        ) !WarpScanOp {655            try loadSpec(ctx);656            var builder = ir.OperationBuilder.init(ctx);657            var state = op_specs.state(@This(), loc);658            state.addOperands(&.{ mask, value });659            state.addTypes(&.{value.type});660            const op = try builder.create(state);661            try setWarpOpAttr(op, ctx, op_kind);662            try setBoolAttr(op, ctx, "inclusive", inclusive);663            return .{ .op = op };664        }665666        pub fn getResult(self: *const WarpScanOp) *ir.Value {667            return self.op.getResult(0).?;668        }669670        pub fn getMask(self: WarpScanOp) *ir.Value {671            return self.op.operands.items[0].value;672        }673674        pub fn getValue(self: WarpScanOp) *ir.Value {675            return self.op.operands.items[1].value;676        }677678        pub fn getOpKind(self: WarpScanOp) ?WarpOpKind {679            return getWarpOpAttr(self.op);680        }681682        pub fn isInclusive(self: WarpScanOp) bool {683            return getBoolAttrValue(self.op, "inclusive");684        }685    };686687    pub const LoadLocalOp: type = loadOp("local");688    pub const StoreLocalOp: type = storeOp("local");689690    pub const LoadGlobalOp: type = loadOp("global");691692    pub const StoreGlobalOp: type = storeOp("global");693694    pub const AtomicGlobalOp = struct {695        op: *ir.Operation,696        pub const operation_spec = op_specs.leaf(.{697            .mnemonic = "atom.global",698            .operands = 3,699            .results = 1,700            .attrs = &.{"kind"},701        });702        pub const operation_name = operation_spec.name;703704        pub fn create(705            ctx: *ir.Context,706            loc: ir.Location,707            kind: choir.dialects.AtomicRmwKind,708            value: *ir.Value,709            memref: *ir.Value,710            index: *ir.Value,711            result_type: ir.Type,712        ) !AtomicGlobalOp {713            try loadSpec(ctx);714            var builder = ir.OperationBuilder.init(ctx);715            var state = op_specs.state(@This(), loc);716            state.addOperands(&.{ value, memref, index });717            state.addTypes(&.{result_type});718            const op = try builder.create(state);719            const kind_attr = try ctx.getDialectAttr("nvptx.atomic_kind", kind.toString());720            try op.setAttr("kind", kind_attr);721            return .{ .op = op };722        }723724        pub fn getKind(self: AtomicGlobalOp) ?choir.dialects.AtomicRmwKind {725            const dialect_attr = self.op.getAttrAs(ir.Attribute.DialectAttr, "kind") orelse return null;726            return choir.dialects.AtomicRmwKind.fromString(dialect_attr.payload);727        }728729        pub fn getValue(self: AtomicGlobalOp) *ir.Value {730            return self.op.operands.items[0].value;731        }732733        pub fn getMemref(self: AtomicGlobalOp) *ir.Value {734            return self.op.operands.items[1].value;735        }736737        pub fn getIndex(self: AtomicGlobalOp) *ir.Value {738            return self.op.operands.items[2].value;739        }740741        pub fn getResult(self: *const AtomicGlobalOp) *ir.Value {742            return self.op.getResult(0).?;743        }744    };745746    pub const AtomicSharedOp = struct {747        op: *ir.Operation,748        pub const operation_spec = op_specs.leaf(.{749            .mnemonic = "atom.shared",750            .operands = 3,751            .results = 1,752            .attrs = &.{"kind"},753        });754        pub const operation_name = operation_spec.name;755756        pub fn create(757            ctx: *ir.Context,758            loc: ir.Location,759            kind: choir.dialects.AtomicRmwKind,760            value: *ir.Value,761            memref: *ir.Value,762            index: *ir.Value,763            result_type: ir.Type,764        ) !AtomicSharedOp {765            try loadSpec(ctx);766            var builder = ir.OperationBuilder.init(ctx);767            var state = op_specs.state(@This(), loc);768            state.addOperands(&.{ value, memref, index });769            state.addTypes(&.{result_type});770            const op = try builder.create(state);771            const kind_attr = try ctx.getDialectAttr("nvptx.atomic_kind", kind.toString());772            try op.setAttr("kind", kind_attr);773            return .{ .op = op };774        }775776        pub fn getKind(self: AtomicSharedOp) ?choir.dialects.AtomicRmwKind {777            const dialect_attr = self.op.getAttrAs(ir.Attribute.DialectAttr, "kind") orelse return null;778            return choir.dialects.AtomicRmwKind.fromString(dialect_attr.payload);779        }780781        pub fn getValue(self: AtomicSharedOp) *ir.Value {782            return self.op.operands.items[0].value;783        }784785        pub fn getMemref(self: AtomicSharedOp) *ir.Value {786            return self.op.operands.items[1].value;787        }788789        pub fn getIndex(self: AtomicSharedOp) *ir.Value {790            return self.op.operands.items[2].value;791        }792793        pub fn getResult(self: *const AtomicSharedOp) *ir.Value {794            return self.op.getResult(0).?;795        }796    };797798    pub const AtomicCasGlobalOp = struct {799        op: *ir.Operation,800        pub const operation_spec = op_specs.leaf(.{801            .mnemonic = "atom.global.cas",802            .operands = 4,803            .results = 1,804        });805        pub const operation_name = operation_spec.name;806807        pub fn create(808            ctx: *ir.Context,809            loc: ir.Location,810            expected: *ir.Value,811            desired: *ir.Value,812            memref: *ir.Value,813            index: *ir.Value,814            result_type: ir.Type,815        ) !AtomicCasGlobalOp {816            try loadSpec(ctx);817            var builder = ir.OperationBuilder.init(ctx);818            var state = op_specs.state(@This(), loc);819            state.addOperands(&.{ expected, desired, memref, index });820            state.addTypes(&.{result_type});821            const op = try builder.create(state);822            return .{ .op = op };823        }824825        pub fn getExpected(self: AtomicCasGlobalOp) *ir.Value {826            return self.op.operands.items[0].value;827        }828829        pub fn getDesired(self: AtomicCasGlobalOp) *ir.Value {830            return self.op.operands.items[1].value;831        }832833        pub fn getMemref(self: AtomicCasGlobalOp) *ir.Value {834            return self.op.operands.items[2].value;835        }836837        pub fn getIndex(self: AtomicCasGlobalOp) *ir.Value {838            return self.op.operands.items[3].value;839        }840841        pub fn getResult(self: *const AtomicCasGlobalOp) *ir.Value {842            return self.op.getResult(0).?;843        }844    };845846    pub const AtomicCasSharedOp = struct {847        op: *ir.Operation,848        pub const operation_spec = op_specs.leaf(.{849            .mnemonic = "atom.shared.cas",850            .operands = 4,851            .results = 1,852        });853        pub const operation_name = operation_spec.name;854855        pub fn create(856            ctx: *ir.Context,857            loc: ir.Location,858            expected: *ir.Value,859            desired: *ir.Value,860            memref: *ir.Value,861            index: *ir.Value,862            result_type: ir.Type,863        ) !AtomicCasSharedOp {864            try loadSpec(ctx);865            var builder = ir.OperationBuilder.init(ctx);866            var state = op_specs.state(@This(), loc);867            state.addOperands(&.{ expected, desired, memref, index });868            state.addTypes(&.{result_type});869            const op = try builder.create(state);870            return .{ .op = op };871        }872873        pub fn getExpected(self: AtomicCasSharedOp) *ir.Value {874            return self.op.operands.items[0].value;875        }876877        pub fn getDesired(self: AtomicCasSharedOp) *ir.Value {878            return self.op.operands.items[1].value;879        }880881        pub fn getMemref(self: AtomicCasSharedOp) *ir.Value {882            return self.op.operands.items[2].value;883        }884885        pub fn getIndex(self: AtomicCasSharedOp) *ir.Value {886            return self.op.operands.items[3].value;887        }888889        pub fn getResult(self: *const AtomicCasSharedOp) *ir.Value {890            return self.op.getResult(0).?;891        }892    };893894    pub const LoadSharedOp: type = loadOp("shared");895896    pub const StoreSharedOp: type = storeOp("shared");897898    pub const SinApproxF32Op = struct {899        op: *ir.Operation,900        pub const operation_spec = op_specs.leaf(.{901            .mnemonic = "sin.approx.f32",902            .operands = 1,903            .results = 1,904        });905        pub const operation_name = operation_spec.name;906907        pub fn create(ctx: *ir.Context, loc: ir.Location, value: *ir.Value) !SinApproxF32Op {908            try loadSpec(ctx);909            var builder = ir.OperationBuilder.init(ctx);910            var state = op_specs.state(@This(), loc);911            state.addOperands(&.{value});912            state.addTypes(&.{value.type});913            const op = try builder.create(state);914            return .{ .op = op };915        }916917        pub fn getValue(self: SinApproxF32Op) *ir.Value {918            return self.op.operands.items[0].value;919        }920921        pub fn getResult(self: *const SinApproxF32Op) *ir.Value {922            return self.op.getResult(0).?;923        }924    };925926    pub const CosApproxF32Op = struct {927        op: *ir.Operation,928        pub const operation_spec = op_specs.leaf(.{929            .mnemonic = "cos.approx.f32",930            .operands = 1,931            .results = 1,932        });933        pub const operation_name = operation_spec.name;934935        pub fn create(ctx: *ir.Context, loc: ir.Location, value: *ir.Value) !CosApproxF32Op {936            try loadSpec(ctx);937            var builder = ir.OperationBuilder.init(ctx);938            var state = op_specs.state(@This(), loc);939            state.addOperands(&.{value});940            state.addTypes(&.{value.type});941            const op = try builder.create(state);942            return .{ .op = op };943        }944945        pub fn getValue(self: CosApproxF32Op) *ir.Value {946            return self.op.operands.items[0].value;947        }948949        pub fn getResult(self: *const CosApproxF32Op) *ir.Value {950            return self.op.getResult(0).?;951        }952    };953954    pub const FmaRnF32Op = struct {955        op: *ir.Operation,956        pub const operation_spec = op_specs.leaf(.{957            .mnemonic = "fma.rn.f32",958            .operands = 3,959            .results = 1,960        });961        pub const operation_name = operation_spec.name;962963        pub fn create(ctx: *ir.Context, loc: ir.Location, a: *ir.Value, b: *ir.Value, c: *ir.Value) !FmaRnF32Op {964            try loadSpec(ctx);965            var builder = ir.OperationBuilder.init(ctx);966            var state = op_specs.state(@This(), loc);967            state.addOperands(&.{ a, b, c });968            state.addTypes(&.{a.type});969            const op = try builder.create(state);970            return .{ .op = op };971        }972973        pub fn getA(self: FmaRnF32Op) *ir.Value {974            return self.op.operands.items[0].value;975        }976977        pub fn getB(self: FmaRnF32Op) *ir.Value {978            return self.op.operands.items[1].value;979        }980981        pub fn getC(self: FmaRnF32Op) *ir.Value {982            return self.op.operands.items[2].value;983        }984985        pub fn getResult(self: *const FmaRnF32Op) *ir.Value {986            return self.op.getResult(0).?;987        }988    };989990    fn loadOp(comptime space: []const u8) type {991        return struct {992            op: *ir.Operation,993            pub const operation_spec = op_specs.leaf(.{994                .mnemonic = "ld." ++ space,995                .operands = 2,996                .results = 1,997            });998            pub const operation_name = operation_spec.name;9991000            pub fn create(1001                ctx: *ir.Context,1002                loc: ir.Location,1003                memref: *ir.Value,1004                index: *ir.Value,1005                result_type: ir.Type,1006            ) !@This() {1007                try loadSpec(ctx);1008                var builder = ir.OperationBuilder.init(ctx);1009                var state = op_specs.state(@This(), loc);1010                state.addOperands(&.{ memref, index });1011                state.addTypes(&.{result_type});1012                const op = try builder.create(state);1013                return .{ .op = op };1014            }10151016            pub fn getResult(self: *const @This()) *ir.Value {1017                return self.op.getResult(0).?;1018            }10191020            pub fn getMemref(self: @This()) *ir.Value {1021                return self.op.operands.items[0].value;1022            }10231024            pub fn getIndex(self: @This()) *ir.Value {1025                return self.op.operands.items[1].value;1026            }1027        };1028    }10291030    fn storeOp(comptime space: []const u8) type {1031        return struct {1032            op: *ir.Operation,1033            pub const operation_spec = op_specs.leaf(.{1034                .mnemonic = "st." ++ space,1035                .operands = 3,1036                .results = 0,1037            });1038            pub const operation_name = operation_spec.name;10391040            pub fn create(1041                ctx: *ir.Context,1042                loc: ir.Location,1043                value: *ir.Value,1044                memref: *ir.Value,1045                index: *ir.Value,1046            ) !@This() {1047                try loadSpec(ctx);1048                var builder = ir.OperationBuilder.init(ctx);1049                var state = op_specs.state(@This(), loc);1050                state.addOperands(&.{ value, memref, index });1051                const op = try builder.create(state);1052                return .{ .op = op };1053            }10541055            pub fn getValue(self: @This()) *ir.Value {1056                return self.op.operands.items[0].value;1057            }10581059            pub fn getMemref(self: @This()) *ir.Value {1060                return self.op.operands.items[1].value;1061            }10621063            pub fn getIndex(self: @This()) *ir.Value {1064                return self.op.operands.items[2].value;1065            }1066        };1067    }10681069    fn loadSpec(ctx: *ir.Context) !void {1070        try ir.dialects.loadDialectSpec(ctx, spec);1071    }10721073    fn setDimensionAttr(op: *ir.Operation, ctx: *ir.Context, dim: Dimension) !void {1074        const attr = try ctx.getDialectAttr("nvptx.dim", dim.toString());1075        try op.setAttr("dim", attr);1076    }10771078    fn getDimensionAttr(op: *const ir.Operation) ?Dimension {1079        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "dim") orelse return null;1080        return Dimension.fromString(dialect_attr.payload);1081    }10821083    fn setShuffleModeAttr(op: *ir.Operation, ctx: *ir.Context, mode: ShuffleMode) !void {1084        const attr = try ctx.getDialectAttr("nvptx.shuffle_mode", mode.toString());1085        try op.setAttr("mode", attr);1086    }10871088    fn getShuffleModeAttr(op: *const ir.Operation) ?ShuffleMode {1089        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "mode") orelse return null;1090        return ShuffleMode.fromString(dialect_attr.payload);1091    }10921093    fn setWarpOpAttr(op: *ir.Operation, ctx: *ir.Context, op_kind: WarpOpKind) !void {1094        const attr = try ctx.getDialectAttr("nvptx.warp_op", op_kind.toString());1095        try op.setAttr("op", attr);1096    }10971098    fn getWarpOpAttr(op: *const ir.Operation) ?WarpOpKind {1099        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "op") orelse return null;1100        return WarpOpKind.fromString(dialect_attr.payload);1101    }11021103    fn setMmaShapeAttr(op: *ir.Operation, ctx: *ir.Context, shape: MmaShape) !void {1104        var buf: [32]u8 = undefined;1105        const shape_str = try shape.toString(buf[0..]);1106        const shape_attr = try ctx.getDialectAttr("nvptx.mma_shape", shape_str);1107        try op.setAttr("shape", shape_attr);1108    }11091110    fn getMmaShapeAttr(op: *const ir.Operation) ?MmaShape {1111        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "shape") orelse return null;1112        return MmaShape.parse(dialect_attr.payload);1113    }11141115    fn setBoolAttr(op: *ir.Operation, ctx: *ir.Context, attr_name: []const u8, value: bool) !void {1116        const bool_attr = try ctx.getBoolAttr(value);1117        try op.setAttr(attr_name, bool_attr);1118    }11191120    fn getBoolAttrValue(op: *const ir.Operation, attr_name: []const u8) bool {1121        const bool_attr = op.getAttrAs(ir.Attribute.BoolAttr, attr_name) orelse return false;1122        return bool_attr.getValue();1123    }1124};