TritonGPUOps¶

`ttg.async_commit_group` (triton::gpu::AsyncCommitGroupOp)¶

异步提交组

语法

operation ::= `ttg.async_commit_group` (`tokens` $inputTokens^)? attr-dict

特质 (Traits): VerifyTensorLayoutsTrait

接口 (Interfaces): InferTypeOpInterface

操作数:¶

操作数	描述
`inputTokens`	async token 类型的可变参数

结果:¶

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

`ttg.async_copy_global_to_local` (triton::gpu::AsyncCopyGlobalToLocalOp)¶

异步地将数据从全局内存复制到本地内存

语法

operation ::= `ttg.async_copy_global_to_local` $src `,` $result (`mask` $mask^)? (`other` $other^)?
              oilist(`cacheModifier` `=` $cache | `evictionPolicy` `=` $evict)
              attr-dict `:` type($src) `->` type($result)

此操作异步地将数据从全局内存复制到本地内存。这类似于 tt.load，但数据被复制到内存描述符指向的本地内存，而不是分布式张量。其余操作数与 tt.load 相同。

特质 (Traits): AttrSizedOperandSegments, VerifyTensorLayoutsTrait

接口 (Interfaces): InferTypeOpInterface

属性：¶

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

操作数：¶

操作数	描述
`src`	ptr 值的秩张量
`result`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`mask`	1 位无符号整数值的张量
`other`	浮点数或浮点数值的秩张量，或整数或整数值的秩张量，或 ptr 或 ptr 值的秩张量，或 ptr

结果：¶

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

`ttg.async_wait` (triton::gpu::AsyncWaitOp)¶

异步等待

语法

operation ::= `ttg.async_wait` ($asyncToken^)? attr-dict

特质 (Traits): VerifyTensorLayoutsTrait

接口 (Interfaces): InferTypeOpInterface

属性 (Attributes):¶

属性	MLIR 类型	描述
`num`	::mlir::IntegerAttr	32位无符号整数属性

操作数:¶

操作数	描述
`asyncToken`	async token 类型的可变参数

结果：¶

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

`ttg.convert_layout` (triton::gpu::ConvertLayoutOp)¶

转换布局

语法

operation ::= `ttg.convert_layout` $src attr-dict `:` type($src) `->` type($result)

特质 (Traits): AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, SameOperandsAndResultShape, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数 (Operands):¶

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

结果：¶

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

`ttg.fp4_to_fp` (triton::gpu::Fp4ToFpOp)¶

将 fp4 (e2m1) 上转换为 fp

语法

operation ::= `ttg.fp4_to_fp` $src attr-dict `:` type($src) `->` type($result)

将以 i8 紧凑表示的 fp4 (e2m1) 上转换为 fp。

i8 的低 4 位表示第一个 fp4 元素，高 4 位表示第二个 fp4 元素。

axis 属性指定了 fp4 元素打包的轴。

特质 (Traits): AlwaysSpeculatableImplTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性：¶

属性	MLIR 类型	描述
`axis`	::mlir::IntegerAttr	32位无符号整数属性

操作数:¶

操作数	描述
`src`	8 位无符号整数值的秩张量

结果：¶

结果	描述
`result`	浮点数值的秩张量

`ttg.global_scratch_alloc` (triton::gpu::GlobalScratchAllocOp)¶

分配全局内存缓冲区

语法

operation ::= `ttg.global_scratch_alloc` attr-dict `:` qualified(type($result))

此操作在全局内存中分配一个对当前程序私有的缓冲区。

特质 (Traits): VerifyTensorLayoutsTrait

属性:¶

属性	MLIR 类型	描述
`nbytes`	::mlir::IntegerAttr	32位无符号整数属性
`alignment`	::mlir::IntegerAttr	32位无符号整数属性

结果:¶

结果	描述
`result`	ptr

`ttg.local_alloc` (triton::gpu::LocalAllocOp)¶

分配张量

语法

operation ::= `ttg.local_alloc` ($src^)? attr-dict `:` functional-type(operands, results)

此操作在共享内存中分配缓冲区，并返回一个包含地址和缓冲区视图的描述符。

显式释放缓冲区是可选的；参见 local_dealloc。

