TritonNvidiaGPUOps

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

对 mbarrier 执行到达 (arrive) 操作

语法

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>

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: MBarrierOpInterface, PredicatedOpInterface

属性：

属性	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))

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: MBarrierOpInterface

属性：

属性	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 `]` (`offsets` `=` `[` $offsets^ `]`)? $result `,` $barrier `,` $pred
              oilist(`cacheModifier` `=` $cache | `evictionPolicy` `=` $evict)
              attr-dict `:` qualified(type($desc)) `,` qualified(type($barrier)) `->` qualified(type($result))

此操作异步地将数据从全局内存拷贝到本地内存。这类似于 tt.load，只是数据被拷贝到内存描述符指向的本地内存，而不是分布式张量。所拷贝的数据取决于 desc 指向的全局内存描述符。

张量模式由描述符类型决定

tt.tensordesc: TILED 模式 - 常规平铺张量内存访问
- 参见: https://docs.nvda.net.cn/cuda/parallel-thread-execution/#tensor-tiled-mode
ttng.tensordesc_im2col: IM2COL 模式 - 用于卷积友好访问模式的 Im2col 模式
- 在 IM2COL 模式下，“coord”是输入张量中的坐标
  - 例如，对于 4D 张量 (NHWC)，“coord”为 [batch_idx, channel_idx, h, w]
- 在 IM2COL 模式下，必须提供额外的 offsets（uint16 值）
  - 对于 3D 张量 (NWC)：1 个偏移量 (offset_w)
  - 对于 4D 张量 (NHWC)：2 个偏移量 (offset_w, offset_h)
  - 对于 5D 张量 (NDHWC)：3 个偏移量 (offset_w, offset_h, offset_d)
  - 一般规则：偏移量数量 = coord.size() - 2
- 参见: https://docs.nvda.net.cn/cuda/parallel-thread-execution/#tensor-im2col-mode

Traits: AttrSizedOperandSegments, VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: MBarrierOpInterface, PredicatedOpInterface, TMALoadLikeOpInterface, TMAOpInterface

属性：

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

操作数：

操作数	描述
`desc`	张量描述符类型
`coord`	可变参数的 32 位无符号整数
`offsets`	16 位无符号整数的可变参数
`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 指向的全局内存描述符。

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: TMAOpInterface, TMAStoreLikeOpInterface

操作数：

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

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

基于描述符异步地从全局内存收集 (gather) 数据到本地内存

语法

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

此操作异步地从全局内存矩阵收集多行数据到本地内存。除了每一行是独立索引的之外，这与 async_tma_copy_global_to_local 类似。

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: MBarrierOpInterface, PredicatedOpInterface, TMALoadLikeOpInterface, TMAOpInterface

属性：

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

操作数：

操作数	描述
`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))

此操作异步地将数据从本地内存拷贝到全局内存，并原子地执行指定的归约类型。原子性处于单个元素的粒度，并且仅暗示宽松的语义。

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: TMAOpInterface, TMAStoreLikeOpInterface

属性：

属性	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)

基于描述符异步地将本地内存中的数据分散 (scatter) 到全局内存

语法

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

ttng.async_tma_scatter 操作异步地将本地内存中多个单独索引的数据行分散到全局内存中。该操作将共享内存中的 2D 张量（通过核心张量平铺 nvmma_shared 布局进行布局）分散到全局内存中给定 y 偏移量的单独索引行中。

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: TMAOpInterface, TMAStoreLikeOpInterface

操作数：

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

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

等待直到所有输入被读取。

语法

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

等待直到所有关联的存储操作的读取操作完成。在写入共享内存之前需要此操作。

Traits: MemWaitOpTrait, VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Attributes:

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

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

发出信号，表明预期有一定数量的字节被拷贝到屏障。

语法

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

此信号告知屏障预期有 size 字节被拷贝。关联的屏障等待操作将阻塞，直到预期的字节数被拷贝完成。

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: MBarrierOpInterface, PredicatedOpInterface

属性：

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

操作数：

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

`ttng.clc_get_program_id` (triton::nvidia_gpu::CLCGetProgramIdOp)

从 CLC 响应获取 CTA ID 坐标

语法

operation ::= `ttng.clc_get_program_id` $clcResult `,` $dim attr-dict `:` type($clcResult) `->` type($result)

解码 CLC 响应以获取已取消集群的第一个 CTA ID 坐标。dim 属性指定维度（0=x，1=y，2=z）。

