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: 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: 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.reinterpret_tensor_descriptor
(triton::ReinterpretTensorDescOp)¶
将指针重新解释为张量描述符
语法
operation ::= `tt.reinterpret_tensor_descriptor` $rawDesc attr-dict `:` qualified(type($rawDesc)) `to` qualified(type($result))
此操作(Op)旨在帮助从无类型原始 TMA 对象过渡到类型化张量描述符对象。理想情况下,一旦 API 完全成熟,我们可以移除此操作。
Traits: AlwaysSpeculatableImplTrait
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
ptr |
结果:¶
结果 |
描述 |
---|---|
|
Triton IR 类型系统中的张量描述符类型 ( |
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
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
可变数量的任何类型 |
tt.addptr
(triton::AddPtrOp)¶
语法
operation ::= `tt.addptr` $ptr `,` $offset attr-dict `:` type($result) `,` type($offset)
Traits: AlwaysSpeculatableImplTrait
, Elementwise
, SameOperandsAndResultEncoding
, SameOperandsAndResultShape
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
ptr 或 ptr 值的 ranked tensor |
|
整数或整数值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
ptr 或 ptr 值的 ranked tensor |
tt.advance
(triton::AdvanceOp)¶
按偏移量推进张量指针
语法
operation ::= `tt.advance` $ptr `,` `[` $offsets `]` attr-dict `:` type($result)
Traits: AlwaysSpeculatableImplTrait
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
ptr |
|
可变数量的 32 位无符号整数 |
结果:¶
结果 |
描述 |
---|---|
|
ptr |
tt.assert
(triton::AssertOp)¶
设备端断言,用于正确性检查,类似于 CUDA
语法
operation ::= `tt.assert` $condition `,` $message attr-dict `:` type($condition)
tt.assert
接受一个条件张量和一个消息字符串。如果条件为 false,则打印消息并中止程序。
Traits: TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: MemoryEffectOpInterface (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
message | ::mlir::StringAttr | 字符串属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
1 位无符号整数或 1 位无符号整数值的 tensor |
tt.atomic_cas
(triton::AtomicCASOp)¶
原子比较和交换
语法
operation ::= `tt.atomic_cas` $sem `,` $scope `,` $ptr `,` $cmp `,` $val attr-dict `:`
functional-type(operands, $result)
在 $ptr 位置比较 $cmp 和数据 $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 值的 ranked tensor |
|
浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
|
浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
tt.atomic_rmw
(triton::AtomicRMWOp)¶
原子读-修改-写
语法
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 值的 ranked tensor |
|
浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
|
1 位无符号整数或 1 位无符号整数值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
tt.bitcast
(triton::BitcastOp)¶
在相同位宽的类型之间进行位转换
语法
operation ::= `tt.bitcast` $src attr-dict `:` type($src) `->` type($result)
Traits: AlwaysSpeculatableImplTrait
, Elementwise
, SameOperandsAndResultEncoding
, SameOperandsAndResultShape
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 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
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
tt.cat
(triton::CatOp)¶
连接 2 个张量
语法
operation ::= `tt.cat` $lhs `,` $rhs attr-dict `:` type($lhs) `->` type($result)
Traits: SameOperandsAndResultElementType
, SameTypeOperands
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
|
浮点数或整数或 ptr 值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
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
Interfaces: ConditionallySpeculatable
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
propagateNan | ::mlir::triton::PropagateNanAttr | 允许的 32 位无符号整数情况:0、65535 |
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor |
|
浮点数或浮点数值的 ranked tensor |
|
浮点数或浮点数值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor |
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 收集操作。
desc_ptr
是指向全局内存中分配的 TMA 描述符的指针。描述符块必须有 1 行,索引必须是 1D 张量。因此,结果是一个包含多行的 2D 张量。
Traits: TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: TT_DescriptorOpInterface
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的张量描述符类型 ( |
|
32 位无符号整数值的 ranked tensor |
|
32 位无符号整数 |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
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 |
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的张量描述符类型 ( |
|
可变数量的 32 位无符号整数 |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
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 |
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的张量描述符类型 ( |
|
浮点数或整数或 ptr 值的 ranked tensor |
|
可变数量的 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 散布操作。
desc_ptr
是指向全局内存中分配的 TMA 描述符的指针。描述符块必须有 1 行,索引必须是 1D 张量。因此,结果是一个包含多行的 2D 张量。
Traits: TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: TT_DescriptorOpInterface
, TT_DescriptorStoreLikeOpInterface
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的张量描述符类型 ( |
|
32 位无符号整数值的 ranked tensor |
|
32 位无符号整数 |
|
浮点数或整数或 ptr 值的 ranked tensor |
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
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的张量描述符类型 ( |
|
浮点数或整数或 ptr 值的 ranked tensor |
|
可变数量的 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。tf32:使用 TC 执行 tf32 操作。tf32x3:实现 3xTF32 技巧。更多信息请参见 F32DotTC.cpp 中的 pass。ieee:不使用 TC,在软件中实现点积。如果 GPU 不具有 Tensor core 或输入不是 f32,则忽略此标志。
Traits: AlwaysSpeculatableImplTrait
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, DotOpInterface
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
inputPrecision | ::mlir::triton::InputPrecisionAttr | 允许的 32 位无符号整数情况:0、1、2 |
maxNumImpreciseAcc | ::mlir::IntegerAttr | 32 位无符号整数属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或整数值的 ranked tensor |
|
浮点数或整数值的 ranked tensor |
|
浮点数或整数值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或整数值的 ranked tensor |
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)
Effects: 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 位无符号整数值的 ranked tensor |
|
浮点数或 8 位无符号整数值的 ranked tensor |
|
浮点数值的 ranked tensor |
|
浮点数或 8 位无符号整数值的 ranked tensor |
|
浮点数或 8 位无符号整数值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数值的 ranked tensor |
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 位无符号整数属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
可变数量的浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
结果:¶
结果 |
描述 |
---|---|
|
可变数量的浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
tt.expand_dims
(triton::ExpandDimsOp)¶
扩展维度
语法
operation ::= `tt.expand_dims` $src attr-dict `:` type($src) `->` type($result)
Traits: AlwaysSpeculatableImplTrait
, SameOperandsAndResultElementType
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
axis | ::mlir::IntegerAttr | 32 位无符号整数属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
tt.experimental_tensormap_create
(triton::ExperimentalTensormapCreateOp)¶
在设备上创建新的 TMA 描述符
语法
operation ::= `tt.experimental_tensormap_create` $desc_ptr `,` $global_address `,`
`[` $box_dim `]` `,`
`[` $global_dim `]` `,`
`[` $global_stride `]` `,`
`[` $element_stride `]`
attr-dict `:` functional-type(operands, results)
Traits: AttrSizedOperandSegments
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: MemoryEffectOpInterface (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{MemoryEffects::Read on ::mlir::triton::GlobalMemory, MemoryEffects::Write on ::mlir::triton::GlobalMemory}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
elem_type | ::mlir::IntegerAttr | 值为非负且最大值为 15 的 32 位无符号整数属性 |
interleave_layout | ::mlir::IntegerAttr | 值为非负且最大值为 2 的 32 位无符号整数属性 |
swizzle_mode | ::mlir::IntegerAttr | 值为非负且最大值为 3 的 32 位无符号整数属性 |
fill_mode | ::mlir::IntegerAttr | 值为非负且最大值为 1 的 32 位无符号整数属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的 Pointer type ( |
|
Triton IR 类型系统中的 Pointer type ( |
|
可变数量的 32 位无符号整数 |
|
可变数量的 32 位无符号整数 |
|
可变数量的 64 位无符号整数 |
|
可变数量的 32 位无符号整数 |
tt.experimental_tensormap_fenceproxy_acquire
(triton::ExperimentalTensormapFenceproxyAcquireOp)¶
在 tensormap 对象上获取围栏
语法
operation ::= `tt.experimental_tensormap_fenceproxy_acquire` $desc_ptr attr-dict `:` qualified(type($desc_ptr))
Traits: TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: MemoryEffectOpInterface (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}
操作数:¶
操作数 |
描述 |
---|---|
|
Triton IR 类型系统中的 Pointer type ( |
tt.extern_elementwise
(triton::ExternElementwiseOp)¶
语法
operation ::= `tt.extern_elementwise` operands attr-dict `:` functional-type(operands, $result)
调用在 $libpath/$libname 中实现的外部函数 $symbol,带 $args,返回 $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 | 布尔属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
可变数量的浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 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
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
rounding | ::mlir::triton::RoundingModeAttr | 允许的 32 位无符号整数情况:0、1 |
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor |
tt.gather
(triton::GatherOp)¶
本地收集操作
语法
operation ::= `tt.gather` $src `[` $indices `]` attr-dict `:`
functional-type(operands, results)
使用索引张量沿着单个指定轴从输入张量收集元素。输出张量具有与索引张量相同的形状。输入张量和索引张量必须具有相同的维度数,并且索引张量中不是收集维度的每个维度不能大于输入张量中的相应维度。
当编译器确定了操作的优化布局时,会设置 efficient_layout
属性,表示不应更改该布局。
Traits: AlwaysSpeculatableImplTrait
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
axis | ::mlir::IntegerAttr | 32 位无符号整数属性 |
efficient_layout | ::mlir::UnitAttr | unit 属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
|
整数值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
tt.get_num_programs
(triton::GetNumProgramsOp)¶
语法
operation ::= `tt.get_num_programs` $axis attr-dict `:` type($result)
Traits: AlwaysSpeculatableImplTrait
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: 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
Interfaces: ConditionallySpeculatable
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
axis | ::mlir::triton::ProgramIDDimAttr | 允许的 32 位无符号整数情况:0、1、2 |
结果:¶
结果 |
描述 |
---|---|
|
32 位无符号整数 |
tt.histogram
(triton::HistogramOp)¶
返回输入的直方图。
语法
operation ::= `tt.histogram` $src attr-dict `:` type($src) `->` type($result)
返回输入张量的直方图。 bin 的数量等于输出张量的维度。每个 bin 的宽度为 1,并且 bin 从 0 开始。
Traits: AlwaysSpeculatableImplTrait
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
整数值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
整数值的 ranked tensor |
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
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
64 位无符号整数或 64 位无符号整数值的 tensor |
结果:¶
结果 |
描述 |
---|---|
|
ptr 或 ptr 值的 ranked tensor |
tt.join
(triton::JoinOp)¶
将两个张量沿一个新的、次要的维度连接
语法
operation ::= `tt.join` $lhs `,` $rhs attr-dict `:` type($lhs) `->` type($result)
例如,如果两个输入张量是 4x8xf32,则返回一个形状为 4x8x2xf32 的张量。
因为 Triton 张量总是具有 2 的幂次方个元素,所以两个输入张量必须具有相同的形状。
Traits: AlwaysSpeculatableImplTrait
, SameTypeOperands
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
|
浮点数或整数或 ptr 值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
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 值的 ranked tensor 或 ptr |
|
1 位无符号整数或 1 位无符号整数值的 ranked tensor |
|
浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
tt.make_range
(triton::MakeRangeOp)¶
生成范围
语法
operation ::= `tt.make_range` attr-dict `:` type($result)
返回一个 1D int32 张量。
值从 $start 到 $end (不包含 $end),步长 = 1
Traits: AlwaysSpeculatableImplTrait
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
start | ::mlir::IntegerAttr | 32 位无符号整数属性 |
end | ::mlir::IntegerAttr | 32 位无符号整数属性 |
结果:¶
结果 |
描述 |
---|---|
|
整数值的 ranked tensor |
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
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
ptr |
|
可变数量的 32 位无符号整数 |
|
可变数量的 64 位无符号整数 |
结果:¶
结果 |
描述 |
---|---|
|
Triton IR 类型系统中的张量描述符类型 ( |
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
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
order | ::mlir::DenseI32ArrayAttr | i32 稠密数组属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
ptr |
|
可变数量的 64 位无符号整数 |
|
可变数量的 64 位无符号整数 |
|
可变数量的 32 位无符号整数 |
结果:¶
结果 |
描述 |
---|---|
|
ptr |
tt.mulhiui
(triton::MulhiUIOp)¶
两个整数的 2N 位乘积中最重要的 N 位
语法
operation ::= `tt.mulhiui` $x `,` $y attr-dict `:` type($x)
两个整数的 2N 位乘积中最重要的 N 位。
Traits: AlwaysSpeculatableImplTrait
, Elementwise
, SameOperandsAndResultType
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
整数或整数值的 ranked tensor |
|
整数或整数值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
整数或整数值的 ranked tensor |
tt.precise_divf
(triton::PreciseDivFOp)¶
浮点类型的精确除法
语法
operation ::= `tt.precise_divf` $x `,` $y attr-dict `:` type($x)
浮点类型的精确除法。
Traits: AlwaysSpeculatableImplTrait
, Elementwise
, SameOperandsAndResultType
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor |
|
浮点数或浮点数值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor |
tt.precise_sqrt
(triton::PreciseSqrtOp)¶
浮点类型的精确平方根
语法
operation ::= `tt.precise_sqrt` $x attr-dict `:` type($x)
浮点类型的精确平方根。
Traits: AlwaysSpeculatableImplTrait
, Elementwise
, SameOperandsAndResultType
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor |
tt.print
(triton::PrintOp)¶
设备端打印,用于调试,类似于 CUDA
语法
operation ::= `tt.print` $prefix attr-dict (`:` $args^ `:` type($args))?
tt.print
接受一个字符串前缀以及任意数量的标量或张量参数,这些参数将被打印。格式会根据参数自动生成。
Traits: SameVariadicOperandSize
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: MemoryEffectOpInterface (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
prefix | ::mlir::StringAttr | 字符串属性 |
hex | ::mlir::BoolAttr | 布尔属性 |
isSigned | ::mlir::DenseI32ArrayAttr | i32 稠密数组属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
可变数量的浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 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
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
ptr 或 ptr 值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
64 位无符号整数或 64 位无符号整数值的 tensor |
tt.reduce
(triton::ReduceOp)¶
使用通用组合算法进行规约
Traits: AlwaysSpeculatableImplTrait
, SameOperandsEncoding
, SameOperandsShape
, SingleBlock
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
axis | ::mlir::IntegerAttr | 32 位无符号整数属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
可变数量的浮点数或整数或 ptr 值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
可变数量的浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
tt.reduce.return
(triton::ReduceReturnOp)¶
规约操作的终止符
语法
operation ::= `tt.reduce.return` $result attr-dict `:` type($result)
Traits: AlwaysSpeculatableImplTrait
, HasParent<ReduceOp>
, ReturnLike
, TensorSizeTrait
, Terminator
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
, RegionBranchTerminatorOpInterface
Effects: 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
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
allow_reorder | ::mlir::UnitAttr | unit 属性 |
efficient_layout | ::mlir::UnitAttr | unit 属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
tt.scan
(triton::ScanOp)¶
使用通用组合算法进行关联扫描
Traits: AlwaysSpeculatableImplTrait
, SameOperandsAndResultEncoding
, SameOperandsAndResultShape
, SingleBlock
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
axis | ::mlir::IntegerAttr | 32 位无符号整数属性 |
reverse | ::mlir::BoolAttr | 布尔属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
可变数量的浮点数或整数或 ptr 值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
可变数量的浮点数或整数或 ptr 值的 ranked tensor |
tt.scan.return
(triton::ScanReturnOp)¶
扫描操作的终止符
语法
operation ::= `tt.scan.return` $result attr-dict `:` type($result)
Traits: AlwaysSpeculatableImplTrait
, HasParent<ScanOp>
, ReturnLike
, TensorSizeTrait
, Terminator
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
, RegionBranchTerminatorOpInterface
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
可变数量的任何类型 |
tt.splat
(triton::SplatOp)¶
扩展
语法
operation ::= `tt.splat` $src attr-dict `:` type($src) `->` type($result)
Traits: AlwaysSpeculatableImplTrait
, SameOperandsAndResultElementType
, SameOperandsAndResultEncoding
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
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
Interfaces: ConditionallySpeculatable
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
|
浮点数或整数或 ptr 值的 ranked tensor |
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 值的 ranked tensor 或 ptr |
|
浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr |
|
1 位无符号整数或 1 位无符号整数值的 ranked tensor |
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] 的寄存器。现在我们将 x 进行 transpose 操作,order=[2,1,0],即将其维度顺序反转。在 TritonGPU 中,我们将为 transpose(x) 的输出选择一个布局,使得 GPU 线程 [i,j,k] 持有 transpose(x) 的元素 [k,j,i]。但这与它之前持有的元素是同一个元素!我们所做的只是“重命名”线程 [i,j,k] 所持有的元素。
“真正的”transpose——即在 GPU 线程之间移动数据——发生在操作之前和/或之后出现的 convertLayout 操作中。
我们这样做是为了您可以链接多个数据移动操作(例如 transpose+reshape+concat),而无需在每次操作后都移动到共享内存。
Traits: AlwaysSpeculatableImplTrait
, InferTypeOpAdaptor
, SameOperandsAndResultElementType
, TensorSizeTrait
, VerifyTensorLayoutsTrait
Interfaces: ConditionallySpeculatable
, InferTypeOpInterface
, NoMemoryEffect (MemoryEffectOpInterface)
, TransposeOpInterface
Effects: MemoryEffects::Effect{}
属性:¶
属性 | MLIR 类型 | 描述 |
---|---|---|
order | ::mlir::DenseI32ArrayAttr | i32 稠密数组属性 |
操作数:¶
操作数 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |
结果:¶
结果 |
描述 |
---|---|
|
浮点数或整数或 ptr 值的 ranked tensor |