TritonNvidiaGPUOps

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

对内存屏障执行到达操作

语法

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

ttng.arrive_barrier 操作对共享内存中的内存屏障对象执行“到达”操作。该操作需要一个至少为1的 count 属性,并按指定计数减少内存屏障的待到达计数。

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

示例

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::IntegerAttr32位无符号整数属性

操作数:

操作数

描述

alloc

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

pred

1位无符号整数

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

一旦所有先前发出的拷贝完成,在内存屏障上到达

语法

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::IntegerAttr32位无符号整数属性

操作数:

操作数

描述

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::IntegerAttr1位无符号整数属性

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

语法

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

特性:VerifyTensorLayoutsTrait

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

异步代理围栏

语法

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

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

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

特性:VerifyTensorLayoutsTrait

属性:

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

操作数:

操作数

描述

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

此操作的存在是为了帮助从无类型原始 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)

使内存屏障跟踪所有先前异步 tcgen5 操作的完成

语法

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

ttng.tc_gen5_commit 是一个异步操作,它使内存屏障对象跟踪所有先前异步 tcgen5 操作的完成。所有异步操作完成后,内存屏障到达操作将以计数1在内存屏障上执行。

如果设置了 two_ctas,则内存屏障也会跟踪所有先前以 two_ctas 设置启动的操作。否则,它会跟踪所有先前未设置 two_ctas 启动的操作。

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

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

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

特性: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 读写。当令牌存在时,它们可用于检查累加器内存上的别名和修改引用。

特性: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 读写。当令牌存在时,它们可用于检查累加器内存上的别名和修改引用。

特性: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::IntegerAttr32位无符号整数属性

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 `:` functional-type(operands, results)

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

此操作降低为 PTX 指令 tcgen05.cp。目前,我们只支持 1CTA 和 warpx4.32x128b 变体。SMEM 中每个 32x128b 块在 4 个 warp 中复制,并存储到 TMEM 的 128 行和 4 列中。此操作的主要用例是将块状标量从 SMEM 拷贝到 TMEM。

输入 SMEM 的形状可以根据用例灵活选择。在最简单的情况下(例如单元测试),源 SMEM 的形状可以是 (32 x num_blocks, 16),目标 TMEM 的形状应为 (128, 16 x num_blocks),用于拷贝 8 位值。对于缩放 GEMM,需要将 rep_m x rep_k 拷贝的 32x128b 块存储在 SMEM 中,其中 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 中,块状标量的 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) 是某些重塑和转置操作的结果。实际上,为了利用原生标量布局和 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 中隐式编码重塑/转置语义。

特性: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 缓冲区上的别名和修改引用。

特性: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 缓冲区上的别名和修改引用。

特性: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::IntegerAttr32位无符号整数属性

操作数:

操作数

描述

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)

在张量图对象上获取围栏

语法

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)

等待内存屏障阶段完成。

语法

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

阻塞程序进程,直到 alloc 中的内存屏障对象完成其当前阶段。

这会降低为使用 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 组点积

语法

operation ::= `ttng.warp_group_dot` $a`,` $b`,` $c (`,` $useC^)? attr-dict
              `:` type($a) `*` 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
maxNumImpreciseAcc::mlir::IntegerAttr32位无符号整数属性
isAsync::mlir::BoolAttr布尔属性

操作数:

操作数

描述

a

TensorOrMemDesc 实例

b

TensorOrMemDesc 实例

c

浮点或整数值的秩张量

useC

1位无符号整数

结果:

结果

描述

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 必须是第一个点积操作的结果。

特性:VerifyTensorLayoutsTrait

接口 (Interfaces): InferTypeOpInterface

属性:

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

操作数:

操作数

描述

inputs

TensorOrMemDesc 实例的变长参数

结果:

结果

描述

outputs

TensorOrMemDesc 实例的变长参数