Traits: AlwaysSpeculatableImplTrait, VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

接口：ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性：

属性	MLIR 类型	描述
`dim`	::mlir::triton::ProgramIDDimAttr	允许的 32 位无符号整数情况：0, 1, 2

操作数：

操作数	描述
`clcResult`	128 位无符号整数

结果：

结果	描述
`result`	32位无符号整数

`ttng.clc_is_canceled` (triton::nvidia_gpu::CLCIsCanceledOp)

检查 CLC 响应是否指示取消成功

语法

operation ::= `ttng.clc_is_canceled` $clcResult attr-dict `:` type($clcResult) `->` type($is_canceled)

解码 CLC 响应以检查集群是否成功取消。如果取消，返回 true，否则返回 false。

Traits: AlwaysSpeculatableImplTrait, VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

接口：ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

Operands:

操作数	描述
`clcResult`	128 位无符号整数

Results:

结果	描述
`is_canceled`	1位无符号整数

`ttng.clc_load_result` (triton::nvidia_gpu::CLCLoadResultOp)

将 CLC 响应从共享内存加载到寄存器

语法

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

将 128 位 CLC 响应从共享内存加载到两个 i64 寄存器中。这允许后续的 is_canceled 和 get_first_ctaid 操作在寄存器上运行，而无需重新读取共享内存。

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

接口 (Interfaces): InferTypeOpInterface

操作数：

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

结果：

结果	描述
`clcResult`	128 位无符号整数

`ttng.clc_try_cancel` (triton::nvidia_gpu::CLCTryCancelOp)

发出 CLC try_cancel 以取消挂起的集群

语法

operation ::= `ttng.clc_try_cancel` $result `,` $mbarrier attr-dict `:` qualified(type($result)) `,` qualified(type($mbarrier))

发出 clusterlaunchcontrol.try_cancel 指令以原子方式取消挂起的集群启动。结果被异步写入结果缓冲区，并在完成时发出 mbarrier 信号。

这用于 Blackwell (SM100+) 上的动态持久内核。

结果缓冲区必须是 16 字节对齐的共享内存。mbarrier 必须是 8 字节对齐的共享内存。

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: MBarrierOpInterface

操作数:

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

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

语法

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

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Attributes:

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

`ttng.cluster_barrier` (triton::nvidia_gpu::ClusterBarrierOp)

在集群范围内同步所有 warp

语法

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

降低为集群到达/等待对。

在 warp 专业化内核中，降低过程将屏障包装在合成的 ttg.warp_specialize 区域中，以便工作 warp 也执行该屏障。此操作不能放置在现有的 ttg.warp_specialize 内。

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Attributes:

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

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

语法

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

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

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

Fence 代理异步

语法

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

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Attributes:

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

`ttng.fence_mbarrier_init_release_cluster` (triton::nvidia_gpu::FenceMBarrierInitReleaseClusterOp)

Fence mbarrier 初始化释放集群

语法

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

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

`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。

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: MBarrierOpInterface

Attributes:

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

Operands:

操作数	描述
`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

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: MBarrierOpInterface

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))

此操作旨在帮助从非类型的原始 TMA 对象转换到类型的张量描述符对象。理想情况下，一旦 API 完全成熟，我们可以删除此操作。

Traits: AlwaysSpeculatableImplTrait, VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数：

操作数	描述
`rawDesc`	ptr

Results:

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

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

使 mbarrier 跟踪所有先前异步 tcgen5 操作的完成

语法

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

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

如果提供了 descs，提交将根据这些描述符的共享布局在 CTA 集群中进行多播。当 tcgen5 MMA 的输入（包括缩放的 MMA 缩放输入）来自使用多播的 TMA 描述符时，应使用此操作。

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

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

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

Traits: AttrSizedOperandSegments, VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: MBarrierOpInterface, PredicatedOpInterface

Operands:

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

`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。

isUnsigned 属性仅在执行整数 MMA 操作时相关。如果为 true，则整数值被视为无符号，否则被视为有符号。

Traits: AttrSizedOperandSegments, VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: DotOpInterface, MBarrierOpInterface, MMAv5OpInterface, MemoryEffectOpInterface, PredicatedOpInterface

Attributes:

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

Operands:

操作数	描述
`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 位无符号整数的变长参数

Results:

