TritonGPUOps¶

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

异步提交组

语法：

operation ::= `ttg.async_commit_group` $inputTokens attr-dict

特性: VerifyTensorLayoutsTrait

接口: InferTypeOpInterface

操作数:¶

操作数	描述
`inputTokens`	异步令牌类型的可变参数

结果:¶

结果	描述
`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相同。

特性: AttrSizedOperandSegments, VerifyTensorLayoutsTrait

接口: InferTypeOpInterface

属性:¶

属性	MLIR类型	描述
`cache`	::mlir::triton::CacheModifierAttr	allowed 32-bit signless integer cases: 1, 2, 3, 4, 5, 6, 7
`evict`	::mlir::triton::EvictionPolicyAttr	allowed 32-bit signless integer cases: 1, 2, 3
`isVolatile`	::mlir::BoolAttr	bool attribute

操作数:¶

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

结果:¶

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

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

异步等待

语法：

operation ::= `ttg.async_wait` $asyncToken attr-dict

特性: VerifyTensorLayoutsTrait

接口: InferTypeOpInterface

属性:¶

属性	MLIR类型	描述
`num`	::mlir::IntegerAttr	32-bit signless integer attribute

操作数:¶

操作数	描述
`asyncToken`	异步令牌类型的可变参数

结果:¶

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

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

转换布局

语法：

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

特性: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, SameOperandsAndResultShape, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数:¶

操作数	描述
`src`	浮点数、整数或指针值的排序张量

结果:¶

结果	描述
`result`	浮点数、整数或指针值的排序张量

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

将fp4 (e2m1) 上转换为fp

语法：

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

将表示为i8s的fp4 (e2m1)上转换为fp。

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

axis 属性指定了 fp4 元素沿哪个轴进行打包。

特性: AlwaysSpeculatableImplTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性:¶

属性	MLIR类型	描述
`axis`	::mlir::IntegerAttr	32-bit signless integer attribute

操作数:¶

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

结果:¶

结果	描述
`result`	浮点值的排序张量

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

分配一个全局内存缓冲区

语法：

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

此操作会在全局内存中分配一个专属于当前程序的缓冲区。

特性: VerifyTensorLayoutsTrait

属性:¶

属性	MLIR类型	描述
`nbytes`	::mlir::IntegerAttr	32-bit signless integer attribute
`alignment`	::mlir::IntegerAttr	32-bit signless integer attribute

结果:¶

结果	描述
`result`	ptr

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

分配张量

语法：

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

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

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

特性: VerifyTensorLayoutsTrait

接口: MemoryEffectOpInterface

属性:¶

属性	MLIR类型	描述
`alignment`	::mlir::IntegerAttr	32-bit signless integer attribute

操作数:¶

操作数	描述
`src`	浮点数、整数或指针值的排序张量

结果:¶

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

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

释放缓冲区

语法：

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

此操作显式释放缓冲区。在此操作之后使用缓冲区是未定义行为。

此操作是可选的。如果您没有显式地释放缓冲区，编译器会假设它在支配所有分配使用后的第一个点被释放。

因为我们假设内存描述符在其使用点后支配的第一个点就已经失效，所以等待内存描述符上异步操作完成的操作（例如ttng.warp_group_dot_wait）也应该将该内存描述符作为操作数。

特性: VerifyTensorLayoutsTrait

操作数:¶

操作数	描述
`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)

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

特性: VerifyTensorLayoutsTrait

操作数:¶

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

结果:¶

结果	描述
`result`	浮点数、整数或指针值的排序张量

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

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

语法：

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

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

特性: VerifyTensorLayoutsTrait

操作数:¶

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

`ttg.memdesc_subview` (triton::gpu::MemDescSubviewOp)¶

获取描述符的子视图。

语法：

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

该操作返回一个新的描述符，表示缓冲区的子视图。它不会影响底层内存。子视图可以进行秩缩减。

例如，假设

输入形状为2x4x16xf16，
输出形状为4x4xf16，且
offsets = [1, 0, 4].

那么在Python语法中，子视图覆盖了input[1][0:4][4:8]。

特性: AlwaysSpeculatableImplTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数:¶

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

结果:¶

结果	描述
`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))

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

特性: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor, SameOperandsAndResultElementType, VerifyTensorLayoutsTrait

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

效果: MemoryEffects::Effect{}

属性:¶

属性	MLIR类型	描述
`order`	::mlir::DenseI32ArrayAttr	i32 dense array attribute

操作数:¶

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

结果:¶

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

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

来自分区区域的隐式终止符

语法：

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

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

TODO: 支持从分区区域返回统一值。

特性: AlwaysSpeculatableImplTrait, HasParent, ReturnLike, Terminator, VerifyTensorLayoutsTrait

接口：ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

效果: MemoryEffects::Effect{}

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

在多个warpgroups上异步执行代码

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

该操作的"default"区域表示当前执行的warp组所执行的代码。该区域允许隐式捕获。该操作包含多个与上层隔离的"partition"区域。这些区域必须隔离，因为它们代表不同的布局域，因为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

特性：AsyncRegions, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, RegionBranchOpInterface

属性:¶

属性	MLIR类型	描述
`partitionNumWarps`	::mlir::DenseI32ArrayAttr	i32 dense array attribute
`warpGroupStartIds`	::mlir::DenseI32ArrayAttr	i32 dense array attribute
`requestedRegisters`	::mlir::DenseI32ArrayAttr	i32 dense array attribute
`actualRegisters`	::mlir::DenseI32ArrayAttr	i32 dense array attribute

操作数:¶

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

结果:¶

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

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

用于ttg.warp_specialize的容器操作

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

特性: HasParent, IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, Terminator, VerifyTensorLayoutsTrait

接口：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>

特性: AlwaysSpeculatableImplTrait, HasParent, ReturnLike, Terminator, VerifyTensorLayoutsTrait

接口：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)¶

属性:¶

操作数:¶

结果:¶

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_dealloc (triton::gpu::LocalDeallocOp)¶

操作数:¶

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

操作数:¶

结果:¶

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

操作数:¶

ttg.memdesc_subview (triton::gpu::MemDescSubviewOp)¶

操作数:¶

结果:¶

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

属性:¶

操作数:¶

结果:¶

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_dealloc` (triton::gpu::LocalDeallocOp)¶

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

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

`ttg.memdesc_subview` (triton::gpu::MemDescSubviewOp)¶

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

`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)¶