TritonAMDGPUOps¶

`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 操作如果被使用，则返回操作前的值，否则该值被隐式丢弃。步幅是连续内存块起始点之间的距离。执行 CAS 时，stride 是每行第一个元素之间以字节计的地址差。编译器在转换为缓冲区操作时会尝试获取 stride，因为它对于优化缓存内存访问非常重要。

特性: SameLoadStoreOperandsAndResultEncoding

属性：¶

属性	MLIR 类型	描述
`sem`	::mlir::triton::MemSemanticAttr	允许的 32 位无符号整数情况：1, 2, 3, 4
`scope`	::mlir::triton::MemSyncScopeAttr	允许的32位无符号整数情况：1, 2, 3

操作数:¶

操作数	描述
`ptr`	ptr
`offsets`	32 位无符号整数值张量
`cmp`	浮点数、整数或 ptr 值的秩张量
`val`	浮点数、整数或 ptr 值的秩张量
`stride`	32位无符号整数

结果:¶

结果	描述
`result`	浮点数、整数或 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` 处加载数据，使用 `$val` 执行 `$rmw_op`，并以指定的内存语义和范围将结果存储到 `$ptr`。原子 RMW 操作如果被使用，则返回操作前的值，否则该值被隐式丢弃。步幅是连续内存块起始点之间的距离。执行 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`	ptr
`offsets`	32 位无符号整数值张量
`value`	浮点数、整数或 ptr 值的秩张量
`stride`	32位无符号整数
`mask`	1 位无符号整数值的分级张量

结果:¶

结果	描述
`result`	浮点数、整数或 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，因为它对于优化缓存内存访问非常重要。

特性: AttrSizedOperandSegments, SameLoadStoreOperandsAndResultEncoding

属性:¶

属性	MLIR 类型	描述
`cache`	::mlir::triton::CacheModifierAttr	允许的32位无符号整数情况：1, 2, 3, 4, 5, 6, 7

操作数:¶

操作数	描述
`ptr`	ptr
`offsets`	32 位无符号整数值张量
`stride`	32位无符号整数
`mask`	1 位无符号整数值的分级张量
`other`	浮点数、整数或 ptr 值的秩张量

结果:¶

结果	描述
`result`	浮点数、整数或 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

属性:¶

属性	MLIR 类型	描述
`cache`	::mlir::triton::CacheModifierAttr	允许的32位无符号整数情况：1, 2, 3, 4, 5, 6, 7

操作数:¶

操作数	描述
`dest`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`ptr`	ptr
`offsets`	32 位无符号整数值张量
`mask`	1 位无符号整数值的分级张量
`other`	浮点数、整数或 ptr 值的秩张量
`stride`	32位无符号整数

结果:¶

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

`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，因为它对于优化缓存内存访问非常重要。

特性: AttrSizedOperandSegments, SameLoadStoreOperandsEncoding

属性:¶

属性	MLIR 类型	描述
`cache`	::mlir::triton::CacheModifierAttr	允许的32位无符号整数情况：1, 2, 3, 4, 5, 6, 7

操作数:¶

操作数	描述
`value`	浮点数、整数或 ptr 值的秩张量
`ptr`	ptr
`offsets`	32 位无符号整数值张量
`stride`	32位无符号整数
`mask`	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 瓦片的逻辑张量形状的子集上对齐。

这确保了连接是一个无操作，意味着不需要在线程之间进行数据重新排列即可组装具有给定形状和布局的目标张量。然而，布局内 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{}

操作数:¶

操作数	描述
`sources`	浮点数、整数或指针值的分级张量可变参数

结果:¶

结果	描述
`result`	任何类型值的分级张量

`amdgpu.cond_barrier` (triton::amdgpu::CondBarrierOp)¶

有条件地设置屏障以同步块中的部分线程

语法

operation ::= `amdgpu.cond_barrier` $pred attr-dict

condBarrierOp 仅当给定参数为真时才设置屏障指令。这提供了一种在块中同步部分线程的方式，故意使执行序列发散。然而，用户应通过使用剩余线程调用 condBarrierOp(true) 来保证所有线程最终会收敛。从概念上讲，这类似于在 if 语句内部设置一个执行屏障。此操作允许我们在适当的时候避免阻塞整个块，以帮助调度。注意：这不会设置任何内存屏障。

操作数:¶

