TritonNvidiaGPUOps¶

`ttng.arrive_barrier` (triton::nvidia_gpu::ArriveBarrierOp)¶

对 mbarrier 执行到达操作

语法

operation ::= `ttng.arrive_barrier` $alloc `,` $count (`,` $pred^)? attr-dict `:` qualified(type($alloc))

ttng.arrive_barrier 操作在共享内存中的 mbarrier 对象上执行“到达”操作。该操作需要一个至少为 1 的 count 属性，并将 mbarrier 的待到达计数减少指定的数量。

该操作接受一个可选的谓词。

示例

ttng.arrive_barrier %barrier, 2 : !ttg.memdesc<1xi64, #shared, #smem, mutable>
ttng.arrive_barrier %barrier, 1, %pred : !ttg.memdesc<1xi64, #shared, #smem, mutable>

特性： VerifyTensorLayoutsTrait

属性：¶

属性	MLIR 类型	描述
`count`	::mlir::IntegerAttr	32位无符号整数属性

操作数:¶

操作数	描述
`alloc`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`pred`	1位无符号整数

`ttng.async_copy_mbarrier_arrive` (triton::nvidia_gpu::AsyncCopyMbarrierArriveOp)¶

一旦所有先前发出的副本完成，就到达 mbarrier

语法

operation ::= `ttng.async_copy_mbarrier_arrive` $barrier attr-dict `:` qualified(type($barrier))

特性： VerifyTensorLayoutsTrait

属性：¶

属性	MLIR 类型	描述
`noIncrement`	::mlir::UnitAttr	单元属性

操作数:¶

操作数	描述
`barrier`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)

`ttng.async_tma_copy_global_to_local` (triton::nvidia_gpu::AsyncTMACopyGlobalToLocalOp)¶

根据描述符将数据从全局内存异步复制到本地内存

语法

operation ::= `ttng.async_tma_copy_global_to_local` $desc `[` $coord `]` $result `,` $barrier `,` $pred
              oilist(`cacheModifier` `=` $cache | `evictionPolicy` `=` $evict)
              attr-dict `:` qualified(type($desc)) `,` qualified(type($barrier)) `->` qualified(type($result))

此操作将数据从全局内存异步复制到本地内存。这类似于 tt.load，不同之处在于数据被复制到由内存描述符指向的本地内存，而不是分布式张量。复制的数据取决于 desc 指向的全局内存描述符。

特性： VerifyTensorLayoutsTrait

属性 (Attributes):¶

属性	MLIR 类型	描述
`cache`	::mlir::triton::CacheModifierAttr	允许的32位无符号整数情况：1, 2, 3, 4, 5, 6, 7
`evict`	::mlir::triton::EvictionPolicyAttr	允许的32位无符号整数情况：1, 2, 3
`isVolatile`	::mlir::BoolAttr	布尔属性

操作数:¶

操作数	描述
`desc`	张量描述符类型（`::mlir::triton::TensorDescType`）在 Triton IR 类型系统中
`coord`	32 位无符号整数的可变参数
`barrier`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`result`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`pred`	1位无符号整数

`ttng.async_tma_copy_local_to_global` (triton::nvidia_gpu::AsyncTMACopyLocalToGlobalOp)¶

根据描述符将数据从本地内存异步复制到全局内存

语法

operation ::= `ttng.async_tma_copy_local_to_global` $desc `[` $coord `]` $src
              attr-dict `:` qualified(type($desc)) `,` qualified(type($src))

此操作将数据从本地内存异步复制到全局内存。这类似于 tt.store，不同之处在于数据从由内存描述符指向的本地内存复制，而不是分布式张量。复制的数据取决于 desc 指向的全局内存描述符。

特性： VerifyTensorLayoutsTrait

操作数:¶

操作数	描述
`desc`	张量描述符类型（`::mlir::triton::TensorDescType`）在 Triton IR 类型系统中
`coord`	32 位无符号整数的可变参数
`src`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)

`ttng.async_tma_gather` (triton::nvidia_gpu::AsyncTMAGatherOp)¶

根据描述符从全局内存异步收集数据到本地内存

语法

operation ::= `ttng.async_tma_gather` $desc `[` $x_offsets `,` $y_offset `]` $result `,` $barrier `,` $pred
              attr-dict `:` type(operands)

此操作异步地将全局内存矩阵中的多行数据收集到本地内存中。这类似于 async_tma_copy_global_to_local，不同之处在于每一行都是独立索引的。

特性： VerifyTensorLayoutsTrait

操作数 (Operands):¶

