TritonOps
tt.addptr (triton::AddPtrOp)
语法
operation ::= `tt.addptr` $ptr `,` $offset attr-dict `:` type($result) `,` type($offset)
特性: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
指针或指针值的秩张量 |
|
整数或整数值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
指针或指针值的秩张量 |
tt.assert (triton::AssertOp)
设备端断言,类似于 CUDA 中的正确性检查
语法
operation ::= `tt.assert` $condition `,` $message attr-dict `:` type($condition)
tt.assert 接受一个条件张量和一条消息字符串。如果条件为假,则打印该消息并终止程序。
特性: TensorSizeTrait, VerifyTensorLayoutsTrait
接口:MemoryEffectOpInterface (MemoryEffectOpInterface)
效果:MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
消息 | ::mlir::StringAttr | 字符串属性 |
操作数:
操作数 |
描述 |
|---|---|
|
1位无符号整数或1位无符号整数值的张量 |
tt.atomic_cas (triton::AtomicCASOp)
原子比较并交换 (Atomic CAS)
语法
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
特性: SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
sem | ::mlir::triton::MemSemanticAttr | 允许的32位无符号整数情况:1, 2, 3, 4 |
scope | ::mlir::triton::MemSyncScopeAttr | 允许的32位无符号整数情况:1, 2, 3 |
操作数:
操作数 |
描述 |
|---|---|
|
指针或指针值的秩张量 |
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
tt.atomic_rmw (triton::AtomicRMWOp)
原子读写改 (Atomic 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 处的旧值
特性: SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: PredicatedOpInterface
属性:
| 属性 | 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 |
操作数:
操作数 |
描述 |
|---|---|
|
指针或指针值的秩张量 |
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
|
1位无符号整数或1位无符号整数值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
tt.bitcast (triton::BitcastOp)
在相同位宽的类型之间进行转换
语法
operation ::= `tt.bitcast` $src attr-dict `:` type($src) `->` type($result)
特性: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
tt.broadcast (triton::BroadcastOp)
广播张量
语法
operation ::= `tt.broadcast` $src attr-dict `:` type($src) `->` type($result)
对于给定的张量,广播将一个或多个尺寸为 1 的维度更改为新的大小,例如 tensor<1x32x1xf32> -> tensor<2x32x4xf32>。不能更改非 1 维度的大小。
特性: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, SameOperandsAndResultEncoding, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
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
特性: TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
被调用者 (callee) | ::mlir::FlatSymbolRefAttr | 平面符号引用属性 |
参数属性 (arg_attrs) | ::mlir::ArrayAttr | 字典属性数组 |
结果属性 (res_attrs) | ::mlir::ArrayAttr | 字典属性数组 |
操作数:
操作数 |
描述 |
|---|---|
|
任意类型的可变参数 |
结果:
结果 |
描述 |
|---|---|
«未命名» |
任意类型的可变参数 |
tt.cat (triton::CatOp)
连接 2 个张量
语法
operation ::= `tt.cat` $lhs `,` $rhs attr-dict `:` type($lhs) `->` type($result)
特性: SameOperandsAndResultElementType, SameTypeOperands, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
|
浮点数、整数或 ptr 值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
tt.clampf (triton::ClampFOp)
浮点类型的钳位 (Clamp) 操作
语法
operation ::= `tt.clampf` $x `,` $min `,` $max `,` `propagateNan` `=` $propagateNan attr-dict `:` type($result)
浮点类型的钳位操作。
该操作接受三个参数:x、min 和 max。它返回一个与 x 形状相同的张量,其值钳位在 [min, max] 范围内。
特性: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
传播 NaN (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 收集操作。
desc_ptr 是指向分配在全局内存中的 TMA 描述符的指针。描述符块必须有 1 行,索引必须是一维张量。因此,结果是一个多行的二维张量。
特性: TensorSizeTrait, VerifyTensorLayoutsTrait
接口: TT_DescriptorLoadLikeOpInterface, TT_DescriptorOpInterface
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
16 位无符号整数或 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 是一个张量描述符对象。目标张量类型和形状必须与描述符匹配,否则结果未定义。
特性: TensorSizeTrait, VerifyTensorLayoutsTrait
接口: TT_DescriptorLoadLikeOpInterface, TT_DescriptorOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
evict | ::mlir::triton::EvictionPolicyAttr | 允许的32位无符号整数情况:1, 2, 3 |
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
可变参数的 32 位无符号整数 |
结果:
结果 |
描述 |
|---|---|
|
浮点数、整数或 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 的形状和类型必须与描述符匹配,否则结果未定义。
特性: TensorSizeTrait, VerifyTensorLayoutsTrait
接口: TT_DescriptorOpInterface, TT_DescriptorStoreLikeOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
类型(kind) | ::mlir::triton::DescriptorReduceKindAttr | 允许的 32 位无符号整数情况: 1, 2, 3, 4, 5, 6, 7, 8 |
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
浮点数、整数或 ptr 值的秩张量 |
|
可变参数的 32 位无符号整数 |
tt.descriptor_scatter (triton::DescriptorScatterOp)
将多行从单个张量分散 (scatter) 到描述符
语法
operation ::= `tt.descriptor_scatter` $desc `[` $x_offsets `,` $y_offset `]` `,` $src
attr-dict `:` type(operands)
tt.descriptor_scatter 操作将在支持的目标上降低为 NVIDIA TMA 分散操作。
desc_ptr 是指向分配在全局内存中的 TMA 描述符的指针。描述符块必须有 1 行,索引必须是一维张量。因此,结果是一个多行的二维张量。
特性: TensorSizeTrait, VerifyTensorLayoutsTrait
接口: TT_DescriptorOpInterface, TT_DescriptorStoreLikeOpInterface
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
16 位无符号整数或 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 的形状和类型必须与描述符匹配,否则结果未定义。
特性: TensorSizeTrait, VerifyTensorLayoutsTrait
接口: TT_DescriptorOpInterface, TT_DescriptorStoreLikeOpInterface
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
浮点数、整数或 ptr 值的秩张量 |
|
可变参数的 32 位无符号整数 |
tt.dot (triton::DotOp)
点积 (Dot)
语法
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:使用带有 tf32 操作的 TC。tf32x3:实现 3xTF32 技巧。更多信息请参阅 F32DotTC.cpp 中的 pass。bf16x3:实现 3xBF16 技巧。更多信息请参阅 F32DotTC.cpp 中的 pass。bf16x6:实现 6xBF16 技巧。更多信息请参阅 F32DotTC.cpp 中的 pass。ieee:不使用 TC,通过软件实现点积。如果 GPU 没有张量核心或输入不是 f32,则此标志被忽略。
特性: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: 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)
缩放点积 (Dot_scaled)
语法
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) 是一个遵循微缩放规范逐块应用缩放的函数。
特性: AlwaysSpeculatableImplTrait, AttrSizedOperandSegments, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: 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 个元素。它接收具体哪些元素是未指定的。
特性: Elementwise, SameOperandsAndResultEncoding, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, MemoryEffectOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
汇编字符串 (asm_string) | ::mlir::StringAttr | 字符串属性 |
约束 (constraints) | ::mlir::StringAttr | 字符串属性 |
纯 (pure) | ::mlir::BoolAttr | 布尔属性 |
打包元素 (packed_element) | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
可变参数,浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
可变参数,浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
tt.expand_dims (triton::ExpandDimsOp)
扩展维度 (Expand_dims)
语法
operation ::= `tt.expand_dims` $src attr-dict `:` type($src) `->` type($result)
特性: 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)
调用在 $libpath/$libname 中实现并带有 $args 的外部函数 $symbol,返回 $libpath/$libname:$symbol($args…)
特性: Elementwise, SameOperandsAndResultEncoding, SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, MemoryEffectOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
库名称 (libname) | ::mlir::StringAttr | 字符串属性 |
库路径 (libpath) | ::mlir::StringAttr | 字符串属性 |
符号 (symbol) | ::mlir::StringAttr | 字符串属性 |
纯 (pure) | ::mlir::BoolAttr | 布尔属性 |
操作数:
操作数 |
描述 |
|---|---|
|
可变参数,浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
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
特性: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
舍入 (rounding) | ::mlir::triton::RoundingModeAttr | 允许的32位无符号整数情况:0, 1 |
操作数:
操作数 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量 |
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}
特性: AffineScope, AutomaticAllocationScope, HasParent<ModuleOp>, IsolatedFromAbove, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: 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.gather (triton::GatherOp)
局部收集 (Gather) 操作
语法
operation ::= `tt.gather` $src `[` $indices `]` attr-dict `:`
functional-type(operands, results)
使用索引张量沿单个指定轴从输入张量中收集元素。输出张量的形状与索引张量相同。输入和索引张量必须具有相同的维度数,并且索引张量中不是收集维度的每个维度不能大于输入张量中的相应维度。
efficient_layout 属性是在编译器确定了操作的优化布局时设置的,表示该布局不应被更改。
特性: 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)
特性: 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)
特性: 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)
返回输入张量的直方图。箱数 (bins) 等于输出张量的维度。每个箱的宽度为 1,且从 0 开始。
特性: 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)
特性: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
64 位无符号整数或 64 位无符号整数值的张量 |
结果:
结果 |
描述 |
|---|---|
|
指针或指针值的秩张量 |
tt.join (triton::JoinOp)
沿新的次要维度连接两个张量
语法
operation ::= `tt.join` $lhs `,` $rhs attr-dict `:` type($lhs) `->` type($result)
例如,如果两个输入张量是 4x8xf32,则返回一个形状为 4x8x2xf32 的张量。
因为 Triton 张量总是具有 2 的幂次个数的元素,所以两个输入张量必须具有相同的形状。
特性: 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)
特性: AttrSizedOperandSegments, SameLoadStoreOperandsAndResultEncoding, SameLoadStoreOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: InferTypeOpInterface, MemoryEffectOpInterface, PredicatedOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
evict | ::mlir::triton::EvictionPolicyAttr | 允许的32位无符号整数情况:1, 2, 3 |
isVolatile | ::mlir::BoolAttr | 布尔属性 |
操作数:
操作数 |
描述 |
|---|---|
|
指针或指针值的秩张量 |
|
1位无符号整数或1位无符号整数值的秩张量 |
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
tt.make_range (triton::MakeRangeOp)
创建范围 (Make range)
语法
operation ::= `tt.make_range` attr-dict `:` type($result)
返回一个一维 int32 张量。
值从 $start 到 $end(不包含),步长 = 1
特性: 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 接受父张量的元信息和块大小,并返回一个描述符对象,该对象可用于从全局内存中的张量进行加载/存储。
特性: AlwaysSpeculatableImplTrait, SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
填充 (padding) | ::mlir::triton::PaddingOptionAttr | 允许的 32 位无符号整数情况: 1, 2 |
操作数:
操作数 |
描述 |
|---|---|
|
ptr |
|
可变参数的 32 位无符号整数 |
|
64 位无符号整数的可变参数 |
结果:
结果 |
描述 |
|---|---|
|
张量描述符类型( |
tt.map_elementwise (triton::MapElementwiseOp)
在张量上映射标量子区域
特性: RecursiveMemoryEffects, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
打包 (pack) | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
浮点数或整数或指针值的秩张量的可变参数 |
结果:
结果 |
描述 |
|---|---|
|
浮点数或整数或指针值的秩张量的可变参数 |
tt.map_elementwise.return (triton::MapElementwiseReturnOp)
映射逐元素操作符的终结符
语法
operation ::= `tt.map_elementwise.return` attr-dict ($result^ `:` type($result))?
特性: 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 位。
特性: 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)
浮点类型的精确除法。
特性: 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)
浮点类型的精确平方根。
特性: 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 接受一个文字字符串前缀和任意数量需要打印的标量或张量参数。格式由参数自动生成。
特性: SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:MemoryEffectOpInterface (MemoryEffectOpInterface)
效果:MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
前缀 (prefix) | ::mlir::StringAttr | 字符串属性 |
十六进制 (hex) | ::mlir::BoolAttr | 布尔属性 |
是否有符号 (isSigned) | ::mlir::DenseI32ArrayAttr | i32 密集数组属性 |
操作数:
操作数 |
描述 |
|---|---|
|
可变参数,浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
tt.ptr_to_int (triton::PtrToIntOp)
将指针转换为 int64
语法
operation ::= `tt.ptr_to_int` $src attr-dict `:` type($src) `->` type($result)
特性: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
指针或指针值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
64 位无符号整数或 64 位无符号整数值的张量 |
tt.reduce (triton::ReduceOp)
使用通用组合算法进行归约
特性: AlwaysSpeculatableImplTrait, SameOperandsEncoding, SameOperandsShape, SingleBlock, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
axis | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
浮点数或整数或指针值的秩张量的可变参数 |
结果:
结果 |
描述 |
|---|---|
|
可变参数,浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
tt.reduce.return (triton::ReduceReturnOp)
归约操作符的终结符
语法
operation ::= `tt.reduce.return` $result attr-dict `:` type($result)
特性: 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,这是对目标布局应因性能原因而保留的提示。编译器仍然可以为了更好的性能而更改它。
特性: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
允许重排 (allow_reorder) | ::mlir::UnitAttr | 单元属性 |
高效布局 (efficient_layout) | ::mlir::UnitAttr | 单元属性 |
操作数:
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
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
}
特性: AlwaysSpeculatableImplTrait, HasParent<FuncOp>, ReturnLike, TensorSizeTrait, Terminator, VerifyTensorLayoutsTrait
接口 (Interfaces): ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
任意类型的可变参数 |
tt.scan (triton::ScanOp)
使用通用组合算法进行关联扫描 (Scan)
特性: AlwaysSpeculatableImplTrait, SameOperandsAndResultEncoding, SameOperandsAndResultShape, SingleBlock, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
axis | ::mlir::IntegerAttr | 32位无符号整数属性 |
反转 (reverse) | ::mlir::BoolAttr | 布尔属性 |
操作数:
操作数 |
描述 |
|---|---|
|
浮点数或整数或指针值的秩张量的可变参数 |
结果:
结果 |
描述 |
|---|---|
|
浮点数或整数或指针值的秩张量的可变参数 |
tt.scan.return (triton::ScanReturnOp)
扫描操作符的终结符
语法
operation ::= `tt.scan.return` $result attr-dict `:` type($result)
特性: 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)
特性: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, SameOperandsAndResultEncoding, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
tt.split (triton::SplitOp)
沿最后一个维度将张量拆分为两个
语法
operation ::= `tt.split` $src attr-dict `:` type($src) `->` type($outLHS)
输入必须是一个最后一个维度大小为 2 的张量。返回两个张量,src[…, 0] 和 src[…, 1]。
例如,如果输入形状为 4x8x2xf32,则返回两个形状为 4x8xf32 的张量。
特性: 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)
特性: SameLoadStoreOperandsEncoding, SameLoadStoreOperandsShape, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: PredicatedOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
evict | ::mlir::triton::EvictionPolicyAttr | 允许的32位无符号整数情况:1, 2, 3 |
忽略 CTA (ignore_cta) | ::mlir::UnitAttr | 单元属性 |
操作数:
操作数 |
描述 |
|---|---|
|
指针或指针值的秩张量 |
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |
|
1位无符号整数或1位无符号整数值的秩张量 |
tt.trans (triton::TransOp)
重排张量的维度
语法
operation ::= `tt.trans` $src attr-dict `:` type($src) `->` type($result)
例如,给定一个形状为 [1,2,4] 的张量 x,transpose(x) 且 order=[2,0,1] 会将张量重排为形状 [4,1,2]。
虽然此操作称为“trans”,但它同时实现了 tl.trans() 和 tl.permute()。(“permute”可能是一个更好的名字,但它被称为“trans”,因为最初它只支持二维张量。)
关于编码的实现说明:
在 TritonGPU 方言(可能还有其他方言)中,为该操作的输出选择了编码,因此从代码生成的角度来看它是无操作 (nop) 的。
例如,假设张量 x 的编码方式使得 GPU 线程 [i,j,k] 拥有的寄存器包含张量的元素 [i,j,k]。现在我们以顺序 [2,1,0] 对 x 进行转置,即我们反转其维度的顺序。在 TritonGPU 中,我们将为转置的输出选择一种布局,使得 GPU 线程 [i,j,k] 拥有转置(x) 的元素 [k,j,i]。但这正是它之前拥有的同一个元素!我们要做的只是“重命名”线程 [i,j,k] 所拥有的元素。
“真正”的转置(即在 GPU 线程之间移动数据)发生在出现在该操作之前和/或之后的操作 convertLayout 中。
我们这样做是为了让您可以链接多个数据移动操作(例如转置+重塑+连接),而无需在每次操作后进入共享内存。
特性: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor, SameOperandsAndResultElementType, TensorSizeTrait, VerifyTensorLayoutsTrait
接口: 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)
特性: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait
接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
浮点数、整数或 ptr 值的秩张量 |
结果:
结果 |
描述 |
|---|---|
|
浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量 |