TritonAMDGPUOps¶
amdgpu.async_tdm_copy_global_to_local (triton::amdgpu::AsyncTDMCopyGlobalToLocalOp)¶
根据描述符异步地将数据从全局内存复制到本地内存
语法
operation ::= `amdgpu.async_tdm_copy_global_to_local` $desc `[` $indices `]` `into` $result `,` $pred
attr-dict `:` qualified(type($desc)) `->` qualified(type($result))
此操作异步地将数据从全局内存复制到本地内存。这类似于 tt.load,但数据被复制到由 result 指向的本地内存,而不是一个分布式张量。复制的数据取决于由 desc 指向的全局内存。将 pred 设置为 false 将禁用复制。此操作不支持共享内存交错(swizzling)。
接口 (Interfaces): InferTypeOpInterface
操作数:¶
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
32位无符号整数的可变参数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
|
1位无符号整数 |
结果:¶
结果 |
描述 |
|---|---|
|
异步令牌类型 |
amdgpu.async_tdm_copy_local_to_global (triton::amdgpu::AsyncTDMCopyLocalToGlobalOp)¶
根据描述符异步地将数据从本地内存复制到全局内存
语法
operation ::= `amdgpu.async_tdm_copy_local_to_global` $desc `[` $indices `]` `from` $src
attr-dict `:` qualified(type($src)) `->` qualified(type($desc))
此操作异步地将数据从本地内存复制到全局内存。这类似于 tt.store,但数据是从由 src 指向的本地内存复制的,而不是一个分布式张量。复制的目标取决于由 desc 指向的全局内存。此操作不支持共享内存填充或交错(swizzling)。
操作数:¶
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
32位无符号整数的可变参数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
amdgpu.async_tdm_wait (triton::amdgpu::AsyncTDMWait)¶
等待直到未完成的 TDM 操作数量小于或等于给定值
语法
operation ::= `amdgpu.async_tdm_wait` $asyncToken attr-dict
此操作会等待,直到未完成的 TDM 操作(包括加载和存储)数量小于或等于给定值。这对于确保数据在使用前已在 LDS (本地数据存储) 中可用是必需的。
接口 (Interfaces): InferTypeOpInterface
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
num | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
异步令牌类型的可变参数 |
结果:¶
结果 |
描述 |
|---|---|
|
异步令牌类型 |
amdgpu.buffer_atomic_cas (triton::amdgpu::BufferAtomicCASOp)¶
原子 CAS 操作,对一个标量基指针和一个张量偏移量进行比较交换
语法
operation ::= `amdgpu.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,因为它对于优化缓存内存访问非常重要。
特性:SameLoadStoreOperandsAndResultEncoding
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
sem | ::mlir::triton::MemSemanticAttr | 允许的32位无符号整数情况:1, 2, 3, 4 |
scope | ::mlir::triton::MemSyncScopeAttr | 允许的32位无符号整数情况:1, 2, 3 |
操作数:¶
操作数 |
描述 |
|---|---|
|
ptr |
|
32位无符号整数值的张量 |
|
浮点数、整数或 ptr 值的秩张量 |
|
浮点数、整数或 ptr 值的秩张量 |
|
32位无符号整数 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
amdgpu.buffer_atomic_rmw (triton::amdgpu::BufferAtomicRMWOp)¶
原子 RMW 操作,对一个标量基指针和一个张量偏移量进行读取、修改和写入
语法
operation ::= `amdgpu.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 处的数据,执行 $rmw_op 操作与 $val,并将结果以指定的内存语义和范围存储到 $ptr。原子 RMW 操作如果被使用,则返回操作前的值,否则该值被隐式丢弃。步幅(Stride)是连续内存块起始位置之间的距离。在执行 RMW 时,stride 是每行第一个元素之间地址的差值(以字节为单位)。编译器在转换为缓冲区操作时会尝试获取 stride,因为它对于优化缓存内存访问非常重要。
特性:AttrSizedOperandSegments, SameLoadStoreOperandsAndResultEncoding
属性:¶
| 属性 | 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 值的秩张量 |
amdgpu.buffer_load (triton::amdgpu::BufferLoadOp)¶
从一个标量基指针和一个张量偏移量加载
语法
operation ::= `amdgpu.buffer_load` $ptr `[` $offsets `]` (`,` $mask^)? (`,` $other^)?
oilist(`cacheModifier` `=` $cache)
(`stride` `=` $stride^)?
attr-dict `:` type($result)
AMD 缓冲区加载操作。缓冲区加载类似于普通加载,但它通过一个标量基指针和一个偏移量张量访问全局内存,而不是通过一个指针张量。其他字段与普通加载类似,即 mask 是一个布尔向量,决定是否应从内存中读取给定元素,而 other 是在 mask[i] == 0 时应在通道 i 上返回的元素。步幅(Stride)是连续内存块起始位置之间的距离。在执行块加载时,stride 是每行第一个元素之间地址的差值(以字节为单位)。编译器在转换为缓冲区操作时会尝试获取 stride,因为它对于优化缓存内存访问非常重要。
特性:AttrSizedOperandSegments, SameLoadStoreOperandsAndResultEncoding
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
操作数:¶
操作数 |
描述 |
|---|---|
|
ptr |
|
32位无符号整数值的张量 |
|
32位无符号整数 |
|
1位无符号整数值的秩张量 |
|
浮点数、整数或 ptr 值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
amdgpu.buffer_load_to_local (triton::amdgpu::BufferLoadToLocalOp)¶
从一个标量基指针和一个张量偏移量加载到共享内存
语法
operation ::= `amdgpu.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 缓冲区加载操作。类似于 amdgpu.buffer_load 操作,但直接写入共享内存而不是寄存器。
特性:AttrSizedOperandSegments
接口 (Interfaces): InferTypeOpInterface
属性 (Attributes):¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
操作数:¶
操作数 |
描述 |
|---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
ptr |
|
32位无符号整数值的张量 |
|
1位无符号整数值的秩张量 |
|
浮点数、整数或 ptr 值的秩张量 |
|
32位无符号整数 |
结果:¶
结果 |
描述 |
|---|---|
|
异步令牌类型 |
amdgpu.buffer_store (triton::amdgpu::BufferStoreOp)¶
存储到标量基指针和一个张量偏移量
语法
operation ::= `amdgpu.buffer_store` $value `,` $ptr `[` $offsets `]` (`,` $mask^)?
oilist(`cacheModifier` `=` $cache)
(`stride` `=` $stride^)?
attr-dict `:` type($value)
AMD 缓冲区存储操作。缓冲区存储类似于普通存储,但它通过一个标量基指针和一个偏移量张量访问全局内存,而不是通过一个指针张量。其他字段与普通存储类似,即 mask 是一个布尔向量,决定是否应将给定元素写入内存,而 value 是在 mask[i] == 1 时应在通道 i 上写入的元素张量。步幅(Stride)是连续内存块起始位置之间的距离。在执行块存储时,stride 是每行第一个元素之间地址的差值(以字节为单位)。编译器在转换为缓冲区操作时会尝试获取 stride,因为它对于优化缓存内存访问非常重要。
特性:AttrSizedOperandSegments, SameLoadStoreOperandsEncoding
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
|
ptr |
|
32位无符号整数值的张量 |
|
32位无符号整数 |
|
1位无符号整数值的秩张量 |
amdgpu.concat (triton::amdgpu::ConcatOp)¶
连接操作
语法
operation ::= `amdgpu.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 = amdgpu.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 = amdgpu.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{}
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点、整数或指针值的秩张量的可变参数 |
结果:¶
结果 |
描述 |
|---|---|
|
任何类型值的秩张量 |
amdgpu.cond_barrier (triton::amdgpu::CondBarrierOp)¶
有条件地设置屏障以同步块中的部分线程
语法
operation ::= `amdgpu.cond_barrier` $pred attr-dict
condBarrierOp 仅在给定参数为 true 时设置屏障指令。这提供了一种同步块中部分线程的方法,有意地使执行序列发散。但是,用户应保证所有线程最终通过调用 condBarrierOp(true) 与剩余线程汇合。从概念上讲,这类似于在 if 语句内设置执行屏障。此操作允许我们在适当的时候避免阻塞整个块,以帮助调度。注意:这不设置任何内存栅栏。
操作数:¶
操作数 |
描述 |
|---|---|
|
1位无符号整数 |
amdgpu.extract_slice (triton::amdgpu::ExtractSliceOp)¶
提取切片操作
语法
operation ::= `amdgpu.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 瓦片(分布在 warps 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 = amdgpu.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 密集数组属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
任何类型值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
任何类型值的秩张量 |
amdgpu.in_thread_transpose (triton::amdgpu::InThreadTransposeOp)¶
对属于每个线程的寄存器值执行转置
语法
operation ::= `amdgpu.in_thread_transpose` $src attr-dict `:` type($src) `->` type($result)
此操作对每个线程中寄存器内的值执行布局转置。具体来说,给定输入布局的分块布局,它会沿着底层线性布局的寄存器维度转置最后两个维度(秩-1和秩-2)。
转换示例
输入布局:分块布局,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 值的秩张量 |
amdgpu.instruction_sched_hint (triton::amdgpu::InstructionSchedHint)¶
一个用于基本块内指令调度提示的占位符操作
语法
operation ::= `amdgpu.instruction_sched_hint` attr-dict
一个用于指令调度提示的占位符操作,应用于占位符操作所在基本块内的指令。此操作主要用于调整 tt.dot 操作产生的主循环内部的指令调度。在较高层级更容易识别点积操作,从而标记预期的调度区域。这些提示操作最终会降级为 LLVM AMDGPU 指令调度原语,旨在控制不同类型的指令(valu/mfma、全局/共享内存等)应如何交错以获得更好的指令级并行性。
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
variant | ::mlir::triton::amdgpu::SchedHintAttr | AMD GPU 的指令调度提示 |
amdgpu.local_load_packed_tranposed (triton::amdgpu::LocalLoadPackedTransposedOp)¶
从共享内存加载一个转置的打包张量到一个分布式张量
语法
operation ::= `amdgpu.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。
特性:LocalLoadTrait
操作数 (Operands):¶
操作数 |
描述 |
|---|---|
|
Triton IR 类型系统中的内存描述符类型 ( |
|
异步令牌类型 |
结果 (Results):¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
amdgpu.masked_load (triton::amdgpu::MaskedLoadOp)¶
带掩码的加载操作
语法
operation ::= `amdgpu.masked_load` $ptr `,` $mask `,` $falseVal
oilist(`cacheModifier` `=` $cache)
(`forceNoAlias` $forceNoAlias^)?
attr-dict `:` functional-type(operands, results)
支持掩码的加载操作。如果掩码为真,则从给定指针加载。与 LLVM 类型一起工作,作为一个简化 LLVM 转换的实用工具操作。
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
forceNoAlias | ::mlir::BoolAttr | 布尔属性 |
操作数 (Operands):¶
操作数 |
描述 |
|---|---|
|
LLVM 指针类型 |
|
1位无符号整数 |
|
LLVM 方言兼容类型 |
结果:¶
结果 |
描述 |
|---|---|
|
LLVM 方言兼容类型 |
amdgpu.masked_store (triton::amdgpu::MaskedStoreOp)¶
带掩码的存储操作
语法
operation ::= `amdgpu.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位无符号整数 |
amdgpu.scaled_upcast_fp4 (triton::amdgpu::ScaledUpcastFp4Op)¶
向上转换 fp4 然后乘以缩放因子
语法
operation ::= `amdgpu.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 类型值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
16位浮点、bfloat16 类型或 32 位浮点值的秩张量 |
amdgpu.scaled_upcast_fp8 (triton::amdgpu::ScaledUpcastFp8Op)¶
向上转换 Fp8 然后乘以缩放因子
语法
operation ::= `amdgpu.scaled_upcast_fp8` $input `scale` $scale attr-dict
`:` type($input) `,` type($scale) `->` type($output)
将 fp8 (e4m3/e5m2) 值向上转换,并与以 BF16 编码的给定 E8M0 缩放因子相乘。这对应于 AMD CDNA4 架构上的 v_cvt_scalef32_* 内在函数。
特性:AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape
接口:ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), UpcastFpOpInterface
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
f8E4M3FN 类型或 f8E5M2 类型值的秩张量 |
|
bfloat16 类型值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
16位浮点、bfloat16 类型或 32 位浮点值的秩张量 |
amdgpu.upcast_mxfp (triton::amdgpu::UpcastMXFPOp)¶
将一个 mxfp 张量转换为 bf16/fp16
语法
operation ::= `amdgpu.upcast_mxfp` $src `,` $scale `fp_type` `=` $fp_type attr-dict `:` type($src) `,` type($scale) `->` type($result)
根据 https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf 计算给定 mxfp 数中编码的 bf16。
特性: AlwaysSpeculatableImplTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
fp_type | ::mlir::triton::ScaleDotElemTypeAttr | 允许的 32 位无符号整数情况:0, 1, 2, 3, 4, 5, 6 |
fastMath | ::mlir::BoolAttr | 布尔属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
|
浮点数、整数或 ptr 值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |