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::FlatSymbolRefAttrflat symbol reference attribute
arg_attrs::mlir::ArrayAttrArray of dictionary attributes
res_attrs::mlir::ArrayAttrArray of dictionary attributes

操作数:

操作数

描述

operands

任意类型的可变参数

结果:

结果

描述

«未命名»

任意类型的可变参数

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, IsolatedFromAbove, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: CallableOpInterface, FunctionOpInterface, OpAsmOpInterface, Symbol

属性:

属性MLIR类型描述
sym_name::mlir::StringAttrstring attribute
function_type::mlir::TypeAttrtype attribute of function type
sym_visibility::mlir::StringAttrstring attribute
arg_attrs::mlir::ArrayAttrArray of dictionary attributes
res_attrs::mlir::ArrayAttrArray of dictionary attributes

tt.reinterpret_tensor_descriptor (triton::ReinterpretTensorDescOp)

将指针重新解释为张量描述符

语法:

operation ::= `tt.reinterpret_tensor_descriptor` $rawDesc attr-dict `:` qualified(type($rawDesc))  `to` qualified(type($result))

该Op的存在是为了帮助从无类型的原始TMA对象过渡到有类型的Tensor描述符对象。 理想情况下,一旦API完全完善后,我们可以移除这个功能。

特性:AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数:

操作数

描述

rawDesc

ptr

结果:

结果

描述

result

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

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, 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

ptr 或 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::StringAttrstring attribute

操作数:

操作数

描述

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::MemSemanticAttrallowed 32-bit signless integer cases: 1, 2, 3, 4
scope::mlir::triton::MemSyncScopeAttrallowed 32-bit signless integer cases: 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::RMWOpAttrallowed 32-bit signless integer cases: 1, 2, 3, 4, 5, 6, 7, 8, 9, 10
sem::mlir::triton::MemSemanticAttrallowed 32-bit signless integer cases: 1, 2, 3, 4
scope::mlir::triton::MemSyncScopeAttrallowed 32-bit signless integer cases: 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

浮点数、整数或指针值的排序张量

结果:

结果

描述

result

浮点数、整数或指针值的排序张量

tt.cat (triton::CatOp)

拼接两个张量

语法:

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

特性:AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, SameTypeOperands, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数:

操作数

描述

lhs

浮点数、整数或指针值的排序张量

rhs

浮点数、整数或指针值的排序张量

结果:

结果

描述

result

浮点数、整数或指针值的排序张量

tt.clampf (triton::ClampFOp)

针对浮点类型的钳位操作

语法:

operation ::= `tt.clampf` $x `,` $min `,` $max `,` `propagateNan` `=` $propagateNan attr-dict `:` type($result)

针对浮点类型的Clamp操作。

该操作接受三个参数:x、min和max。它返回一个与x形状相同的张量,其值被限制在[min, max]范围内。

特性: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性:

属性MLIR类型描述
propagateNan::mlir::triton::PropagateNanAttrallowed 32-bit signless integer cases: 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

操作数:

操作数

描述

描述

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

x_offsets

32位无符号整数值的排序张量

y_offset

32位无符号整数

结果:

结果

描述

result

浮点数、整数或指针值的排序张量

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

属性:

属性MLIR类型描述
cache::mlir::triton::CacheModifierAttrallowed 32-bit signless integer cases: 1, 2, 3, 4, 5, 6, 7
evict::mlir::triton::EvictionPolicyAttrallowed 32-bit signless integer cases: 1, 2, 3

操作数:

操作数

描述

描述

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

indices

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

结果:

结果

描述

result

浮点数、整数或指针值的排序张量

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行且索引必须是一维张量。 因此,结果是一个包含多行的二维张量。

特性: TensorSizeTrait, VerifyTensorLayoutsTrait

操作数:

操作数

描述

描述

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

x_offsets

32位无符号整数值的排序张量

y_offset

32位无符号整数

src

浮点数、整数或指针值的排序张量

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

操作数:

操作数

描述

描述

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

src

浮点数、整数或指针值的排序张量

indices

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: 使用张量核心执行tf32运算。 tf32x3: 实现3xTF32技巧。更多信息请参阅F32DotTC.cpp中的实现 ieee: 不使用张量核心,在软件中实现点积运算。 如果GPU没有张量核心或输入不是f32,则忽略此标志。

特性:AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

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

效果: MemoryEffects::Effect{}

属性:

属性MLIR类型描述
inputPrecision::mlir::triton::InputPrecisionAttrallowed 32-bit signless integer cases: 0, 1, 2
maxNumImpreciseAcc::mlir::IntegerAttr32-bit signless integer attribute

