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{}

操作数:

操作数

描述

ptr

指针或指针值的秩张量

偏移量

整数或整数值的秩张量

结果:

结果

描述

result

指针或指针值的秩张量

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

操作数:

操作数

描述

ptr

指针或指针值的秩张量

cmp

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

val

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

结果:

结果

描述

result

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

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

操作数:

操作数

描述

ptr

指针或指针值的秩张量

val

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

mask

1位无符号整数或1位无符号整数值的秩张量

结果:

结果

描述

result

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

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{}

操作数:

操作数

描述

src

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

结果:

结果

描述

result

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

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{}

操作数:

操作数

描述

src

浮点数、整数或 ptr 值的秩张量

结果:

结果

描述

result

浮点数、整数或 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{}

操作数:

操作数

描述

左操作数 (lhs)

浮点数、整数或 ptr 值的秩张量

右操作数 (rhs)

浮点数、整数或 ptr 值的秩张量

结果:

结果

描述

result

浮点数、整数或 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

操作数:

操作数

描述

x

浮点数或浮点值的秩张量

最小值 (min)

浮点数或浮点值的秩张量

最大值 (max)

浮点数或浮点值的秩张量

结果:

结果

描述

result

浮点数或浮点值的秩张量

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

操作数:

操作数

描述

desc

张量描述符类型(::mlir::triton::TensorDescType)在 Triton IR 类型系统中

x_offsets

16 位无符号整数或 32 位无符号整数值的秩张量

y_offset

32位无符号整数

结果:

结果

描述

result

浮点数、整数或 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

操作数:

操作数

描述

desc

张量描述符类型(::mlir::triton::TensorDescType)在 Triton IR 类型系统中

indices

可变参数的 32 位无符号整数

结果:

结果

描述

result

浮点数、整数或 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

操作数:

操作数

描述

desc

张量描述符类型(::mlir::triton::TensorDescType)在 Triton IR 类型系统中

src

浮点数、整数或 ptr 值的秩张量

indices

可变参数的 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

操作数:

操作数

描述

desc

张量描述符类型(::mlir::triton::TensorDescType)在 Triton IR 类型系统中

x_offsets

16 位无符号整数或 32 位无符号整数值的秩张量

y_offset

32位无符号整数

src

浮点数、整数或 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

操作数:

操作数

描述

desc

张量描述符类型(::mlir::triton::TensorDescType)在 Triton IR 类型系统中

src

浮点数、整数或 ptr 值的秩张量

indices

可变参数的 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::IntegerAttr32位无符号整数属性

操作数:

操作数

描述

a

浮点数或整数值的秩张量

b

浮点数或整数值的秩张量

c

浮点数或整数值的秩张量

结果:

结果

描述

d

浮点数或整数值的秩张量

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

操作数:

操作数

描述

a

浮点数或 8 位无符号整数值的秩张量

b

浮点数或 8 位无符号整数值的秩张量

c

浮点值的秩张量

a_scale

浮点数或 8 位无符号整数值的秩张量

b_scale

浮点数或 8 位无符号整数值的秩张量

结果:

结果

描述

d

浮点值的秩张量

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::IntegerAttr32位无符号整数属性

操作数:

操作数

描述

参数 (args)

可变参数,浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

结果:

结果

描述

result

可变参数,浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

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::IntegerAttr32位无符号整数属性

操作数:

操作数

描述

src

浮点数、整数或 ptr 值的秩张量

结果:

结果

描述

result

浮点数、整数或 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布尔属性

操作数:

操作数

描述

源张量 (srcs)

可变参数,浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

结果:

结果

描述

result

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

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

操作数:

操作数

描述

src

浮点数或浮点值的秩张量

结果:

结果

描述

result

浮点数或浮点值的秩张量

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::IntegerAttr32位无符号整数属性
高效布局 (efficient_layout)::mlir::UnitAttr单元属性

操作数:

操作数

描述

src

浮点数、整数或 ptr 值的秩张量

indices

整数值的秩张量

结果:

结果

描述

result

浮点数、整数或 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

结果:

结果

描述

result

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

结果:

结果

描述

result

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{}

操作数:

操作数

描述

src

整数值的秩张量

mask

1位无符号整数或1位无符号整数值的秩张量

结果:

结果

描述

result

整数值的秩张量

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{}

操作数:

操作数

描述

src

64 位无符号整数或 64 位无符号整数值的张量

结果:

结果

描述

result

指针或指针值的秩张量

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{}

操作数:

操作数

描述

左操作数 (lhs)

浮点数、整数或 ptr 值的秩张量

右操作数 (rhs)

浮点数、整数或 ptr 值的秩张量