结果	描述
`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)) 如果设置了 $two_ctas，操作将在两个连续的 CTA 上执行矩阵乘法，它将读取跨越两个 CTA 分布的数据，如果操作是同步的，它将同步两个 CTA。如果 is_async 为 false，操作同步执行。在这种情况下，屏障操作数一定不能存在。否则，如果给定了屏障，操作将触发对其的提交/到达。在屏障等待后，读取结果将是安全的。

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

Traits: AttrSizedOperandSegments, VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: DotOpInterface, MBarrierOpInterface, MMAv5OpInterface, MemoryEffectOpInterface, PredicatedOpInterface

Attributes:

属性	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
`two_ctas`	::mlir::UnitAttr	单元属性
`multicast`	::mlir::UnitAttr	单元属性
`is_async`	::mlir::UnitAttr	单元属性

Operands:

操作数	描述
`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 位无符号整数的变长参数

Results:

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

`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)

Traits: AttrSizedOperandSegments, VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

接口：MemoryEffectOpInterface (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{MemoryEffects::Read on ::mlir::triton::GlobalMemory, MemoryEffects::Write on ::mlir::triton::GlobalMemory}

Attributes:

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

Operands:

操作数	描述
`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))

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

接口：MemoryEffectOpInterface (MemoryEffectOpInterface)

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

Operands:

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

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

分配张量内存

语法

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

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

显式取消分配缓冲区是可选的；请参阅 local_dealloc。

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

接口： MemoryEffectOpInterface

操作数：

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

结果:

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

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

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

语法

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

连续存储在 SMEM 中的 2D 块被拷贝到目标地址指定的 TMEM 中。此操作降低为 PTX 指令 tcgen05.cp。这支持写入缩放 tmem 布局以及默认 tmem 布局。目前，写入 tmem 缩放布局时的语义不同。

对于默认布局，拷贝不会更改源内存描述符和目标内存描述符之间的逻辑元素。

对于缩放布局：SMEM 中的每个 32x128b 块都在 4 个 warp 上复制，并存储到 TMEM 的 128 行和 4 列中。此操作的主要用例是将阻塞缩放从 SMEM 拷贝到 TMEM。

根据用例，SMEM 的形状可以灵活选择。在最简单的情况下（例如单元测试），源 SMEM 的形状可以是 (32 x num_blocks, 16)，目标 TMEM 的形状应为 (128, 16 x num_blocks)，用于拷贝 8 位值。对于缩放 GEMM，需要将 32x128b 块的 rep_m x rep_k 拷贝存储在 SMEM 中，其中 rep_m = BLOCK_M / 128，rep_k = BLOCK_K / scale_vec_size / 4，并且 scale_vec_size = 32（对于 MXFP）。概念上，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 中，阻塞缩放的 TMEM memdesc 必须具有以下形式

其形状必须是 (BLOCK_MN, BLOCK_K / scale_vec_size)，代表阻塞缩放的逻辑形状。
它必须附加 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) 是某些重塑 (reshape) 和转置 (transpose) 操作的结果。实际上，为了利用本机缩放布局和 TMEM 拷贝操作，用户需要在将缩放输入 dot_scaled 之前执行 scales5D.trans(0, 3, 2, 1, 4).reshape(BLOCK_M, BLOCK_K // scale_vec_size)。当我们在 IR 中使用 tmem_copy 时，此类重塑和转置操作将被删除。但是，它们在寄存器上导致的逻辑形状变化现在被理解为合并到了 tmem_copy 本身中。理想情况下，我们会将寄存器上完成的重塑/转置提升到 SMEM memdesc 上，使 tmem_copy 成为直接的 2D 拷贝操作：(BLOCK_MN, BLOCK_K / scale_vec_size) -> (BLOCK_MN, BLOCK_K / scale_vec_size)。在 memdesc 上没有此类操作的情况下，我们诉诸于隐式地在 tmem_copy 中编码重塑/转置语义。

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Operands:

操作数	描述
`src`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`dst`	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) (`,` type($red)^)?

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

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

可选归约修饰符：当指定 redOp 时，加载操作还会沿输入的 N 维度执行元素归约，并产生第二个结果张量 red。对于形状为 [M, N] 的输入，归约结果的形状为 [M]，包含 N 维度的每个“切片”的一个归约值。

目前仅限于 f32 元素类型。