src 操作数是分配缓冲区的可选初始化器。它必须与缓冲区的元素类型相同。如果未指定 src，返回的缓冲区必须是可变的。

特质 (Traits): VerifyTensorLayoutsTrait

接口： MemoryEffectOpInterface

属性 (Attributes):¶

属性	MLIR 类型	描述
`alignment`	::mlir::IntegerAttr	32位无符号整数属性

操作数:¶

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

结果：¶

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

`ttg.local_barrier` (triton::gpu::LocalBarrierOp)¶

同步 CTA 中所有线程的执行以及共享内存的读/写操作。

语法

operation ::= `ttg.local_barrier` attr-dict

local_barrier 操作同步 CTA 中所有线程的执行以及共享内存和寄存器之间的所有操作。它用于协调 CTA 中线程之间的通信。

此操作会等待，直到 CTA 中的所有线程都到达一个 local_barrier，并且这些线程在该操作之前对共享内存和寄存器所做的操作对 CTA 中的所有线程可见。

通过在访问相同内存的线程之间使用 local_barrier 进行同步，可以避免数据冒险。

local_barrier 操作不提供对全局内存的同步保证。

特质 (Traits): VerifyTensorLayoutsTrait

`ttg.local_dealloc` (triton::gpu::LocalDeallocOp)¶

释放缓冲区

语法

operation ::= `ttg.local_dealloc` $src attr-dict `:` qualified(type($src))

此操作显式地释放一个缓冲区。在此操作之后使用该缓冲区是未定义的。

此操作是可选的。如果你不显式释放缓冲区，编译器会假定它在后支配所有 alloc 用法的第一个点被释放。

因为我们假定一个内存描述符在其所有用法的后支配点即为死亡，所以等待内存描述符上异步操作完成的操作（例如 ttng.warp_group_dot_wait）也应该将内存描述符作为操作数。

特质 (Traits): VerifyTensorLayoutsTrait

操作数 (Operands):¶

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

`ttg.local_load` (triton::gpu::LocalLoadOp)¶

将缓冲区从本地内存加载到分布式张量

语法

operation ::= `ttg.local_load` $src (`token` $token^)? attr-dict `:` qualified(type($src)) `->` type($result)

从本地内存描述符中将一个张量加载到一个分布式张量中。

特质 (Traits): LocalLoadTrait, VerifyTensorLayoutsTrait

操作数:¶

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

结果：¶

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

`ttg.local_store` (triton::gpu::LocalStoreOp)¶

将分布式张量存储到本地内存的缓冲区中

语法

operation ::= `ttg.local_store` $src `,` $dst attr-dict `:` type($src) `->` qualified(type($dst))

将一个分布式张量存储到本地内存的缓冲区中。

特质 (Traits): VerifyTensorLayoutsTrait

操作数 (Operands):¶

操作数	描述
`src`	浮点数、整数或 ptr 值的秩张量
`dst`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)

`ttg.mask` (triton::gpu::MaskOp)¶

用于流水线操作的掩码操作

特质 (Traits): SingleBlock, VerifyTensorLayoutsTrait

操作数：¶

操作数	描述
`pred`	1位无符号整数

结果：¶

结果	描述
`result`	任意类型的可变参数

`ttg.mask.return` (triton::gpu::MaskReturnOp)¶

掩码操作符的终止符

语法

operation ::= `ttg.mask.return` $result attr-dict `:` type($result)

特质 (Traits): AlwaysSpeculatableImplTrait, HasParent<MaskOp>, ReturnLike, Terminator, VerifyTensorLayoutsTrait

接口 (Interfaces): ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

效果: MemoryEffects::Effect{}

操作数：¶

操作数	描述
`result`	任意类型的可变参数

`ttg.memdesc_index` (triton::gpu::MemDescIndexOp)¶

获取描述符的子视图。

语法

operation ::= `ttg.memdesc_index` $src `[` $index `]` attr-dict `:` qualified(type($src)) `->` qualified(type($result))

此操作返回一个指向输入描述符沿第 0 维的第 i 个元素的新描述符。

它不影响底层内存。

例如，假设

