# TritonNvidiaGPUOps ### `ttng.arrive_barrier` (triton::nvidia_gpu::ArriveBarrierOp) _Perform the arrive operation on an mbarrier_ Syntax: ``` operation ::= `ttng.arrive_barrier` $alloc `,` $count (`,` $pred^)? attr-dict `:` qualified(type($alloc)) ``` The `ttng.arrive_barrier` operation performs the "arrive" operation on an mbarrier object in shared memory. The operation requires a `count` attribute of at least 1 and decreases the pending arrival count of the mbarrier by the specified count. When `multicastCTA` is non-zero, the arrive is multicast across the cluster. Each bit set in the mask identifies a CTA ID dimension to multicast along. CTA IDs `a` and `b` belong to the same equivalence class iff `a & ~multicastCTA == b & ~multicastCTA`; all CTAs in a class multicast to each other. Multicast requires `numCTAs > 1`, `0 < multicastCTA <= numCTAs - 1`, and the barrier must have the identity CGA layout `[[1], [2], ...]`. The default value of `multicastCTA` is 0 (no multicast). The operation accepts an optional predicate and `fromCTA` attribute. The latter is a CTA-ID basis-selection mask: set bits are preserved and unset bits are zeroed when selecting the CTA that routes the arrival to each per-CTA barrier. A non-identity `fromCTA` cannot be combined with multicast. Example: ```mlir ttng.arrive_barrier %barrier, 2 : !ttg.memdesc<1xi64, #shared, #smem, mutable> ttng.arrive_barrier %barrier, 1, %pred : !ttg.memdesc<1xi64, #shared, #smem, mutable> ttng.arrive_barrier %barrier, 1 {multicastCTA = 1 : i32} : !ttg.memdesc<2xi64, #shared, #smem, mutable> ``` Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MBarrierOpInterface`, `PredicatedOpInterface` #### Attributes:
AttributeMLIR TypeDescription
count::mlir::IntegerAttr32-bit signless integer attribute
fromCTA::mlir::IntegerAttr32-bit signless integer attribute
multicastCTA::mlir::IntegerAttr32-bit signless integer attribute
#### Operands: | Operand | Description | | :-----: | ----------- | | `alloc` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `pred` | 1-bit signless integer | ### `ttng.async_copy_mbarrier_arrive` (triton::nvidia_gpu::AsyncCopyMbarrierArriveOp) _Arrive on mbarrier once all previously issued copies are completed_ Syntax: ``` operation ::= `ttng.async_copy_mbarrier_arrive` $barrier attr-dict `:` qualified(type($barrier)) ``` Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MBarrierOpInterface` #### Attributes:
AttributeMLIR TypeDescription
noIncrement::mlir::UnitAttrunit attribute
#### Operands: | Operand | Description | | :-----: | ----------- | | `barrier` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | ### `ttng.async_shared_store` (triton::nvidia_gpu::AsyncSharedStoreOp) _Store a distributed tensor to shared memory asynchronously_ Syntax: ``` operation ::= `ttng.async_shared_store` $src `,` $dst `,` $mbarrier attr-dict `:` type($src) `->` qualified(type($dst)) `,` qualified(type($mbarrier)) ``` Store a distributed tensor into shared memory using PTX st.async.shared. The store completion decrements the transaction count of `mbarrier`. This op requires a CTA cluster with at least two CTAs. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MBarrierOpInterface` #### Operands: | Operand | Description | | :-----: | ----------- | | `src` | ranked tensor of floating-point or integer or ptr values | | `dst` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `mbarrier` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | ### `ttng.async_tma_copy_global_to_local` (triton::nvidia_gpu::AsyncTMACopyGlobalToLocalOp) _Copy data based on descriptor from global memory to local memory asynchronously_ Syntax: ``` operation ::= `ttng.async_tma_copy_global_to_local` $desc `[` $coord `]` (`offsets` `=` `[` $offsets^ `]`)? $result `,` $barrier `,` $pred oilist(`cacheModifier` `=` $cache | `evictionPolicy` `=` $evict) attr-dict `:` qualified(type($desc)) `,` qualified(type($barrier)) `->` qualified(type($result)) ``` This operation copies data from global memory to local memory asynchronously. This is analogous to tt.load except the data is copied to local memory pointed to by the memory descriptor instead of to a distributed tensor. The copied data depends on the global memory descriptor pointed to by `desc`. The tensor mode is determined by the descriptor type: - tt.tensordesc: TILED mode - Regular tiled tensor memory access - See: https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-tiled-mode - ttng.tensordesc_im2col: IM2COL mode - Im2col mode for convolution-friendly access patterns - In IM2COL mode, 'coord' is the coordinate in the input tensor - For example, for a 4D tensor (NHWC), 'coord' is [batch_idx, channel_idx, h, w] - In IM2COL mode, additional `offsets` must be provided (uint16 values) - For 3D tensors (NWC): 1 offset (offset_w) - For 4D tensors (NHWC): 2 offsets (offset_w, offset_h) - For 5D tensors (NDHWC): 3 offsets (offset_w, offset_h, offset_d) - General rule: number of offsets = coord.size() - 2 - See: https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-im2col-mode Traits: `AttrSizedOperandSegments`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MBarrierOpInterface`, `PredicatedOpInterface`, `TMALoadLikeOpInterface`, `TMAOpInterface` #### Attributes:
AttributeMLIR TypeDescription
multicast::mlir::UnitAttrunit attribute
cache::mlir::triton::CacheModifierAttrallowed 32-bit signless integer cases: 1, 2, 3, 4, 5, 6, 7
evict::mlir::triton::EvictionPolicyAttrallowed 32-bit signless integer cases: 1, 2, 3
isVolatile::mlir::BoolAttrbool attribute
#### Operands: | Operand | Description | | :-----: | ----------- | | `desc` | tensor descriptor type | | `coord` | variadic of 32-bit signless integer | | `offsets` | variadic of 16-bit signless integer | | `barrier` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `result` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `pred` | 1-bit signless integer | ### `ttng.async_tma_copy_local_to_global` (triton::nvidia_gpu::AsyncTMACopyLocalToGlobalOp) _Copy data based on descriptor from local memory to global memory asynchronously_ Syntax: ``` operation ::= `ttng.async_tma_copy_local_to_global` $desc `[` $coord `]` $src attr-dict `:` qualified(type($desc)) `,` qualified(type($src)) ``` This operation copies data from local memory to global memory asynchronously. This is analogous to tt.store except the data is copied from local memory pointed to by the memory descriptor instead of from a distributed tensor. The copied data depends on the global memory descriptor pointed to by `desc`. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `TMAOpInterface`, `TMAStoreLikeOpInterface` #### Operands: | Operand | Description | | :-----: | ----------- | | `desc` | Tensor descriptor type (`::mlir::triton::TensorDescType`) in Triton IR type system | | `coord` | variadic of 32-bit signless integer | | `src` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | ### `ttng.async_tma_gather` (triton::nvidia_gpu::AsyncTMAGatherOp) _Gather data based on descriptor from global memory to local memory asynchronously_ Syntax: ``` operation ::= `ttng.async_tma_gather` $desc `[` $x_offsets `,` $y_offset `]` $result `,` $barrier `,` $pred attr-dict `:` type(operands) ``` This operation gathers multiple rows of data from a global memory matrix to local memory asynchronously. This is similar to async_tma_copy_global_to_local except that each row is indexed independently. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MBarrierOpInterface`, `PredicatedOpInterface`, `TMALoadLikeOpInterface`, `TMAOpInterface` #### Attributes:
AttributeMLIR TypeDescription
multicast::mlir::UnitAttrunit attribute
#### Operands: | Operand | Description | | :-----: | ----------- | | `desc` | Tensor descriptor type (`::mlir::triton::TensorDescType`) in Triton IR type system | | `x_offsets` | ranked tensor of 32-bit signless integer values | | `y_offset` | 32-bit signless integer | | `barrier` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `result` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `pred` | 1-bit signless integer | ### `ttng.async_tma_reduce` (triton::nvidia_gpu::AsyncTMAReduceOp) _Reduce result in gmem based on a TMA descriptor_ Syntax: ``` operation ::= `ttng.async_tma_reduce` $kind `,` $desc `[` $coord `]` $src attr-dict `:` qualified(type($desc)) `,` qualified(type($src)) ``` This operation copies data from local memory to global memory asynchronously, and atomically performs the specified reduction kind. Atomicity is at the granularity of individual elements, and only relaxed semantics are implied. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `TMAOpInterface`, `TMAStoreLikeOpInterface` #### Attributes:
AttributeMLIR TypeDescription
kind::mlir::triton::DescriptorReduceKindAttrallowed 32-bit signless integer cases: 1, 2, 3, 4, 5, 6, 7, 8
#### Operands: | Operand | Description | | :-----: | ----------- | | `desc` | Tensor descriptor type (`::mlir::triton::TensorDescType`) in Triton IR type system | | `coord` | variadic of 32-bit signless integer | | `src` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | ### `ttng.async_tma_scatter` (triton::nvidia_gpu::AsyncTMAScatterOp) _Scatter data from local memory into global memory based on a descriptor asynchronously_ Syntax: ``` operation ::= `ttng.async_tma_scatter` $desc `[` $x_offsets `,` $y_offset `]` $src attr-dict `:` type(operands) ``` The `ttng.async_tma_scatter` operation scatters multiple separately-indexed rows of data from local memory into global memory asynchronously. The operation scatters a 2D tensor in shared memory, laid out using the core tensor tiles nvmma_shared layout, into separately indexed rows in global memory at a given `y` offset. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `TMAOpInterface`, `TMAStoreLikeOpInterface` #### Operands: | Operand | Description | | :-----: | ----------- | | `desc` | Tensor descriptor type (`::mlir::triton::TensorDescType`) in Triton IR type system | | `x_offsets` | ranked tensor of 32-bit signless integer values | | `y_offset` | 32-bit signless integer | | `src` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | ### `ttng.async_tma_store_wait` (triton::nvidia_gpu::TMAStoreWaitOp) _Wait for pending TMA store operations._ Syntax: ``` operation ::= `ttng.async_tma_store_wait` attr-dict ``` Wait for the associated store operations to complete. When `read_only` is set, only wait until their reads from shared memory have completed. This is needed before shared memory can be written to again. Traits: `MemWaitOpTrait`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` #### Attributes:
AttributeMLIR TypeDescription
pendings::mlir::IntegerAttr32-bit signless integer attribute
read_only::mlir::UnitAttrunit attribute
### `ttng.barrier_expect` (triton::nvidia_gpu::BarrierExpectOp) _Signal a barrier of an expected number of bytes to be copied._ Syntax: ``` operation ::= `ttng.barrier_expect` $alloc `,` $size attr-dict `,` $pred `:` qualified(type($alloc)) ``` This signals the barrier that `size` bytes are expected to be copied. The associated barrier wait will block until the expected number of bytes has been copied. The optional `fromCTA` attribute is a CTA-ID basis-selection mask: set bits are preserved and unset bits are zeroed when selecting the CTA that routes the arrival to each per-CTA barrier. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MBarrierOpInterface`, `PredicatedOpInterface` #### Attributes:
AttributeMLIR TypeDescription
size::mlir::IntegerAttr32-bit signless integer attribute
fromCTA::mlir::IntegerAttr32-bit signless integer attribute
#### Operands: | Operand | Description | | :-----: | ----------- | | `alloc` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `pred` | 1-bit signless integer | ### `ttng.clc_get_program_id` (triton::nvidia_gpu::CLCGetProgramIdOp) _Get program ID coordinate from CLC response_ Syntax: ``` operation ::= `ttng.clc_get_program_id` $clcResult `,` $dim attr-dict `:` type($clcResult) `->` type($result) ``` Decodes the CLC response to get the program ID coordinate of the canceled cluster. The dim attribute specifies which dimension (0=x, 1=y, 2=z). Traits: `AlwaysSpeculatableImplTrait`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `ConditionallySpeculatable`, `InferTypeOpInterface`, `NoMemoryEffect (MemoryEffectOpInterface)` Effects: `MemoryEffects::Effect{}` #### Attributes:
AttributeMLIR TypeDescription
dim::mlir::triton::ProgramIDDimAttrallowed 32-bit signless integer cases: 0, 1, 2
#### Operands: | Operand | Description | | :-----: | ----------- | | `clcResult` | 128-bit signless integer | #### Results: | Result | Description | | :----: | ----------- | | `result` | 32-bit signless integer | ### `ttng.clc_is_canceled` (triton::nvidia_gpu::CLCIsCanceledOp) _Check if CLC response indicates successful cancellation_ Syntax: ``` operation ::= `ttng.clc_is_canceled` $clcResult attr-dict `:` type($clcResult) `->` type($is_canceled) ``` Decodes the CLC response to check if a cluster was successfully canceled. Returns true if canceled, false otherwise. Traits: `AlwaysSpeculatableImplTrait`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `ConditionallySpeculatable`, `InferTypeOpInterface`, `NoMemoryEffect (MemoryEffectOpInterface)` Effects: `MemoryEffects::Effect{}` #### Operands: | Operand | Description | | :-----: | ----------- | | `clcResult` | 128-bit signless integer | #### Results: | Result | Description | | :----: | ----------- | | `is_canceled` | 1-bit signless integer | ### `ttng.clc_load_result` (triton::nvidia_gpu::CLCLoadResultOp) _Load CLC response from shared memory into registers_ Syntax: ``` operation ::= `ttng.clc_load_result` $src attr-dict `:` qualified(type($src)) `->` type($clcResult) ``` Loads the 128-bit CLC response from shared memory into an i128 value. This allows subsequent clc_is_canceled and clc_get_program_id operations to operate on registers without re-reading shared memory. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `InferTypeOpInterface` #### Operands: | Operand | Description | | :-----: | ----------- | | `src` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | #### Results: | Result | Description | | :----: | ----------- | | `clcResult` | 128-bit signless integer | ### `ttng.clc_try_cancel` (triton::nvidia_gpu::CLCTryCancelOp) _Issue CLC try_cancel to cancel a pending cluster_ Syntax: ``` operation ::= `ttng.clc_try_cancel` $result `,` $mbarrier attr-dict `:` qualified(type($result)) `,` qualified(type($mbarrier)) ``` Issues a clusterlaunchcontrol.try_cancel instruction to atomically cancel a pending cluster launch. The result is written asynchronously to the result buffer and the mbarrier is signaled on completion. This is used for dynamic persistent kernels on Blackwell (SM100+). The result buffer must be a 16-byte aligned `2xi64` shared memory memdesc. The mbarrier must be 8-byte aligned shared memory. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MBarrierOpInterface` #### Operands: | Operand | Description | | :-----: | ----------- | | `result` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `mbarrier` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | ### `ttng.clc_try_cancel_sync` (triton::nvidia_gpu::CLCTryCancelSyncOp) _Request cancellation of a pending cluster_ Syntax: ``` operation ::= `ttng.clc_try_cancel_sync` attr-dict `:` type($response) ``` Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` #### Results: | Result | Description | | :----: | ----------- | | `response` | ranked tensor of floating-point or integer or ptr values | ### `ttng.cluster_barrier` (triton::nvidia_gpu::ClusterBarrierOp) _Synchronize all warps at cluster scope_ Syntax: ``` operation ::= `ttng.cluster_barrier` attr-dict ``` Lowers to a cluster arrive/wait pair. In warp-specialized kernels, lowering wraps the barrier in a synthetic `ttg.warp_specialize` region so worker warps also execute a barrier outside existing warp-specialized code. A barrier inside `ttg.warp_specialize` lowers through a cluster-scoped mbarrier. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` #### Attributes:
AttributeMLIR TypeDescription
relaxed::mlir::BoolAttrbool attribute
### `ttng.fence_async_shared` (triton::nvidia_gpu::FenceAsyncSharedOp) _Fence proxy async_ Syntax: ``` operation ::= `ttng.fence_async_shared` attr-dict ``` Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` #### Attributes:
AttributeMLIR TypeDescription
bCluster::mlir::BoolAttrbool attribute
### `ttng.fence_mbarrier_init_release_cluster` (triton::nvidia_gpu::FenceMBarrierInitReleaseClusterOp) _Fence mbarrier init release.cluster_ Syntax: ``` operation ::= `ttng.fence_mbarrier_init_release_cluster` attr-dict ``` Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` ### `ttng.init_barrier` (triton::nvidia_gpu::InitBarrierOp) _Initialize a barrier in the given shared memory allocation._ Syntax: ``` operation ::= `ttng.init_barrier` $alloc `,` $count attr-dict `:` qualified(type($alloc)) ``` Initializes a shared memory allocation with mbarrier information. `alloc` is a descriptor to the shared memory allocation. `count` is the number of arrives expected by the barrier. This lowers to PTX mbarrier.init.shared::cta.b64. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MBarrierOpInterface` #### Attributes:
AttributeMLIR TypeDescription
count::mlir::IntegerAttr32-bit signless integer attribute
#### Operands: | Operand | Description | | :-----: | ----------- | | `alloc` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | ### `ttng.inval_barrier` (triton::nvidia_gpu::InvalBarrierOp) _Invalidate a barrier allocation._ Syntax: ``` operation ::= `ttng.inval_barrier` $alloc attr-dict `:` qualified(type($alloc)) ``` Invalidate a barrier allocation so that it can be re-used. According to PTX spec this has to be done before any reuse of the memory used by mbarrier. https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-inval Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MBarrierOpInterface` #### Operands: | Operand | Description | | :-----: | ----------- | | `alloc` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | ### `ttng.packed_arith` (triton::nvidia_gpu::PackedArithOp) _Apply elementwise arithmetic using a packed PTX instruction_ Syntax: ``` operation ::= `ttng.packed_arith` $op_kind $operands attr-dict `:` functional-type($operands, $result) ``` Applies packed x2, mixed-precision, and FP8/FP4/UE8M0 x4 arithmetic. Instruction types, modifiers, and grouping are inferred from tensor types. FP4 uses lower-nibble-first i8 values along the inferred halved dimension. Traits: `AlwaysSpeculatableImplTrait`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `ConditionallySpeculatable`, `NoMemoryEffect (MemoryEffectOpInterface)` Effects: `MemoryEffects::Effect{}` #### Attributes:
AttributeMLIR TypeDescription
op_kind::mlir::triton::nvidia_gpu::PackedArithOpKindAttrpacked arithmetic operation
#### Operands: | Operand | Description | | :-----: | ----------- | | `operands` | variadic of ranked tensor of packed arithmetic operand element values | #### Results: | Result | Description | | :----: | ----------- | | `result` | ranked tensor of packed arithmetic result element values | ### `ttng.reinterpret_tensor_descriptor` (triton::nvidia_gpu::ReinterpretTensorDescOp) _Reinterpret a pointer as a tensor descriptor_ Syntax: ``` operation ::= `ttng.reinterpret_tensor_descriptor` $rawDesc attr-dict `:` qualified(type($rawDesc)) `to` qualified(type($result)) ``` This op exists to help the transition from untyped raw TMA objects to typed tensor descriptor objects. Ideally, we can remove this once the APIs are fully fleshed out. Traits: `AlwaysSpeculatableImplTrait`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `ConditionallySpeculatable`, `NoMemoryEffect (MemoryEffectOpInterface)` Effects: `MemoryEffects::Effect{}` #### Operands: | Operand | Description | | :-----: | ----------- | | `rawDesc` | ptr | #### Results: | Result | Description | | :----: | ----------- | | `result` | Tensor descriptor type (`::mlir::triton::TensorDescType`) in Triton IR type system | ### `ttng.tc_gen5_commit` (triton::nvidia_gpu::TCGen5CommitOp) _Make an mbarrier track completion of all prior async tcgen5 ops_ Syntax: ``` operation ::= `ttng.tc_gen5_commit` $barrier (`,` $pred^)? (`descs` $descs^)? attr-dict `:` qualified(type($barrier)) (`,` qualified(type($descs))^)? ``` The `ttng.tc_gen5_commit` is an asynchronous operation that makes the mbarrier object track the completion of all prior asynchronous tcgen5 operations. Upon completion of all asynchronous operations, the mbarrier arrive operation is performed on the mbarrier with a count of 1. If `descs` are provided, the commit will be multicast across the CTA cluster based on the shared layouts of those descriptors. This should be used when the inputs to the tcgen5 MMA, including scaled-MMA scale inputs, come from TMA descriptors using multicast. Note that the completion mechanisms are guaranteed to occur sequentially in the order the commit operations were issued. This means, for example: ```mlir ttng.tmem_copy ttng.tc_gen5_mma ttng.tc_gen5_commit %barrierA ttng.tc_gen5_commit %barrierB ``` `%barrierA` tracks the completion of the previous TMEM copy and MMA operations, but since the commit groups are sequential, the arrive-on operation on `%barrierA` is guaranteed to be performed before the arrive-on operation on `%barrierB`, even though its commit group is empty. Traits: `AttrSizedOperandSegments`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MBarrierOpInterface`, `PredicatedOpInterface` #### Operands: | Operand | Description | | :-----: | ----------- | | `barrier` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `pred` | 1-bit signless integer | | `descs` | variadic of memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | ### `ttng.tc_gen5_mma` (triton::nvidia_gpu::TCGen5MMAOp) _Block-level op mapping to Tensor Core Gen5 MMA_ Syntax: ``` operation ::= `ttng.tc_gen5_mma` $a `,` $b `,` $d `` custom($acc_dep, type($token)) `,` $useD`,` $pred `` custom($barriers, $barrier_preds) attr-dict `:` qualified(type($a)) `,` qualified(type($b)) `,` qualified(type($d)) (`,` qualified(type($barriers))^)? ``` $d += matrix_multiply($a, $b). If `is_async` is false, the op executes synchronously. The barrier operands must not be present in that case. Otherwise, if a barrier is given, the op will trigger a commit/arrive on it. The result will be safe to read after a barrier wait. If `$two_ctas` is set, the op will execute a matmul across two contiguous CTAs and read the data distributed across the two CTAs. It will also synchronize both CTAs if the op is synchronous. This operation takes and produces an optional token to indicate TMEM read and write on its accumulator operand. When the tokens are present, they can be used to check aliasing and modref on the accumulator memory. The `isUnsigned` attribute is only relevant when performing an integer MMA operation. If true, the integer values are treated as unsigned, otherwise they are treated as signed. Traits: `AttrSizedOperandSegments`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `DotOpInterface`, `MBarrierOpInterface`, `MMAv5OpInterface`, `MemoryEffectOpInterface`, `PredicatedOpInterface` #### Attributes:
AttributeMLIR TypeDescription
is_async::mlir::UnitAttrunit attribute
two_ctas::mlir::UnitAttrunit attribute
multicast::mlir::UnitAttrunit attribute
is_unsigned::mlir::UnitAttrunit attribute
#### Operands: | Operand | Description | | :-----: | ----------- | | `a` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `b` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `d` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `acc_dep` | async token type | | `useD` | 1-bit signless integer | | `pred` | 1-bit signless integer | | `barriers` | variadic of memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `barrier_preds` | variadic of 1-bit signless integer | #### Results: | Result | Description | | :----: | ----------- | | `token` | async token type | ### `ttng.tc_gen5_mma_scaled` (triton::nvidia_gpu::TCGen5MMAScaledOp) _Block-level op mapping to Tensor Core Gen5 MMA_ Syntax: ``` operation ::= `ttng.tc_gen5_mma_scaled` $a `,` $b `,` $d `` custom($acc_dep, type($token)) `,` $a_scale `,` $b_scale `,` $useD `,` $pred `lhs` `=` $a_type `rhs` `=` $b_type `` custom($barriers, $barrier_preds) attr-dict `:` qualified(type($a)) `,` qualified(type($b)) `,` qualified(type($d)) `,` qualified(type($a_scale)) `,` qualified(type($b_scale)) (`,` qualified(type($barriers))^)? ``` $d += matrix_multiply(scale($a, $a_scale), scale($b, $b_scale)) If `$two_ctas` is set, the op will execute a matmul across two contiguous CTAs, read the data distributed across the two CTAs, and synchronize both CTAs if the op is synchronous. If `is_async` is false, the op executes synchronously. The barrier operands must not be present in that case. Otherwise, if a barrier is given, the op will trigger a commit/arrive on it. The result will be safe to read after a barrier wait. This operation takes and produces an optional token to indicate TMEM read and write on its accumulator operand. When the tokens are present, they can be used to check aliasing and modref on the accumulator memory. Traits: `AttrSizedOperandSegments`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `DotOpInterface`, `MBarrierOpInterface`, `MMAv5OpInterface`, `MemoryEffectOpInterface`, `PredicatedOpInterface` #### Attributes:
AttributeMLIR TypeDescription
a_type::mlir::triton::ScaleDotElemTypeAttrallowed 32-bit signless integer cases: 0, 1, 2, 3, 4, 5, 6
b_type::mlir::triton::ScaleDotElemTypeAttrallowed 32-bit signless integer cases: 0, 1, 2, 3, 4, 5, 6
two_ctas::mlir::UnitAttrunit attribute
multicast::mlir::UnitAttrunit attribute
is_async::mlir::UnitAttrunit attribute
#### Operands: | Operand | Description | | :-----: | ----------- | | `a` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `b` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `d` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `acc_dep` | async token type | | `a_scale` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `b_scale` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `useD` | 1-bit signless integer | | `pred` | 1-bit signless integer | | `barriers` | variadic of memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `barrier_preds` | variadic of 1-bit signless integer | #### Results: | Result | Description | | :----: | ----------- | | `token` | async token type | ### `ttng.tensormap_create` (triton::nvidia_gpu::TensormapCreateOp) _Create a new TMA descriptor on device_ Syntax: ``` operation ::= `ttng.tensormap_create` $desc_ptr `,` $global_address `,` `[` $box_dim `]` `,` `[` $global_dim `]` `,` `[` $global_stride `]` `,` `[` $element_stride `]` attr-dict `:` functional-type(operands, results) ``` Traits: `AttrSizedOperandSegments`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MemoryEffectOpInterface (MemoryEffectOpInterface)` Effects: `MemoryEffects::Effect{MemoryEffects::Read on ::mlir::triton::GlobalMemory, MemoryEffects::Write on ::mlir::triton::GlobalMemory}` #### Attributes:
AttributeMLIR TypeDescription
elem_type::mlir::IntegerAttr32-bit signless integer attribute whose value is non-negative whose maximum value is 15
interleave_layout::mlir::IntegerAttr32-bit signless integer attribute whose value is non-negative whose maximum value is 2
swizzle_mode::mlir::IntegerAttr32-bit signless integer attribute whose value is non-negative whose maximum value is 3
fill_mode::mlir::IntegerAttr32-bit signless integer attribute whose value is non-negative whose maximum value is 1
#### Operands: | Operand | Description | | :-----: | ----------- | | `desc_ptr` | Pointer type (`::mlir::triton::PointerType`) in Triton IR type system | | `global_address` | Pointer type (`::mlir::triton::PointerType`) in Triton IR type system | | `box_dim` | variadic of 32-bit signless integer | | `global_dim` | variadic of 32-bit signless integer | | `global_stride` | variadic of 64-bit signless integer | | `element_stride` | variadic of 32-bit signless integer | ### `ttng.tensormap_fenceproxy_acquire` (triton::nvidia_gpu::TensormapFenceproxyAcquireOp) _Acquire fence on a tensormap object_ Syntax: ``` operation ::= `ttng.tensormap_fenceproxy_acquire` $desc_ptr attr-dict `:` qualified(type($desc_ptr)) ``` Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MemoryEffectOpInterface (MemoryEffectOpInterface)` Effects: `MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}` #### Operands: | Operand | Description | | :-----: | ----------- | | `desc_ptr` | Pointer type (`::mlir::triton::PointerType`) in Triton IR type system | ### `ttng.tmem_alloc` (triton::nvidia_gpu::TMEMAllocOp) _Allocate tensor memory_ Syntax: ``` operation ::= `ttng.tmem_alloc` ($src^)? attr-dict `:` functional-type(operands, results) ``` This operation allocates a buffer in tensor memory and returns a descriptor containing the address and a view of the buffer. This is similar to ttg.local_alloc except the buffer is allocated in tensor memory. Explicitly deallocating a buffer is optional; see local_dealloc. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MemoryEffectOpInterface` #### Operands: | Operand | Description | | :-----: | ----------- | | `src` | ranked tensor of floating-point or integer or ptr values | #### Results: | Result | Description | | :----: | ----------- | | `result` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `token` | async token type | ### `ttng.tmem_copy` (triton::nvidia_gpu::TMEMCopyOp) _Initiate an asynchronous copy operation from shared memory to tensor memory._ Syntax: ``` operation ::= `ttng.tmem_copy` $src `,` $dst attr-dict `:` qualified(type(operands)) ``` 2D blocks stored contiguously in SMEM are copied into TMEM as specified by the destination address. This op lowers to the PTX instruction tcgen05.cp. This supports writing to the scales TMEM layout as well as the default TMEM layout. For the default layout, each contiguous TMEM row segment must contain at least 128 bits, the minimum width supported by tcgen05.cp. Currently, the semantics are different when writing to the TMEM scale layout. In the case of the default layout, the copy doesn't change the logical elements between the source and destination memdesc. In the case of the scale layout: Each 32x128b block in SMEM is duplicated over 4 warps and stored into 128 rows and 4 columns of TMEM. The primary use case of this op is to copy blocked scales from SMEM to TMEM. The shape of the input SMEM can be flexibly chosen depending on use cases. In the simplest case (e.g. unit test), the source SMEM can be of shape (32 x num_blocks, 16), and the destination TMEM should be of shape (128, 16 x num_blocks), for copying 8-bit values. For scaled GEMM, rep_m x rep_k copies of a 32x128b block need to be stored in SMEM, where rep_m = BLOCK_M / 128, rep_k = BLOCK_K / scale_vec_size / 4, and scale_vec_size = 32 for MXFP. Conceptually, the SMEM is organized in a high-dimensional layout, (rep_m, rep_k, 32, 4, 4B). Some axes can be flattened into one to reduce the rank of the load. For example, the following patterns are supported: * (rep_m, rep_k * 32 x 4 x 4B), 2D scale load with cp.async * (rep_m, rep_k, 32, 16B), 4D scale load with TMA * (rep_m, rep_k, 32, 4, 4B), 5D scale load with cp.async Since rep_m blocks are not contiguous in SMEM, this axis cannot be flattened into inner ones. In Triton, the TMEM memdesc for blocked scales must be of the following form: * Its shape must be (BLOCK_MN, BLOCK_K / scale_vec_size), representing the logical shape of blocked scales. * It must be attached with `tensor_memory_scales_encoding` to indicate the chunk-based layout and its duplication over 4 warps. In contrast, the source SMEM must be in the explicit chunk-based layout as described above. Therefore, the IR might look like this: %0 = ttng.tmem_alloc : () -> !ttg.memdesc<128x4xi8, #tmem_scales, #ttng.tensor_memory> ttng.tmem_copy %1, %0 : !ttg.memdesc<1x1x32x4x4xi8, #shared1, #smem>, !ttg.memdesc<128x4xi8, #tmem_scales, #ttng.tensor_memory> We interpret the semantics of this copy operation as follows. The chunk-based layout in SMEM implies that the logical shape (BLOCK_MN, BLOCK_K / scale_vec_size) in TMEM is the result of certain reshape and transpose operations. In practice, to take advantage of the native scale layout and the TMEM copy op, users need to do `scales5D.trans(0, 3, 2, 1, 4).reshape(BLOCK_M, BLOCK_K // scale_vec_size)` before feeding scales into dot_scaled. When we use tmem_copy in the IR, such reshape and transpose operations are removed. But the change in the logical shape they have caused in registers is now understood to be incorporated into tmem_copy itself. Ideally, we would lift reshape / transpose done on registers onto the SMEM memdesc, making tmem_copy a straightforward 2D copy operation: (BLOCK_MN, BLOCK_K / scale_vec_size) -> (BLOCK_MN, BLOCK_K / scale_vec_size). In the absence of such operations on memdesc, we resort to implicitly encoding the reshape/transpose semantics in tmem_copy. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` #### Operands: | Operand | Description | | :-----: | ----------- | | `src` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `dst` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | ### `ttng.tmem_load` (triton::nvidia_gpu::TMEMLoadOp) _Load a buffer from tensor memory into a distributed tensor_ Syntax: ``` operation ::= `ttng.tmem_load` $src `` custom($dep, type($token)) attr-dict `:` qualified(type($src)) `->` type($result) (`,` type($red)^)? ``` This is similar to ttg.local_load except the result layout is restricted to only a few possibilities. Therefore, we cannot combine this op with layout conversions like local_load. This operation takes and produces an optional token to indicate TMEM read on its source operand. When the tokens are present, they can be used to check aliasing and modref on the TMEM buffer. Optional reduction modifier: When `redOp` is specified, the load operation additionally performs an element-wise reduction along the N-dimension of the input and produces a second result tensor `red`. For an input of shape `[M, N]`, the reduced result has shape `[M]`, containing one reduced value per "slice" of the N-dimension. Currently restricted to f32 element type. - redOp: Specifies the reduction operation (MIN or MAX) to apply along the N-dimension. When set, the `red` result must be present. - abs: When true, applies absolute value to each element before performing the reduction. Only valid when `redOp` is specified. - NaN: When true, the reduction propagates NaN values (if any input element in a slice is NaN, the corresponding reduced value is NaN). When false, NaN values are ignored during reduction. Only valid when `redOp` is specified. Example: Input in TMEM of shape [M=2, N=4]: [[ 1.0, 3.0, 2.0, 4.0], [-5.0, 1.0, 8.0, 2.0]] With redOp=MAX: result = [[ 1.0, 3.0, 2.0, 4.0], // unchanged [-5.0, 1.0, 8.0, 2.0]] red = [4.0, 8.0] // max along N per row With redOp=MIN, abs=true: red = [1.0, 1.0] // min of |values| per row This operation lowers to hardware-accelerated reduction via the PTX tcgen05.ld.red instruction on supported architectures, e.g. Blackwell Ultra. Traits: `AttrSizedResultSegments`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` #### Attributes:
AttributeMLIR TypeDescription
redOp::mlir::triton::nvidia_gpu::TMEMLoadReduceModifierAttrallowed 32-bit signless integer cases: 1, 2
abs::mlir::BoolAttrbool attribute
NaN::mlir::BoolAttrbool attribute
#### Operands: | Operand | Description | | :-----: | ----------- | | `src` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `dep` | async token type | #### Results: | Result | Description | | :----: | ----------- | | `result` | ranked tensor of floating-point or integer or ptr values | | `token` | async token type | | `red` | ranked tensor of floating-point or integer or ptr values | ### `ttng.tmem_store` (triton::nvidia_gpu::TMEMStoreOp) _Store a distributed tensor into a buffer in tensor memory_ Syntax: ``` operation ::= `ttng.tmem_store` $src `,` $dst `` custom($dep, type($token)) `,` $pred attr-dict `:` type($src) `->` qualified(type($dst)) ``` This is similar to ttg.local_store except the source layout is restricted to only a few possibilities. This operation takes and produces an optional token to indicate TMEM write on its destination operand. When the tokens are present, they can be used to check aliasing and modref on the TMEM buffer. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `PredicatedOpInterface` #### Operands: | Operand | Description | | :-----: | ----------- | | `dst` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `dep` | async token type | | `src` | ranked tensor of floating-point or integer or ptr values | | `pred` | 1-bit signless integer | #### Results: | Result | Description | | :----: | ----------- | | `token` | async token type | ### `ttng.tmem_subslice` (triton::nvidia_gpu::TMEMSubSliceOp) _Take a subslice of a tensor memory allocation_ Syntax: ``` operation ::= `ttng.tmem_subslice` $src attr-dict `:` qualified(type($src)) `->` qualified(type($result)) ``` This operation takes a subslice of a tensor memory allocation and returns a new descriptor containing the address and a view of the subslice. This is similar to ttg.memdesc_subslice and can slice either physical-layout dimension or the leading pipeline dimension of an ordinary descriptor. Scale descriptors support their two physical-layout dimensions. Traits: `AlwaysSpeculatableImplTrait`, `MemDescViewTrait`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `ConditionallySpeculatable`, `NoMemoryEffect (MemoryEffectOpInterface)` Effects: `MemoryEffects::Effect{}` #### Attributes:
AttributeMLIR TypeDescription
offset::mlir::IntegerAttr32-bit signless integer attribute
dim::mlir::IntegerAttr32-bit signless integer attribute
#### Operands: | Operand | Description | | :-----: | ----------- | | `src` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | #### Results: | Result | Description | | :----: | ----------- | | `result` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | ### `ttng.tmem_wait` (triton::nvidia_gpu::TMEMWaitOp) _Wait for tensor memory loads or stores to complete_ Syntax: ``` operation ::= `ttng.tmem_wait` $kind attr-dict ``` Waits for all prior tensor memory loads or stores issued by the executing threads, as selected by `kind`. All threads in each executing warp must execute this operation. It does not synchronize different warps or wait for asynchronous MMA or tensor memory copy operations. Lowers to `tcgen05.wait::ld.sync.aligned` or `tcgen05.wait::st.sync.aligned`. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` #### Attributes:
AttributeMLIR TypeDescription
kind::mlir::triton::nvidia_gpu::TMEMWaitKindAttrtensor memory wait kind
### `ttng.wait_barrier` (triton::nvidia_gpu::WaitBarrierOp) _Wait until the mbarrier phase completes._ Syntax: ``` operation ::= `ttng.wait_barrier` $alloc `,` $phase (`,` $pred^)? (`deps` $deps^)? attr-dict `:` qualified(type($alloc)) (`,` type($deps)^)? ``` Blocks the program progress until the mbarrier object in `alloc` completes its current phase. This lowers a waitloop using PTX instruction mbarrier.try_wait.parity.shared::cta.b64. Accepts an optional list of memory dependencies. If present, it is assumed that any of the dependencies may be accessed until the barrier completes. The barrier behavior is described here: https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-asynchronous-copy-completion-mechanisms Traits: `AttrSizedOperandSegments`, `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `MBarrierOpInterface`, `PredicatedOpInterface` #### Operands: | Operand | Description | | :-----: | ----------- | | `alloc` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `phase` | 32-bit signless integer | | `pred` | 1-bit signless integer | | `deps` | variadic of memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | ### `ttng.warp_group_dot` (triton::nvidia_gpu::WarpGroupDotOp) _Warp group dot_ Syntax: ``` operation ::= `ttng.warp_group_dot` $a`,` $b`,` $c (`,` $useC^)? attr-dict `:` type($a) `*` qualified(type($b)) `->` type($d) ``` $d = matrix_multiply($a, $b) + $c. For docs on InputPrecisionAttr, see TT_DotOp Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `DotOpInterface`, `InferTypeOpInterface`, `MemoryEffectOpInterface` #### Attributes:
AttributeMLIR TypeDescription
inputPrecision::mlir::triton::InputPrecisionAttrallowed 32-bit signless integer cases: 0, 1, 2, 3, 4
maxNumImpreciseAcc::mlir::IntegerAttr32-bit signless integer attribute
isAsync::mlir::BoolAttrbool attribute
#### Operands: | Operand | Description | | :-----: | ----------- | | `a` | TensorOrMemDesc instance | | `b` | memory descriptor type (`::mlir::triton::gpu::MemDescType`) in Triton IR type system | | `c` | ranked tensor of floating-point or integer values | | `useC` | 1-bit signless integer | #### Results: | Result | Description | | :----: | ----------- | | `d` | ranked tensor of floating-point or integer values | ### `ttng.warp_group_dot_wait` (triton::nvidia_gpu::WarpGroupDotWaitOp) _Warp group dot wait_ Syntax: ``` operation ::= `ttng.warp_group_dot_wait` $inputs attr-dict `:` type($inputs) ``` Waits until there are $pendings or fewer outstanding async dot operations. $inputs must be the tensors corresponding to the async dot ops that we're waiting on. For example, if there are N pending async dot ops and we wait until one remains pending, then $inputs must include the results of the first N - 1 dot ops. The `warpGroupLocal` attribute specifies that we only wait for the local warp group, and if num_warps > 4 then shared memory inputs may still be in use by other warp groups. This is useful when only the result register values need to be ready and the shared memory inputs will not be overwritten. Traits: `VerifyMemDescLayoutsTrait`, `VerifyTensorLayoutsTrait` Interfaces: `InferTypeOpInterface` #### Attributes:
AttributeMLIR TypeDescription
pendings::mlir::IntegerAttr32-bit signless integer attribute
warpGroupLocal::mlir::UnitAttrunit attribute
#### Operands: | Operand | Description | | :-----: | ----------- | | `inputs` | variadic of TensorOrMemDesc instance | #### Results: | Result | Description | | :----: | ----------- | | `outputs` | variadic of TensorOrMemDesc instance |