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

特性:TensorSizeTrait, VerifyTensorLayoutsTrait

接口: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}

特性:AffineScope, AutomaticAllocationScope, HasParent<ModuleOp>, IsolatedFromAbove, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: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
}

特性:AlwaysSpeculatableImplTrait, HasParent<FuncOp>, ReturnLike, TensorSizeTrait, Terminator, VerifyTensorLayoutsTrait

接口:ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

效果: MemoryEffects::Effect{}

操作数:

操作数

描述

srcs

任意类型的可变参数

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

指针或指针值的秩张量

offset

整数或整数值的秩张量

结果:

结果

描述

result

指针或指针值的秩张量

tt.advance (triton::AdvanceOp)

按偏移量推进张量指针

语法

operation ::= `tt.advance` $ptr `,` `[` $offsets `]` attr-dict `:` type($result)

特性:AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数:

操作数

描述

ptr

ptr

offsets

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

结果:

结果

描述

result

ptr

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 类型描述
message::mlir::StringAttr字符串属性

操作数:

操作数

描述

condition

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

tt.atomic_cas (triton::AtomicCASOp)

原子比较并交换

语法

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

特性: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)

原子读-改-写

语法

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

属性:

属性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.cat (triton::CatOp)

连接2个张量

语法

operation ::= `tt.cat` $lhs `,` $rhs attr-dict `:` type($lhs) `->` type($result)

特性:SameOperandsAndResultElementType, SameTypeOperands, TensorSizeTrait, VerifyTensorLayoutsTrait

接口:NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数 (Operands):

操作数

描述

lhs

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

rhs

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

结果:

结果

描述

result

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

tt.clampf (triton::ClampFOp)

浮点类型钳位操作

语法

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 类型描述
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行,并且索引必须是1D张量。因此,结果是多行的2D张量。

特性:TensorSizeTrait, VerifyTensorLayoutsTrait

接口:TT_DescriptorOpInterface

操作数:

操作数

描述

desc

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

x_offsets

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_DescriptorOpInterface

属性:

属性MLIR 类型描述
cache::mlir::triton::CacheModifierAttr允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7
evict::mlir::triton::EvictionPolicyAttr允许的32位无符号整数情况:1, 2, 3

操作数 (Operands):

操作数

描述

desc

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

indices

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

结果 (Results):

结果

描述

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

操作数 (Operands):

操作数

描述

desc

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

src

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

indices

可变数量的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张量。

特性:TensorSizeTrait, VerifyTensorLayoutsTrait

接口:TT_DescriptorOpInterface, TT_DescriptorStoreLikeOpInterface

操作数:

操作数

描述

desc

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

x_offsets

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)

点积

语法

operation ::= `tt.dot` $a`,` $b`,` $c (`,` `inputPrecision` `=` $inputPrecision^)? attr-dict `:`
              type($a) `*` type($b) `->` type($d)

$d = 矩阵乘法($a, $b) + $c. $inputPrecision 描述了当输入为 f32 时如何使用 TC。它可以是以下之一:tf32, tf32x3, ieee。tf32: 使用 TC 进行 tf32 操作。tf32x3: 实现 3xTF32 技巧。更多信息请参阅 F32DotTC.cpp 中的 pass。ieee: 不使用 TC,在软件中实现点积。如果 GPU 没有 Tensor cores 或输入不是 f32,则此标志被忽略。

特性:AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

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

效果: MemoryEffects::Effect{}

属性:

属性MLIR 类型描述
inputPrecision::mlir::triton::InputPrecisionAttr允许的32位无符号整数情况:0, 1, 2
maxNumImpreciseAcc::mlir::IntegerAttr32位无符号整数属性

操作数:

操作数

描述

a

浮点数或整数值的秩张量

b

浮点数或整数值的秩张量

c

浮点数或整数值的秩张量

结果:

结果

描述

d

浮点数或整数值的秩张量

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 = 矩阵乘法(缩放($a, $a_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)

扩展维度

语法

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 中实现的外部函数 $symbol,参数为 $args,返回 $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.gather (triton::GatherOp)

本地收集操作

语法

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)

返回输入张量的直方图。bin 的数量等于输出张量的维度。每个 bin 的宽度为1,且 bin 从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

属性:

属性MLIR 类型描述
boundaryCheck::mlir::DenseI32ArrayAttri32 密集数组属性
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

指针或指针值的秩张量或指针

mask

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

other

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

结果:

结果

描述

result

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

tt.make_range (triton::MakeRangeOp)

生成范围

语法

operation ::= `tt.make_range` attr-dict `:` type($result)

返回一个1D 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{}

操作数:

操作数

描述

base

ptr

shape

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

strides

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

结果:

结果

描述

result

张量描述符类型(::mlir::triton::TensorDescType)在 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>>

特性:AlwaysSpeculatableImplTrait, SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性:

属性MLIR 类型描述
order::mlir::DenseI32ArrayAttri32 密集数组属性

操作数:

操作数

描述

base

ptr

shape

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

strides

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

offsets

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

结果:

结果

描述

result

ptr

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

接口: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

接口: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.scan (triton::ScanOp)

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

特性: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

接口: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

属性:

属性MLIR 类型描述
boundaryCheck::mlir::DenseI32ArrayAttri32 密集数组属性
cache::mlir::triton::CacheModifierAttr允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7
evict::mlir::triton::EvictionPolicyAttr允许的32位无符号整数情况:1, 2, 3

操作数:

操作数

描述

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]。现在我们用 order [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

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