输入形状为 2x4x16xf16，
输出形状为 4x16xf16，并且
index = 1。那么输出描述符等价于 input[1]，其中 input 是逻辑张量。

特质 (Traits): AlwaysSpeculatableImplTrait, MemDescViewTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数：¶

操作数	描述
`src`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)
`index`	32位无符号整数

结果：¶

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

`ttg.memdesc_reinterpret` (triton::gpu::MemDescReinterpretOp)¶

将内存描述符重新解释为不同的类型和形状

语法

operation ::= `ttg.memdesc_reinterpret` $src attr-dict `:` qualified(type($src)) `->` qualified(type($result))

ttg.memdesc_reinterpret 操作将一个内存描述符重新解释为具有不同形状和元素类型的描述符。由于内存描述符缺少步长信息，此操作仅在原始内存描述符是连续的情况下有效。

特质 (Traits): AlwaysSpeculatableImplTrait, MemDescViewTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数：¶

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

结果：¶

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

`ttg.memdesc_reshape` (triton::gpu::MemDescReshapeOp)¶

为新形状创建描述符

语法

operation ::= `ttg.memdesc_reshape` $src attr-dict `:` qualified(type($src)) `->` qualified(type($result))

此操作返回一个新描述符，表示底层缓冲区的重塑视图。这不影响内存。

特质 (Traits): AlwaysSpeculatableImplTrait, MemDescViewTrait, SameOperandsAndResultElementType, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数 (Operands):¶

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

结果 (Results):¶

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

`ttg.memdesc_subslice` (triton::gpu::MemDescSubsliceOp)¶

获取描述符的子视图。

语法

operation ::= `ttg.memdesc_subslice` $src `[` custom<Offsets>($offsets) `]` attr-dict `:` qualified(type($src))
              `->` qualified(type($result))

此操作返回一个表示逻辑张量子视图的新描述符。它不影响底层内存。

例如，假设

输入形状为 32x16xf16，
输出形状为 8x16xf16，并且
offsets = [2, 1]。那么在 Python 语法中，子视图覆盖了 input[2:8+2, 1:16+1]，其中 input 是逻辑张量。

偏移量必须大于或等于张量的 tile（或零）。

特质 (Traits): AlwaysSpeculatableImplTrait, MemDescViewTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性：¶

属性	MLIR 类型	描述
`offsets`	::mlir::DenseI32ArrayAttr	i32 密集数组属性

操作数 (Operands):¶

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

结果：¶

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

`ttg.memdesc_trans` (triton::gpu::MemDescTransOp)¶

转置描述符

语法

operation ::= `ttg.memdesc_trans` $src attr-dict `:` qualified(type($src)) `->` qualified(type($result))

此操作返回一个表示缓冲区转置视图的新描述符。

特质 (Traits): AlwaysSpeculatableImplTrait, InferTypeOpAdaptor, MemDescViewTrait, SameOperandsAndResultElementType, VerifyTensorLayoutsTrait

接口 (Interfaces): ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface), TransposeOpInterface

效果: MemoryEffects::Effect{}

属性：¶

属性	MLIR 类型	描述
`order`	::mlir::DenseI32ArrayAttr	i32 密集数组属性

操作数：¶

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

结果：¶

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

`ttg.predicate_stage` (triton::gpu::PredicateStageOp)¶

流水线阶段谓词

语法

operation ::= `ttg.predicate_stage` $iv `,` $ub `,` $step `maxStage` $maxStage `stage` $stage attr-dict `:` type($iv) `->` type($result)

特质 (Traits): AlwaysSpeculatableImplTrait, VerifyTensorLayoutsTrait

接口：ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性：¶

属性	MLIR 类型	描述
`maxStage`	::mlir::IntegerAttr	32位无符号整数属性
`stage`	::mlir::IntegerAttr	32位无符号整数属性

操作数：¶

操作数	描述
`iv`	无符号整数或索引
`ub`	无符号整数或索引
`step`	无符号整数或索引

结果：¶

结果	描述
`result`	1位无符号整数

`ttg.warp_return` (triton::gpu::WarpReturnOp)¶

分区区域的隐式终止符

语法

operation ::= `ttg.warp_return` attr-dict

