TritonOps¶
tt.call (triton::CallOp)¶
调用操作
语法
operation ::= `tt.call` $callee `(` $operands `)` attr-dict `:` functional-type($operands, results)
tt.call 操作表示对一个函数的直接调用,该函数与调用位于相同的符号作用域内。调用的操作数和结果类型必须与指定的函数类型匹配。被调用者被编码为一个名为“callee”的符号引用属性。
示例
%2 = tt.call @my_add(%0, %1) : (f32, f32) -> f32
Traits: TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
callee | ::mlir::FlatSymbolRefAttr | 扁平符号引用属性 |
arg_attrs | ::mlir::ArrayAttr | 字典属性数组 |
res_attrs | ::mlir::ArrayAttr | 字典属性数组 |
操作数:¶
操作数 |
描述 |
|---|---|
|
任意类型的可变参数 |
结果:¶
结果 |
描述 |
|---|---|
«未命名» |
任意类型的可变参数 |
tt.func (triton::FuncOp)¶
一个带有名称的操作,其中包含一个单独的 SSACFG 区域
函数内的操作不能隐式捕获在函数外定义的值,即函数是 IsolatedFromAbove(与上层隔离)的。所有外部引用必须使用函数参数或建立符号连接的属性(例如,通过像 SymbolRefAttr 这样的字符串属性按名称引用的符号)。外部函数声明(用于引用在其他模块中声明的函数)没有函数体。虽然 MLIR 文本形式为函数参数提供了很好的内联语法,但它们在内部表示为区域中第一个块的“块参数”。
在函数参数、结果或函数本身的属性字典中,只能指定方言属性名称。
示例
// External function definitions.
tt.func @abort()
tt.func @scribble(i32, i64, memref<? x 128 x f32, #layout_map0>) -> f64
// A function that returns its argument twice:
tt.func @count(%x: i64) -> (i64, i64)
attributes {fruit: "banana"} {
return %x, %x: i64, i64
}
// A function with an argument attribute
tt.func @example_fn_arg(%x: i32 {swift.self = unit})
// A function with a result attribute
tt.func @example_fn_result() -> (f64 {dialectName.attrName = 0 : i64})
// A function with an attribute
tt.func @example_fn_attr() attributes {dialectName.attrName = false}
Traits: AffineScope, AutomaticAllocationScope, HasParent<ModuleOp>, IsolatedFromAbove, TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: ArgAndResultAttrsOpInterface, CallableOpInterface, FunctionOpInterface, OpAsmOpInterface, Symbol
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
sym_name | ::mlir::StringAttr | 字符串属性 |
function_type | ::mlir::TypeAttr | 函数类型的类型属性 |
sym_visibility | ::mlir::StringAttr | 字符串属性 |
arg_attrs | ::mlir::ArrayAttr | 字典属性数组 |
res_attrs | ::mlir::ArrayAttr | 字典属性数组 |
tt.return (triton::ReturnOp)¶
函数返回操作
语法
operation ::= `tt.return` attr-dict ($srcs^ `:` type($srcs))?
tt.return 操作表示函数内的返回操作。该操作接受可变数量的操作数,并且不产生任何结果。操作数的数量和类型必须与包含该操作的函数的签名相匹配。
示例
tt.func @foo() : (i32, f8) {
...
tt.return %0, %1 : i32, f8
}
Traits: AlwaysSpeculatableImplTrait, HasParent<FuncOp>, ReturnLike, TensorSizeTrait, Terminator, VerifyTensorLayoutsTrait
接口 (Interfaces): ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
任意类型的可变参数 |
tt.addptr (triton::AddPtrOp)¶
语法
operation ::= `tt.addptr` $ptr `,` $offset attr-dict `:` type($result) `,` type($offset)
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
ptr 或 ptr 值的秩张量 |
|
整数或整数值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
ptr 或 ptr 值的秩张量 |
tt.advance (triton::AdvanceOp)¶
按偏移量前进张量指针
语法
operation ::= `tt.advance` $ptr `,` `[` $offsets `]` attr-dict `:` type($result)
Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
ptr |
|
可变参数的 32 位无符号整数 |
结果:¶
结果 |
描述 |
|---|---|
|
ptr |
tt.assert (triton::AssertOp)¶
设备端断言,类似于 CUDA 中用于正确性检查的断言
语法
operation ::= `tt.assert` $condition `,` $message attr-dict `:` type($condition)
tt.assert 接受一个条件张量和一个消息字符串。如果条件为假,则打印消息并中止程序。
Traits: TensorSizeTrait, VerifyTensorLayoutsTrait
接口:MemoryEffectOpInterface (MemoryEffectOpInterface)
效果:MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
message | ::mlir::StringAttr | 字符串属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
1 位无符号整数或 1 位无符号整数值的张量 |
tt.atomic_cas (triton::AtomicCASOp)¶
原子比较并交换(CAS)
语法
operation ::= `tt.atomic_cas` $sem `,` $scope `,` $ptr `,` $cmp `,` $val attr-dict `:`
functional-type(operands, $result)
将 $cmp 与位于 $ptr 的数据 $old 进行比较,
如果 $old == $cmp,则将 $val 存储到 $ptr,
否则将 $old 存储到 $ptr,
返回 $old
Traits: SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
sem | ::mlir::triton::MemSemanticAttr | 允许的32位无符号整数情况:1, 2, 3, 4 |
scope | ::mlir::triton::MemSyncScopeAttr | 允许的32位无符号整数情况:1, 2, 3 |
操作数:¶
操作数 |
描述 |
|---|---|
|
ptr 或 ptr 值的秩张量 |
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |
tt.atomic_rmw (triton::AtomicRMWOp)¶
原子读-改-写(RMW)
语法
operation ::= `tt.atomic_rmw` $atomic_rmw_op `,` $sem `,` $scope `,` $ptr `,` $val (`,` $mask^)? attr-dict `:`
functional-type(operands, $result)
加载 $ptr 处的数据,与 $val 执行 $rmw_op 操作,并将结果存储回 $ptr。
返回 $ptr 处的旧值
Traits: SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
属性:¶
| 属性 | 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 值的秩张量 |
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |
|
1 位无符号整数或 1 位无符号整数值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |
tt.bitcast (triton::BitcastOp)¶
在相同位宽的类型之间进行转换
语法
operation ::= `tt.bitcast` $src attr-dict `:` type($src) `->` type($result)
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |
tt.broadcast (triton::BroadcastOp)¶
广播一个张量
语法
operation ::= `tt.broadcast` $src attr-dict `:` type($src) `->` type($result)
对于给定的张量,广播操作会将一个或多个大小为 1 的维度更改为新的大小,例如 tensor<1x32x1xf32> -> tensor<2x32x4xf32>。您不能更改非 1 维度的大小。
Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, SameOperandsAndResultEncoding, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
tt.cat (triton::CatOp)¶
连接 2 个张量
语法
operation ::= `tt.cat` $lhs `,` $rhs attr-dict `:` type($lhs) `->` type($result)
Traits: SameOperandsAndResultElementType, SameTypeOperands, TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数 (Operands):¶
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
|
浮点数、整数或 ptr 值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
tt.clampf (triton::ClampFOp)¶
浮点类型的截断操作
语法
operation ::= `tt.clampf` $x `,` $min `,` $max `,` `propagateNan` `=` $propagateNan attr-dict `:` type($result)
浮点类型的截断操作。
该操作接受三个参数:x、min 和 max。它返回一个与 x 形状相同的张量,其值被截断到 [min, max] 范围内。
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
propagateNan | ::mlir::triton::PropagateNanAttr | 允许的 32 位无符号整数情况:0, 65535 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量 |
|
浮点数或浮点值的秩张量 |
|
浮点数或浮点值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量 |
tt.descriptor_gather (triton::DescriptorGatherOp)¶
从描述符中收集多行到一个张量中
语法
operation ::= `tt.descriptor_gather` $desc `[` $x_offsets `,` $y_offset `]`
attr-dict `:` functional-type(operands, results)
tt.descriptor_gather 操作将在支持它的目标上被降级为 NVIDIA TMA gather 操作。
desc_ptr 是一个指向在全局内存中分配的 TMA 描述符的指针。描述符块必须有 1 行,并且索引必须是 1D 张量。因此,结果是一个多行的 2D 张量。
Traits: TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: TT_DescriptorOpInterface
操作数:¶
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
32 位无符号整数值的秩张量 |
|
32位无符号整数 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
tt.descriptor_load (triton::DescriptorLoadOp)¶
从描述符加载
语法
operation ::= `tt.descriptor_load` $desc `[` $indices `]`
oilist(
`cacheModifier` `=` $cache |
`evictionPolicy` `=` $evict
)
attr-dict `:` qualified(type($desc)) `->` type($result)
此操作将在支持它的目标上被降级为 Nvidia TMA 加载操作。desc 是一个张量描述符对象。目标张量的类型和形状必须与描述符匹配,否则结果是未定义的。
Traits: TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: TT_DescriptorOpInterface
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
evict | ::mlir::triton::EvictionPolicyAttr | 允许的32位无符号整数情况:1, 2, 3 |
操作数 (Operands):¶
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
可变参数的 32 位无符号整数 |
结果 (Results):¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
tt.descriptor_reduce (triton::DescriptorReduceOp)¶
基于描述符执行规约存储操作
语法
operation ::= `tt.descriptor_reduce` $kind `,` $desc `[` $indices `]` `,` $src
attr-dict `:` qualified(type($desc)) `,` type($src)
此操作将在支持它的目标上被降级为 Nvidia TMA 存储操作。desc 是一个张量描述符对象。src 的形状和类型必须与描述符匹配,否则结果是未定义的。
Traits: TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: TT_DescriptorOpInterface, TT_DescriptorStoreLikeOpInterface
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
kind | ::mlir::triton::DescriptorReduceKindAttr | 允许的 32 位无符号整数情况:1, 2, 3, 4, 5, 6, 7, 8 |
操作数 (Operands):¶
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
浮点数、整数或 ptr 值的秩张量 |
|
可变参数的 32 位无符号整数 |
tt.descriptor_scatter (triton::DescriptorScatterOp)¶
将单个张量中的多行分散到一个描述符中
语法
operation ::= `tt.descriptor_scatter` $desc `[` $x_offsets `,` $y_offset `]` `,` $src
attr-dict `:` type(operands)
tt.descriptor_scatter 操作将在支持它的目标上被降级为 NVIDIA TMA scatter 操作。
desc_ptr 是一个指向在全局内存中分配的 TMA 描述符的指针。描述符块必须有 1 行,并且索引必须是 1D 张量。因此,结果是一个多行的 2D 张量。
Traits: TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: TT_DescriptorOpInterface, TT_DescriptorStoreLikeOpInterface
操作数:¶
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
32 位无符号整数值的秩张量 |
|
32位无符号整数 |
|
浮点数、整数或 ptr 值的秩张量 |
tt.descriptor_store (triton::DescriptorStoreOp)¶
基于描述符存储值
语法
operation ::= `tt.descriptor_store` $desc `[` $indices `]` `,` $src
attr-dict `:` qualified(type($desc)) `,` type($src)
此操作将在支持它的目标上被降级为 Nvidia TMA 存储操作。desc 是一个张量描述符对象。src 的形状和类型必须与描述符匹配,否则结果是未定义的。
Traits: TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: TT_DescriptorOpInterface, TT_DescriptorStoreLikeOpInterface
操作数:¶
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
浮点数、整数或 ptr 值的秩张量 |
|
可变参数的 32 位无符号整数 |
tt.dot (triton::DotOp)¶
点积
语法
operation ::= `tt.dot` $a`,` $b`,` $c (`,` `inputPrecision` `=` $inputPrecision^)? attr-dict `:`
type($a) `*` type($b) `->` type($d)
$d = matrix_multiply($a, $b) + $c。$inputPrecision 描述了当输入为 f32 时如何使用 TC。它可以是以下之一:tf32、tf32x3、ieee、bf16x3、bf16x6。tf32:使用 TC 和 tf32 操作。tf32x3:实现 3xTF32 技巧。更多信息请参见 F32DotTC.cpp 中的 pass。bf16x3:实现 3xBF16 技巧。更多信息请参见 F32DotTC.cpp 中的 pass。bf16x6:实现 6xBF16 技巧。更多信息请参见 F32DotTC.cpp 中的 pass。ieee:不使用 TC,在软件中实现点积。如果 GPU 没有 Tensor Core 或输入不是 f32,则此标志被忽略。
Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable, DotOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
inputPrecision | ::mlir::triton::InputPrecisionAttr | 允许的 32 位无符号整数情况:0, 1, 2, 3, 4 |
maxNumImpreciseAcc | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点或整数值的秩张量 |
|
浮点或整数值的秩张量 |
|
浮点或整数值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点或整数值的秩张量 |
tt.dot_scaled (triton::DotScaledOp)¶
缩放点积
语法
operation ::= `tt.dot_scaled` $a (`scale` $a_scale^)? `,` $b (`scale` $b_scale^)? `,` $c
`lhs` `=` $a_elem_type `rhs` `=` $b_elem_type attr-dict
`:` type($a) (`,` type($a_scale)^)? `*` type($b) (`,` type($b_scale)^)? `->` type($d)
$d = matrix_multiply(scale($a, $a_scale), scale($b, $b_scale)) + $c。其中 scale(x, s) 是一个函数,它根据微缩放规范对每个块应用缩放。
Traits: AlwaysSpeculatableImplTrait, AttrSizedOperandSegments, TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable, DotOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
a_elem_type | ::mlir::triton::ScaleDotElemTypeAttr | 允许的 32 位无符号整数情况:0, 1, 2, 3, 4, 5, 6 |
b_elem_type | ::mlir::triton::ScaleDotElemTypeAttr | 允许的 32 位无符号整数情况:0, 1, 2, 3, 4, 5, 6 |
fastMath | ::mlir::BoolAttr | 布尔属性 |
lhs_k_pack | ::mlir::BoolAttr | 布尔属性 |
rhs_k_pack | ::mlir::BoolAttr | 布尔属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点或 8 位无符号整数值的秩张量 |
|
浮点或 8 位无符号整数值的秩张量 |
|
浮点值的秩张量 |
|
浮点或 8 位无符号整数值的秩张量 |
|
浮点或 8 位无符号整数值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点值的秩张量 |
tt.elementwise_inline_asm (triton::ElementwiseInlineAsmOp)¶
对一组打包元素应用逐元素操作的内联汇编。
语法
operation ::= `tt.elementwise_inline_asm` $asm_string attr-dict ($args^ `:` type($args))? `->` type($result)
运行一个内联汇编块以生成一个或多个张量。
汇编块一次被给予 packed_element 个元素。具体接收哪些元素是未指定的。
Traits: Elementwise, SameOperandsAndResultEncoding, TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable, MemoryEffectOpInterface
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
asm_string | ::mlir::StringAttr | 字符串属性 |
constraints | ::mlir::StringAttr | 字符串属性 |
pure | ::mlir::BoolAttr | 布尔属性 |
packed_element | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量、或整数或整数值的秩张量、或 ptr 或 ptr 值的秩张量、或 ptr 的可变参数 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量、或整数或整数值的秩张量、或 ptr 或 ptr 值的秩张量、或 ptr 的可变参数 |
tt.expand_dims (triton::ExpandDimsOp)¶
扩展维度
语法
operation ::= `tt.expand_dims` $src attr-dict `:` type($src) `->` type($result)
Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
axis | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
tt.extern_elementwise (triton::ExternElementwiseOp)¶
语法
operation ::= `tt.extern_elementwise` operands attr-dict `:` functional-type(operands, $result)
使用 $args 调用在 $libpath/$libname 中实现的外部函数 $symbol,返回 $libpath/$libname:$symbol($args…)
Traits: Elementwise, SameOperandsAndResultEncoding, SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable, MemoryEffectOpInterface
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
libname | ::mlir::StringAttr | 字符串属性 |
libpath | ::mlir::StringAttr | 字符串属性 |
symbol | ::mlir::StringAttr | 字符串属性 |
pure | ::mlir::BoolAttr | 布尔属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量、或整数或整数值的秩张量、或 ptr 或 ptr 值的秩张量、或 ptr 的可变参数 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |
tt.fp_to_fp (triton::FpToFpOp)¶
自定义类型的浮点转换
语法
operation ::= `tt.fp_to_fp` $src attr-dict (`,` `rounding` `=` $rounding^)? `:` type($src) `->` type($result)
用于自定义类型(F8)和非默认舍入模式的浮点转换。
F8 <-> FP16、BF16、FP32、FP64
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
rounding | ::mlir::triton::RoundingModeAttr | 允许的32位无符号整数情况:0, 1 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量 |
tt.gather (triton::GatherOp)¶
局部收集操作
语法
operation ::= `tt.gather` $src `[` $indices `]` attr-dict `:`
functional-type(operands, results)
使用索引张量沿指定的单个轴从输入张量中收集元素。输出张量的形状与索引张量相同。输入张量和索引张量的维度数必须相同,并且索引张量中非收集维度的每个维度都不能大于输入张量中对应的维度。
当编译器确定了该操作的优化布局时,会设置 efficient_layout 属性,表示不应更改该布局。
Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
axis | ::mlir::IntegerAttr | 32位无符号整数属性 |
efficient_layout | ::mlir::UnitAttr | 单元属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
|
整数值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
tt.get_num_programs (triton::GetNumProgramsOp)¶
语法
operation ::= `tt.get_num_programs` $axis attr-dict `:` type($result)
Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
axis | ::mlir::triton::ProgramIDDimAttr | 允许的 32 位无符号整数情况:0, 1, 2 |
结果:¶
结果 |
描述 |
|---|---|
|
32位无符号整数 |
tt.get_program_id (triton::GetProgramIdOp)¶
语法
operation ::= `tt.get_program_id` $axis attr-dict `:` type($result)
Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
axis | ::mlir::triton::ProgramIDDimAttr | 允许的 32 位无符号整数情况:0, 1, 2 |
结果:¶
结果 |
描述 |
|---|---|
|
32位无符号整数 |
tt.histogram (triton::HistogramOp)¶
返回输入的直方图。
语法
operation ::= `tt.histogram` $src (`,` $mask^)? attr-dict `:` type($src) `->` type($result)
返回输入张量的直方图。箱子的数量等于输出张量的维度。每个箱子的宽度为 1,并且箱子从 0 开始。
Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
整数值的秩张量 |
|
1 位无符号整数或 1 位无符号整数值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
整数值的秩张量 |
tt.int_to_ptr (triton::IntToPtrOp)¶
将 int64 转换为指针
语法
operation ::= `tt.int_to_ptr` $src attr-dict `:` type($src) `->` type($result)
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
64 位无符号整数或 64 位无符号整数值的张量 |
结果:¶
结果 |
描述 |
|---|---|
|
ptr 或 ptr 值的秩张量 |
tt.join (triton::JoinOp)¶
沿一个新的、次要维度连接两个张量
语法
operation ::= `tt.join` $lhs `,` $rhs attr-dict `:` type($lhs) `->` type($result)
例如,如果两个输入张量是 4x8xf32,则返回一个形状为 4x8x2xf32 的张量。
因为 Triton 张量总是具有 2 的幂次个元素,所以两个输入张量必须具有相同的形状。
Traits: AlwaysSpeculatableImplTrait, SameTypeOperands, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
|
浮点数、整数或 ptr 值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
tt.load (triton::LoadOp)¶
从指针张量或张量指针加载
语法
operation ::= `tt.load` $ptr (`,` $mask^)? (`,` $other^)?
oilist(
`cacheModifier` `=` $cache |
`evictionPolicy` `=` $evict
)
attr-dict `:` type($ptr)
Traits: AttrSizedOperandSegments, SameLoadStoreOperandsAndResultEncoding, SameLoadStoreOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: InferTypeOpInterface, MemoryEffectOpInterface
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
boundaryCheck | ::mlir::DenseI32ArrayAttr | i32 密集数组属性 |
padding | ::mlir::triton::PaddingOptionAttr | 允许的 32 位无符号整数情况:1, 2 |
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
evict | ::mlir::triton::EvictionPolicyAttr | 允许的32位无符号整数情况:1, 2, 3 |
isVolatile | ::mlir::BoolAttr | 布尔属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
ptr 或 ptr 值的秩张量,或 ptr |
|
1 位无符号整数或 1 位无符号整数值的秩张量 |
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |
tt.make_range (triton::MakeRangeOp)¶
创建范围
语法
operation ::= `tt.make_range` attr-dict `:` type($result)
返回一个 1D int32 张量。
值范围从 $start 到 $end(不含),步长 = 1
Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
start | ::mlir::IntegerAttr | 32位无符号整数属性 |
end | ::mlir::IntegerAttr | 32位无符号整数属性 |
结果:¶
结果 |
描述 |
|---|---|
|
整数值的秩张量 |
tt.make_tensor_descriptor (triton::MakeTensorDescOp)¶
创建一个带有父张量元信息和块大小的张量描述符类型
语法
operation ::= `tt.make_tensor_descriptor` $base `,` `[` $shape `]` `,` `[` $strides `]` attr-dict `:` type($base) `,` type($result)
tt.make_tensor_descriptor 接受父张量的元信息和块大小,并返回一个描述符对象,该对象可用于从全局内存中的张量加载/存储。
Traits: AlwaysSpeculatableImplTrait, SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
padding | ::mlir::triton::PaddingOptionAttr | 允许的 32 位无符号整数情况:1, 2 |
操作数:¶
操作数 |
描述 |
|---|---|
|
ptr |
|
可变参数的 32 位无符号整数 |
|
64 位无符号整数的可变参数 |
结果:¶
结果 |
描述 |
|---|---|
|
张量描述符类型( |
tt.make_tensor_ptr (triton::MakeTensorPtrOp)¶
创建一个带有父张量元信息和指定块的张量指针类型
语法
operation ::= `tt.make_tensor_ptr` $base `,` `[` $shape `]` `,` `[` $strides `]` `,` `[` $offsets `]` attr-dict `:` type($result)
tt.make_tensor_ptr 接受父张量的元信息和块张量,然后返回指向块张量的指针,例如返回一个 tt.ptr<tensor<8x8xf16>> 类型。
Traits: AlwaysSpeculatableImplTrait, SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
order | ::mlir::DenseI32ArrayAttr | i32 密集数组属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
ptr |
|
64 位无符号整数的可变参数 |
|
64 位无符号整数的可变参数 |
|
可变参数的 32 位无符号整数 |
结果:¶
结果 |
描述 |
|---|---|
|
ptr |
tt.map_elementwise (triton::MapElementwiseOp)¶
将标量子区域映射到张量上
Traits: RecursiveMemoryEffects, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
pack | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点、整数或 ptr 值的秩张量的可变参数 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点、整数或 ptr 值的秩张量的可变参数 |
tt.map_elementwise.return (triton::MapElementwiseReturnOp)¶
map elementwise 运算符的终止符
语法
operation ::= `tt.map_elementwise.return` attr-dict ($result^ `:` type($result))?
Traits: AlwaysSpeculatableImplTrait, HasParent<MapElementwiseOp>, ReturnLike, TensorSizeTrait, Terminator, VerifyTensorLayoutsTrait
接口 (Interfaces): ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
任意类型的可变参数 |
tt.mulhiui (triton::MulhiUIOp)¶
两个整数的 2N 位乘积的最高 N 位
语法
operation ::= `tt.mulhiui` $x `,` $y attr-dict `:` type($x)
两个整数的 2N 位乘积的最高 N 位。
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
整数或整数值的秩张量 |
|
整数或整数值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
整数或整数值的秩张量 |
tt.precise_divf (triton::PreciseDivFOp)¶
浮点类型的精确除法
语法
operation ::= `tt.precise_divf` $x `,` $y attr-dict `:` type($x)
浮点类型的精确除法。
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量 |
|
浮点数或浮点值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量 |
tt.precise_sqrt (triton::PreciseSqrtOp)¶
浮点类型的精确平方根
语法
operation ::= `tt.precise_sqrt` $x attr-dict `:` type($x)
浮点类型的精确平方根。
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量 |
tt.print (triton::PrintOp)¶
设备端打印,类似于 CUDA 中用于调试的打印
语法
operation ::= `tt.print` $prefix attr-dict (`:` $args^ `:` type($args))?
tt.print 接受一个字面字符串前缀和任意数量的应打印的标量或张量参数。格式会根据参数自动生成。
Traits: SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:MemoryEffectOpInterface (MemoryEffectOpInterface)
效果:MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
prefix | ::mlir::StringAttr | 字符串属性 |
hex | ::mlir::BoolAttr | 布尔属性 |
isSigned | ::mlir::DenseI32ArrayAttr | i32 密集数组属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量、或整数或整数值的秩张量、或 ptr 或 ptr 值的秩张量、或 ptr 的可变参数 |
tt.ptr_to_int (triton::PtrToIntOp)¶
将指针转换为 int64
语法
operation ::= `tt.ptr_to_int` $src attr-dict `:` type($src) `->` type($result)
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
ptr 或 ptr 值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
64 位无符号整数或 64 位无符号整数值的张量 |
tt.reduce (triton::ReduceOp)¶
使用通用组合算法进行规约
Traits: AlwaysSpeculatableImplTrait, SameOperandsEncoding, SameOperandsShape, SingleBlock, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
axis | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点、整数或 ptr 值的秩张量的可变参数 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量、或整数或整数值的秩张量、或 ptr 或 ptr 值的秩张量、或 ptr 的可变参数 |
tt.reduce.return (triton::ReduceReturnOp)¶
reduce 运算符的终止符
语法
operation ::= `tt.reduce.return` $result attr-dict `:` type($result)
Traits: AlwaysSpeculatableImplTrait, HasParent<ReduceOp>, ReturnLike, TensorSizeTrait, Terminator, VerifyTensorLayoutsTrait
接口 (Interfaces): ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
任意类型的可变参数 |
tt.reshape (triton::ReshapeOp)¶
将一个张量重新解释为不同的形状。如果设置了属性,它可能会改变元素的顺序。
语法
operation ::= `tt.reshape` $src (`allow_reorder` $allow_reorder^)? (`efficient_layout` $efficient_layout^)? attr-dict `:` type($src) `->` type($result)
将一个张量重新解释为不同的形状。
如果设置了 allow_reorder,编译器可以自由更改元素的顺序以生成更高效的代码。
如果设置了 efficient_layout,这是一个提示,表明为了性能原因应保留目标布局。编译器仍然可以为了更好的性能而更改它。
Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
allow_reorder | ::mlir::UnitAttr | 单元属性 |
efficient_layout | ::mlir::UnitAttr | 单元属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
tt.scan (triton::ScanOp)¶
使用通用组合算法进行关联扫描
Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultEncoding, SameOperandsAndResultShape, SingleBlock, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
axis | ::mlir::IntegerAttr | 32位无符号整数属性 |
reverse | ::mlir::BoolAttr | 布尔属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点、整数或 ptr 值的秩张量的可变参数 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点、整数或 ptr 值的秩张量的可变参数 |
tt.scan.return (triton::ScanReturnOp)¶
scan 运算符的终止符
语法
operation ::= `tt.scan.return` $result attr-dict `:` type($result)
Traits: AlwaysSpeculatableImplTrait, HasParent<ScanOp>, ReturnLike, TensorSizeTrait, Terminator, VerifyTensorLayoutsTrait
接口 (Interfaces): ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
任意类型的可变参数 |
tt.splat (triton::SplatOp)¶
Splat
语法
operation ::= `tt.splat` $src attr-dict `:` type($src) `->` type($result)
Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, SameOperandsAndResultEncoding, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
tt.split (triton::SplitOp)¶
将一个张量沿其最后一个维度拆分为两个
语法
operation ::= `tt.split` $src attr-dict `:` type($src) `->` type($outLHS)
输入必须是一个最后一个维度大小为 2 的张量。返回两个张量,src[…, 0] 和 src[…, 1]。
例如,如果输入形状是 4x8x2xf32,则返回两个形状为 4x8xf32 的张量。
Traits: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
|
浮点数、整数或 ptr 值的秩张量 |
tt.store (triton::StoreOp)¶
通过指针张量或张量指针进行存储
语法
operation ::= `tt.store` $ptr `,` $value (`,` $mask^)?
oilist(`cacheModifier` `=` $cache | `evictionPolicy` `=` $evict)
attr-dict `:` type($ptr)
Traits: SameLoadStoreOperandsEncoding, SameLoadStoreOperandsShape, TensorSizeTrait, VerifyTensorLayoutsTrait
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
boundaryCheck | ::mlir::DenseI32ArrayAttr | i32 密集数组属性 |
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
evict | ::mlir::triton::EvictionPolicyAttr | 允许的32位无符号整数情况:1, 2, 3 |
操作数:¶
操作数 |
描述 |
|---|---|
|
ptr 或 ptr 值的秩张量,或 ptr |
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |
|
1 位无符号整数或 1 位无符号整数值的秩张量 |
tt.trans (triton::TransOp)¶
重排张量的维度
语法
operation ::= `tt.trans` $src attr-dict `:` type($src) `->` type($result)
例如,给定一个形状为 [1,2,4] 的张量 x,使用 order=[2,0,1] 的 transpose(x) 会将张量重排为形状 [4,1,2]。
尽管此操作名为“trans”,但它同时实现了 tl.trans() 和 tl.permute()。(“permute”可能是一个更好的名字,但它被称为“trans”是因为最初它只支持 2D 张量。)
关于编码的实现说明:¶
在 TritonGPU 方言(可能还有其他方言)中,为此操作的输出选择了一种编码,使其从代码生成的角度来看是一个空操作(nop)。
例如,假设张量 x 有一个编码,使得 GPU 线程 [i,j,k] 的寄存器中包含张量的元素 [i,j,k]。现在我们用 order [2,1,0] 对 x 进行转置,即我们反转其维度的顺序。在 TritonGPU 中,我们会为转置的输出选择一个布局,使得 GPU 线程 [i,j,k] 拥有 transpose(x) 的元素 [k,j,i]。但这与它之前拥有的元素是相同的!我们所做的只是“重命名”了线程 [i,j,k] 拥有的元素。
“真正的”转置——即在 GPU 线程之间移动数据——发生在操作之前和/或之后的 convertLayout 操作中。
我们这样做是为了让您可以链接多个数据移动操作(例如 transpose+reshape+concat),而无需在每次操作后都访问共享内存。
Traits: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor, SameOperandsAndResultElementType, TensorSizeTrait, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface), TransposeOpInterface
效果: MemoryEffects::Effect{}
属性:¶
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
order | ::mlir::DenseI32ArrayAttr | i32 密集数组属性 |
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
tt.unsplat (triton::UnsplatOp)¶
将只有一个元素的张量转换为标量
语法
operation ::= `tt.unsplat` $src attr-dict `:` type($src)
Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
结果:¶
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或 ptr 或 ptr 值的秩张量,或 ptr |