操作数	描述
`pred`	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 是一个无操作，意味着不需要在线程之间进行数据重新排列即可提取具有给定形状和布局的目标张量。

+——-+——-+ | 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 = 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 密集数组属性

操作数 (Operands):¶

操作数	描述
`source`	任何类型值的分级张量

结果:¶

结果	描述
`result`	任何类型值的分级张量

`amdgpu.in_thread_transpose` (triton::amdgpu::InThreadTransposeOp)¶

对属于每个线程的寄存器值执行转置

语法

operation ::= `amdgpu.in_thread_transpose` $src attr-dict `:` type($src) `->` type($result)

此操作对每个线程寄存器中的值执行布局转置。具体来说，给定输入布局的块状布局，它会沿底层线性布局的寄存器维度转置最后两个维度（rank-1 和 rank-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 高效合并加载，并在 HBM 和共享内存顺序不同以及目标 AMD 硬件不原生支持此转置的情况下，紧随其后地向量化写入共享内存。这是 ttg.convert_layout 的一个特定变体，在降低到 LLVM 时将转换为 ttg.convert_layout。我们不希望此转换被优化掉，因为我们需要在从 HBM 加载后和写入共享内存之前显式地实例化指令以在每个线程内执行转置。

特性: AlwaysSpeculatableImplTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数:¶

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

结果:¶

结果	描述
`result`	浮点数、整数或 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

操作数:¶

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

结果:¶

结果	描述
`result`	浮点数、整数或 ptr 值的秩张量

`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	布尔属性

操作数:¶

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

结果:¶

结果	描述
`result`	浮点数、整数或 ptr 值的秩张量

TritonAMDGPUOps¶

amdgpu.buffer_atomic_cas (triton::amdgpu::BufferAtomicCASOp)¶

属性：¶

操作数:¶

结果:¶

amdgpu.buffer_atomic_rmw (triton::amdgpu::BufferAtomicRMWOp)¶

属性:¶

操作数:¶

结果:¶

amdgpu.buffer_load (triton::amdgpu::BufferLoadOp)¶

属性:¶

操作数:¶

结果:¶

amdgpu.buffer_load_to_local (triton::amdgpu::BufferLoadToLocalOp)¶

属性:¶

操作数:¶

结果:¶

amdgpu.buffer_store (triton::amdgpu::BufferStoreOp)¶

属性:¶

操作数:¶

amdgpu.concat (triton::amdgpu::ConcatOp)¶

操作数:¶

结果:¶

amdgpu.cond_barrier (triton::amdgpu::CondBarrierOp)¶

操作数:¶

amdgpu.extract_slice (triton::amdgpu::ExtractSliceOp)¶

属性:¶

操作数 (Operands):¶

结果:¶

amdgpu.in_thread_transpose (triton::amdgpu::InThreadTransposeOp)¶

操作数:¶

结果:¶

amdgpu.instruction_sched_hint (triton::amdgpu::InstructionSchedHint)¶

属性:¶

amdgpu.local_load_packed_tranposed (triton::amdgpu::LocalLoadPackedTransposedOp)¶

操作数:¶

结果:¶

amdgpu.upcast_mxfp (triton::amdgpu::UpcastMXFPOp)¶

属性:¶

操作数:¶

结果:¶

`amdgpu.buffer_atomic_cas` (triton::amdgpu::BufferAtomicCASOp)¶

`amdgpu.buffer_atomic_rmw` (triton::amdgpu::BufferAtomicRMWOp)¶

`amdgpu.buffer_load` (triton::amdgpu::BufferLoadOp)¶

`amdgpu.buffer_load_to_local` (triton::amdgpu::BufferLoadToLocalOp)¶

`amdgpu.buffer_store` (triton::amdgpu::BufferStoreOp)¶

`amdgpu.concat` (triton::amdgpu::ConcatOp)¶

`amdgpu.cond_barrier` (triton::amdgpu::CondBarrierOp)¶

`amdgpu.extract_slice` (triton::amdgpu::ExtractSliceOp)¶

`amdgpu.in_thread_transpose` (triton::amdgpu::InThreadTransposeOp)¶

`amdgpu.instruction_sched_hint` (triton::amdgpu::InstructionSchedHint)¶

`amdgpu.local_load_packed_tranposed` (triton::amdgpu::LocalLoadPackedTransposedOp)¶

`amdgpu.upcast_mxfp` (triton::amdgpu::UpcastMXFPOp)¶