redOp：指定沿 N 维度应用的归约操作（MIN 或 MAX）。设置后，red 结果必须存在。
abs：当为 true 时，在执行归约之前对每个元素应用绝对值。仅在指定 redOp 时有效。
NaN：当为 true 时，归约传播 NaN 值（如果切片中的任何输入元素是 NaN，则相应的归约值是 NaN）。当为 false 时，归约期间忽略 NaN 值。仅在指定 redOp 时有效。

示例：TMEM 中的输入，形状[M=2, N=4]: [[ 1.0, 3.0, 2.0, 4.0], [-5.0, 1.0, 8.0, 2.0]]

使用 redOp=MAX: 结果 = [[ 1.0, 3.0, 2.0, 4.0], // 不变 [-5.0, 1.0, 8.0, 2.0]] red = [4.0, 8.0] // 每行沿 N 的最大值

使用 redOp=MIN, abs=true: red = [1.0, 1.0] // 每行 |值| 的最小值

此操作在支持的架构（例如 Blackwell Ultra）上通过 PTX tcgen05.ld.red 指令降低为硬件加速归约。

Traits: AttrSizedResultSegments, VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Attributes:

属性	MLIR 类型	描述
`redOp`	::mlir::triton::nvidia_gpu::TMEMLoadReduceModifierAttr	允许的 32 位无符号整数情况：1, 2
`abs`	::mlir::BoolAttr	布尔属性
`NaN`	::mlir::BoolAttr	布尔属性

Operands:

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

Results:

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

`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。

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: PredicatedOpInterface

操作数:

操作数	描述
`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 唯一能做的事情。

Traits: AlwaysSpeculatableImplTrait, MemDescViewTrait, VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性：

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

Operands:

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

Results:

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

`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::cta.b64 降低一个等待循环。

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

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

Traits: AttrSizedOperandSegments, VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: MBarrierOpInterface, PredicatedOpInterface

Operands:

操作数	描述
`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 组点积

语法

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

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

Interfaces: DotOpInterface, InferTypeOpInterface, MemoryEffectOpInterface

Attributes:

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

Operands:

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

Results:

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

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

Warp 组点积等待

语法

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

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

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

Traits: VerifyMemDescLayoutsTrait, VerifyTensorLayoutsTrait

接口 (Interfaces): InferTypeOpInterface

Attributes:

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

Operands:

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

Results:

结果	描述
`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)

属性：

操作数：

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.async_tma_store_wait (triton::nvidia_gpu::TMAStoreWaitOp)

Attributes:

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

属性：

操作数：

ttng.clc_get_program_id (triton::nvidia_gpu::CLCGetProgramIdOp)

属性：

操作数：

结果：

ttng.clc_is_canceled (triton::nvidia_gpu::CLCIsCanceledOp)

Operands:

Results:

ttng.clc_load_result (triton::nvidia_gpu::CLCLoadResultOp)

操作数：

结果：

ttng.clc_try_cancel (triton::nvidia_gpu::CLCTryCancelOp)

操作数:

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

Attributes:

ttng.cluster_barrier (triton::nvidia_gpu::ClusterBarrierOp)

Attributes:

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

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

Attributes:

ttng.fence_mbarrier_init_release_cluster (triton::nvidia_gpu::FenceMBarrierInitReleaseClusterOp)

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

Attributes:

Operands:

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

Operands:

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

操作数：

Results:

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

Operands:

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

Attributes:

Operands:

Results:

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

Attributes:

Operands:

Results:

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

Attributes:

Operands:

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

Operands:

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

操作数：

结果:

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

Operands:

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

Attributes:

Operands:

Results:

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

操作数:

结果:

`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.async_tma_store_wait` (triton::nvidia_gpu::TMAStoreWaitOp)

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

`ttng.clc_get_program_id` (triton::nvidia_gpu::CLCGetProgramIdOp)

`ttng.clc_is_canceled` (triton::nvidia_gpu::CLCIsCanceledOp)

`ttng.clc_load_result` (triton::nvidia_gpu::CLCLoadResultOp)

`ttng.clc_try_cancel` (triton::nvidia_gpu::CLCTryCancelOp)

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

`ttng.cluster_barrier` (triton::nvidia_gpu::ClusterBarrierOp)

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

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

`ttng.fence_mbarrier_init_release_cluster` (triton::nvidia_gpu::FenceMBarrierInitReleaseClusterOp)

`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.tensormap_create` (triton::nvidia_gpu::TensormapCreateOp)

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

`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.wait_barrier` (triton::nvidia_gpu::WaitBarrierOp)

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

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