ttg.warp_return 操作是结束 ttg.warp_specialize 操作的分区区域的隐式终止符。它没有操作数，因为这些区域不能返回任何东西。

待办事项：支持从分区区域返回统一值。

特质 (Traits): AlwaysSpeculatableImplTrait, HasParent<WarpSpecializePartitionsOp>, ReturnLike, Terminator, VerifyTensorLayoutsTrait

接口 (Interfaces): ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

效果: MemoryEffects::Effect{}

`ttg.warp_specialize` (triton::gpu::WarpSpecializeOp)¶

在多个 warp 组上异步执行代码

ttg.warp_specialize 操作表示在不同的 warp 组上同时执行不同的代码。一个 warp 组是一组 2 的幂次方的 warp，其 warp 数量可以与外部区域中的 warp 数量不同。

该操作的“默认”区域表示当前执行的 warp 组执行的代码。该区域允许隐式捕获。该操作包含多个“分区”区域，这些区域与上层隔离。它们必须被隔离，因为这些区域代表不同的布局域，因为 warp 的数量不同。

从语义上讲，每个区域的执行对每个 warp 组同时开始，并且所有 warp 组在该操作结束时汇合。

示例

%0 = ttg.warp_specialize(%a, %b)
default {
  %out = some_operation(%a) // implicit capture of `%a`
  ttg.warp_yield %out : i32
}
partition0(%arg0: i32, %arg1: i32) num_warps(8) {
  some_async_dispatch(%arg0, %arg1)
  ttg.warp_return
}
partition1(%arg0: i32, %arg1: i32) num_warps(1) {
  some_async_dispatch(%arg0, %arg1)
  ttg.warp_return
} : (i32, i32) -> i32

特质 (Traits): AsyncRegions, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, VerifyTensorLayoutsTrait

接口 (Interfaces): ConditionallySpeculatable, RegionBranchOpInterface

属性：¶

属性	MLIR 类型	描述
`partitionNumWarps`	::mlir::DenseI32ArrayAttr	i32 密集数组属性
`warpGroupStartIds`	::mlir::DenseI32ArrayAttr	i32 密集数组属性
`requestedRegisters`	::mlir::DenseI32ArrayAttr	i32 密集数组属性
`actualRegisters`	::mlir::DenseI32ArrayAttr	i32 密集数组属性

操作数：¶

操作数	描述
`explicitCaptures`	任意类型的可变参数

结果：¶

结果	描述
`defaultPassthrough`	任意类型的可变参数

`ttg.warp_specialize.partitions` (triton::gpu::WarpSpecializePartitionsOp)¶

ttg.warp_specialize 的容器操作

由于 MLIR 要求整个操作与上层隔离，此操作包含了 ttg.warp_specialize 中实际与上层隔离的区域。

特质 (Traits): HasParent<WarpSpecializeOp>, IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, Terminator, VerifyTensorLayoutsTrait

接口 (Interfaces): ConditionallySpeculatable

`ttg.warp_yield` (triton::gpu::WarpYieldOp)¶

从 ttg.warp_specialize 的默认区域产生

语法

operation ::= `ttg.warp_yield` ($values^)? attr-dict (`:` type($values)^)?

ttg.warp_yield 操作是 ttg.warp_specialize 操作的“默认”区域的终止符。其操作数作为 ttg.warp_specialize 操作的 SSA 结果被透明地传递。

示例

ttg.warp_yield %a, %b : i32, tensor<32xbf16, #blocked>

特质 (Traits): AlwaysSpeculatableImplTrait, HasParent<WarpSpecializeOp>, ReturnLike, Terminator, VerifyTensorLayoutsTrait

接口 (Interfaces): ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

效果: MemoryEffects::Effect{}

操作数：¶

操作数	描述
`values`	任意类型的可变参数

TritonGPUOps¶

ttg.async_commit_group (triton::gpu::AsyncCommitGroupOp)¶

操作数:¶

结果:¶

ttg.async_copy_global_to_local (triton::gpu::AsyncCopyGlobalToLocalOp)¶

属性：¶

操作数：¶

结果：¶

ttg.async_wait (triton::gpu::AsyncWaitOp)¶

属性 (Attributes):¶

操作数:¶

结果：¶