操作数	描述
`desc`	张量描述符类型（`::mlir::triton::TensorDescType`）在 Triton IR 类型系统中
`x_offsets`	32 位无符号整数值的秩张量
`y_offset`	32位无符号整数
`barrier`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`result`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`pred`	1位无符号整数

`ttng.async_tma_reduce` (triton::nvidia_gpu::AsyncTMAReduceOp)¶

基于 TMA 描述符在 gmem 中归约结果

语法

operation ::= `ttng.async_tma_reduce` $kind `,` $desc `[` $coord `]` $src
              attr-dict `:` qualified(type($desc)) `,` qualified(type($src))

此操作将数据从本地内存异步复制到全局内存，并原子地执行指定的归约类型。原子性以单个元素为粒度，并且只隐含松散语义。

特性： VerifyTensorLayoutsTrait

接口：MemoryEffectOpInterface (MemoryEffectOpInterface)

效果： MemoryEffects::Effect{MemoryEffects::Read on ::mlir::triton::GlobalMemory, MemoryEffects::Write on ::mlir::triton::GlobalMemory}

属性:¶

属性	MLIR 类型	描述
`kind`	::mlir::triton::DescriptorReduceKindAttr	允许的 32 位无符号整数情况：1, 2, 3, 4, 5, 6, 7, 8

操作数:¶

操作数	描述
`desc`	张量描述符类型（`::mlir::triton::TensorDescType`）在 Triton IR 类型系统中
`coord`	32 位无符号整数的可变参数
`src`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)

`ttng.async_tma_scatter` (triton::nvidia_gpu::AsyncTMAScatterOp)¶

基于描述符异步地将数据从本地内存散布到全局内存中

语法

operation ::= `ttng.async_tma_scatter` $desc `[` $x_offsets `,` $y_offset `]` $src
              attr-dict `:` type(operands)

ttng.async_tma_scatter 操作将多个单独索引的数据行从本地内存异步散布到全局内存中。该操作在共享内存中散布一个 2D 张量，该张量由核心张量瓦片 nvmma_shared 布局，散布到全局内存中给定 y 偏移量的单独索引行中。

特性： VerifyTensorLayoutsTrait

操作数:¶

操作数	描述
`desc`	张量描述符类型（`::mlir::triton::TensorDescType`）在 Triton IR 类型系统中
`x_offsets`	32 位无符号整数值的秩张量
`y_offset`	32位无符号整数
`src`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)

`ttng.barrier_expect` (triton::nvidia_gpu::BarrierExpectOp)¶

向屏障发出信号，表示预期要复制的字节数。

语法

operation ::= `ttng.barrier_expect` $alloc `,` $size attr-dict `,` $pred `:` qualified(type($alloc))

这向屏障发出信号，表示预期要复制 size 字节。相关的屏障等待将阻塞，直到预期的字节数被复制完毕。

特性： VerifyTensorLayoutsTrait

属性：¶

属性	MLIR 类型	描述
`size`	::mlir::IntegerAttr	32位无符号整数属性

操作数：¶

操作数	描述
`alloc`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`pred`	1位无符号整数

`ttng.cluster_arrive` (triton::nvidia_gpu::ClusterArriveOp)¶

语法

operation ::= `ttng.cluster_arrive` attr-dict

特性： VerifyTensorLayoutsTrait

属性：¶

属性	MLIR 类型	描述
`relaxed`	::mlir::IntegerAttr	1 位无符号整数属性

`ttng.cluster_wait` (triton::nvidia_gpu::ClusterWaitOp)¶

语法

operation ::= `ttng.cluster_wait` attr-dict

特性： VerifyTensorLayoutsTrait

`ttng.fence_async_shared` (triton::nvidia_gpu::FenceAsyncSharedOp)¶

Fence proxy async

语法

operation ::= `ttng.fence_async_shared` attr-dict

特性： VerifyTensorLayoutsTrait

属性 (Attributes):¶

属性	MLIR 类型	描述
`bCluster`	::mlir::BoolAttr	布尔属性

`ttng.init_barrier` (triton::nvidia_gpu::InitBarrierOp)¶

在给定的共享内存分配中初始化一个屏障。

语法

operation ::= `ttng.init_barrier` $alloc `,` $count attr-dict `:` qualified(type($alloc))

使用 mbarrier 信息初始化一个共享内存分配。 alloc 是共享内存分配的描述符。 count 是屏障预期的到达次数。

这会降级为 PTX mbarrier.init.shared::cta.b64。

特性： VerifyTensorLayoutsTrait

属性：¶

属性	MLIR 类型	描述
`count`	::mlir::IntegerAttr	32位无符号整数属性

操作数：¶

操作数	描述
`alloc`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)

`ttng.inval_barrier` (triton::nvidia_gpu::InvalBarrierOp)¶

使屏障分配无效。

语法

operation ::= `ttng.inval_barrier` $alloc attr-dict `:` qualified(type($alloc))

使屏障分配无效，以便可以重新使用。根据 PTX 规范，在重新使用 mbarrier 使用的内存之前必须执行此操作。

https://docs.nvda.net.cn/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-inval

特性： VerifyTensorLayoutsTrait

操作数 (Operands):¶

操作数	描述
`alloc`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)

`ttng.reinterpret_tensor_descriptor` (triton::nvidia_gpu::ReinterpretTensorDescOp)¶

将指针重新解释为张量描述符

语法

operation ::= `ttng.reinterpret_tensor_descriptor` $rawDesc attr-dict `:` qualified(type($rawDesc))  `to` qualified(type($result))

此 Op 的存在是为了帮助从无类型的原始 TMA 对象过渡到有类型的张量描述符对象。理想情况下，一旦 API 完全充实，我们就可以移除它。

特性： AlwaysSpeculatableImplTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数:¶

操作数	描述
`rawDesc`	ptr

结果:¶

结果	描述
`result`	张量描述符类型（`::mlir::triton::TensorDescType`）在 Triton IR 类型系统中

`ttng.tc_gen5_commit` (triton::nvidia_gpu::TCGen5CommitOp)¶

让 mbarrier 跟踪所有先前异步 tcgen5 操作的完成情况

语法

operation ::= `ttng.tc_gen5_commit` $barrier (`,` $pred^)? attr-dict `:` qualified(type($barrier))

ttng.tc_gen5_commit 是一个异步操作，它使 mbarrier 对象跟踪所有先前异步 tcgen5 操作的完成情况。在所有异步操作完成后，将对 mbarrier 执行计数为 1 的 mbarrier 到达操作。

如果设置了 two_ctas，那么 mbarrier 也会跟踪所有先前用 two_ctas 设置启动的操作。否则，它会跟踪所有先前没有设置 two_ctas 启动的操作。

请注意，完成机制保证按提交操作发出的顺序顺序发生。这意味着，例如

ttng.tmem_copy
ttng.tc_gen5_mma
ttng.tc_gen5_commit %barrierA
ttng.tc_gen5_commit %barrierB

%barrierA 跟踪先前 TMEM 复制和 MMA 操作的完成情况，但由于提交组是顺序的，因此保证在 %barrierB 上的到达操作之前执行 %barrierA 上的到达操作，即使其提交组为空。

特性： VerifyTensorLayoutsTrait

属性：¶

属性	MLIR 类型	描述
`two_ctas`	::mlir::UnitAttr	单元属性

操作数 (Operands):¶

操作数	描述
`barrier`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`pred`	1位无符号整数

`ttng.tc_gen5_mma` (triton::nvidia_gpu::TCGen5MMAOp)¶

映射到 tensorcore gen5 mma 的块级操作

语法

operation ::= `ttng.tc_gen5_mma` $a `,` $b `,` $d `` custom<Token>($acc_dep, type($token)) `,` $useD`,`
              $pred `` custom<BarriersAndPreds>($barriers, $barrier_preds)
              attr-dict `:` qualified(type($a)) `,` qualified(type($b)) `,`
              qualified(type($d)) (`,` qualified(type($barriers))^)?

$d += matrix_multiply($a, $b)。如果 is_async 为 false，则该操作同步执行。在这种情况下，屏障操作数必须不存在。否则，如果给定了屏障，该操作将触发对其的提交/到达。在屏障等待之后，结果将是安全的。如果设置了 $two_ctas，该操作将跨越两个连续的 CTA 执行矩阵乘法，它将读取分布在两个 CTA 上的数据，并在操作是同步的情况下同步两个 CTA。

此操作接受并生成一个可选的令牌，以指示对其累加器操作数的 TMEM 读写。当令牌存在时，它们可用于检查累加器内存上的别名和 modref。

特性： AttrSizedOperandSegments, VerifyTensorLayoutsTrait

接口： DotOpInterface, MMAv5OpInterface, MemoryEffectOpInterface

属性：¶

属性	MLIR 类型	描述
`is_async`	::mlir::UnitAttr	单元属性
`two_ctas`	::mlir::UnitAttr	单元属性

操作数：¶

