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	允许的 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

特性: VerifyTensorLayoutsTrait

接口: InferTypeOpInterface

属性:¶

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

操作数:¶

操作数	描述
`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`	浮点值或整数值或 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 元素打包的轴。

特性: 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))

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

特性: 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，则返回的缓冲区必须是可变的。

特性: VerifyTensorLayoutsTrait

接口: MemoryEffectOpInterface

属性:¶

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

操作数:¶

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

结果:¶

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

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

释放缓冲区

语法

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

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

此操作是可选的。如果不显式释放缓冲区，编译器假定它在所有使用 alloc 的后支配点中最先出现的点被释放。

因为我们假定 memdesc 在其使用后的后支配点中最先出现的点无效，所以等待 memdesc 上的异步操作完成的操作（例如 ttng.warp_group_dot_wait）也应将 memdesc 作为操作数。

特性: 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`	浮点值或整数值或 ptr 值的分级张量

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

将分布式张量存储到局部内存的缓冲区中

语法

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

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

特性: VerifyTensorLayoutsTrait

操作数:¶

操作数	描述
`src`	浮点值或整数值或 ptr 值的分级张量
`dst`	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))

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

特性: AlwaysSpeculatableImplTrait, MemDescViewTrait, SameOperandsAndResultElementType, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数:¶

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

结果:¶

结果	描述
`result`	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, MemDescViewTrait, 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, MemDescViewTrait, SameOperandsAndResultElementType, VerifyTensorLayoutsTrait

接口: 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)

特性: 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 操作的分区区域的隐式终止符。由于这些区域不能返回任何内容，因此它没有操作数。

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

特性: AlwaysSpeculatableImplTrait, HasParent<WarpSpecializePartitionsOp>, ReturnLike, Terminator, VerifyTensorLayoutsTrait

接口: 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 组在该操作结束时汇合。

示例

%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 密集数组属性
`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 实际的与上方隔离的区域。

特性: HasParent<WarpSpecializeOp>, IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, Terminator, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable

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

从 ttg.warp_specialize 的默认区域 yield

语法

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<WarpSpecializeOp>, 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_reshape (triton::gpu::MemDescReshapeOp)¶

操作数:¶

结果:¶

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

操作数:¶

结果:¶

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

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

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

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

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

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