ttg.convert_layout (triton::gpu::ConvertLayoutOp)¶

操作数 (Operands):¶

结果：¶

ttg.fp4_to_fp (triton::gpu::Fp4ToFpOp)¶

属性：¶

操作数:¶

结果：¶

ttg.global_scratch_alloc (triton::gpu::GlobalScratchAllocOp)¶

属性:¶

结果:¶

ttg.local_alloc (triton::gpu::LocalAllocOp)¶

属性 (Attributes):¶

操作数:¶

结果：¶

ttg.local_barrier (triton::gpu::LocalBarrierOp)¶

ttg.local_dealloc (triton::gpu::LocalDeallocOp)¶

操作数 (Operands):¶

ttg.local_load (triton::gpu::LocalLoadOp)¶

操作数:¶

结果：¶

ttg.local_store (triton::gpu::LocalStoreOp)¶

操作数 (Operands):¶

ttg.mask (triton::gpu::MaskOp)¶

操作数：¶

结果：¶

ttg.mask.return (triton::gpu::MaskReturnOp)¶

操作数：¶

ttg.memdesc_index (triton::gpu::MemDescIndexOp)¶

操作数：¶

结果：¶

ttg.memdesc_reinterpret (triton::gpu::MemDescReinterpretOp)¶

操作数：¶

结果：¶

ttg.memdesc_reshape (triton::gpu::MemDescReshapeOp)¶

操作数 (Operands):¶

结果 (Results):¶

ttg.memdesc_subslice (triton::gpu::MemDescSubsliceOp)¶

属性：¶

操作数 (Operands):¶

结果：¶

ttg.memdesc_trans (triton::gpu::MemDescTransOp)¶

属性：¶

操作数：¶

结果：¶

ttg.predicate_stage (triton::gpu::PredicateStageOp)¶

属性：¶

操作数：¶

结果：¶

ttg.warp_return (triton::gpu::WarpReturnOp)¶

ttg.warp_specialize (triton::gpu::WarpSpecializeOp)¶

属性：¶

操作数：¶

结果：¶

ttg.warp_specialize.partitions (triton::gpu::WarpSpecializePartitionsOp)¶

ttg.warp_yield (triton::gpu::WarpYieldOp)¶

操作数：¶

`ttg.async_commit_group` (triton::gpu::AsyncCommitGroupOp)¶

`ttg.async_copy_global_to_local` (triton::gpu::AsyncCopyGlobalToLocalOp)¶

`ttg.async_wait` (triton::gpu::AsyncWaitOp)¶

`ttg.convert_layout` (triton::gpu::ConvertLayoutOp)¶

`ttg.fp4_to_fp` (triton::gpu::Fp4ToFpOp)¶

`ttg.global_scratch_alloc` (triton::gpu::GlobalScratchAllocOp)¶

`ttg.local_alloc` (triton::gpu::LocalAllocOp)¶

`ttg.local_barrier` (triton::gpu::LocalBarrierOp)¶

`ttg.local_dealloc` (triton::gpu::LocalDeallocOp)¶

`ttg.local_load` (triton::gpu::LocalLoadOp)¶

`ttg.local_store` (triton::gpu::LocalStoreOp)¶

`ttg.mask` (triton::gpu::MaskOp)¶

`ttg.mask.return` (triton::gpu::MaskReturnOp)¶

`ttg.memdesc_index` (triton::gpu::MemDescIndexOp)¶

`ttg.memdesc_reinterpret` (triton::gpu::MemDescReinterpretOp)¶

`ttg.memdesc_reshape` (triton::gpu::MemDescReshapeOp)¶

`ttg.memdesc_subslice` (triton::gpu::MemDescSubsliceOp)¶

`ttg.memdesc_trans` (triton::gpu::MemDescTransOp)¶

`ttg.predicate_stage` (triton::gpu::PredicateStageOp)¶

`ttg.warp_return` (triton::gpu::WarpReturnOp)¶

`ttg.warp_specialize` (triton::gpu::WarpSpecializeOp)¶

`ttg.warp_specialize.partitions` (triton::gpu::WarpSpecializePartitionsOp)¶

`ttg.warp_yield` (triton::gpu::WarpYieldOp)¶