操作数:

操作数

描述

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 = 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::ScaleDotElemTypeAttrallowed 32-bit signless integer cases: 0, 1, 2, 3, 4, 5, 6
b_elem_type::mlir::triton::ScaleDotElemTypeAttrallowed 32-bit signless integer cases: 0, 1, 2, 3, 4, 5, 6
fastMath::mlir::BoolAttrbool attribute
lhs_k_pack::mlir::BoolAttrbool attribute
rhs_k_pack::mlir::BoolAttrbool attribute

操作数:

操作数

描述

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)

运行一个内联汇编块以生成一个或多个张量。

asm块每次接收packed_element元素。具体接收哪些元素是未指定的。

特性:Elementwise, SameOperandsAndResultEncoding, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, MemoryEffectOpInterface

属性:

属性MLIR类型描述
asm_string::mlir::StringAttrstring attribute
constraints::mlir::StringAttrstring attribute
pure::mlir::BoolAttrbool attribute
packed_element::mlir::IntegerAttr32-bit signless integer attribute

操作数:

操作数

描述

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-bit signless integer attribute

操作数:

操作数

描述

src

浮点数、整数或指针值的排序张量

结果:

结果

描述

result

浮点数、整数或指针值的排序张量

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)

特性: AttrSizedOperandSegments, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: MemoryEffectOpInterface (MemoryEffectOpInterface)

效果:MemoryEffects::Effect{MemoryEffects::Read on ::mlir::triton::GlobalMemory, MemoryEffects::Write on ::mlir::triton::GlobalMemory}

属性:

属性MLIR类型描述
elem_type::mlir::IntegerAttr32-bit signless integer attribute whose value is non-negative whose maximum value is 15
interleave_layout::mlir::IntegerAttr32-bit signless integer attribute whose value is non-negative whose maximum value is 2
swizzle_mode::mlir::IntegerAttr32-bit signless integer attribute whose value is non-negative whose maximum value is 3
fill_mode::mlir::IntegerAttr32-bit signless integer attribute whose value is non-negative whose maximum value is 1

操作数:

操作数

描述

desc_ptr

Triton IR类型系统中的指针类型(::mlir::triton::PointerType)

global_address

Triton IR类型系统中的指针类型(::mlir::triton::PointerType)

box_dim

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

global_dim

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

global_stride

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

element_stride

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

tt.experimental_tensormap_fenceproxy_acquire (triton::ExperimentalTensormapFenceproxyAcquireOp)

在tensormap对象上获取fence

语法:

operation ::= `tt.experimental_tensormap_fenceproxy_acquire` $desc_ptr attr-dict `:` qualified(type($desc_ptr))

特性: TensorSizeTrait, VerifyTensorLayoutsTrait

接口: MemoryEffectOpInterface (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}

操作数:

操作数

描述

desc_ptr

Triton IR类型系统中的指针类型(::mlir::triton::PointerType)

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::StringAttrstring attribute
libpath::mlir::StringAttrstring attribute
symbol::mlir::StringAttrstring attribute
pure::mlir::BoolAttrbool attribute

操作数:

操作数

描述

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::RoundingModeAttrallowed 32-bit signless integer cases: 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-bit signless integer attribute
efficient_layout::mlir::UnitAttrunit attribute

操作数:

操作数

描述

src

浮点数、整数或指针值的排序张量

indices

整数值的排序张量

结果:

结果

描述

result

浮点数、整数或指针值的排序张量

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::ProgramIDDimAttrallowed 32-bit signless integer cases: 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::ProgramIDDimAttrallowed 32-bit signless integer cases: 0, 1, 2

结果:

结果

描述

result

32位无符号整数

tt.histogram (triton::HistogramOp)

返回输入数据的直方图。

语法:

operation ::= `tt.histogram` $src attr-dict `:` type($src) `->` type($result)

返回输入张量的直方图。箱子的数量等于输出张量的维度。每个箱子的宽度为1,且箱子从0开始。

特性:AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数:

操作数

描述

src

整数值的排序张量

结果:

结果

描述

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

浮点数、整数或指针值的排序张量

rhs

浮点数、整数或指针值的排序张量

结果:

结果

描述

result

浮点数、整数或指针值的排序张量

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 dense array attribute
padding::mlir::triton::PaddingOptionAttrallowed 32-bit signless integer cases: 1, 2
cache::mlir::triton::CacheModifierAttrallowed 32-bit signless integer cases: 1, 2, 3, 4, 5, 6, 7
evict::mlir::triton::EvictionPolicyAttrallowed 32-bit signless integer cases: 1, 2, 3
isVolatile::mlir::BoolAttrbool attribute

