TritonAMDGPUOps
amdg.arrive_barrier (triton::amdgpu::ArriveBarrierOp)
在 mbarrier 上执行到达(arrive)操作
语法
operation ::= `amdg.arrive_barrier` $alloc `,` $count attr-dict `:` qualified(type($alloc)) `->` type($result)
在共享内存中的 mbarrier 对象上执行“arrive”操作。该操作需要至少为 1 的 count 属性,并将 mbarrier 的挂起到达计数减去指定的计数值。如果挂起计数达到零,相位会发生变化(以循环方式递减),并使用初始计数值重新加载挂起计数。返回“arrive”操作之前 mbarrier 对象的相位奇偶性(0 为偶,1 为奇)。
示例
ttag.arrive_barrier %barrier, 2 : !ttg.memdesc<1xi64, #shared, #smem, mutable>
接口: InferTypeOpInterface, MBarrierOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
count | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
结果:
结果 |
描述 |
|---|---|
|
32位无符号整数 |
amdg.async_copy_local_to_global (triton::amdgpu::AsyncCopyLocalToGlobalOp)
异步将数据从本地内存复制到全局内存
语法
operation ::= `amdg.async_copy_local_to_global` $src `,` $dst (`mask` $mask^)?
oilist(`cacheModifier` `=` $cache | `evictionPolicy` `=` $evict)
attr-dict `:` qualified(type($src)) `->` type($dst)
此操作异步地将数据从本地内存复制到全局内存。这与 tt.store 类似,只是数据是从内存描述符指向的本地内存中复制,而不是从分布式张量中复制。连续性(contiguity)是指在给定布局和掩码下,可以在单个向量中存储的最大元素数量。即使无法基于 IR 证明对齐,这也允许操作使用 async_copy_local_to_global。
接口: InferTypeOpInterface, PredicatedOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
evict | ::mlir::triton::EvictionPolicyAttr | 允许的32位无符号整数情况:1, 2, 3 |
contiguity | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
指针值的秩张量(ranked tensor) |
|
1 位无符号整数值的张量 |
结果:
结果 |
描述 |
|---|---|
|
异步令牌类型 |
amdg.async_copy_mbarrier_arrive (triton::amdgpu::AsyncCopyMbarrierArriveOp)
在所有先前发出的拷贝完成后,在 mbarrier 上到达
语法
operation ::= `amdg.async_copy_mbarrier_arrive` $barrier attr-dict `:` qualified(type($barrier))
执行“async arrive”操作:当所有先前的异步加载到 LDS(特别是,非 TDM)的操作完成后,将挂起计数减 1。指令本身是异步的;它立即返回。递减屏障挂起计数。用于递减的更新值固定为 1。如果挂起计数变为零,相位会发生变化(以循环方式递减),并使用初始计数值重新加载挂起计数。
接口: MBarrierOpInterface
操作数:
操作数 |
描述 |
|---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
amdg.async_tdm_copy_global_to_local (triton::amdgpu::AsyncTDMCopyGlobalToLocalOp)
基于描述符异步将数据从全局内存复制到本地内存
语法
operation ::= `amdg.async_tdm_copy_global_to_local` $desc `[` $indices `]` `into` $result `,` `pred` `=` $pred (`,` `barrier` `=` $barrier^)?
attr-dict `:` qualified(type($desc)) (`,` qualified(type($barrier))^)? `->` qualified(type($result))
此操作异步地将数据从全局内存复制到本地内存。这与 tt.load 类似,只是数据被复制到由 result 指向的本地内存中,而不是分布式张量。复制的数据取决于 desc 指向的全局内存。将 pred 设置为 false 将禁用复制。此操作不支持共享内存交错(swizzling)。该操作还可以接受一个可选的 64 位 LDS 屏障地址,在这种情况下,它会发送一个“LDS 原子到达”信号来表示完成。
特性(Traits): AttrSizedOperandSegments
接口: InferTypeOpInterface, MBarrierOpInterface, PredicatedOpInterface, TDMOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
可变参数的 32 位无符号整数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
32位无符号整数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
结果:
结果 |
描述 |
|---|---|
|
异步令牌类型 |
amdg.async_tdm_copy_local_to_global (triton::amdgpu::AsyncTDMCopyLocalToGlobalOp)
基于描述符异步将数据从本地内存复制到全局内存
语法
operation ::= `amdg.async_tdm_copy_local_to_global` $desc `[` $indices `]` `from` $src (`,` `barrier` `=` $barrier^)?
attr-dict `:` qualified(type($src)) (`,` qualified(type($barrier))^)? `->` qualified(type($desc))
此操作异步地将数据从本地内存复制到全局内存。这与 tt.store 类似,只是数据是从 src 指向的本地内存中复制,而不是从分布式张量中复制。复制的目标取决于 desc 指向的全局内存。此操作不支持共享内存填充或交错(swizzling)。该操作还可以接受一个可选的 64 位 LDS 屏障地址,在这种情况下,它会发送一个“LDS 原子到达”信号来表示完成。
特性(Traits): AttrSizedOperandSegments
接口: MBarrierOpInterface, TDMOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
可变参数的 32 位无符号整数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
Triton IR 类型系统中的内存描述符类型 ( |
amdg.async_tdm_gather (triton::amdgpu::AsyncTDMGatherOp)
异步从全局内存的非连续行收集数据到本地内存
语法
operation ::= `amdg.async_tdm_gather` $desc `[` $src_row_indices `,` $src_col_offset `]` `to` $dst `,` `pred` `=` $pred (`,` `barrier` `=` $barrier^)?
attr-dict `:` qualified(type($src_row_indices)) `,` qualified(type($dst)) (`,` qualified(type($barrier))^)? `->` qualified(type($desc))
此操作使用 TDM 收集模式异步地从全局内存的非连续行收集数据到本地内存。与从连续内存读取的常规 async_tdm_copy_global_to_local 不同,此操作使用 src_row_indices 来指定要从全局内存中读取哪些行。
描述符必须是二维的。src_row_indices 指定要从全局内存中读取哪些行。src_row_indices 的元素类型决定了索引大小
I16:16 位索引,每条指令最多 16 行
I32:32 位索引,每条指令最多 8 行。如果需要更多行,将发出多条 TDM 指令。
src_col_offset 指定所有收集行在源张量中的起始列。
接口: InferTypeOpInterface, MBarrierOpInterface, PredicatedOpInterface, TDMOpInterface
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
16 位无符号整数或 32 位无符号整数值的张量 |
|
32位无符号整数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
32位无符号整数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
结果:
结果 |
描述 |
|---|---|
|
异步令牌类型 |
amdg.async_tdm_intrinsic_wait (triton::amdgpu::AsyncTDMIntrinsicWait)
等待,直到未完成的 TDM 内在函数(intrinsic)数量小于或等于给定值
语法
operation ::= `amdg.async_tdm_intrinsic_wait` ($asyncToken^)? attr-dict
此操作会一直等待,直到未完成的 TDM 内在函数(汇编指令)的数量小于或等于给定值。这对于确保数据在被使用之前在 LDS(加载、收集)或 HBM(存储、散布)中可用是必要的。
与计算 IR 操作的 AsyncTDMWait 不同,此操作计算已发出的 TDM 汇编指令(例如,tensor_load_to_lds,tensor_store_from_lds)的实际数量,这是降低到 LLVM 所必需的。
特性(Traits): MemWaitOpTrait
接口 (Interfaces): InferTypeOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
count | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
异步令牌类型的变参 |
结果:
结果 |
描述 |
|---|---|
|
异步令牌类型 |
amdg.async_tdm_scatter (triton::amdgpu::AsyncTDMScatterOp)
异步将数据从本地内存散布到全局内存的非连续行
语法
operation ::= `amdg.async_tdm_scatter` $desc `[` $dst_row_indices `,` $dst_col_offset `]` `from` $src (`,` `barrier` `=` $barrier^)?
attr-dict `:` qualified(type($dst_row_indices)) `,` qualified(type($src)) (`,` qualified(type($barrier))^)? `->` qualified(type($desc))
此操作使用 TDM 散布模式异步地将数据从本地内存散布到全局内存的非连续行。与复制到连续内存的常规 async_tdm_copy_local_to_global 不同,此操作使用 dst_row_indices 来指定要写入全局内存的哪些行。
描述符必须是二维的。dst_row_indices 指定要写入全局内存的哪些行。dst_row_indices 的元素类型决定了索引大小
I16:16 位索引,每条指令最多 16 行
I32:32 位索引,每条指令最多 8 行。如果需要更多行,将发出多条 TDM 指令。
dst_col_offset 指定所有散布行在目标张量中的起始列。
接口: InferTypeOpInterface, MBarrierOpInterface, TDMOpInterface
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
16 位无符号整数或 32 位无符号整数值的张量 |
|
32位无符号整数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
Triton IR 类型系统中的内存描述符类型 ( |
结果:
结果 |
描述 |
|---|---|
|
异步令牌类型 |
amdg.async_tdm_wait (triton::amdgpu::AsyncTDMWait)
等待,直到未完成的 TDM 操作数量小于或等于给定值
语法
operation ::= `amdg.async_tdm_wait` $asyncToken attr-dict
此操作会一直等待,直到未完成的异步 TDM 操作数量小于或等于给定值。这对于确保数据在被使用之前在 LDS(加载、收集)或 HBM(存储、散布)中可用是必要的。
此操作计算 TDM IR 操作的数量(AsyncTDMCopyGlobalToLocalOp, AsyncTDMCopyLocalToGlobalOp, AsyncTDMScatterOp, AsyncTDMGatherOp)。
特性(Traits): MemWaitOpTrait
接口 (Interfaces): InferTypeOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
num | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
异步令牌类型的变参 |
结果:
结果 |
描述 |
|---|---|
|
异步令牌类型 |
amdg.async_wait (triton::amdgpu::AsyncWaitOp)
等待,直到未完成的异步内在函数(intrinsics)数量小于或等于给定值
语法
operation ::= `amdg.async_wait` ($asyncToken^)? attr-dict
类似于 ttg.async_wait,但此操作等待的是未完成的异步指令/内在函数的数量(这是降低到 AMD 后端的 LLVM 所必需的),而不是等待未完成的 ttg.async_commit_groups。
特性(Traits): MemWaitOpTrait
接口 (Interfaces): InferTypeOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
num_inst | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
异步令牌类型的变参 |
结果:
结果 |
描述 |
|---|---|
|
异步令牌类型 |
amdg.buffer_atomic_cas (triton::amdgpu::BufferAtomicCASOp)
原子 CAS 操作,对标量基地址和张量偏移量执行比较并交换
语法
operation ::= `amdg.buffer_atomic_cas` $sem `,` $scope `,` $cmp `,` $val `,` $ptr `[` $offsets `]`
(`stride` `=` $stride^)?
attr-dict `:` type($result)
AMD 缓冲区原子 CAS 操作。缓冲区原子操作与普通原子操作类似,但它们通过标量基地址和偏移量张量访问全局内存,而不是通过指针张量。类似于 TT_AtomicCASOp:缓冲区原子 CAS 操作加载 $ptr 处的数据,如果 $ptr 处的值等于 $cmp,则以指定的内存语义和作用域原子地将 $val 存储到 $ptr。如果使用了原子 CAS 操作的返回值,则返回预操作值,否则该值被隐式丢弃。步长(stride)是连续内存块开始处之间的距离。执行 CAS 时,stride 是每一行第一个元素之间的地址差(以字节为单位)。编译器在转换为缓冲区操作时会尝试获取 stride,因为它对于优化缓存内存访问非常重要。
特性(Traits): SameLoadStoreOperandsAndResultEncoding
接口: BufferOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
sem | ::mlir::triton::MemSemanticAttr | 允许的32位无符号整数情况:1, 2, 3, 4 |
scope | ::mlir::triton::MemSyncScopeAttr | 允许的32位无符号整数情况:1, 2, 3 |
操作数:
操作数 |
描述 |
|---|---|
|
ptr |
|
32 位无符号整数值的张量 |
|
浮点数、整数或 ptr 值的秩张量 |
|
浮点数、整数或 ptr 值的秩张量 |
|
32位无符号整数 |
结果:
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
amdg.buffer_atomic_rmw (triton::amdgpu::BufferAtomicRMWOp)
原子 RMW 操作,对标量基地址和张量偏移量执行读取、修改和写入
语法
operation ::= `amdg.buffer_atomic_rmw` $atomic_rmw_op `,` $sem `,` $scope `,` $value `,` $ptr `[` $offsets `]` (`,` $mask^)?
(`stride` `=` $stride^)?
attr-dict `:` type($result)
AMD 缓冲区原子 RMW 操作。缓冲区原子操作与普通原子操作类似,但它们通过标量基地址和偏移量张量访问全局内存,而不是通过指针张量。与其它缓冲区操作类似,mask 是一个布尔向量,决定是否应该对给定元素执行原子 RMW 操作。mask[i] == 0 的元素会被丢弃(即,原子操作不会被执行)。类似于 TT_AtomicRMWOp:缓冲区原子 RMW 操作加载 $ptr 处的数据,使用 $val 执行 $rmw_op,并以指定的内存语义和作用域将结果存储到 $ptr。如果使用了原子 RMW 操作的返回值,则返回预操作值,否则该值被隐式丢弃。步长(stride)是连续内存块开始处之间的距离。执行 RMW 时,stride 是每一行第一个元素之间的地址差(以字节为单位)。编译器在转换为缓冲区操作时会尝试获取 stride,因为它对于优化缓存内存访问非常重要。
特性(Traits): AttrSizedOperandSegments, SameLoadStoreOperandsAndResultEncoding
接口: BufferOpInterface, PredicatedOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
atomic_rmw_op | ::mlir::triton::RMWOpAttr | 允许的 32 位无符号整数值:1, 2, 3, 4, 5, 6, 7, 8, 9, 10 |
sem | ::mlir::triton::MemSemanticAttr | 允许的32位无符号整数情况:1, 2, 3, 4 |
scope | ::mlir::triton::MemSyncScopeAttr | 允许的32位无符号整数情况:1, 2, 3 |
操作数:
操作数 |
描述 |
|---|---|
|
ptr |
|
32 位无符号整数值的张量 |
|
浮点数、整数或 ptr 值的秩张量 |
|
32位无符号整数 |
|
1 位无符号整数值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
amdg.buffer_load (triton::amdgpu::BufferLoadOp)
从标量基地址和张量偏移量加载
语法
operation ::= `amdg.buffer_load` $ptr `[` $offsets `]` (`,` $mask^)? (`,` $other^)?
oilist(`cacheModifier` `=` $cache)
(`stride` `=` $stride^)?
attr-dict `:` type($result)
AMD 缓冲区加载操作。缓冲区加载类似于普通加载,但它通过标量基地址和偏移量张量访问全局内存,而不是通过指针张量。其他字段与普通加载类似,即 mask 是一个布尔向量,决定是否应该从内存中读取给定元素;other 是当 mask[i] == 0 时 lane i 应返回的元素。步长(stride)是连续内存块开始处之间的距离。执行块加载时,stride 是每一行第一个元素之间的地址差(以字节为单位)。编译器在转换为缓冲区操作时会尝试获取 stride,因为它对于优化缓存内存访问非常重要。连续性(contiguity)是指在给定布局和掩码下,可以在单个向量中加载的最大元素数量。即使无法基于 IR 证明对齐,这也允许使用 buffer_load。
特性(Traits): AttrSizedOperandSegments, SameLoadStoreOperandsAndResultEncoding
接口: BufferOpInterface, PredicatedOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
contiguity | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
ptr |
|
32 位无符号整数值的张量 |
|
32位无符号整数 |
|
1 位无符号整数值的秩张量 |
|
浮点数、整数或 ptr 值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
amdg.buffer_load_to_local (triton::amdgpu::BufferLoadToLocalOp)
从标量基地址和张量偏移量加载到共享内存
语法
operation ::= `amdg.buffer_load_to_local` $ptr `[` $offsets `]` (`mask` `=` $mask^)? (`other` `=` $other^)? (`stride` `=` $stride^)?
oilist(`cacheModifier` `=` $cache) `into` $dest
attr-dict `:` type($ptr) `[` type($offsets) `]` type($other) `->` type($dest)
AMD 缓冲区加载操作。类似于 amdg.buffer_load 操作,但它直接写入共享内存而不是寄存器。连续性是指在给定布局和掩码下,可以在单个向量中加载的最大元素数量。即使无法基于 IR 证明对齐,这也允许使用 buffer_load_to_local。
特性(Traits): AttrSizedOperandSegments
接口: BufferOpInterface, InferTypeOpInterface, PredicatedOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
contiguity | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
ptr |
|
32 位无符号整数值的张量 |
|
1 位无符号整数值的秩张量 |
|
浮点数、整数或 ptr 值的秩张量 |
|
32位无符号整数 |
结果:
结果 |
描述 |
|---|---|
|
异步令牌类型 |
amdg.buffer_store (triton::amdgpu::BufferStoreOp)
存储到标量基地址和张量偏移量
语法
operation ::= `amdg.buffer_store` $value `,` $ptr `[` $offsets `]` (`,` $mask^)?
oilist(`cacheModifier` `=` $cache)
(`stride` `=` $stride^)?
attr-dict `:` type($value)
AMD 缓冲区存储操作。缓冲区存储类似于普通存储,但它通过标量基地址和偏移量张量访问全局内存,而不是通过指针张量。其他字段与普通存储类似,即 mask 是一个布尔向量,决定是否应该将给定元素写入内存;value 是当 mask[i] == 1 时 lane i 应写入的元素张量。步长(stride)是连续内存块开始处之间的距离。执行块存储时,stride 是每一行第一个元素之间的地址差(以字节为单位)。编译器在转换为缓冲区操作时会尝试获取 stride,因为它对于优化缓存内存访问非常重要。连续性是指在给定布局和掩码下,可以在单个向量中加载的最大元素数量。即使无法基于 IR 证明对齐,这也允许使用 buffer_store。
特性(Traits): AttrSizedOperandSegments, SameLoadStoreOperandsEncoding
接口: BufferOpInterface, PredicatedOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
contiguity | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
|
ptr |
|
32 位无符号整数值的张量 |
|
32位无符号整数 |
|
1 位无符号整数值的秩张量 |
amdg.cluster_barrier_arrive (triton::amdgpu::ClusterBarrierArriveOp)
在集群屏障上到达
语法
operation ::= `amdg.cluster_barrier_arrive` attr-dict
表明集群已到达屏障,用于同步集群内的 CTA。
参见 ClusterBarrierWaitOp,了解如何等待已到达的集群屏障。
amdg.cluster_barrier_wait (triton::amdgpu::ClusterBarrierWaitOp)
等待集群屏障
语法
operation ::= `amdg.cluster_barrier_wait` attr-dict
等待同一集群的所有 CTA 到达集群屏障。到达和等待操作必须成对出现。在到达之前等待,或者在没有相应等待的情况下多次到达,将导致未定义行为。
amdg.concat (triton::amdgpu::ConcatOp)
拼接(Concat)操作
语法
operation ::= `amdg.concat` $sources attr-dict `:` type($sources) `->` type($result)
“concat”操作将一系列源 n 维张量组合成一个更大的目标张量。
所有源张量必须具有相同的形状、元素类型和编码。拼接维度由用户提供的源形状和目标形状推断得出。例如,两个 64x128 形状的张量可以产生 128x128 的目标形状(表示沿维度 0 拼接);或者产生 64x256(表示沿维度 1 拼接)。
通常,作为操作参数传递的源张量可以以多种方式排列成最终形状。例如,给定四个 64x64 形状的张量:concat s0<64x64>, s1<64x64>, s2<64x64>, s3<64x64> -> <128x128>
它们可以在结果张量中以不同的配置进行布局
s0 s1 2) s0 s2 s2 s3 s1 s3
从逻辑张量的角度来看,源张量被视为张量中张量的元素。换句话说,输入张量的一维数组在概念上被重塑为 n 维网格。此操作的语义假定行优先顺序(或其 n 维泛化),这意味着最快变化的维度首先填充,最慢变化的维度最后填充。在上面的示例中,这对应于布局 1)。
源张量和目标张量在 CTA 瓦片级别必须具有相同的线性布局。也就是说,输入维度的所有基向量必须匹配(寄存器输入维度除外)。寄存器基必须在定义单个 CTA 瓦片逻辑张量形状的子集上对齐。
这确保了拼接是一个空操作(no-op),意味着不需要在线程之间进行数据重排来组装具有给定形状和布局的目标张量。然而,布局内 CTA 瓦片的顺序不需要在源布局和目标布局之间匹配。正确处理这一点是该操作降低逻辑的责任。
此操作旨在直接作用于逻辑张量,避免了复杂的布局重新解释或重塑的需要。例如,tt.join 操作仅支持沿最内层维度进行拼接,并要求结果的最内层维度为每个线程提供 2 个元素,分布在寄存器中。相比之下,此 concat 操作对拼接维度或维度大小没有限制。
sources:输入张量的列表。
示例 1
#blocked = #ttg.blocked<{sizePerThread = [1, 8],
threadsPerWarp = [8, 8], warpsPerCTA = [4, 1], order = [1, 0]}>
%0 = amdg.concat %arg0, %arg1: tensor<32x64xf32, #blocked>,tensor<32x64xf32, #blocked>,
-> tensor<64x64xf32, #blocked>
示例 2
#src_layout = #ttg.linear<{register=[[0, 1], [0, 2], [0, 8], [0, 16], [0, 64], [64, 0]], lane=[[1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [0, 4]], warp=[[0, 32], [32, 0]], block=[]}>
#dst_layout = #ttg.linear<{register=[[0, 1], [0, 2], [0, 8], [0, 16], [0, 64], [0, 128], [64, 0], [128, 0]], lane=[[1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [0, 4]], warp=[[0, 32], [32, 0]], block=[]}>
%0 = amdg.concat %arg0, %arg1, %arg2, %arg3 : tensor<128x128xf16, #src_layout>, tensor<128x128xf16, #src_layout>, tensor<128x128xf16, #src_layout>,
tensor<128x128xf16, #src_layout> -> tensor<256x256xf16, #dst_layout>
特性: AlwaysSpeculatableImplTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
浮点数、整数或指针值的秩张量的变参 |
结果:
结果 |
描述 |
|---|---|
|
任何类型值的秩张量 |
amdg.cond_barrier (triton::amdgpu::CondBarrierOp)
有条件地设置屏障以同步块中的部分线程
语法
operation ::= `amdg.cond_barrier` $pred attr-dict
condBarrierOp 仅在给定参数为真时才设置屏障指令。这提供了一种同步块中部分线程的方法,有意地使执行序列发生分叉。但是,用户应保证所有线程最终通过使用剩余的线程调用 condBarrierOp(true) 来汇合。概念上,这类似于在 if 语句中设置执行屏障。此操作允许我们在合适时避免阻塞整个块,以帮助调度。注意:这不会设置任何内存栅栏(fence)。
操作数:
操作数 |
描述 |
|---|---|
|
1位无符号整数 |
amdg.extract_slice (triton::amdgpu::ExtractSliceOp)
提取切片操作
语法
operation ::= `amdg.extract_slice` $source $static_offsets attr-dict `:` type($source) `to` type($result)
“extract_slice”操作能够从寄存器中的张量提取切片。
“extract_slice”操作支持以下参数
source:要创建视图张量的基张量
offsets:要在基张量中创建视图的偏移量
在分布式布局中,张量被划分为 CTA 瓦片。CTA 瓦片代表分布在工作组内所有线程和 warp 中的最小连续张量部分。ExtractSlice 操作提取张量中属于 CTA 瓦片倍数的部分。
源张量和目标张量在 CTA 瓦片级别必须具有匹配的线性布局。这确保了 extract_slice 是一个空操作(no-op),意味着不需要在线程之间进行数据重排来提取具有给定形状和布局的目标张量。
+——-+——-+ | W0 | W1 | | | | | + | + | | W2 | W3 | <– 单个 CTA 瓦片(分布在 warp W0-W3 上) | | | | + | + | | | | +——-+——-+ | 源张量 提取切片 | . +————–+ | . | W0 | W1 | | . | | | | | + | + | | | W2 | W3 | | | | | | | + | + | | | | | | +——-+——+ | | W0 | W1 | | | | | | | + | + | | | W2 W3 | | | | | | | + | + | | | | | | +————–+
此操作旨在直接作用于逻辑张量,避免了复杂的布局重新解释或重塑的需要。例如,tt.split 操作仅支持沿最内层维度进行拆分,并要求结果的最内层维度为每个线程提供 2 个元素,分布在寄存器中。相比之下,extract_slice 操作对提取维度或维度大小没有限制。
示例 1
#blocked = #ttg.blocked<{sizePerThread = [1, 8],
threadsPerWarp = [4, 16], warpsPerCTA = [4, 1], order = [0, 1]}>
#blocked1 = #ttg.blocked<{sizePerThread = [1, 8],
threadsPerWarp = [16, 4], warpsPerCTA = [4, 1], order = [0, 1]}>
%1 = ttg.convert_layout %0 : tensor<128x128xf16, #blocked>
-> tensor<128x128xf16, #blocked1>
// create a slice of base tensor %1 with static offsets
%2 = amdg.extract_slice %0 [0, 0] :
tensor<128x128xf16, #blocked1> to tensor<128x32xf16, #blocked1>
示例 1 展示了“extract_slice”操作的使用方法。在此示例中,创建了一个 128x32 的新切片。“extract_slice”作用于期望的切片在 CTA 瓦片上具有与源张量相同布局的张量。“%0”不能直接进行切片,因为生成的切片不满足此条件。因此,它需要转换为适合切片的布局。“#blocked1”布局适用于此,因为它保持 sizePerThread 不变,从而保持合并属性不变。为了利用 warp 中的所有线程,此新布局的“threadsPerWarp”设置为 [16,4]。在执行“extract_slice”之前进行此布局转换,确保切片仍能高效地利用所有线程。切片的大小由结果类型决定。
特性: AlwaysSpeculatableImplTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
static_offsets | ::mlir::DenseI64ArrayAttr | i64 稠密数组属性 |
操作数:
操作数 |
描述 |
|---|---|
|
任何类型值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
任何类型值的秩张量 |
amdg.in_thread_transpose (triton::amdgpu::InThreadTransposeOp)
对属于每个线程的寄存器值执行转置
语法
operation ::= `amdg.in_thread_transpose` $src attr-dict `:` type($src) `->` type($result)
此操作对每个线程寄存器中的值执行布局转置。具体来说,给定输入布局的 blocked 布局,它会沿着底层线性布局的寄存器维度转置最后两个维度(rank-1 和 rank-2)。
转换示例
输入布局:blocked 布局,sizePerThread=[2, 2],order=[0, 1]。其线性布局寄存器基 = [[1, 0], [2, 0], [0, 1], [0, 2]]
输出布局:与输入相同的线程和 warp 基,寄存器基 = [[0, 1], [0, 2], [1, 0], [2, 0]]
此操作使得在 HBM 和共享内存顺序不同,且目标 AMD 硬件不支持原生转置时,能够从 HBM 进行高效的合并加载,并随后向量化写入共享内存。这是 ttg.convert_layout 的一个特定变体,在降低到 LLVM 时将转换为 ttg.convert_layout。我们不希望这种转换被优化掉,因为我们需要在从 HBM 加载后、写入共享内存之前,显式地实现转置指令。
特性: AlwaysSpeculatableImplTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
amdg.init_barrier (triton::amdgpu::InitBarrierOp)
在给定的共享内存分配中初始化屏障。
语法
operation ::= `amdg.init_barrier` $alloc `,` $count attr-dict `:` qualified(type($alloc))
使用 mbarrier 信息初始化共享内存分配。alloc 是共享内存分配的描述符。count 是屏障预期的到达次数。
接口: MBarrierOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
count | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
amdg.instruction_sched_hint (triton::amdgpu::InstructionSchedHint)
基本块内指令调度提示的占位符操作
语法
operation ::= `amdg.instruction_sched_hint` attr-dict
应用于占位符操作所在基本块内指令的指令调度提示的占位符操作。此操作主要旨在用于调整 tt.dot 操作结果主循环内的指令调度。在高层次上更容易识别 dot 操作,从而标记预期的调度区域。提示操作最终降低为 LLVM AMDGPU 指令调度原语,旨在控制不同类型的指令(valu/mfma,全局/共享内存等)如何交错以实现更好的指令级并行性。
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
variant | ::mlir::triton::amdgpu::SchedHintAttr | AMD GPU 的指令调度提示 |
amdg.local_load_packed_tranposed (triton::amdgpu::LocalLoadPackedTransposedOp)
将转置的打包张量从共享内存加载到分布式张量中
语法
operation ::= `amdg.local_load_packed_tranposed` $src (`token` $token^)? attr-dict `:` qualified(type($src)) `->` type($result)
需要共享内存中的 M/N 打包和 M/N 连续张量,并产生寄存器中的 K 打包 K 连续张量。打包的改变将通过使 M/N 维度加倍并使 K 维度减半来改变张量的形状。例如,如果 A 在共享内存中是 16x64,则此操作的结果将是 32x32。
特性(Traits): LocalLoadTrait
操作数:
操作数 |
描述 |
|---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
异步令牌类型 |
结果:
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
amdg.masked_load (triton::amdgpu::MaskedLoadOp)
带掩码的加载操作
语法
operation ::= `amdg.masked_load` $ptr `,` $mask `,` $falseVal (`,` $multicastMask^)?
oilist(`cacheModifier` `=` $cache)
(`forceNoAlias` $forceNoAlias^)?
attr-dict `:` functional-type(operands, results)
支持掩码和多播(multicast)的加载操作。如果掩码为真,则从给定的指针加载。作为简化 LLVM 转换的实用操作,使用 LLVM 类型。在支持多播的架构上,multicastMask 指定集群中的哪些 CTA 请求相同的数据。这允许硬件高效地将数据广播到集群中的多个 CTA。
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
forceNoAlias | ::mlir::BoolAttr | 布尔属性 |
操作数:
操作数 |
描述 |
|---|---|
|
LLVM 指针类型 |
|
1位无符号整数 |
|
LLVM 方言兼容类型 |
|
16 位无符号整数 |
结果:
结果 |
描述 |
|---|---|
|
LLVM 方言兼容类型 |
amdg.masked_store (triton::amdgpu::MaskedStoreOp)
带掩码的存储操作
语法
operation ::= `amdg.masked_store` $ptr `,` $value `,` $mask
oilist(`cacheModifier` `=` $cache)
(`forceNoAlias` $forceNoAlias^)?
attr-dict `:` type(operands)
支持掩码的存储操作。如果掩码为真,则从给定的指针存储。作为简化 LLVM 转换的实用操作,使用 LLVM 类型。
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
forceNoAlias | ::mlir::BoolAttr | 布尔属性 |
操作数:
操作数 |
描述 |
|---|---|
|
LLVM 指针类型 |
|
LLVM 方言兼容类型 |
|
1位无符号整数 |
amdg.memory_counter_wait (triton::amdgpu::MemoryCounterWaitOp)
等待指定的硬件计数器
语法
operation ::= `amdg.memory_counter_wait` oilist( `load` `(` $load `)` | `store` `(` $store `)` | `ds` `(` $ds `)` ) attr-dict
在继续之前,等待指定的计数器小于或等于提供的值。
在不同的架构上,计数器可以降低到不同的指令,包括钳位到某些硬件支持的最大值或将多个计数器合并为一个。
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
load | ::mlir::IntegerAttr | 32位无符号整数属性 |
store | ::mlir::IntegerAttr | 32位无符号整数属性 |
ds | ::mlir::IntegerAttr | 32位无符号整数属性 |
amdg.scaled_upcast_fp4 (triton::amdgpu::ScaledUpcastFp4Op)
向上转换 fp4 并乘以缩放因子
语法
operation ::= `amdg.scaled_upcast_fp4` $input `scale` $scale attr-dict
`:` type($input) `,` type($scale) `->` type($output)
将打包为 i8 值的 fp4 (e2m1) 值向上转换,并乘以以 BF16 编码的给定 E8M0 缩放因子。这映射到 AMD CDNA4 架构上的 v_cvt_scalef32_* 内在函数。
i8 的低 4 位表示第一个 fp4 元素,高 4 位表示第二个 fp4 元素。
axis 属性指定了 fp4 元素的打包轴。
特性: AlwaysSpeculatableImplTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), UpcastFpOpInterface
效果: MemoryEffects::Effect{}
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
axis | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
8 位无符号整数值的秩张量 |
|
bfloat16 类型或 8 位无符号整数值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
16 位浮点数或 bfloat16 类型值的秩张量 |
amdg.scaled_upcast_fp8 (triton::amdgpu::ScaledUpcastFp8Op)
向上转换 Fp8 并乘以缩放因子
语法
operation ::= `amdg.scaled_upcast_fp8` $input `scale` $scale attr-dict
`:` type($input) `,` type($scale) `->` type($output)
将 fp8 (e4m3/e5m2) 值向上转换,并乘以以 BF16 编码的给定 E8M0 缩放因子。这映射到 AMD CDNA4 架构上的 v_cvt_scalef32_* 内在函数。
特性(Traits): AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), UpcastFpOpInterface
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
f8E4M3FN 类型或 f8E5M2 类型值的秩张量 |
|
bfloat16 类型或 8 位无符号整数值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
16 位浮点数或 bfloat16 类型值的秩张量 |
amdg.tdm_prefetch (triton::amdgpu::TDMPrefetchOp)
基于 TDM 描述符将数据从全局内存预取到 L2。
语法
operation ::= `amdg.tdm_prefetch` $desc `[` $indices `]` `,` $pred `,` `speculative` `=` $speculative
(`returnOffsets` $returnOffsets^)?
attr-dict `:` qualified(type($desc))
(`->` type($maybeOffsets)^)?
此操作将数据从全局内存预取到 L2。它类似于 AsyncTDMCopyGlobalToLocalOp,但它不将数据复制到本地内存,而是仅将数据预取到 L2 缓存中。推测性预取可以生成更高效的汇编,因为它们不需要越界检查。但是,如果虚拟地址转换尚未在 CU 级别缓存,它们会被硬件丢弃。
接口: InferTypeOpInterface, MemoryEffectOpInterface (MemoryEffectOpInterface), PredicatedOpInterface
效果(Effects): MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::amd::L2Cache}
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
speculative | ::mlir::BoolAttr | 布尔属性 |
returnOffsets | ::mlir::UnitAttr | 单元属性 |
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
可变参数的 32 位无符号整数 |
|
1位无符号整数 |
结果:
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
amdg.update_tensor_descriptor (triton::amdgpu::UpdateTensorDescriptorOp)
更新 TDM 描述符的选定字段;产生一个新的描述符 SSA 值
语法
operation ::= `amdg.update_tensor_descriptor` $desc
oilist(
`add_offsets` `=` `[` $add_offsets `]`
| `set_bounds` `=` `[` $set_bounds `]`
| `dest` `=` $dest `:` qualified(type($dest))
| `pred` `=` $pred
| `barrier` `=` $barrier `:` qualified(type($barrier))
)
attr-dict `:` qualified(type($desc))
返回一个新的 TDM 描述符,其中包含重写的选定字段。每个参数都是独立可选的;仅写入调用者指定的字段。其它所有内容均继承自输入描述符。
参数 |
类型(kind) |
对描述符的影响 |
|---|---|---|
|
增量式 |
|
|
重写 |
|
|
重写 |
|
|
重写 |
|
|
重写 |
|
add_offsets 是增量式的(以元素为单位的增量)并且不触及 tensor_dim — 对于越界(OOB)正确的循环,在循环外设置一次 set_bounds(内部瓦片看到循环不变的 tensor_dim,可以提升),或者在必须触发越界时,按迭代设置 set_bounds。
示例
// K-loop interior: bump tile position only
%d1 = amdg.update_tensor_descriptor %d, add_offsets=[%c0, %BLOCK_K]
: !tt.tensordesc<64x64xf16, #shared>
// Prologue: position at first tile + wire LDS and barrier
%d1 = amdg.update_tensor_descriptor %d,
add_offsets=[%pid_m_off, %c0],
dest=%a_shared, barrier=%a_bar
: !tt.tensordesc<64x64xf16, #shared>
// Peel epilogue: install real OOB extent for the partial last tile
%d1 = amdg.update_tensor_descriptor %d, set_bounds=[%m_remain, %k_remain]
: !tt.tensordesc<64x64xf16, #shared>
特性(Traits): AlwaysSpeculatableImplTrait, AttrSizedOperandSegments
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
可变参数的 32 位无符号整数 |
|
可变参数的 32 位无符号整数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
32位无符号整数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
结果:
结果 |
描述 |
|---|---|
|
张量描述符类型( |
amdg.wait_barrier (triton::amdgpu::WaitBarrierOp)
等待,直到 mbarrier 相位完成。
语法
operation ::= `amdg.wait_barrier` $alloc `,` $phase attr-dict `:` qualified(type($alloc))
阻塞程序进度,直到 alloc 中的 mbarrier 对象完成其当前相位。
接口: MBarrierOpInterface
操作数:
操作数 |
描述 |
|---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
32位无符号整数 |