操作数	描述
`a`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`b`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`d`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`acc_dep`	异步令牌类型
`useD`	1位无符号整数
`pred`	1位无符号整数
`barriers`	Triton IR 类型系统中内存描述符类型 (`::mlir::triton::gpu::MemDescType`) 的可变参数
`barrier_preds`	1 位无符号整数的可变参数

结果：¶

结果	描述
`token`	异步令牌类型

`ttng.tc_gen5_mma_scaled` (triton::nvidia_gpu::TCGen5MMAScaledOp)¶

映射到 tensorcore gen5 mma 的块级操作

语法

operation ::= `ttng.tc_gen5_mma_scaled` $a `,` $b `,` $d `` custom<Token>($acc_dep, type($token)) `,` $a_scale `,`
              $b_scale `,` $useD `,` $pred `lhs` `=` $a_type `rhs` `=` $b_type
              `` custom<BarriersAndPreds>($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($lhs, $lhs_scale), scale(rlhs, $rhs_scale)) 如果 is_async 为 false，则该操作同步执行。在这种情况下，屏障操作数必须不存在。否则，如果给定了屏障，该操作将触发对其的提交/到达。在屏障等待之后，结果将是安全的。

此操作接受并生成一个可选的令牌，以指示对其累加器操作数的 TMEM 读写。当令牌存在时，它们可用于检查累加器内存上的别名和 modref。

特性： AttrSizedOperandSegments, VerifyTensorLayoutsTrait

接口： DotOpInterface, MMAv5OpInterface, MemoryEffectOpInterface

属性：¶

属性	MLIR 类型	描述
`a_type`	::mlir::triton::ScaleDotElemTypeAttr	允许的 32 位无符号整数情况：0, 1, 2, 3, 4, 5, 6
`b_type`	::mlir::triton::ScaleDotElemTypeAttr	允许的 32 位无符号整数情况：0, 1, 2, 3, 4, 5, 6
`is_async`	::mlir::UnitAttr	单元属性

操作数:¶

操作数	描述
`a`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`b`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`d`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`acc_dep`	异步令牌类型
`a_scale`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`b_scale`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`useD`	1位无符号整数
`pred`	1位无符号整数
`barriers`	Triton IR 类型系统中内存描述符类型 (`::mlir::triton::gpu::MemDescType`) 的可变参数
`barrier_preds`	1 位无符号整数的可变参数

结果:¶

结果	描述
`token`	异步令牌类型

`ttng.async_tma_store_wait` (triton::nvidia_gpu::TMAStoreWaitOp)¶

等到所有输入都读取完毕。

语法

operation ::= `ttng.async_tma_store_wait` attr-dict

等到所有相关联的存储操作的读取操作都完成。在可以写入共享内存之前，这是必需的。

特性： VerifyTensorLayoutsTrait

属性：¶

属性	MLIR 类型	描述
`pendings`	::mlir::IntegerAttr	32位无符号整数属性

`ttng.tmem_alloc` (triton::nvidia_gpu::TMEMAllocOp)¶

分配张量内存

语法

operation ::= `ttng.tmem_alloc` ($src^)? attr-dict `:` functional-type(operands, results)

此操作在张量内存中分配缓冲区，并返回一个包含地址和缓冲区视图的描述符。这类似于 ttg.local_alloc，不同之处在于缓冲区是在张量内存中分配的。

显式释放缓冲区是可选的；请参阅 local_dealloc。

特性： VerifyTensorLayoutsTrait

接口： MemoryEffectOpInterface

操作数 (Operands):¶

操作数	描述
`src`	浮点数、整数或 ptr 值的秩张量

结果 (Results):¶

结果	描述
`result`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`token`	异步令牌类型

`ttng.tmem_copy` (triton::nvidia_gpu::TMEMCopyOp)¶

启动从共享内存到张量内存的异步复制操作。

语法

operation ::= `ttng.tmem_copy` $src `,` $dst (`,` $barrier^)? attr-dict `:` qualified(type(operands))

连续存储在 SMEM 中的 2D 块被复制到 TMEM 中，由目标地址指定。可以通过等待可选的屏障来观察复制的完成情况。如果此操作与 MMA 操作一起使用，一个屏障可以用来等待复制和 MMA。我们不需要在 MMA 之前等待复制完成，因为 tcgen05.cp 后跟 tcgen05.mma 保证按此顺序执行。

此操作会降级为 PTX 指令 tcgen05.cp。这支持写入 scales tmem 布局以及默认 tmem 布局。目前，写入 tmem scale 布局的语义是不同的。

在默认布局的情况下，复制不会改变源和目标 memdesc 之间的逻辑元素。

