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::IntegerAttr | 32位无符号整数属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
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 | 单元属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
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 | 布尔属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
张量描述符类型( |
|
32位无符号整数的变长参数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
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
操作数:¶
操作数 |
描述 |
---|---|
|
张量描述符类型( |
|
32位无符号整数的变长参数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
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):¶
操作数 |
描述 |
---|---|
|
张量描述符类型( |
|
32位无符号整数值的秩张量 |
|
32位无符号整数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
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 |
操作数:¶
操作数 |
描述 |
---|---|
|
张量描述符类型( |
|
32位无符号整数的变长参数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
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
操作数:¶
操作数 |
描述 |
---|---|
|
张量描述符类型( |
|
32位无符号整数值的秩张量 |
|
32位无符号整数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
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位无符号整数属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
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.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::IntegerAttr | 32位无符号整数属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
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):¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
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{}
操作数:¶
操作数 |
描述 |
---|---|
|
ptr |
结果:¶
结果 |
描述 |
---|---|
|
张量描述符类型( |
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):¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
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 | 单元属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
异步令牌类型 |
|
1位无符号整数 |
|
1位无符号整数 |
|
Triton IR 类型系统中内存描述符类型 ( |
|
1位无符号整数的变长参数 |
结果:¶
结果 |
描述 |
---|---|
|
异步令牌类型 |
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 | 单元属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
异步令牌类型 |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
1位无符号整数 |
|
1位无符号整数 |
|
Triton IR 类型系统中内存描述符类型 ( |
|
1位无符号整数的变长参数 |
结果:¶
结果 |
描述 |
---|---|
|
异步令牌类型 |
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):¶
操作数 |
描述 |
---|---|
|
浮点数、整数或 ptr 值的秩张量 |
结果 (Results):¶
结果 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
异步令牌类型 |
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
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
Triton IR 类型系统中的内存描述符类型 ( |
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):¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
异步令牌类型 |
结果:¶
结果 |
描述 |
---|---|
|
浮点数、整数或 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 缓冲区上的别名和修改引用。
特性:VerifyTensorLayoutsTrait
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
异步令牌类型 |
|
浮点数、整数或 ptr 值的秩张量 |
|
1位无符号整数 |
结果:¶
结果 |
描述 |
---|---|
|
异步令牌类型 |
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位无符号整数属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
结果:¶
结果 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
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 位无符号整数属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的指针类型 ( |
|
Triton IR 类型系统中的指针类型 ( |
|
32位无符号整数的变长参数 |
|
32位无符号整数的变长参数 |
|
64位无符号整数的变长参数 |
|
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}
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的指针类型 ( |
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
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
32位无符号整数 |
|
1位无符号整数 |
|
Triton IR 类型系统中内存描述符类型 ( |
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::IntegerAttr | 32位无符号整数属性 |
isAsync | ::mlir::BoolAttr | 布尔属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
TensorOrMemDesc 实例 |
|
TensorOrMemDesc 实例 |
|
浮点或整数值的秩张量 |
|
1位无符号整数 |
结果:¶
结果 |
描述 |
---|---|
|
浮点或整数值的秩张量 |
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::IntegerAttr | 32位无符号整数属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
TensorOrMemDesc 实例的变长参数 |
结果:¶
结果 |
描述 |
---|---|
|
TensorOrMemDesc 实例的变长参数 |