结果:

结果

描述

result

浮点数、整数或 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布尔属性

操作数:

操作数

描述

ptr

指针或指针值的秩张量

mask

1位无符号整数或1位无符号整数值的秩张量

other

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

结果:

结果

描述

result

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

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::IntegerAttr32位无符号整数属性
结束 (end)::mlir::IntegerAttr32位无符号整数属性

结果:

结果

描述

result

整数值的秩张量

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

操作数:

操作数

描述

基址 (base)

ptr

形状 (shape)

可变参数的 32 位无符号整数

步长 (strides)

64 位无符号整数的可变参数

结果:

结果

描述

result

张量描述符类型(::mlir::triton::TensorDescType)在 Triton IR 类型系统中

tt.map_elementwise (triton::MapElementwiseOp)

在张量上映射标量子区域

特性: RecursiveMemoryEffects, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait

属性:

属性MLIR 类型描述
打包 (pack)::mlir::IntegerAttr32位无符号整数属性

操作数:

操作数

描述

源张量 (srcs)

浮点数或整数或指针值的秩张量的可变参数

结果:

结果

描述

result

浮点数或整数或指针值的秩张量的可变参数

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{}

操作数:

操作数

描述

result

任意类型的可变参数

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{}

操作数:

操作数

描述

x

整数或整数值的秩张量

y

整数或整数值的秩张量

结果:

结果

描述

result

整数或整数值的秩张量

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{}

操作数:

操作数

描述

x

浮点数或浮点值的秩张量

y

浮点数或浮点值的秩张量

结果:

结果

描述

result

浮点数或浮点值的秩张量

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{}

操作数:

操作数

描述

x

浮点数或浮点值的秩张量

结果:

结果

描述

result

浮点数或浮点值的秩张量

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::DenseI32ArrayAttri32 密集数组属性

操作数:

操作数

描述

参数 (args)

可变参数,浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

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{}

操作数:

操作数

描述

src

指针或指针值的秩张量

结果:

结果

描述

result

64 位无符号整数或 64 位无符号整数值的张量

tt.reduce (triton::ReduceOp)

使用通用组合算法进行归约

特性: AlwaysSpeculatableImplTrait, SameOperandsEncoding, SameOperandsShape, SingleBlock, TensorSizeTrait, VerifyTensorLayoutsTrait

接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性:

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

操作数:

操作数

描述

源张量 (srcs)

浮点数或整数或指针值的秩张量的可变参数

结果:

结果

描述

result

可变参数,浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

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{}

操作数:

操作数

描述

result

任意类型的可变参数

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单元属性

操作数:

操作数

描述

src

浮点数、整数或 ptr 值的秩张量

结果:

结果

描述

result

浮点数、整数或 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{}

操作数:

操作数

描述

源张量 (srcs)

任意类型的可变参数

tt.scan (triton::ScanOp)

使用通用组合算法进行关联扫描 (Scan)

特性: AlwaysSpeculatableImplTrait, SameOperandsAndResultEncoding, SameOperandsAndResultShape, SingleBlock, TensorSizeTrait, VerifyTensorLayoutsTrait

接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性:

属性MLIR 类型描述
axis::mlir::IntegerAttr32位无符号整数属性
反转 (reverse)::mlir::BoolAttr布尔属性

操作数:

操作数

描述

源张量 (srcs)

浮点数或整数或指针值的秩张量的可变参数

结果:

结果

描述

result

浮点数或整数或指针值的秩张量的可变参数

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{}

操作数:

操作数

描述

result

任意类型的可变参数

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{}

操作数:

操作数

描述

src

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

结果:

结果

描述

result

浮点数、整数或 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{}

操作数:

操作数

描述

src

浮点数、整数或 ptr 值的秩张量

结果:

结果

描述

左出 (outLHS)

浮点数、整数或 ptr 值的秩张量

右出 (outRHS)

浮点数、整数或 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单元属性

操作数:

操作数

描述

ptr

指针或指针值的秩张量

值 (value)

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量

mask

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::DenseI32ArrayAttri32 密集数组属性

操作数:

操作数

描述

src

浮点数、整数或 ptr 值的秩张量

结果:

结果

描述

result

浮点数、整数或 ptr 值的秩张量

tt.unsplat (triton::UnsplatOp)

将具有单个元素的张量转换为标量

语法

operation ::= `tt.unsplat` $src attr-dict `:` type($src)

特性: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

接口:ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数:

操作数

描述

src

浮点数、整数或 ptr 值的秩张量

结果:

结果

描述

result

浮点数或浮点值的秩张量,或整数或整数值的秩张量,或指针或指针值的秩张量