在 scale 布局的情况下：SMEM 中的每个 32x128b 块在 4 个 warp 上重复，并存储到 TMEM 的 128 行和 4 列中。此操作的主要用例是将分块的 scale 从 SMEM 复制到 TMEM。

输入 SMEM 的形状可以根据用例灵活选择。在最简单的情况下（例如，单元测试），源 SMEM 的形状可以是 (32 x num_blocks, 16)，目标 TMEM 的形状应该是 (128, 16 x num_blocks)，用于复制 8 位值。对于缩放 GEMM，需要在 SMEM 中存储 32x128b 块的 rep_m x rep_k 个副本，其中 rep_m = BLOCK_M / 128，rep_k = BLOCK_K / scale_vec_size / 4，对于 MXFP，scale_vec_size = 32。从概念上讲，SMEM 是以高维布局组织的，(rep_m, rep_k, 32, 4, 4B)。某些轴可以展平为一个，以减少加载的秩。例如，支持以下模式

(rep_m, rep_k * 32 x 4 x 4B)，使用 cp.async 进行 2D 缩放加载
(rep_m, rep_k, 32, 16B)，使用 TMA 进行 4D 缩放加载
(rep_m, rep_k, 32, 4, 4B)，使用 cp.async 进行 5D 缩放加载，因为 rep_m 块在 SMEM 中不连续，所以此轴不能展平到内轴。

在 Triton 中，用于分块 scale 的 TMEM memdesc 必须具有以下形式

其形状必须是 (BLOCK_MN, BLOCK_K / scale_vec_size)，表示分块 scale 的逻辑形状。
它必须附加 tensor_memory_scales_encoding 以指示基于块的布局及其在 4 个 warp 上的重复。

相比之下，源 SMEM 必须是如上所述的显式基于块的布局。因此 IR 可能如下所示

ttng.tmem_copy %1, %0 : (!ttg.memdesc<1x1x32x4x4xi8, #shared1, #smem>, !ttg.memdesc<128x4xi8, #tmem_scales, #ttng.tensor_memory>) -> ()

我们将此复制操作的语义解释如下。SMEM 中的基于块的布局意味着 TMEM 中的逻辑形状 (BLOCK_MN, BLOCK_K / scale_vec_size) 是某些重塑和转置操作的结果。在实践中，为了利用原生 scale 布局和 TMEM 复制操作，用户需要在将 scale 输入 dot_scaled 之前执行 scales5D.trans(0, 3, 2, 1, 4).reshape(BLOCK_M, BLOCK_K // scale_vec_size)。当我们在 IR 中使用 tmem_copy 时，会移除此类重塑和转置操作。但它们对寄存器造成的逻辑形状变化现在被理解为已并入 tmem_copy 本身。理想情况下，我们会将对寄存器执行的 reshape / transpose 提升到 SMEM memdesc 上，使 tmem_copy 成为一个直接的 2D 复制操作：(BLOCK_MN, BLOCK_K / scale_vec_size) -> (BLOCK_MN, BLOCK_K / scale_vec_size)。在没有对 memdesc 进行此类操作的情况下，我们诉诸于在 tmem_copy 中隐式编码 reshape/transpose 语义。

特性： VerifyTensorLayoutsTrait

操作数：¶

操作数	描述
`src`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`dst`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`barrier`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)

`ttng.tmem_load` (triton::nvidia_gpu::TMEMLoadOp)¶

将缓冲区从张量内存加载到分布式张量中

语法

operation ::= `ttng.tmem_load` $src `` custom<Token>($dep, type($token))
              attr-dict `:` qualified(type($src)) `->` type($result)

这类似于 ttg.local_load，但结果布局仅限于少数几种可能性。因此，我们不能像 local_load 那样将此操作与任何转换布局结合起来。

此操作接受并生成一个可选的令牌，以指示对其源操作数的 TMEM 读取。当令牌存在时，它们可用于检查 TMEM 缓冲区上的别名和 modref。

特性： VerifyTensorLayoutsTrait

操作数 (Operands):¶

操作数	描述
`src`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`dep`	异步令牌类型

结果：¶

结果	描述
`result`	浮点数、整数或 ptr 值的秩张量
`token`	异步令牌类型

`ttng.tmem_store` (triton::nvidia_gpu::TMEMStoreOp)¶

将分布式张量存储到张量内存的缓冲区中

语法

operation ::= `ttng.tmem_store` $src `,` $dst `` custom<Token>($dep, type($token)) `,` $pred
              attr-dict `:` type($src) `->` qualified(type($dst))

这类似于 ttg.local_store，但源布局仅限于少数几种可能性。

此操作接受并生成一个可选的令牌，以指示对其源操作数的 TMEM 写入。当令牌存在时，它们可用于检查 TMEM 缓冲区上的别名和 modref。

特性： VerifyTensorLayoutsTrait

操作数：¶

操作数	描述
`dst`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`dep`	异步令牌类型
`src`	浮点数、整数或 ptr 值的秩张量
`pred`	1位无符号整数

结果：¶

结果	描述
`token`	异步令牌类型

`ttng.tmem_subslice` (triton::nvidia_gpu::TMEMSubSliceOp)¶

获取张量内存分配的子切片

语法

operation ::= `ttng.tmem_subslice` $src attr-dict `:` qualified(type($src)) `->` qualified(type($result))

此操作获取张量内存分配的子切片，并返回一个包含地址和子切片视图的新描述符。这类似于 ttg.memdesc_subslice，但我们只能沿 2D memdesc 的内部维度进行切片，因为这是我们唯一可以为 TMem 做的。

特性： AlwaysSpeculatableImplTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性：¶

属性	MLIR 类型	描述
`N`	::mlir::IntegerAttr	32位无符号整数属性

操作数：¶

操作数	描述
`src`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)

