tiny.choir.dialects.gpu.dialect.GpuDialect

1pub const GpuDialect = struct {2    pub const name = "gpu";3    const op_specs = ir.dialects.opSpec.dialect(@This());4    pub const spec = ir.dialects.dialectSpec(@This(), .{5        .types = ir.dialects.typeNames(type_specs),6    });7    const symbol_table_trait = ir.dialects.trait(ir.traits.SymbolTable);89    const func_symbol_vtable = interfaces.SymbolOpInterface.VTable{10        .getSymbolName = getFuncSymbolName,11        .setSymbolName = setFuncSymbolName,12        .isDeclaration = isFuncDeclaration,13    };1415    const type_specs = struct {16        pub const tma_desc = type_names.tma_desc;17        pub const mbarrier = type_names.mbarrier;18        pub const sampled_texture = type_names.sampled_texture;19    };2021    pub const StageInputOp = stage.StageInputOp;22    pub const StageOutputOp = stage.StageOutputOp;23    pub const PositionOp = stage.PositionOp;24    pub const FragCoordOp = stage.FragCoordOp;25    pub const VertexIndexOp = stage.VertexIndexOp;26    pub const InstanceIndexOp = stage.InstanceIndexOp;27    pub const FrontFacingOp = stage.FrontFacingOp;28    pub const SampledTextureOp = stage.SampledTextureOp;29    pub const SampleOp = stage.SampleOp;30    pub const SampleLodOp = stage.SampleLodOp;31    pub const DpdxOp = stage.DpdxOp;32    pub const DpdyOp = stage.DpdyOp;33    pub const FwidthOp = stage.FwidthOp;34    pub const PushConstantOp = stage.PushConstantOp;35    pub const UniformOp = stage.UniformOp;3637    pub const ModuleOp = struct {38        op: *ir.Operation,3940        pub const operation_spec = op_specs.define(.{41            .mnemonic = "module",42            .operands = 0,43            .results = 0,44            .regions = .{"body"},45            .successors = 0,46            .dynamic_traits = .{symbol_table_trait},47        });48        pub const operation_name = operation_spec.name;4950        pub fn create(ctx: *ir.Context, loc: ir.Location) !ModuleOp {51            try loadSpec(ctx);52            var builder = ir.OperationBuilder.init(ctx);53            var state = ir.Operation.State.init(operation_name, loc);54            var body = ir.context.initRegion(ctx);55            defer body.deinit();56            var body_builder = ir.OperationBuilder.init(ctx);57            _ = try body_builder.createBlock(&body, &.{}, &.{});58            var regions = [_]*ir.Region{&body};59            state.addRegionBodies(&regions);6061            const op = try builder.create(state);62            errdefer op.erase();6364            return .{ .op = op };65        }6667        pub fn getBody(self: ModuleOp) *ir.Region {68            return self.op.getRegion(0).?;69        }7071        pub fn getBodyBlock(self: ModuleOp) *ir.Block {72            return self.getBody().getEntryBlock().?;73        }74    };7576    pub const FuncOp = struct {77        op: *ir.Operation,7879        pub const operation_spec = op_specs.define(.{80            .mnemonic = "func",81            .operands = 0,82            .regions = .{"body"},83            .successors = 0,84            .attrs = &.{ "kernel", ir.SymbolTable.symbol_attr_names.sym_visibility },85            .required_attrs = &.{"sym_name"},86            .interfaces = &.{87                interfaces.SymbolOpInterface.entry(&func_symbol_vtable),88                effects.EffectOpInterface.entryFor(.{ .facts = &.{.{ .region = .{89                    .index = 0,90                    .execution = .latent,91                    .may_diverge = false,92                    .captures = false,93                } }} }),94            },95        });96        pub const operation_name = operation_spec.name;9798        pub fn create(99            ctx: *ir.Context,100            loc: ir.Location,101            func_name: []const u8,102            input_types: []const ir.Type,103            result_types: []const ir.Type,104        ) !FuncOp {105            try loadSpec(ctx);106            var builder = ir.OperationBuilder.init(ctx);107            var state = ir.Operation.State.init(operation_name, loc);108            state.addTypes(result_types);109            var body = ir.context.initRegion(ctx);110            defer body.deinit();111            var body_builder = ir.OperationBuilder.init(ctx);112            _ = try body_builder.createBlockWithLoc(&body, input_types, loc);113            var regions = [_]*ir.Region{&body};114            state.addRegionBodies(&regions);115116            const op = try builder.create(state);117            errdefer op.erase();118119            const name_attr = try func.FuncDialect.getSymNameAttr(ctx, func_name);120            try op.setAttr("sym_name", name_attr);121122            return .{ .op = op };123        }124125        pub fn createKernel(126            ctx: *ir.Context,127            loc: ir.Location,128            kernel_name: []const u8,129            input_types: []const ir.Type,130        ) !FuncOp {131            var func_op = try create(ctx, loc, kernel_name, input_types, &.{});132            errdefer func_op.op.erase();133            const kernel_attr = try func.FuncDialect.getKernelAttr(ctx);134            try func_op.op.setAttr("kernel", kernel_attr);135            return func_op;136        }137138        pub fn getName(self: FuncOp) ?[]const u8 {139            return ir.SymbolTable.getSymbolName(self.op);140        }141142        pub fn isKernel(self: FuncOp) bool {143            return self.op.getAttr("kernel") != null;144        }145146        pub fn getBody(self: FuncOp) *ir.Region {147            return self.op.getRegion(0).?;148        }149150        pub fn getEntryBlock(self: FuncOp) *ir.Block {151            return self.getBody().getEntryBlock().?;152        }153154        pub fn getArguments(self: FuncOp) []*ir.Value {155            return self.getEntryBlock().arguments.items;156        }157158        pub fn getNumArguments(self: FuncOp) usize {159            return self.getEntryBlock().arguments.items.len;160        }161162        pub fn getArgument(self: FuncOp, index: usize) *ir.Value {163            return self.getEntryBlock().arguments.items[index];164        }165166        pub fn getResultTypes(self: FuncOp) []const ir.Type {167            return self.op.getResultTypes();168        }169170        pub fn getNumResults(self: FuncOp) usize {171            return self.op.results.items.len;172        }173    };174175    pub const YieldOp = struct {176        op: *ir.Operation,177178        pub const operation_spec = op_specs.terminator(.{ .mnemonic = "yield" });179        pub const operation_name = operation_spec.name;180181        pub fn create(182            ctx: *ir.Context,183            loc: ir.Location,184            operands: []const *ir.Value,185        ) !YieldOp {186            try loadSpec(ctx);187            var builder = ir.OperationBuilder.init(ctx);188            var state = ir.Operation.State.init(operation_name, loc);189            state.addOperands(operands);190            const op = try builder.create(state);191            errdefer op.erase();192            return .{ .op = op };193        }194195        pub fn getOperands(self: YieldOp) []const *ir.Value {196            return self.op.getOperandValues();197        }198    };199200    pub const LaunchOp = struct {201        op: *ir.Operation,202203        const dim_attr_keys = struct {204            pub const grid_x = "grid_x";205            pub const grid_y = "grid_y";206            pub const grid_z = "grid_z";207            pub const block_x = "block_x";208            pub const block_y = "block_y";209            pub const block_z = "block_z";210        };211212        pub const operation_spec = op_specs.leaf(.{213            .mnemonic = "launch",214            .interfaces = &.{gpuEffects(.launch, &.{}, &.{})},215            .operands = ir.dialects.shape.any(),216            .results = 0,217            .required_attrs = &.{218                dim_attr_keys.block_x,219                dim_attr_keys.block_y,220                dim_attr_keys.block_z,221                dim_attr_keys.grid_x,222                dim_attr_keys.grid_y,223                dim_attr_keys.grid_z,224                "kernel",225                "num_kernel_args",226            },227        });228        pub const operation_name = operation_spec.name;229230        fn setDimAttr(op: *ir.Operation, ctx: *ir.Context, key: []const u8, value: u32) !void {231            const attr = try ctx.getI64Attr(@intCast(value));232            try op.setAttr(key, attr);233        }234235        fn getDimAttr(op: *const ir.Operation, key: []const u8) ?u32 {236            const int_attr = op.getAttrAs(ir.Attribute.IntegerAttr, key) orelse return null;237            const raw = int_attr.getUnsignedValue();238            if (raw > std.math.maxInt(u32)) return null;239            return @intCast(raw);240        }241242        pub fn create(243            ctx: *ir.Context,244            loc: ir.Location,245            kernel_name: []const u8,246            kernel_args: []const *ir.Value,247            grid_dim: [3]u32,248            block_dim: [3]u32,249        ) !LaunchOp {250            try loadSpec(ctx);251            var builder = ir.OperationBuilder.init(ctx);252            var state = ir.Operation.State.init(operation_name, loc);253254            var all_operands: std.ArrayList(*ir.Value) = .empty;255            const allocator = ir.context.transientAllocator(ctx);256            defer all_operands.deinit(allocator);257            for (kernel_args) |arg| {258                try all_operands.append(allocator, arg);259            }260            state.addOperands(all_operands.items);261262            const op = try builder.create(state);263            errdefer op.erase();264265            const kernel_attr = try func.FuncDialect.getSymNameAttr(ctx, kernel_name);266            try op.setAttr("kernel", kernel_attr);267268            var buf: [16]u8 = undefined;269            const num_args_str = try ir.format.intPayload(buf[0..], kernel_args.len);270            const num_args_attr = try ctx.getDialectAttr("gpu.num_kernel_args", num_args_str);271            try op.setAttr("num_kernel_args", num_args_attr);272273            try setDimAttr(op, ctx, dim_attr_keys.grid_x, grid_dim[0]);274            try setDimAttr(op, ctx, dim_attr_keys.grid_y, grid_dim[1]);275            try setDimAttr(op, ctx, dim_attr_keys.grid_z, grid_dim[2]);276            try setDimAttr(op, ctx, dim_attr_keys.block_x, block_dim[0]);277            try setDimAttr(op, ctx, dim_attr_keys.block_y, block_dim[1]);278            try setDimAttr(op, ctx, dim_attr_keys.block_z, block_dim[2]);279280            return .{ .op = op };281        }282283        pub fn getKernelName(self: LaunchOp) ?[]const u8 {284            if (self.op.getAttr("kernel")) |attr| {285                return func.FuncDialect.getSymNameValue(attr);286            }287            return null;288        }289290        pub fn getNumKernelArgs(self: LaunchOp) usize {291            const dialect_attr = self.op.getAttrAs(ir.Attribute.DialectAttr, "num_kernel_args") orelse return 0;292            return std.fmt.parseInt(usize, dialect_attr.payload, 10) catch 0;293        }294295        pub fn getKernelArgs(self: LaunchOp) []const *ir.Value {296            const num_args = self.getNumKernelArgs();297            return self.op.getOperandValues()[0..num_args];298        }299300        pub fn getGridDim(self: LaunchOp) ?[3]u32 {301            const gx = getDimAttr(self.op, dim_attr_keys.grid_x) orelse return null;302            const gy = getDimAttr(self.op, dim_attr_keys.grid_y) orelse return null;303            const gz = getDimAttr(self.op, dim_attr_keys.grid_z) orelse return null;304            return .{ gx, gy, gz };305        }306307        pub fn getBlockDim(self: LaunchOp) ?[3]u32 {308            const bx = getDimAttr(self.op, dim_attr_keys.block_x) orelse return null;309            const by = getDimAttr(self.op, dim_attr_keys.block_y) orelse return null;310            const bz = getDimAttr(self.op, dim_attr_keys.block_z) orelse return null;311            return .{ bx, by, bz };312        }313    };314315    pub const ThreadIdxOp = struct {316        op: *ir.Operation,317318        pub const operation_spec = dimIndexSpec("thread_idx");319        pub const operation_name = operation_spec.name;320321        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !ThreadIdxOp {322            const arith = choir.dialects.arith;323            try loadSpec(ctx);324            var builder = ir.OperationBuilder.init(ctx);325            const index_type = try arith.ArithDialect.getIndexType(ctx);326            var state = ir.Operation.State.init(operation_name, loc);327            state.addTypes(&.{index_type});328329            const op = try builder.create(state);330            errdefer op.erase();331            try setDimensionAttr(op, ctx, dim);332333            return .{ .op = op };334        }335336        pub fn getResult(self: *const ThreadIdxOp) *ir.Value {337            return self.op.getResult(0).?;338        }339340        pub fn getDimension(self: ThreadIdxOp) ?Dimension {341            return getDimensionAttr(self.op);342        }343    };344345    pub const BlockIdxOp = struct {346        op: *ir.Operation,347348        pub const operation_spec = dimIndexSpec("block_idx");349        pub const operation_name = operation_spec.name;350351        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !BlockIdxOp {352            const arith = choir.dialects.arith;353            try loadSpec(ctx);354            var builder = ir.OperationBuilder.init(ctx);355            const index_type = try arith.ArithDialect.getIndexType(ctx);356            var state = ir.Operation.State.init(operation_name, loc);357            state.addTypes(&.{index_type});358359            const op = try builder.create(state);360            errdefer op.erase();361            try setDimensionAttr(op, ctx, dim);362363            return .{ .op = op };364        }365366        pub fn getResult(self: *const BlockIdxOp) *ir.Value {367            return self.op.getResult(0).?;368        }369370        pub fn getDimension(self: BlockIdxOp) ?Dimension {371            return getDimensionAttr(self.op);372        }373    };374375    pub const BlockDimOp = struct {376        op: *ir.Operation,377378        pub const operation_spec = dimIndexSpec("block_dim");379        pub const operation_name = operation_spec.name;380381        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !BlockDimOp {382            const arith = choir.dialects.arith;383            try loadSpec(ctx);384            var builder = ir.OperationBuilder.init(ctx);385            const index_type = try arith.ArithDialect.getIndexType(ctx);386            var state = ir.Operation.State.init(operation_name, loc);387            state.addTypes(&.{index_type});388389            const op = try builder.create(state);390            errdefer op.erase();391            try setDimensionAttr(op, ctx, dim);392393            return .{ .op = op };394        }395396        pub fn getResult(self: *const BlockDimOp) *ir.Value {397            return self.op.getResult(0).?;398        }399400        pub fn getDimension(self: BlockDimOp) ?Dimension {401            return getDimensionAttr(self.op);402        }403    };404405    pub const GridDimOp = struct {406        op: *ir.Operation,407408        pub const operation_spec = dimIndexSpec("grid_dim");409        pub const operation_name = operation_spec.name;410411        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !GridDimOp {412            const arith = choir.dialects.arith;413            try loadSpec(ctx);414            var builder = ir.OperationBuilder.init(ctx);415            const index_type = try arith.ArithDialect.getIndexType(ctx);416            var state = ir.Operation.State.init(operation_name, loc);417            state.addTypes(&.{index_type});418419            const op = try builder.create(state);420            errdefer op.erase();421            try setDimensionAttr(op, ctx, dim);422423            return .{ .op = op };424        }425426        pub fn getResult(self: *const GridDimOp) *ir.Value {427            return self.op.getResult(0).?;428        }429430        pub fn getDimension(self: GridDimOp) ?Dimension {431            return getDimensionAttr(self.op);432        }433    };434435    pub const GlobalIdxOp = struct {436        op: *ir.Operation,437438        pub const operation_spec = dimIndexSpec("global_idx");439        pub const operation_name = operation_spec.name;440441        pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !GlobalIdxOp {442            const arith = choir.dialects.arith;443            try loadSpec(ctx);444            var builder = ir.OperationBuilder.init(ctx);445            const index_type = try arith.ArithDialect.getIndexType(ctx);446            var state = ir.Operation.State.init(operation_name, loc);447            state.addTypes(&.{index_type});448449            const op = try builder.create(state);450            errdefer op.erase();451            try setDimensionAttr(op, ctx, dim);452453            return .{ .op = op };454        }455456        pub fn getResult(self: *const GlobalIdxOp) *ir.Value {457            return self.op.getResult(0).?;458        }459460        pub fn getDimension(self: GlobalIdxOp) ?Dimension {461            return getDimensionAttr(self.op);462        }463    };464465    pub const LaneIdOp = struct {466        op: *ir.Operation,467468        pub const operation_spec = indexSpec("lane_id");469        pub const operation_name = operation_spec.name;470471        pub fn create(ctx: *ir.Context, loc: ir.Location) !LaneIdOp {472            const arith = choir.dialects.arith;473            try loadSpec(ctx);474            var builder = ir.OperationBuilder.init(ctx);475            const index_type = try arith.ArithDialect.getIndexType(ctx);476            var state = ir.Operation.State.init(operation_name, loc);477            state.addTypes(&.{index_type});478479            const op = try builder.create(state);480            errdefer op.erase();481            return .{ .op = op };482        }483484        pub fn getResult(self: *const LaneIdOp) *ir.Value {485            return self.op.getResult(0).?;486        }487    };488489    pub const WarpIdOp = struct {490        op: *ir.Operation,491492        pub const operation_spec = indexSpec("warp_id");493        pub const operation_name = operation_spec.name;494495        pub fn create(ctx: *ir.Context, loc: ir.Location) !WarpIdOp {496            const arith = choir.dialects.arith;497            try loadSpec(ctx);498            var builder = ir.OperationBuilder.init(ctx);499            const index_type = try arith.ArithDialect.getIndexType(ctx);500            var state = ir.Operation.State.init(operation_name, loc);501            state.addTypes(&.{index_type});502503            const op = try builder.create(state);504            errdefer op.erase();505            return .{ .op = op };506        }507508        pub fn getResult(self: *const WarpIdOp) *ir.Value {509            return self.op.getResult(0).?;510        }511    };512513    pub const BarrierOp = struct {514        op: *ir.Operation,515516        pub const operation_spec = op_specs.leaf(.{517            .mnemonic = "barrier",518            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},519            .operands = 0,520            .results = 0,521            .required_attrs = &.{"scope"},522        });523        pub const operation_name = operation_spec.name;524525        pub fn create(ctx: *ir.Context, loc: ir.Location, scope: Scope) !BarrierOp {526            try loadSpec(ctx);527            var builder = ir.OperationBuilder.init(ctx);528            const state = ir.Operation.State.init(operation_name, loc);529530            const op = try builder.create(state);531            errdefer op.erase();532            try setScopeAttr(op, ctx, scope);533534            return .{ .op = op };535        }536537        pub fn getScope(self: BarrierOp) ?Scope {538            return getScopeAttr(self.op);539        }540    };541542    pub const FenceOp = struct {543        op: *ir.Operation,544545        pub const operation_spec = op_specs.leaf(.{546            .mnemonic = "fence",547            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},548            .operands = 0,549            .results = 0,550            .required_attrs = &.{ "ordering", "scope" },551        });552        pub const operation_name = operation_spec.name;553554        pub fn create(ctx: *ir.Context, loc: ir.Location, scope: Scope, ordering: MemoryOrder) !FenceOp {555            try loadSpec(ctx);556            var builder = ir.OperationBuilder.init(ctx);557            const state = ir.Operation.State.init(operation_name, loc);558559            const op = try builder.create(state);560            errdefer op.erase();561            try setScopeAttr(op, ctx, scope);562            try setOrderingAttr(op, ctx, ordering);563564            return .{ .op = op };565        }566567        pub fn getScope(self: FenceOp) ?Scope {568            return getScopeAttr(self.op);569        }570571        pub fn getOrdering(self: FenceOp) ?MemoryOrder {572            return getOrderingAttr(self.op);573        }574    };575576    pub const MemcpyAsyncOp = struct {577        op: *ir.Operation,578579        pub const operation_spec = op_specs.leaf(.{580            .mnemonic = "memcpy_async",581            .interfaces = &.{gpuEffects(.launch, &.{0}, &.{1})},582            .operands = ir.dialects.shape.between(3, 5),583            .operand_names = .{ "src", "dst", "num_bytes", "stream", "event" },584            .results = 0,585            .operand_segments = ir.dialects.segments.operands(.{586                1,587                1,588                1,589                ir.dialects.shape.atMost(1),590                ir.dialects.shape.atMost(1),591            }),592        });593        pub const operation_name = operation_spec.name;594595        pub fn create(596            ctx: *ir.Context,597            loc: ir.Location,598            src: *ir.Value,599            dst: *ir.Value,600            num_bytes: *ir.Value,601            stream: ?*ir.Value,602            event: ?*ir.Value,603        ) !MemcpyAsyncOp {604            try loadSpec(ctx);605            var builder = ir.OperationBuilder.init(ctx);606            var state = ir.Operation.State.init(operation_name, loc);607608            var operands: [5]*ir.Value = undefined;609            var operand_count: usize = 0;610            operands[operand_count] = src;611            operand_count += 1;612            operands[operand_count] = dst;613            operand_count += 1;614            operands[operand_count] = num_bytes;615            operand_count += 1;616617            if (stream) |stream_val| {618                operands[operand_count] = stream_val;619                operand_count += 1;620            }621622            if (event) |event_val| {623                operands[operand_count] = event_val;624                operand_count += 1;625            }626627            state.addOperands(operands[0..operand_count]);628            const op = try builder.create(state);629            errdefer op.erase();630            const segment_sizes = [_]usize{631                1,632                1,633                1,634                if (stream != null) 1 else 0,635                if (event != null) 1 else 0,636            };637            try ir.dialects.setOperandSegmentSizes(operation_spec, op, &segment_sizes);638639            return .{ .op = op };640        }641642        pub fn getSrc(self: MemcpyAsyncOp) *ir.Value {643            return ir.dialects.operand(operation_spec, self.op, "src");644        }645646        pub fn getDst(self: MemcpyAsyncOp) *ir.Value {647            return ir.dialects.operand(operation_spec, self.op, "dst");648        }649650        pub fn getNumBytes(self: MemcpyAsyncOp) *ir.Value {651            return ir.dialects.operand(operation_spec, self.op, "num_bytes");652        }653654        pub fn getStream(self: MemcpyAsyncOp) ?*ir.Value {655            return ir.dialects.operandSegmentValue(operation_spec, self.op, "stream");656        }657658        pub fn getEvent(self: MemcpyAsyncOp) ?*ir.Value {659            return ir.dialects.operandSegmentValue(operation_spec, self.op, "event");660        }661    };662663    pub const TmaCreateDescriptorOp = struct {664        op: *ir.Operation,665666        pub const operation_spec = op_specs.leaf(.{667            .mnemonic = "tma.create_descriptor",668            .operands = .{ "tensor", "box_shape" },669            .results = .{"descriptor"},670        });671        pub const operation_name = operation_spec.name;672673        pub fn create(674            ctx: *ir.Context,675            loc: ir.Location,676            tensor: *ir.Value,677            box_shape: *ir.Value,678        ) !TmaCreateDescriptorOp {679            try loadSpec(ctx);680            var builder = ir.OperationBuilder.init(ctx);681            var state = ir.Operation.State.init(operation_name, loc);682            state.addOperands(&.{ tensor, box_shape });683684            const desc_type = try getTmaDescriptorType(ctx);685            state.addTypes(&.{desc_type});686687            const op = try builder.create(state);688            errdefer op.erase();689            return .{ .op = op };690        }691692        pub fn getResult(self: *const TmaCreateDescriptorOp) *ir.Value {693            return self.op.getResult(0).?;694        }695696        pub fn getTensor(self: TmaCreateDescriptorOp) *ir.Value {697            return self.op.operands.items[0].value;698        }699700        pub fn getBoxShape(self: TmaCreateDescriptorOp) *ir.Value {701            return self.op.operands.items[1].value;702        }703    };704705    pub const TmaLoadOp = struct {706        op: *ir.Operation,707708        pub const operation_spec = noResultSpec("tma.load", .{ "descriptor", "shared_mem", "barrier", "coords" });709        pub const operation_name = operation_spec.name;710711        pub fn create(712            ctx: *ir.Context,713            loc: ir.Location,714            desc: *ir.Value,715            shmem: *ir.Value,716            mbarrier: *ir.Value,717            coords: *ir.Value,718        ) !TmaLoadOp {719            try loadSpec(ctx);720            var builder = ir.OperationBuilder.init(ctx);721            var state = ir.Operation.State.init(operation_name, loc);722            state.addOperands(&.{ desc, shmem, mbarrier, coords });723724            const op = try builder.create(state);725            errdefer op.erase();726            return .{ .op = op };727        }728729        pub fn getDescriptor(self: TmaLoadOp) *ir.Value {730            return self.op.operands.items[0].value;731        }732733        pub fn getSharedMem(self: TmaLoadOp) *ir.Value {734            return self.op.operands.items[1].value;735        }736737        pub fn getBarrier(self: TmaLoadOp) *ir.Value {738            return self.op.operands.items[2].value;739        }740741        pub fn getCoords(self: TmaLoadOp) *ir.Value {742            return self.op.operands.items[3].value;743        }744    };745746    pub const TmaCommitGroupOp = struct {747        op: *ir.Operation,748749        pub const operation_spec = noResultSpec("tma.commit_group", 0);750        pub const operation_name = operation_spec.name;751752        pub fn create(ctx: *ir.Context, loc: ir.Location) !TmaCommitGroupOp {753            try loadSpec(ctx);754            var builder = ir.OperationBuilder.init(ctx);755            const state = ir.Operation.State.init(operation_name, loc);756757            const op = try builder.create(state);758            errdefer op.erase();759            return .{ .op = op };760        }761    };762763    pub const TmaWaitGroupOp = struct {764        op: *ir.Operation,765766        pub const operation_spec = op_specs.leaf(.{767            .mnemonic = "tma.wait_group",768            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},769            .operands = 0,770            .results = 0,771            .required_attrs = &.{"count"},772        });773        pub const operation_name = operation_spec.name;774775        pub fn create(ctx: *ir.Context, loc: ir.Location, count: i64) !TmaWaitGroupOp {776            try loadSpec(ctx);777            var builder = ir.OperationBuilder.init(ctx);778            const state = ir.Operation.State.init(operation_name, loc);779780            const op = try builder.create(state);781            errdefer op.erase();782            try setI64Attr(op, ctx, "count", count);783            return .{ .op = op };784        }785786        pub fn getCount(self: TmaWaitGroupOp) ?i64 {787            return getI64AttrValue(self.op, "count");788        }789    };790791    pub const ShflSyncOp = struct {792        op: *ir.Operation,793794        pub const operation_spec = op_specs.leaf(.{795            .mnemonic = "shfl_sync",796            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},797            .operands = .{ "mask", "src", "lane_or_delta" },798            .results = .{"result"},799            .required_attrs = &.{"mode"},800        });801        pub const operation_name = operation_spec.name;802803        pub fn create(804            ctx: *ir.Context,805            loc: ir.Location,806            mode: ShuffleMode,807            mask: *ir.Value,808            src: *ir.Value,809            lane_or_delta: *ir.Value,810        ) !ShflSyncOp {811            try loadSpec(ctx);812            var builder = ir.OperationBuilder.init(ctx);813            var state = ir.Operation.State.init(operation_name, loc);814            state.addOperands(&.{ mask, src, lane_or_delta });815            state.addTypes(&.{src.type});816817            const op = try builder.create(state);818            errdefer op.erase();819820            const mode_attr = try ctx.getDialectAttr("gpu.shuffle_mode", mode.toString());821            try op.setAttr("mode", mode_attr);822823            return .{ .op = op };824        }825826        pub fn getResult(self: *const ShflSyncOp) *ir.Value {827            return self.op.getResult(0).?;828        }829830        pub fn getMask(self: ShflSyncOp) *ir.Value {831            return self.op.operands.items[0].value;832        }833834        pub fn getSrc(self: ShflSyncOp) *ir.Value {835            return self.op.operands.items[1].value;836        }837838        pub fn getLaneOrDelta(self: ShflSyncOp) *ir.Value {839            return self.op.operands.items[2].value;840        }841842        pub fn getMode(self: ShflSyncOp) ?ShuffleMode {843            const dialect_attr = self.op.getAttrAs(ir.Attribute.DialectAttr, "mode") orelse return null;844            return ShuffleMode.fromString(dialect_attr.payload);845        }846    };847848    pub const AllSyncOp = struct {849        op: *ir.Operation,850851        pub const operation_spec = op_specs.leaf(.{852            .mnemonic = "all_sync",853            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},854            .operands = .{ "mask", "predicate" },855            .results = .{"result"},856        });857        pub const operation_name = operation_spec.name;858859        pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !AllSyncOp {860            const arith = choir.dialects.arith;861            try loadSpec(ctx);862            var builder = ir.OperationBuilder.init(ctx);863            const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);864            var state = ir.Operation.State.init(operation_name, loc);865            state.addOperands(&.{ mask, pred });866            state.addTypes(&.{bool_type});867868            const op = try builder.create(state);869            errdefer op.erase();870            return .{ .op = op };871        }872873        pub fn getResult(self: *const AllSyncOp) *ir.Value {874            return self.op.getResult(0).?;875        }876877        pub fn getMask(self: AllSyncOp) *ir.Value {878            return self.op.operands.items[0].value;879        }880881        pub fn getPredicate(self: AllSyncOp) *ir.Value {882            return self.op.operands.items[1].value;883        }884    };885886    pub const AnySyncOp = struct {887        op: *ir.Operation,888889        pub const operation_spec = op_specs.leaf(.{890            .mnemonic = "any_sync",891            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},892            .operands = .{ "mask", "predicate" },893            .results = .{"result"},894        });895        pub const operation_name = operation_spec.name;896897        pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !AnySyncOp {898            const arith = choir.dialects.arith;899            try loadSpec(ctx);900            var builder = ir.OperationBuilder.init(ctx);901            const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);902            var state = ir.Operation.State.init(operation_name, loc);903            state.addOperands(&.{ mask, pred });904            state.addTypes(&.{bool_type});905906            const op = try builder.create(state);907            errdefer op.erase();908            return .{ .op = op };909        }910911        pub fn getResult(self: *const AnySyncOp) *ir.Value {912            return self.op.getResult(0).?;913        }914915        pub fn getMask(self: AnySyncOp) *ir.Value {916            return self.op.operands.items[0].value;917        }918919        pub fn getPredicate(self: AnySyncOp) *ir.Value {920            return self.op.operands.items[1].value;921        }922    };923924    pub const BallotSyncOp = struct {925        op: *ir.Operation,926927        pub const operation_spec = op_specs.leaf(.{928            .mnemonic = "ballot_sync",929            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},930            .operands = .{ "mask", "predicate" },931            .results = .{"result"},932        });933        pub const operation_name = operation_spec.name;934935        pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !BallotSyncOp {936            const arith = choir.dialects.arith;937            try loadSpec(ctx);938            var builder = ir.OperationBuilder.init(ctx);939            const i32_type = try arith.ArithDialect.getI32Type(ctx);940            var state = ir.Operation.State.init(operation_name, loc);941            state.addOperands(&.{ mask, pred });942            state.addTypes(&.{i32_type});943944            const op = try builder.create(state);945            errdefer op.erase();946            return .{ .op = op };947        }948949        pub fn getResult(self: *const BallotSyncOp) *ir.Value {950            return self.op.getResult(0).?;951        }952953        pub fn getMask(self: BallotSyncOp) *ir.Value {954            return self.op.operands.items[0].value;955        }956957        pub fn getPredicate(self: BallotSyncOp) *ir.Value {958            return self.op.operands.items[1].value;959        }960    };961962    pub const WarpReduceOp = struct {963        op: *ir.Operation,964965        pub const operation_spec = op_specs.leaf(.{966            .mnemonic = "warp_reduce",967            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},968            .operands = .{ "mask", "value" },969            .results = .{"result"},970            .required_attrs = &.{"op"},971        });972        pub const operation_name = operation_spec.name;973974        pub fn create(975            ctx: *ir.Context,976            loc: ir.Location,977            op_kind: WarpOpKind,978            mask: *ir.Value,979            value: *ir.Value,980        ) !WarpReduceOp {981            try loadSpec(ctx);982            var builder = ir.OperationBuilder.init(ctx);983            var state = ir.Operation.State.init(operation_name, loc);984            state.addOperands(&.{ mask, value });985            state.addTypes(&.{value.type});986987            const op = try builder.create(state);988            errdefer op.erase();989            try setWarpOpAttr(op, ctx, op_kind);990991            return .{ .op = op };992        }993994        pub fn getResult(self: *const WarpReduceOp) *ir.Value {995            return self.op.getResult(0).?;996        }997998        pub fn getMask(self: WarpReduceOp) *ir.Value {999            return self.op.operands.items[0].value;1000        }10011002        pub fn getValue(self: WarpReduceOp) *ir.Value {1003            return self.op.operands.items[1].value;1004        }10051006        pub fn getOpKind(self: WarpReduceOp) ?WarpOpKind {1007            return getWarpOpAttr(self.op);1008        }1009    };10101011    pub const WarpScanOp = struct {1012        op: *ir.Operation,10131014        pub const operation_spec = op_specs.leaf(.{1015            .mnemonic = "warp_scan",1016            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},1017            .operands = .{ "mask", "value" },1018            .results = .{"result"},1019            .required_attrs = &.{ "inclusive", "op" },1020        });1021        pub const operation_name = operation_spec.name;10221023        pub fn create(1024            ctx: *ir.Context,1025            loc: ir.Location,1026            op_kind: WarpOpKind,1027            inclusive: bool,1028            mask: *ir.Value,1029            value: *ir.Value,1030        ) !WarpScanOp {1031            try loadSpec(ctx);1032            var builder = ir.OperationBuilder.init(ctx);1033            var state = ir.Operation.State.init(operation_name, loc);1034            state.addOperands(&.{ mask, value });1035            state.addTypes(&.{value.type});10361037            const op = try builder.create(state);1038            errdefer op.erase();1039            try setWarpOpAttr(op, ctx, op_kind);1040            try setBoolAttr(op, ctx, "inclusive", inclusive);10411042            return .{ .op = op };1043        }10441045        pub fn getResult(self: *const WarpScanOp) *ir.Value {1046            return self.op.getResult(0).?;1047        }10481049        pub fn getMask(self: WarpScanOp) *ir.Value {1050            return self.op.operands.items[0].value;1051        }10521053        pub fn getValue(self: WarpScanOp) *ir.Value {1054            return self.op.operands.items[1].value;1055        }10561057        pub fn getOpKind(self: WarpScanOp) ?WarpOpKind {1058            return getWarpOpAttr(self.op);1059        }10601061        pub fn isInclusive(self: WarpScanOp) bool {1062            return getBoolAttrValue(self.op, "inclusive");1063        }1064    };10651066    pub const MatchAnyOp = struct {1067        op: *ir.Operation,10681069        pub const operation_spec = op_specs.leaf(.{1070            .mnemonic = "match_any",1071            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},1072            .operands = .{"value"},1073            .results = .{"mask"},1074        });1075        pub const operation_name = operation_spec.name;10761077        pub fn create(ctx: *ir.Context, loc: ir.Location, value: *ir.Value) !MatchAnyOp {1078            const arith = choir.dialects.arith;1079            try loadSpec(ctx);1080            var builder = ir.OperationBuilder.init(ctx);1081            const i32_type = try arith.ArithDialect.getI32Type(ctx);1082            var state = ir.Operation.State.init(operation_name, loc);1083            state.addOperands(&.{value});1084            state.addTypes(&.{i32_type});10851086            const op = try builder.create(state);1087            errdefer op.erase();1088            return .{ .op = op };1089        }10901091        pub fn getResult(self: *const MatchAnyOp) *ir.Value {1092            return self.op.getResult(0).?;1093        }10941095        pub fn getValue(self: MatchAnyOp) *ir.Value {1096            return self.op.operands.items[0].value;1097        }1098    };10991100    pub const MatchAllOp = struct {1101        op: *ir.Operation,11021103        pub const operation_spec = op_specs.leaf(.{1104            .mnemonic = "match_all",1105            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},1106            .operands = .{"value"},1107            .results = .{ "mask", "all_equal" },1108        });1109        pub const operation_name = operation_spec.name;11101111        pub fn create(ctx: *ir.Context, loc: ir.Location, value: *ir.Value) !MatchAllOp {1112            const arith = choir.dialects.arith;1113            try loadSpec(ctx);1114            var builder = ir.OperationBuilder.init(ctx);1115            const i32_type = try arith.ArithDialect.getI32Type(ctx);1116            const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);1117            var state = ir.Operation.State.init(operation_name, loc);1118            state.addOperands(&.{value});1119            state.addTypes(&.{ i32_type, bool_type });11201121            const op = try builder.create(state);1122            errdefer op.erase();1123            return .{ .op = op };1124        }11251126        pub fn getMaskResult(self: *const MatchAllOp) *ir.Value {1127            return self.op.getResult(0).?;1128        }11291130        pub fn getAllEqualResult(self: *const MatchAllOp) *ir.Value {1131            return self.op.getResult(1).?;1132        }11331134        pub fn getValue(self: MatchAllOp) *ir.Value {1135            return self.op.operands.items[0].value;1136        }1137    };11381139    pub const ActiveMaskOp = struct {1140        op: *ir.Operation,11411142        pub const operation_spec = op_specs.leaf(.{1143            .mnemonic = "active_mask",1144            .interfaces = &.{gpuEffects(.state_observe, &.{}, &.{})},1145            .operands = 0,1146            .results = .{"mask"},1147        });1148        pub const operation_name = operation_spec.name;11491150        pub fn create(ctx: *ir.Context, loc: ir.Location) !ActiveMaskOp {1151            const arith = choir.dialects.arith;1152            try loadSpec(ctx);1153            var builder = ir.OperationBuilder.init(ctx);1154            const i32_type = try arith.ArithDialect.getI32Type(ctx);1155            var state = ir.Operation.State.init(operation_name, loc);1156            state.addTypes(&.{i32_type});11571158            const op = try builder.create(state);1159            errdefer op.erase();1160            return .{ .op = op };1161        }11621163        pub fn getResult(self: *const ActiveMaskOp) *ir.Value {1164            return self.op.getResult(0).?;1165        }1166    };11671168    pub const SyncWarpOp = struct {1169        op: *ir.Operation,11701171        pub const operation_spec = noResultSpec("sync_warp", .{"mask"});1172        pub const operation_name = operation_spec.name;11731174        pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value) !SyncWarpOp {1175            try loadSpec(ctx);1176            var builder = ir.OperationBuilder.init(ctx);1177            var state = ir.Operation.State.init(operation_name, loc);1178            state.addOperands(&.{mask});11791180            const op = try builder.create(state);1181            errdefer op.erase();1182            return .{ .op = op };1183        }11841185        pub fn getMask(self: SyncWarpOp) *ir.Value {1186            return self.op.operands.items[0].value;1187        }1188    };11891190    pub const mma_sync_a_count = 4;1191    pub const mma_sync_b_count = 2;1192    pub const mma_sync_acc_count = 4;11931194    pub const MmaSyncOp = struct {1195        op: *ir.Operation,11961197        pub const operation_spec = op_specs.leaf(.{1198            .mnemonic = "mma_sync",1199            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},1200            .operands = .{ "a0", "a1", "a2", "a3", "b0", "b1", "c0", "c1", "c2", "c3" },1201            .results = .{ "d0", "d1", "d2", "d3" },1202            .required_attrs = &.{"shape"},1203        });1204        pub const operation_name = operation_spec.name;12051206        pub fn create(1207            ctx: *ir.Context,1208            loc: ir.Location,1209            a: [mma_sync_a_count]*ir.Value,1210            b: [mma_sync_b_count]*ir.Value,1211            c: [mma_sync_acc_count]*ir.Value,1212            shape: MmaShape,1213        ) !MmaSyncOp {1214            try loadSpec(ctx);1215            var builder = ir.OperationBuilder.init(ctx);1216            var state = ir.Operation.State.init(operation_name, loc);1217            state.addOperands(&.{ a[0], a[1], a[2], a[3], b[0], b[1], c[0], c[1], c[2], c[3] });1218            state.addTypes(&.{ c[0].type, c[1].type, c[2].type, c[3].type });12191220            const op = try builder.create(state);1221            errdefer op.erase();1222            try setMmaShapeAttr(op, ctx, shape);1223            return .{ .op = op };1224        }12251226        pub fn getA(self: MmaSyncOp, index: usize) *ir.Value {1227            return self.op.operands.items[index].value;1228        }12291230        pub fn getB(self: MmaSyncOp, index: usize) *ir.Value {1231            return self.op.operands.items[mma_sync_a_count + index].value;1232        }12331234        pub fn getC(self: MmaSyncOp, index: usize) *ir.Value {1235            return self.op.operands.items[mma_sync_a_count + mma_sync_b_count + index].value;1236        }12371238        pub fn getD(self: *const MmaSyncOp, index: usize) *ir.Value {1239            return self.op.getResult(index).?;1240        }12411242        pub fn getShape(self: MmaSyncOp) ?MmaShape {1243            return getMmaShapeAttr(self.op);1244        }1245    };12461247    pub const CpAsyncSharedOp = struct {1248        op: *ir.Operation,12491250        pub const operation_spec = op_specs.leaf(.{1251            .mnemonic = "cp_async_shared",1252            .interfaces = &.{gpuEffects(.launch, &.{2}, &.{0})},1253            .operands = .{ "dst", "dst_index", "src", "src_index" },1254            .results = 0,1255            .required_attrs = &.{"bytes"},1256        });1257        pub const operation_name = operation_spec.name;12581259        pub fn create(1260            ctx: *ir.Context,1261            loc: ir.Location,1262            dst: *ir.Value,1263            dst_index: *ir.Value,1264            src: *ir.Value,1265            src_index: *ir.Value,1266            bytes: u32,1267        ) !CpAsyncSharedOp {1268            try loadSpec(ctx);1269            var builder = ir.OperationBuilder.init(ctx);1270            var state = ir.Operation.State.init(operation_name, loc);1271            state.addOperands(&.{ dst, dst_index, src, src_index });1272            const op = try builder.create(state);1273            errdefer op.erase();1274            const bytes_attr = try ctx.getI64Attr(@intCast(bytes));1275            try op.setAttr("bytes", bytes_attr);1276            return .{ .op = op };1277        }12781279        pub fn getDst(self: CpAsyncSharedOp) *ir.Value {1280            return ir.dialects.operand(operation_spec, self.op, "dst");1281        }12821283        pub fn getDstIndex(self: CpAsyncSharedOp) *ir.Value {1284            return ir.dialects.operand(operation_spec, self.op, "dst_index");1285        }12861287        pub fn getSrc(self: CpAsyncSharedOp) *ir.Value {1288            return ir.dialects.operand(operation_spec, self.op, "src");1289        }12901291        pub fn getSrcIndex(self: CpAsyncSharedOp) *ir.Value {1292            return ir.dialects.operand(operation_spec, self.op, "src_index");1293        }12941295        pub fn getBytes(self: CpAsyncSharedOp) ?u32 {1296            const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "bytes") orelse return null;1297            const raw = int_attr.getUnsignedValue();1298            if (raw > std.math.maxInt(u32)) return null;1299            return @intCast(raw);1300        }1301    };13021303    pub const CpAsyncCommitOp = struct {1304        op: *ir.Operation,13051306        pub const operation_spec = op_specs.leaf(.{1307            .mnemonic = "cp_async_commit",1308            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},1309            .operands = 0,1310            .results = 0,1311        });1312        pub const operation_name = operation_spec.name;13131314        pub fn create(ctx: *ir.Context, loc: ir.Location) !CpAsyncCommitOp {1315            try loadSpec(ctx);1316            var builder = ir.OperationBuilder.init(ctx);1317            const state = ir.Operation.State.init(operation_name, loc);1318            const op = try builder.create(state);1319            errdefer op.erase();1320            return .{ .op = op };1321        }1322    };13231324    pub const CpAsyncWaitOp = struct {1325        op: *ir.Operation,13261327        pub const operation_spec = op_specs.leaf(.{1328            .mnemonic = "cp_async_wait",1329            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},1330            .operands = 0,1331            .results = 0,1332            .required_attrs = &.{"groups"},1333        });1334        pub const operation_name = operation_spec.name;13351336        pub fn create(ctx: *ir.Context, loc: ir.Location, groups: u32) !CpAsyncWaitOp {1337            try loadSpec(ctx);1338            var builder = ir.OperationBuilder.init(ctx);1339            const state = ir.Operation.State.init(operation_name, loc);1340            const op = try builder.create(state);1341            errdefer op.erase();1342            const groups_attr = try ctx.getI64Attr(@intCast(groups));1343            try op.setAttr("groups", groups_attr);1344            return .{ .op = op };1345        }13461347        pub fn getGroups(self: CpAsyncWaitOp) ?u32 {1348            const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "groups") orelse return null;1349            const raw = int_attr.getUnsignedValue();1350            if (raw > std.math.maxInt(u32)) return null;1351            return @intCast(raw);1352        }1353    };13541355    pub const AtomicLoadOp = struct {1356        op: *ir.Operation,13571358        pub const operation_spec = op_specs.leaf(.{1359            .mnemonic = "atomic_load",1360            .interfaces = &.{gpuEffects(.synchronize, &.{0}, &.{})},1361            .operands = .{ "memref", "index" },1362            .results = .{"value"},1363            .required_attrs = &.{"ordering"},1364        });1365        pub const operation_name = operation_spec.name;13661367        pub fn create(1368            ctx: *ir.Context,1369            loc: ir.Location,1370            memref: *ir.Value,1371            index: *ir.Value,1372            result_type: ir.Type,1373            ordering: MemoryOrder,1374        ) !AtomicLoadOp {1375            try loadSpec(ctx);1376            var builder = ir.OperationBuilder.init(ctx);1377            var state = ir.Operation.State.init(operation_name, loc);1378            state.addOperands(&.{ memref, index });1379            state.addTypes(&.{result_type});13801381            const op = try builder.create(state);1382            errdefer op.erase();1383            try setOrderingAttr(op, ctx, ordering);1384            return .{ .op = op };1385        }13861387        pub fn getResult(self: *const AtomicLoadOp) *ir.Value {1388            return self.op.getResult(0).?;1389        }13901391        pub fn getMemref(self: AtomicLoadOp) *ir.Value {1392            return self.op.operands.items[0].value;1393        }13941395        pub fn getIndex(self: AtomicLoadOp) *ir.Value {1396            return self.op.operands.items[1].value;1397        }13981399        pub fn getOrdering(self: AtomicLoadOp) ?MemoryOrder {1400            return getOrderingAttr(self.op);1401        }1402    };14031404    pub const AtomicStoreOp = struct {1405        op: *ir.Operation,14061407        pub const operation_spec = op_specs.leaf(.{1408            .mnemonic = "atomic_store",1409            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{1})},1410            .operands = .{ "value", "memref", "index" },1411            .results = 0,1412            .required_attrs = &.{"ordering"},1413        });1414        pub const operation_name = operation_spec.name;14151416        pub fn create(1417            ctx: *ir.Context,1418            loc: ir.Location,1419            value: *ir.Value,1420            memref: *ir.Value,1421            index: *ir.Value,1422            ordering: MemoryOrder,1423        ) !AtomicStoreOp {1424            try loadSpec(ctx);1425            var builder = ir.OperationBuilder.init(ctx);1426            var state = ir.Operation.State.init(operation_name, loc);1427            state.addOperands(&.{ value, memref, index });14281429            const op = try builder.create(state);1430            errdefer op.erase();1431            try setOrderingAttr(op, ctx, ordering);1432            return .{ .op = op };1433        }14341435        pub fn getValue(self: AtomicStoreOp) *ir.Value {1436            return self.op.operands.items[0].value;1437        }14381439        pub fn getMemref(self: AtomicStoreOp) *ir.Value {1440            return self.op.operands.items[1].value;1441        }14421443        pub fn getIndex(self: AtomicStoreOp) *ir.Value {1444            return self.op.operands.items[2].value;1445        }14461447        pub fn getOrdering(self: AtomicStoreOp) ?MemoryOrder {1448            return getOrderingAttr(self.op);1449        }1450    };14511452    pub const AtomicAddOp = struct {1453        op: *ir.Operation,14541455        pub const operation_spec = op_specs.leaf(.{1456            .mnemonic = "atomic_add",1457            .interfaces = &.{gpuEffects(.synchronize, &.{0}, &.{0})},1458            .operands = .{ "memref", "index", "value" },1459            .results = .{"old_value"},1460            .attrs = &.{"scope"},1461            .required_attrs = &.{"ordering"},1462        });1463        pub const operation_name = operation_spec.name;14641465        pub fn create(1466            ctx: *ir.Context,1467            loc: ir.Location,1468            memref: *ir.Value,1469            index: *ir.Value,1470            val: *ir.Value,1471            ordering: MemoryOrder,1472            scope: ?Scope,1473        ) !AtomicAddOp {1474            try loadSpec(ctx);1475            var builder = ir.OperationBuilder.init(ctx);1476            var state = ir.Operation.State.init(operation_name, loc);1477            state.addOperands(&.{ memref, index, val });1478            state.addTypes(&.{val.type});14791480            const op = try builder.create(state);1481            errdefer op.erase();1482            try setOrderingAttr(op, ctx, ordering);1483            if (scope) |scope_value| {1484                try setScopeAttr(op, ctx, scope_value);1485            }1486            return .{ .op = op };1487        }14881489        pub fn getResult(self: *const AtomicAddOp) *ir.Value {1490            return self.op.getResult(0).?;1491        }14921493        pub fn getMemref(self: AtomicAddOp) *ir.Value {1494            return self.op.operands.items[0].value;1495        }14961497        pub fn getIndex(self: AtomicAddOp) *ir.Value {1498            return self.op.operands.items[1].value;1499        }15001501        pub fn getVal(self: AtomicAddOp) *ir.Value {1502            return self.op.operands.items[2].value;1503        }15041505        pub fn getOrdering(self: AtomicAddOp) ?MemoryOrder {1506            return getOrderingAttr(self.op);1507        }15081509        pub fn getScope(self: AtomicAddOp) ?Scope {1510            return getScopeAttr(self.op);1511        }1512    };15131514    pub const AtomicMaxOp = struct {1515        op: *ir.Operation,15161517        pub const operation_spec = op_specs.leaf(.{1518            .mnemonic = "atomic_max",1519            .interfaces = &.{gpuEffects(.synchronize, &.{0}, &.{0})},1520            .operands = .{ "memref", "index", "value" },1521            .results = .{"old_value"},1522        });1523        pub const operation_name = operation_spec.name;15241525        pub fn create(1526            ctx: *ir.Context,1527            loc: ir.Location,1528            memref: *ir.Value,1529            index: *ir.Value,1530            val: *ir.Value,1531        ) !AtomicMaxOp {1532            try loadSpec(ctx);1533            var builder = ir.OperationBuilder.init(ctx);1534            var state = ir.Operation.State.init(operation_name, loc);1535            state.addOperands(&.{ memref, index, val });1536            state.addTypes(&.{val.type});15371538            const op = try builder.create(state);1539            errdefer op.erase();1540            return .{ .op = op };1541        }15421543        pub fn getResult(self: *const AtomicMaxOp) *ir.Value {1544            return self.op.getResult(0).?;1545        }15461547        pub fn getMemref(self: AtomicMaxOp) *ir.Value {1548            return self.op.operands.items[0].value;1549        }15501551        pub fn getIndex(self: AtomicMaxOp) *ir.Value {1552            return self.op.operands.items[1].value;1553        }15541555        pub fn getVal(self: AtomicMaxOp) *ir.Value {1556            return self.op.operands.items[2].value;1557        }1558    };15591560    pub const AtomicCasOp = struct {1561        op: *ir.Operation,15621563        pub const operation_spec = op_specs.leaf(.{1564            .mnemonic = "atomic_cas",1565            .interfaces = &.{gpuEffects(.synchronize, &.{0}, &.{0})},1566            .operands = .{ "memref", "index", "expected", "desired" },1567            .results = .{"old_value"},1568            .attrs = &.{"scope"},1569            .required_attrs = &.{"ordering"},1570        });1571        pub const operation_name = operation_spec.name;15721573        pub fn create(1574            ctx: *ir.Context,1575            loc: ir.Location,1576            memref: *ir.Value,1577            index: *ir.Value,1578            expected: *ir.Value,1579            desired: *ir.Value,1580            ordering: MemoryOrder,1581            scope: ?Scope,1582        ) !AtomicCasOp {1583            try loadSpec(ctx);1584            var builder = ir.OperationBuilder.init(ctx);1585            var state = ir.Operation.State.init(operation_name, loc);1586            state.addOperands(&.{ memref, index, expected, desired });1587            state.addTypes(&.{expected.type});15881589            const op = try builder.create(state);1590            errdefer op.erase();1591            try setOrderingAttr(op, ctx, ordering);1592            if (scope) |scope_value| {1593                try setScopeAttr(op, ctx, scope_value);1594            }1595            return .{ .op = op };1596        }15971598        pub fn getResult(self: *const AtomicCasOp) *ir.Value {1599            return self.op.getResult(0).?;1600        }16011602        pub fn getMemref(self: AtomicCasOp) *ir.Value {1603            return self.op.operands.items[0].value;1604        }16051606        pub fn getIndex(self: AtomicCasOp) *ir.Value {1607            return self.op.operands.items[1].value;1608        }16091610        pub fn getExpected(self: AtomicCasOp) *ir.Value {1611            return self.op.operands.items[2].value;1612        }16131614        pub fn getDesired(self: AtomicCasOp) *ir.Value {1615            return self.op.operands.items[3].value;1616        }16171618        pub fn getOrdering(self: AtomicCasOp) ?MemoryOrder {1619            return getOrderingAttr(self.op);1620        }16211622        pub fn getScope(self: AtomicCasOp) ?Scope {1623            return getScopeAttr(self.op);1624        }1625    };16261627    fn loadSpec(ctx: *ir.Context) !void {1628        try ir.dialects.loadDialectSpec(ctx, spec);1629    }16301631    fn dimIndexSpec(comptime mnemonic: []const u8) ir.dialects.OperationSpec {1632        return op_specs.leaf(.{1633            .mnemonic = mnemonic,1634            .interfaces = &.{gpuEffects(.state_observe, &.{}, &.{})},1635            .operands = 0,1636            .results = .{"index"},1637            .required_attrs = &.{"dim"},1638        });1639    }16401641    fn indexSpec(comptime mnemonic: []const u8) ir.dialects.OperationSpec {1642        return op_specs.leaf(.{1643            .mnemonic = mnemonic,1644            .interfaces = &.{gpuEffects(.state_observe, &.{}, &.{})},1645            .operands = 0,1646            .results = .{"index"},1647        });1648    }16491650    fn noResultSpec(comptime mnemonic: []const u8, comptime operands: anytype) ir.dialects.OperationSpec {1651        return op_specs.leaf(.{1652            .mnemonic = mnemonic,1653            .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},1654            .operands = operands,1655            .results = 0,1656        });1657    }16581659    fn getFuncSymbolName(op_ptr: *const anyopaque) ?[]const u8 {1660        const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));1661        if (op.getAttr("sym_name")) |attr| {1662            return func.FuncDialect.getSymNameValue(attr);1663        }1664        return null;1665    }16661667    fn setFuncSymbolName(op_ptr: *const anyopaque, symbol_name: []const u8) anyerror!void {1668        const op: *ir.Operation = @ptrCast(@alignCast(@constCast(op_ptr)));1669        try op.setAttr("sym_name", try func.FuncDialect.getSymNameAttr(op.getContext(), symbol_name));1670    }16711672    fn isFuncDeclaration(_: *const anyopaque) bool {1673        return false;1674    }16751676    pub fn getTmaDescriptorType(ctx: *ir.Context) !ir.Type {1677        try loadSpec(ctx);1678        return ctx.getDialectTypeFromName(type_names.tma_desc);1679    }16801681    pub fn getMBarrierType(ctx: *ir.Context) !ir.Type {1682        try loadSpec(ctx);1683        return ctx.getDialectTypeFromName(type_names.mbarrier);1684    }16851686    fn setDimensionAttr(op: *ir.Operation, ctx: *ir.Context, dim: Dimension) !void {1687        const dim_attr = try ctx.getDialectAttr("gpu.dim", dim.toString());1688        try op.setAttr("dim", dim_attr);1689    }16901691    fn getDimensionAttr(op: *const ir.Operation) ?Dimension {1692        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "dim") orelse return null;1693        return Dimension.fromString(dialect_attr.payload);1694    }16951696    fn setScopeAttr(op: *ir.Operation, ctx: *ir.Context, scope: Scope) !void {1697        const scope_attr = try ctx.getDialectAttr("gpu.scope", scope.toString());1698        try op.setAttr("scope", scope_attr);1699    }17001701    fn getScopeAttr(op: *const ir.Operation) ?Scope {1702        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "scope") orelse return null;1703        return Scope.fromString(dialect_attr.payload);1704    }17051706    fn setOrderingAttr(op: *ir.Operation, ctx: *ir.Context, ordering: MemoryOrder) !void {1707        const order_attr = try ctx.getDialectAttr("gpu.ordering", ordering.toString());1708        try op.setAttr("ordering", order_attr);1709    }17101711    fn getOrderingAttr(op: *const ir.Operation) ?MemoryOrder {1712        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "ordering") orelse return null;1713        return MemoryOrder.fromString(dialect_attr.payload);1714    }17151716    fn setBoolAttr(op: *ir.Operation, ctx: *ir.Context, attr_name: []const u8, value: bool) !void {1717        const bool_attr = try ctx.getBoolAttr(value);1718        try op.setAttr(attr_name, bool_attr);1719    }17201721    fn getBoolAttrValue(op: *const ir.Operation, attr_name: []const u8) bool {1722        const bool_attr = op.getAttrAs(ir.Attribute.BoolAttr, attr_name) orelse return false;1723        return bool_attr.getValue();1724    }17251726    fn setI64Attr(op: *ir.Operation, ctx: *ir.Context, attr_name: []const u8, value: i64) !void {1727        const int_attr = try ctx.getI64Attr(value);1728        try op.setAttr(attr_name, int_attr);1729    }17301731    fn getI64AttrValue(op: *const ir.Operation, attr_name: []const u8) ?i64 {1732        const int_attr = op.getAttrAs(ir.Attribute.IntegerAttr, attr_name) orelse return null;1733        return int_attr.getValue();1734    }17351736    fn setWarpOpAttr(op: *ir.Operation, ctx: *ir.Context, op_kind: WarpOpKind) !void {1737        const op_attr = try ctx.getDialectAttr("gpu.warp_op", op_kind.toString());1738        try op.setAttr("op", op_attr);1739    }17401741    fn getWarpOpAttr(op: *const ir.Operation) ?WarpOpKind {1742        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "op") orelse return null;1743        return WarpOpKind.fromString(dialect_attr.payload);1744    }17451746    fn setMmaShapeAttr(op: *ir.Operation, ctx: *ir.Context, shape: MmaShape) !void {1747        var buf: [32]u8 = undefined;1748        const shape_str = try shape.toString(buf[0..]);1749        const shape_attr = try ctx.getDialectAttr("gpu.mma_shape", shape_str);1750        try op.setAttr("shape", shape_attr);1751    }17521753    fn getMmaShapeAttr(op: *const ir.Operation) ?MmaShape {1754        const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "shape") orelse return null;1755        return MmaShape.parse(dialect_attr.payload);1756    }1757};