操作数:

操作数

描述

ptr

ptr 或 ptr 值的排名张量或 ptr

mask

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

other

浮点数或浮点值的排名张量或整数或整数值的排名张量或指针或指针值的排名张量或指针

结果:

结果

描述

result

浮点数或浮点值的排名张量或整数或整数值的排名张量或指针或指针值的排名张量或指针

tt.make_range (triton::MakeRangeOp)

生成范围

语法:

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

返回一个一维int32张量。

数值范围从$start到$end(不包含$end),步长=1

特性:AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性:

属性MLIR类型描述
start::mlir::IntegerAttr32-bit signless integer attribute
end::mlir::IntegerAttr32-bit signless integer attribute

结果:

结果

描述

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

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

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>

特性: AlwaysSpeculatableImplTrait, SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性:

属性MLIR类型描述
order::mlir::DenseI32ArrayAttri32 dense array attribute

操作数:

操作数

描述

base

ptr

shape

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

strides

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

offsets

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

结果:

结果

描述

result

ptr

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::StringAttrstring attribute
hex::mlir::BoolAttrbool attribute
isSigned::mlir::DenseI32ArrayAttri32 dense array attribute

操作数:

操作数

描述

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-bit signless integer attribute

操作数:

操作数

描述

srcs

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

结果:

结果

描述

result

可变数量的浮点值或浮点值的分级张量,或整数值或整数值的分级张量,或指针或指针值的分级张量或指针

tt.reduce.return (triton::ReduceReturnOp)

reduce操作符的终止条件

语法:

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

特性: AlwaysSpeculatableImplTrait, HasParent, 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::UnitAttrunit attribute
efficient_layout::mlir::UnitAttrunit attribute

操作数:

操作数

描述

src

浮点数、整数或指针值的排序张量

结果:

结果

描述

result

浮点数、整数或指针值的排序张量

tt.scan (triton::ScanOp)

使用通用组合算法的关联扫描

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

接口: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

属性:

属性MLIR类型描述
axis::mlir::IntegerAttr32-bit signless integer attribute
reverse::mlir::BoolAttrbool attribute

操作数:

操作数

描述

srcs

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

结果:

结果

描述

result

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

tt.scan.return (triton::ScanReturnOp)

扫描操作符的终止器

语法:

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

特性: AlwaysSpeculatableImplTrait, HasParent, ReturnLike, TensorSizeTrait, Terminator, VerifyTensorLayoutsTrait

接口:ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

效果: MemoryEffects::Effect{}

操作数:

操作数

描述

result

任意类型的可变参数

tt.splat (triton::SplatOp)

飞溅效果

语法:

operation ::= `tt.splat` $src attr-dict `:` type($src) `->` type($result)

特性: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, SameOperandsAndResultEncoding, TensorSizeTrait, VerifyTensorLayoutsTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数:

操作数

描述

src

浮点数或浮点值的排名张量或整数或整数值的排名张量或指针或指针值的排名张量或指针

结果:

结果

描述

result

浮点数、整数或指针值的排序张量

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

浮点数、整数或指针值的排序张量

结果:

结果

描述

outLHS

浮点数、整数或指针值的排序张量

outRHS

浮点数、整数或指针值的排序张量

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 dense array attribute
cache::mlir::triton::CacheModifierAttrallowed 32-bit signless integer cases: 1, 2, 3, 4, 5, 6, 7
evict::mlir::triton::EvictionPolicyAttrallowed 32-bit signless integer cases: 1, 2, 3

操作数:

操作数

描述

ptr

ptr 或 ptr 值的排名张量或 ptr

value

浮点数或浮点值的排名张量或整数或整数值的排名张量或指针或指针值的排名张量或指针

mask

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

tt.trans (triton::TransOp)

重新排列张量的维度

语法:

operation ::= `tt.trans` $src attr-dict `:` type($src) `->` type($result)

例如,给定一个形状为[1,2,4]的张量x,使用order=[2,0,1]进行transpose(x)操作会将张量重新排列为形状[4,1,2]。

虽然这个操作被称为"trans",但它同时实现了tl.trans()和tl.permute()功能。("permute"可能是更合适的名称,但由于最初它仅支持2D张量,所以被称为"trans"。)

关于编码实现的注意事项:

在TritonGPU方言(以及其他方言中可能也是),为此操作的输出选择了一种编码,因此从代码生成的角度来看它是一个无操作。

例如,假设张量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 dense array attribute

操作数:

操作数

描述

src

浮点数、整数或指针值的排序张量

结果:

结果

描述

result

浮点数、整数或指针值的排序张量