结果：¶

结果	描述
`result`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)

`ttng.tensormap_create` (triton::nvidia_gpu::TensormapCreateOp)¶

在设备上创建新的 TMA 描述符

语法

operation ::= `ttng.tensormap_create` $desc_ptr `,` $global_address `,`
              `[` $box_dim `]` `,`
              `[` $global_dim `]` `,`
              `[` $global_stride `]` `,`
              `[` $element_stride `]`
              attr-dict `:` functional-type(operands, results)

特性： AttrSizedOperandSegments, VerifyTensorLayoutsTrait

接口：MemoryEffectOpInterface (MemoryEffectOpInterface)

效果： MemoryEffects::Effect{MemoryEffects::Read on ::mlir::triton::GlobalMemory, MemoryEffects::Write on ::mlir::triton::GlobalMemory}

属性：¶

属性	MLIR 类型	描述
`elem_type`	::mlir::IntegerAttr	值为非负且最大值为 15 的 32 位无符号整数属性
`interleave_layout`	::mlir::IntegerAttr	值为非负且最大值为 2 的 32 位无符号整数属性
`swizzle_mode`	::mlir::IntegerAttr	值为非负且最大值为 3 的 32 位无符号整数属性
`fill_mode`	::mlir::IntegerAttr	值为非负且最大值为 1 的 32 位无符号整数属性

操作数：¶

操作数	描述
`desc_ptr`	Triton IR 类型系统中的指针类型 (`::mlir::triton::PointerType`)
`global_address`	Triton IR 类型系统中的指针类型 (`::mlir::triton::PointerType`)
`box_dim`	32 位无符号整数的可变参数
`global_dim`	32 位无符号整数的可变参数
`global_stride`	64 位无符号整数的可变参数
`element_stride`	32 位无符号整数的可变参数

`ttng.tensormap_fenceproxy_acquire` (triton::nvidia_gpu::TensormapFenceproxyAcquireOp)¶

在 tensormap 对象上获取 fence

语法

operation ::= `ttng.tensormap_fenceproxy_acquire` $desc_ptr attr-dict `:` qualified(type($desc_ptr))

特性： VerifyTensorLayoutsTrait

接口：MemoryEffectOpInterface (MemoryEffectOpInterface)

效果：MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}

操作数：¶

操作数	描述
`desc_ptr`	Triton IR 类型系统中的指针类型 (`::mlir::triton::PointerType`)

`ttng.wait_barrier` (triton::nvidia_gpu::WaitBarrierOp)¶

等到 mbarrier 阶段完成。

语法

operation ::= `ttng.wait_barrier` $alloc `,` $phase (`,` $pred^)? (`deps` $deps^)?
              attr-dict `:` qualified(type($alloc)) (`,` type($deps)^)?

阻塞程序进程，直到 alloc 中的 mbarrier 对象完成其当前阶段。

这会使用 PTX 指令 mbarrier.try_wait.parity.shared.b64 降级等待循环。

接受可选的内存列表。如果存在，则假定在屏障完成之前可能会访问任何依赖项。

屏障行为在此处描述：https://docs.nvda.net.cn/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-asynchronous-copy-completion-mechanisms

特性： AttrSizedOperandSegments, VerifyTensorLayoutsTrait

操作数：¶

