TritonGPUOps¶

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

异步提交组

语法

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

特性 (Traits): VerifyTensorLayoutsTrait

接口 (Interfaces): 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 相同。

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

操作数 (Operands):¶

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

结果 (Results):¶

结果	描述
`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`	可变参数的异步令牌类型

结果 (Results):¶

结果	描述
`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 值的秩张量

结果 (Results):¶

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

`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。

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

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

特性 (Traits): AlwaysSpeculatableImplTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性 (Attributes):¶

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

操作数:¶

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

结果 (Results):¶

结果	描述
`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 值的秩张量

结果 (Results):¶

结果	描述
`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）也应将该内存描述符作为操作数。

特性 (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`	异步令牌类型

结果 (Results):¶

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

用于流水线的 Mask 操作

特性 (Traits): SingleBlock, VerifyTensorLayoutsTrait

操作数 (Operands):¶

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

结果 (Results):¶

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

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

mask 操作符的终止器

语法

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

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

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

效果: MemoryEffects::Effect{}

操作数 (Operands):¶

操作数	描述
`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，并且
索引 = 1。则输出描述符等同于 input[1]，其中 input 是逻辑张量。

当输入秩为1（即 shape=[k]）时，输出形状将为 shape=[1]。

特性 (Traits): AlwaysSpeculatableImplTrait, MemDescViewTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数 (Operands):¶

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

结果 (Results):¶

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

操作数 (Operands):¶

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

结果 (Results):¶

结果	描述
`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，并且
偏移量 = [2, 1]。则在 Python 语法中，子视图覆盖 input[2:8+2, 1:16+1]，其中 input 是逻辑张量。

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

特性 (Traits): AlwaysSpeculatableImplTrait, MemDescViewTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性 (Attributes):¶

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

操作数 (Operands):¶

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

结果 (Results):¶

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

属性 (Attributes):¶

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

操作数 (Operands):¶

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

结果 (Results):¶

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

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

效果: MemoryEffects::Effect{}

属性 (Attributes):¶

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

操作数 (Operands):¶

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

结果 (Results):¶

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

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

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

语法

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

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

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

特性 (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的幂个 warps 组成的小组，其 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

属性 (Attributes):¶

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

操作数 (Operands):¶

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

结果 (Results):¶

结果	描述
`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 的默认区域中 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>

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

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

效果: MemoryEffects::Effect{}

操作数 (Operands):¶

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

TritonGPUOps¶

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

操作数:¶

结果:¶

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

属性：¶

操作数 (Operands):¶

结果 (Results):¶

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

属性 (Attributes):¶

操作数:¶

结果 (Results):¶

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

操作数 (Operands):¶

结果 (Results):¶

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

属性 (Attributes):¶

操作数:¶

结果 (Results):¶

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

属性:¶

结果:¶

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

属性 (Attributes):¶

操作数:¶

结果 (Results):¶

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

操作数 (Operands):¶

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

操作数:¶

结果 (Results):¶

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

操作数 (Operands):¶

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

操作数 (Operands):¶

结果 (Results):¶

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

操作数 (Operands):¶

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

操作数 (Operands):¶

结果 (Results):¶

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

操作数 (Operands):¶

结果 (Results):¶

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

操作数 (Operands):¶

结果 (Results):¶

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

属性 (Attributes):¶

操作数 (Operands):¶

结果 (Results):¶

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

属性 (Attributes):¶

操作数 (Operands):¶

结果 (Results):¶

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

属性 (Attributes):¶

操作数 (Operands):¶

结果 (Results):¶

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

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

属性 (Attributes):¶

操作数 (Operands):¶

结果 (Results):¶

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

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

操作数 (Operands):¶

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