操作数	描述
`alloc`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`phase`	32位无符号整数
`pred`	1位无符号整数
`deps`	Triton IR 类型系统中内存描述符类型 (`::mlir::triton::gpu::MemDescType`) 的可变参数

`ttng.warp_group_dot` (triton::nvidia_gpu::WarpGroupDotOp)¶

Warp group 点积

语法

operation ::= `ttng.warp_group_dot` $a`,` $b`,` $c (`,` $useC^)? attr-dict
              `:` type($a) `*` qualified(type($b)) `->` type($d)

$d = matrix_multiply($a, $b) + $c。有关 InputPrecisionAttr 的文档，请参见 TT_DotOp

特性： VerifyTensorLayoutsTrait

接口： DotOpInterface, InferTypeOpInterface, MemoryEffectOpInterface

属性：¶

属性	MLIR 类型	描述
`inputPrecision`	::mlir::triton::InputPrecisionAttr	允许的 32 位无符号整数情况：0, 1, 2, 3, 4
`maxNumImpreciseAcc`	::mlir::IntegerAttr	32位无符号整数属性
`isAsync`	::mlir::BoolAttr	布尔属性

操作数：¶

操作数	描述
`a`	TensorOrMemDesc 实例
`b`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`c`	浮点或整数值的秩张量
`useC`	1位无符号整数

结果：¶

结果	描述
`d`	浮点或整数值的秩张量

`ttng.warp_group_dot_wait` (triton::nvidia_gpu::WarpGroupDotWaitOp)¶

Warp group 点积等待

语法

operation ::= `ttng.warp_group_dot_wait` $inputs attr-dict `:` type($inputs)

等到有 $pendings 个或更少的未完成异步点积操作。

$inputs 必须是与我们正在等待的异步点积操作相对应的张量。例如，如果有 N 个待处理的异步点积操作，我们调用 warp_group_dot_wait 1，那么 $inputs 必须是第一个点积操作的结果。

特性： VerifyTensorLayoutsTrait

接口 (Interfaces): InferTypeOpInterface

属性：¶

属性	MLIR 类型	描述
`pendings`	::mlir::IntegerAttr	32位无符号整数属性

操作数：¶

操作数	描述
`inputs`	TensorOrMemDesc 实例的可变参数

结果：¶

结果	描述
`outputs`	TensorOrMemDesc 实例的可变参数

TritonNvidiaGPUOps¶

ttng.arrive_barrier (triton::nvidia_gpu::ArriveBarrierOp)¶

属性：¶

操作数:¶

ttng.async_copy_mbarrier_arrive (triton::nvidia_gpu::AsyncCopyMbarrierArriveOp)¶

属性：¶

操作数:¶

ttng.async_tma_copy_global_to_local (triton::nvidia_gpu::AsyncTMACopyGlobalToLocalOp)¶

属性 (Attributes):¶

操作数:¶

ttng.async_tma_copy_local_to_global (triton::nvidia_gpu::AsyncTMACopyLocalToGlobalOp)¶

操作数:¶

ttng.async_tma_gather (triton::nvidia_gpu::AsyncTMAGatherOp)¶

操作数 (Operands):¶

ttng.async_tma_reduce (triton::nvidia_gpu::AsyncTMAReduceOp)¶

属性:¶

操作数:¶

ttng.async_tma_scatter (triton::nvidia_gpu::AsyncTMAScatterOp)¶

操作数:¶

ttng.barrier_expect (triton::nvidia_gpu::BarrierExpectOp)¶

属性：¶

操作数：¶

ttng.cluster_arrive (triton::nvidia_gpu::ClusterArriveOp)¶

属性：¶

ttng.cluster_wait (triton::nvidia_gpu::ClusterWaitOp)¶

ttng.fence_async_shared (triton::nvidia_gpu::FenceAsyncSharedOp)¶

属性 (Attributes):¶

ttng.init_barrier (triton::nvidia_gpu::InitBarrierOp)¶

属性：¶

操作数：¶

ttng.inval_barrier (triton::nvidia_gpu::InvalBarrierOp)¶

操作数 (Operands):¶

ttng.reinterpret_tensor_descriptor (triton::nvidia_gpu::ReinterpretTensorDescOp)¶

操作数:¶

结果:¶

ttng.tc_gen5_commit (triton::nvidia_gpu::TCGen5CommitOp)¶

属性：¶

操作数 (Operands):¶

ttng.tc_gen5_mma (triton::nvidia_gpu::TCGen5MMAOp)¶

属性：¶

操作数：¶

结果：¶

ttng.tc_gen5_mma_scaled (triton::nvidia_gpu::TCGen5MMAScaledOp)¶

属性：¶

操作数:¶

结果:¶

ttng.async_tma_store_wait (triton::nvidia_gpu::TMAStoreWaitOp)¶

属性：¶

ttng.tmem_alloc (triton::nvidia_gpu::TMEMAllocOp)¶

操作数 (Operands):¶

结果 (Results):¶

ttng.tmem_copy (triton::nvidia_gpu::TMEMCopyOp)¶

操作数：¶

ttng.tmem_load (triton::nvidia_gpu::TMEMLoadOp)¶

操作数 (Operands):¶

结果：¶

ttng.tmem_store (triton::nvidia_gpu::TMEMStoreOp)¶

操作数：¶

结果：¶

ttng.tmem_subslice (triton::nvidia_gpu::TMEMSubSliceOp)¶

属性：¶

操作数：¶

结果：¶

ttng.tensormap_create (triton::nvidia_gpu::TensormapCreateOp)¶

属性：¶

操作数：¶

ttng.tensormap_fenceproxy_acquire (triton::nvidia_gpu::TensormapFenceproxyAcquireOp)¶

操作数：¶

ttng.wait_barrier (triton::nvidia_gpu::WaitBarrierOp)¶

操作数：¶

ttng.warp_group_dot (triton::nvidia_gpu::WarpGroupDotOp)¶

属性：¶

操作数：¶

结果：¶

ttng.warp_group_dot_wait (triton::nvidia_gpu::WarpGroupDotWaitOp)¶

属性：¶

操作数：¶

结果：¶

`ttng.arrive_barrier` (triton::nvidia_gpu::ArriveBarrierOp)¶

`ttng.async_copy_mbarrier_arrive` (triton::nvidia_gpu::AsyncCopyMbarrierArriveOp)¶

`ttng.async_tma_copy_global_to_local` (triton::nvidia_gpu::AsyncTMACopyGlobalToLocalOp)¶

`ttng.async_tma_copy_local_to_global` (triton::nvidia_gpu::AsyncTMACopyLocalToGlobalOp)¶

`ttng.async_tma_gather` (triton::nvidia_gpu::AsyncTMAGatherOp)¶

`ttng.async_tma_reduce` (triton::nvidia_gpu::AsyncTMAReduceOp)¶

`ttng.async_tma_scatter` (triton::nvidia_gpu::AsyncTMAScatterOp)¶

`ttng.barrier_expect` (triton::nvidia_gpu::BarrierExpectOp)¶

`ttng.cluster_arrive` (triton::nvidia_gpu::ClusterArriveOp)¶

`ttng.cluster_wait` (triton::nvidia_gpu::ClusterWaitOp)¶

`ttng.fence_async_shared` (triton::nvidia_gpu::FenceAsyncSharedOp)¶

`ttng.init_barrier` (triton::nvidia_gpu::InitBarrierOp)¶

`ttng.inval_barrier` (triton::nvidia_gpu::InvalBarrierOp)¶

`ttng.reinterpret_tensor_descriptor` (triton::nvidia_gpu::ReinterpretTensorDescOp)¶

`ttng.tc_gen5_commit` (triton::nvidia_gpu::TCGen5CommitOp)¶

`ttng.tc_gen5_mma` (triton::nvidia_gpu::TCGen5MMAOp)¶

`ttng.tc_gen5_mma_scaled` (triton::nvidia_gpu::TCGen5MMAScaledOp)¶

`ttng.async_tma_store_wait` (triton::nvidia_gpu::TMAStoreWaitOp)¶

`ttng.tmem_alloc` (triton::nvidia_gpu::TMEMAllocOp)¶

`ttng.tmem_copy` (triton::nvidia_gpu::TMEMCopyOp)¶

`ttng.tmem_load` (triton::nvidia_gpu::TMEMLoadOp)¶

`ttng.tmem_store` (triton::nvidia_gpu::TMEMStoreOp)¶

`ttng.tmem_subslice` (triton::nvidia_gpu::TMEMSubSliceOp)¶

`ttng.tensormap_create` (triton::nvidia_gpu::TensormapCreateOp)¶

`ttng.tensormap_fenceproxy_acquire` (triton::nvidia_gpu::TensormapFenceproxyAcquireOp)¶

`ttng.wait_barrier` (triton::nvidia_gpu::WaitBarrierOp)¶

`ttng.warp_group_dot` (triton::nvidia_gpu::WarpGroupDotOp)¶

`ttng.warp_group_dot_wait` (triton::nvidia_gpu::WarpGroupDotWaitOp)¶