TritonOps

tt.call (triton::CallOp)

调用操作

语法

operation ::= `tt.call` $callee `(` $operands `)` attr-dict `:` functional-type($operands, results)

tt.call 操作表示对调用所在符号范围内的函数的直接调用。调用的操作数和结果类型必须与指定的函数类型匹配。被调用者编码为名为“callee”的符号引用属性。

示例

%2 = tt.call @my_add(%0, %1) : (f32, f32) -> f32

Traits: TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: CallOpInterface, SymbolUserOpInterface

属性:

属性MLIR 类型描述
callee::mlir::FlatSymbolRefAttr平面符号引用属性
arg_attrs::mlir::ArrayAttr字典属性数组
res_attrs::mlir::ArrayAttr字典属性数组

操作数:

操作数

描述

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}

Traits: AffineScope, AutomaticAllocationScope, HasParent<ModuleOp>, IsolatedFromAbove, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: CallableOpInterface, FunctionOpInterface, OpAsmOpInterface, Symbol

属性:

属性MLIR 类型描述
sym_name::mlir::StringAttr字符串属性
function_type::mlir::TypeAttr函数类型的类型属性
sym_visibility::mlir::StringAttr字符串属性
arg_attrs::mlir::ArrayAttr字典属性数组
res_attrs::mlir::ArrayAttr字典属性数组

tt.reinterpret_tensor_descriptor (triton::ReinterpretTensorDescOp)

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

语法

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

此操作(Op)旨在帮助从无类型原始 TMA 对象过渡到类型化张量描述符对象。理想情况下,一旦 API 完全成熟,我们可以移除此操作。

Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

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
}

Traits: AlwaysSpeculatableImplTrait, HasParent<FuncOp>, ReturnLike, TensorSizeTrait, Terminator, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

srcs

可变数量的任何类型

tt.addptr (triton::AddPtrOp)

语法

operation ::= `tt.addptr` $ptr `,` $offset attr-dict `:` type($result) `,` type($offset)

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

ptr

ptr 或 ptr 值的 ranked tensor

offset

整数或整数值的 ranked tensor

结果:

结果

描述

result

ptr 或 ptr 值的 ranked tensor

tt.advance (triton::AdvanceOp)

按偏移量推进张量指针

语法

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

Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

ptr

ptr

offsets

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

结果:

结果

描述

result

ptr

tt.assert (triton::AssertOp)

设备端断言,用于正确性检查,类似于 CUDA

语法

operation ::= `tt.assert` $condition `,` $message attr-dict `:` type($condition)

tt.assert 接受一个条件张量和一个消息字符串。如果条件为 false,则打印消息并中止程序。

Traits: TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: MemoryEffectOpInterface (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}

属性:

属性MLIR 类型描述
message::mlir::StringAttr字符串属性

操作数:

操作数

描述

condition

1 位无符号整数或 1 位无符号整数值的 tensor

tt.atomic_cas (triton::AtomicCASOp)

原子比较和交换

语法

operation ::= `tt.atomic_cas` $sem `,` $scope `,` $ptr `,` $cmp `,` $val attr-dict `:`
              functional-type(operands, $result)

在 $ptr 位置比较 $cmp 和数据 $old,

如果 $old == $cmp,则将 $val 存储到 $ptr,

否则将 $old 存储到 $ptr,

返回 $old

Traits: SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait

属性:

属性MLIR 类型描述
sem::mlir::triton::MemSemanticAttr允许的 32 位无符号整数情况:1、2、3、4
scope::mlir::triton::MemSyncScopeAttr允许的 32 位无符号整数情况:1、2、3

操作数:

操作数

描述

ptr

ptr 或 ptr 值的 ranked tensor

cmp

浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

val

浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

结果:

结果

描述

result

浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

tt.atomic_rmw (triton::AtomicRMWOp)

原子读-修改-写

语法

operation ::= `tt.atomic_rmw` $atomic_rmw_op `,` $sem `,` $scope `,` $ptr `,` $val (`,` $mask^)?  attr-dict `:`
              functional-type(operands, $result)

加载 $ptr 位置的数据,使用 $val 执行 $rmw_op,并将结果存储到 $ptr。

返回 $ptr 位置的旧值

Traits: SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait

属性:

属性MLIR 类型描述
atomic_rmw_op::mlir::triton::RMWOpAttr允许的 32 位无符号整数情况:1、2、3、4、5、6、7、8、9、10
sem::mlir::triton::MemSemanticAttr允许的 32 位无符号整数情况:1、2、3、4
scope::mlir::triton::MemSyncScopeAttr允许的 32 位无符号整数情况:1、2、3

操作数:

操作数

描述

ptr

ptr 或 ptr 值的 ranked tensor

val

浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

mask

1 位无符号整数或 1 位无符号整数值的 ranked tensor

结果:

结果

描述

result

浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

tt.bitcast (triton::BitcastOp)

在相同位宽的类型之间进行位转换

语法

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

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

src

浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

结果:

结果

描述

result

浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

tt.broadcast (triton::BroadcastOp)

广播张量

语法

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

对于给定的张量,广播将一个或多个大小为 1 的维度更改为新的大小,例如 tensor<1x32x1xf32> -> tensor<2x32x4xf32>。您不能更改非 1 维度的大小。

Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, SameOperandsAndResultEncoding, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

src

浮点数或整数或 ptr 值的 ranked tensor

结果:

结果

描述

result

浮点数或整数或 ptr 值的 ranked tensor

tt.cat (triton::CatOp)

连接 2 个张量

语法

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

Traits: SameOperandsAndResultElementType, SameTypeOperands, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

lhs

浮点数或整数或 ptr 值的 ranked tensor

rhs

浮点数或整数或 ptr 值的 ranked tensor

结果:

结果

描述

result

浮点数或整数或 ptr 值的 ranked tensor

tt.clampf (triton::ClampFOp)

浮点类型的截断操作

语法

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

浮点类型的截断操作。

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

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

属性MLIR 类型描述
propagateNan::mlir::triton::PropagateNanAttr允许的 32 位无符号整数情况:0、65535

操作数:

操作数

描述

x

浮点数或浮点数值的 ranked tensor

min

浮点数或浮点数值的 ranked tensor

max

浮点数或浮点数值的 ranked tensor

结果:

结果

描述

result

浮点数或浮点数值的 ranked tensor

tt.descriptor_gather (triton::DescriptorGatherOp)

从描述符中收集多行到单个张量

语法

operation ::= `tt.descriptor_gather` $desc `[` $x_offsets `,` $y_offset `]`
              attr-dict `:` functional-type(operands, results)

tt.descriptor_gather 操作将在支持它的目标上降低为 NVIDIA TMA 收集操作。

desc_ptr 是指向全局内存中分配的 TMA 描述符的指针。描述符块必须有 1 行,索引必须是 1D 张量。因此,结果是一个包含多行的 2D 张量。

Traits: TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: TT_DescriptorOpInterface

操作数:

操作数

描述

desc

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

x_offsets

32 位无符号整数值的 ranked tensor

y_offset

32 位无符号整数

结果:

结果

描述

result

浮点数或整数或 ptr 值的 ranked tensor

tt.descriptor_load (triton::DescriptorLoadOp)

从描述符加载

语法

operation ::= `tt.descriptor_load` $desc `[` $indices `]`
              oilist(
              `cacheModifier` `=` $cache |
              `evictionPolicy` `=` $evict
              )
              attr-dict `:` qualified(type($desc)) `->` type($result)

此操作将在支持它的目标上降低为 Nvidia TMA 加载操作。desc 是一个张量描述符对象。目标张量类型和形状必须与描述符匹配,否则结果是未定义的。

Traits: TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: TT_DescriptorOpInterface

属性:

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

操作数:

操作数

描述

desc

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

indices

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

结果:

结果

描述

result

浮点数或整数或 ptr 值的 ranked tensor

tt.descriptor_reduce (triton::DescriptorReduceOp)

根据描述符执行规约存储操作

语法

operation ::= `tt.descriptor_reduce` $kind `,` $desc `[` $indices `]` `,` $src
              attr-dict `:` qualified(type($desc)) `,` type($src)

此操作将在支持它的目标上降低为 Nvidia TMA 存储操作。desc 是一个张量描述符对象。src 的形状和类型必须与描述符匹配,否则结果是未定义的。

Traits: TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: TT_DescriptorOpInterface, TT_DescriptorStoreLikeOpInterface

属性:

属性MLIR 类型描述
kind::mlir::triton::DescriptorReduceKindAttr允许的 32 位无符号整数情况:1、2、3、4、5、6、7、8

操作数:

操作数

描述

desc

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

src

浮点数或整数或 ptr 值的 ranked tensor

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 张量。

Traits: TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: TT_DescriptorOpInterface, TT_DescriptorStoreLikeOpInterface

操作数:

操作数

描述

desc

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

x_offsets

32 位无符号整数值的 ranked tensor

y_offset

32 位无符号整数

src

浮点数或整数或 ptr 值的 ranked tensor

tt.descriptor_store (triton::DescriptorStoreOp)

根据描述符存储值

语法

operation ::= `tt.descriptor_store` $desc `[` $indices `]` `,` $src
              attr-dict `:` qualified(type($desc)) `,` type($src)

此操作将在支持它的目标上降低为 Nvidia TMA 存储操作。desc 是一个张量描述符对象。src 的形状和类型必须与描述符匹配,否则结果是未定义的。

Traits: TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: TT_DescriptorOpInterface, TT_DescriptorStoreLikeOpInterface

操作数:

操作数

描述

desc

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

src

浮点数或整数或 ptr 值的 ranked tensor

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:使用 TC 执行 tf32 操作。tf32x3:实现 3xTF32 技巧。更多信息请参见 F32DotTC.cpp 中的 pass。ieee:不使用 TC,在软件中实现点积。如果 GPU 不具有 Tensor core 或输入不是 f32,则忽略此标志。

Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, DotOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

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

操作数:

操作数

描述

a

浮点数或整数值的 ranked tensor

b

浮点数或整数值的 ranked tensor

c

浮点数或整数值的 ranked tensor

结果:

结果

描述

d

浮点数或整数值的 ranked tensor

tt.dot_scaled (triton::DotScaledOp)

缩放点积

语法

operation ::= `tt.dot_scaled` $a (`scale` $a_scale^)? `,` $b (`scale` $b_scale^)? `,` $c
              `lhs` `=` $a_elem_type `rhs` `=` $b_elem_type attr-dict
              `:` type($a) (`,` type($a_scale)^)? `*` type($b) (`,` type($b_scale)^)? `->` type($d)

$d = matrix_multiply(scale($a, $a_scale), scale($b, $b_scale)) + $c。其中 scale(x, s) 是一个根据微缩放规范对每个块应用缩放的函数。

Traits: AlwaysSpeculatableImplTrait, AttrSizedOperandSegments, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, DotOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

属性MLIR 类型描述
a_elem_type::mlir::triton::ScaleDotElemTypeAttr允许的 32 位无符号整数情况:0、1、2、3、4、5、6
b_elem_type::mlir::triton::ScaleDotElemTypeAttr允许的 32 位无符号整数情况:0、1、2、3、4、5、6
fastMath::mlir::BoolAttr布尔属性
lhs_k_pack::mlir::BoolAttr布尔属性
rhs_k_pack::mlir::BoolAttr布尔属性

操作数:

操作数

描述

a

浮点数或 8 位无符号整数值的 ranked tensor

b

浮点数或 8 位无符号整数值的 ranked tensor

c

浮点数值的 ranked tensor

a_scale

浮点数或 8 位无符号整数值的 ranked tensor

b_scale

浮点数或 8 位无符号整数值的 ranked tensor

结果:

结果

描述

d

浮点数值的 ranked tensor

tt.elementwise_inline_asm (triton::ElementwiseInlineAsmOp)

应用逐元素操作到一组打包元素的内联汇编。

语法

operation ::= `tt.elementwise_inline_asm` $asm_string attr-dict ($args^ `:` type($args))? `->` type($result)

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

汇编块每次给定 packed_element 个元素。它接收哪些元素是未指定的。

Traits: Elementwise, SameOperandsAndResultEncoding, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, MemoryEffectOpInterface

属性:

属性MLIR 类型描述
asm_string::mlir::StringAttr字符串属性
constraints::mlir::StringAttr字符串属性
pure::mlir::BoolAttr布尔属性
packed_element::mlir::IntegerAttr32 位无符号整数属性

操作数:

操作数

描述

args

可变数量的浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

结果:

结果

描述

result

可变数量的浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

tt.expand_dims (triton::ExpandDimsOp)

扩展维度

语法

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

Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

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

操作数:

操作数

描述

src

浮点数或整数或 ptr 值的 ranked tensor

结果:

结果

描述

result

浮点数或整数或 ptr 值的 ranked tensor

tt.experimental_tensormap_create (triton::ExperimentalTensormapCreateOp)

在设备上创建新的 TMA 描述符

语法

operation ::= `tt.experimental_tensormap_create` $desc_ptr `,` $global_address `,`
              `[` $box_dim `]` `,`
              `[` $global_dim `]` `,`
              `[` $global_stride `]` `,`
              `[` $element_stride `]`
              attr-dict `:` functional-type(operands, results)

Traits: AttrSizedOperandSegments, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: MemoryEffectOpInterface (MemoryEffectOpInterface)

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

属性:

属性MLIR 类型描述
elem_type::mlir::IntegerAttr值为非负且最大值为 15 的 32 位无符号整数属性
interleave_layout::mlir::IntegerAttr值为非负且最大值为 2 的 32 位无符号整数属性
swizzle_mode::mlir::IntegerAttr值为非负且最大值为 3 的 32 位无符号整数属性
fill_mode::mlir::IntegerAttr值为非负且最大值为 1 的 32 位无符号整数属性

操作数:

操作数

描述

desc_ptr

Triton IR 类型系统中的 Pointer type (::mlir::triton::PointerType)

global_address

Triton IR 类型系统中的 Pointer type (::mlir::triton::PointerType)

box_dim

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

global_dim

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

global_stride

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

element_stride

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

tt.experimental_tensormap_fenceproxy_acquire (triton::ExperimentalTensormapFenceproxyAcquireOp)

在 tensormap 对象上获取围栏

语法

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

Traits: TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: MemoryEffectOpInterface (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}

操作数:

操作数

描述

desc_ptr

Triton IR 类型系统中的 Pointer type (::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…)

Traits: Elementwise, SameOperandsAndResultEncoding, SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, MemoryEffectOpInterface

属性:

属性MLIR 类型描述
libname::mlir::StringAttr字符串属性
libpath::mlir::StringAttr字符串属性
symbol::mlir::StringAttr字符串属性
pure::mlir::BoolAttr布尔属性

操作数:

操作数

描述

srcs

可变数量的浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

结果:

结果

描述

result

浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

tt.fp_to_fp (triton::FpToFpOp)

自定义类型的浮点转换

语法

operation ::= `tt.fp_to_fp` $src attr-dict  (`,` `rounding` `=` $rounding^)? `:` type($src) `->` type($result)

自定义类型 (F8) 和非默认舍入模式的浮点转换。

F8 <-> FP16, BF16, FP32, FP64

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

属性MLIR 类型描述
rounding::mlir::triton::RoundingModeAttr允许的 32 位无符号整数情况:0、1

操作数:

操作数

描述

src

浮点数或浮点数值的 ranked tensor

结果:

结果

描述

result

浮点数或浮点数值的 ranked tensor

tt.gather (triton::GatherOp)

本地收集操作

语法

operation ::= `tt.gather` $src `[` $indices `]` attr-dict `:`
              functional-type(operands, results)

使用索引张量沿着单个指定轴从输入张量收集元素。输出张量具有与索引张量相同的形状。输入张量和索引张量必须具有相同的维度数,并且索引张量中不是收集维度的每个维度不能大于输入张量中的相应维度。

当编译器确定了操作的优化布局时,会设置 efficient_layout 属性,表示不应更改该布局。

Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

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

操作数:

操作数

描述

src

浮点数或整数或 ptr 值的 ranked tensor

indices

整数值的 ranked tensor

结果:

结果

描述

result

浮点数或整数或 ptr 值的 ranked tensor

tt.get_num_programs (triton::GetNumProgramsOp)

语法

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

Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

属性MLIR 类型描述
axis::mlir::triton::ProgramIDDimAttr允许的 32 位无符号整数情况:0、1、2

结果:

结果

描述

result

32 位无符号整数

tt.get_program_id (triton::GetProgramIdOp)

语法

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

Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

属性MLIR 类型描述
axis::mlir::triton::ProgramIDDimAttr允许的 32 位无符号整数情况:0、1、2

结果:

结果

描述

result

32 位无符号整数

tt.histogram (triton::HistogramOp)

返回输入的直方图。

语法

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

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

Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

src

整数值的 ranked tensor

结果:

结果

描述

result

整数值的 ranked tensor

tt.int_to_ptr (triton::IntToPtrOp)

将 int64 转换为指针

语法

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

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

src

64 位无符号整数或 64 位无符号整数值的 tensor

结果:

结果

描述

result

ptr 或 ptr 值的 ranked tensor

tt.join (triton::JoinOp)

将两个张量沿一个新的、次要的维度连接

语法

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

例如,如果两个输入张量是 4x8xf32,则返回一个形状为 4x8x2xf32 的张量。

因为 Triton 张量总是具有 2 的幂次方个元素,所以两个输入张量必须具有相同的形状。

Traits: AlwaysSpeculatableImplTrait, SameTypeOperands, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

lhs

浮点数或整数或 ptr 值的 ranked tensor

rhs

浮点数或整数或 ptr 值的 ranked tensor

结果:

结果

描述

result

浮点数或整数或 ptr 值的 ranked tensor

tt.load (triton::LoadOp)

从指针张量或张量指针加载

语法

operation ::= `tt.load` $ptr (`,` $mask^)? (`,` $other^)?
              oilist(
              `cacheModifier` `=` $cache |
              `evictionPolicy` `=` $evict
              )
              attr-dict `:` type($ptr)

Traits: AttrSizedOperandSegments, SameLoadStoreOperandsAndResultEncoding, SameLoadStoreOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: InferTypeOpInterface, MemoryEffectOpInterface

属性:

属性MLIR 类型描述
boundaryCheck::mlir::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

ptr 或 ptr 值的 ranked tensor 或 ptr

mask

1 位无符号整数或 1 位无符号整数值的 ranked tensor

other

浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

结果:

结果

描述

result

浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

tt.make_range (triton::MakeRangeOp)

生成范围

语法

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

返回一个 1D int32 张量。

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

Traits: AlwaysSpeculatableImplTrait, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

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

结果:

结果

描述

result

整数值的 ranked tensor

tt.make_tensor_descriptor (triton::MakeTensorDescOp)

创建一个包含父张量和块大小元信息的张量描述符类型

语法

operation ::= `tt.make_tensor_descriptor` $base `,` `[` $shape `]` `,` `[` $strides `]` attr-dict `:` type($base) `,` type($result)

tt.make_tensor_descriptor 接受父张量和块大小的元信息,并返回一个描述符对象,该对象可用于从全局内存中的张量进行加载/存储。

Traits: AlwaysSpeculatableImplTrait, SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

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<tensor<8x8xf16>>

Traits: AlwaysSpeculatableImplTrait, SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

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

操作数:

操作数

描述

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 位。

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

x

整数或整数值的 ranked tensor

y

整数或整数值的 ranked tensor

结果:

结果

描述

result

整数或整数值的 ranked tensor

tt.precise_divf (triton::PreciseDivFOp)

浮点类型的精确除法

语法

operation ::= `tt.precise_divf` $x `,` $y attr-dict `:` type($x)

浮点类型的精确除法。

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

x

浮点数或浮点数值的 ranked tensor

y

浮点数或浮点数值的 ranked tensor

结果:

结果

描述

result

浮点数或浮点数值的 ranked tensor

tt.precise_sqrt (triton::PreciseSqrtOp)

浮点类型的精确平方根

语法

operation ::= `tt.precise_sqrt` $x attr-dict `:` type($x)

浮点类型的精确平方根。

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

x

浮点数或浮点数值的 ranked tensor

结果:

结果

描述

result

浮点数或浮点数值的 ranked tensor

tt.print (triton::PrintOp)

设备端打印,用于调试,类似于 CUDA

语法

operation ::= `tt.print` $prefix attr-dict (`:` $args^ `:` type($args))?

tt.print 接受一个字符串前缀以及任意数量的标量或张量参数,这些参数将被打印。格式会根据参数自动生成。

Traits: SameVariadicOperandSize, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: MemoryEffectOpInterface (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{MemoryEffects::Write on ::mlir::triton::GlobalMemory}

属性:

属性MLIR 类型描述
prefix::mlir::StringAttr字符串属性
hex::mlir::BoolAttr布尔属性
isSigned::mlir::DenseI32ArrayAttri32 稠密数组属性

操作数:

操作数

描述

args

可变数量的浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

tt.ptr_to_int (triton::PtrToIntOp)

将指针转换为 int64

语法

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

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultEncoding, SameOperandsAndResultShape, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

src

ptr 或 ptr 值的 ranked tensor

结果:

结果

描述

result

64 位无符号整数或 64 位无符号整数值的 tensor

tt.reduce (triton::ReduceOp)

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

Traits: AlwaysSpeculatableImplTrait, SameOperandsEncoding, SameOperandsShape, SingleBlock, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

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

操作数:

操作数

描述

srcs

可变数量的浮点数或整数或 ptr 值的 ranked tensor

结果:

结果

描述

result

可变数量的浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

tt.reduce.return (triton::ReduceReturnOp)

规约操作的终止符

语法

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

Traits: AlwaysSpeculatableImplTrait, HasParent<ReduceOp>, ReturnLike, TensorSizeTrait, Terminator, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

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,这是一个提示,表示应出于性能原因保留目标布局。编译器仍然可以自由地更改它以获得更好的性能。

Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

属性MLIR 类型描述
allow_reorder::mlir::UnitAttrunit 属性
efficient_layout::mlir::UnitAttrunit 属性

操作数:

操作数

描述

src

浮点数或整数或 ptr 值的 ranked tensor

结果:

结果

描述

result

浮点数或整数或 ptr 值的 ranked tensor

tt.scan (triton::ScanOp)

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

Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultEncoding, SameOperandsAndResultShape, SingleBlock, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

属性:

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

操作数:

操作数

描述

srcs

可变数量的浮点数或整数或 ptr 值的 ranked tensor

结果:

结果

描述

result

可变数量的浮点数或整数或 ptr 值的 ranked tensor

tt.scan.return (triton::ScanReturnOp)

扫描操作的终止符

语法

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

Traits: AlwaysSpeculatableImplTrait, HasParent<ScanOp>, ReturnLike, TensorSizeTrait, Terminator, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

result

可变数量的任何类型

tt.splat (triton::SplatOp)

扩展

语法

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

Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType, SameOperandsAndResultEncoding, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

src

浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

结果:

结果

描述

result

浮点数或整数或 ptr 值的 ranked tensor

tt.split (triton::SplitOp)

将张量沿其最后一个维度分割成两部分

语法

operation ::= `tt.split` $src attr-dict `:` type($src) `->` type($outLHS)

输入必须是一个其最后一个维度大小为 2 的张量。返回两个张量,src[…, 0] 和 src[…, 1]。

例如,如果输入形状是 4x8x2xf32,则返回两个形状为 4x8xf32 的张量。

Traits: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

操作数:

操作数

描述

src

浮点数或整数或 ptr 值的 ranked tensor

结果:

结果

描述

outLHS

浮点数或整数或 ptr 值的 ranked tensor

outRHS

浮点数或整数或 ptr 值的 ranked tensor

tt.store (triton::StoreOp)

通过指针张量或张量指针存储

语法

operation ::= `tt.store` $ptr `,` $value (`,` $mask^)?
              oilist(`cacheModifier` `=` $cache | `evictionPolicy` `=` $evict)
              attr-dict `:` type($ptr)

Traits: SameLoadStoreOperandsEncoding, SameLoadStoreOperandsShape, TensorSizeTrait, VerifyTensorLayoutsTrait

属性:

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

操作数:

操作数

描述

ptr

ptr 或 ptr 值的 ranked tensor 或 ptr

value

浮点数或浮点数值的 ranked tensor 或整数或整数值的 ranked tensor 或 ptr 或 ptr 值的 ranked tensor 或 ptr

mask

1 位无符号整数或 1 位无符号整数值的 ranked tensor

tt.trans (triton::TransOp)

重新排列张量的维度

语法

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

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

虽然此操作名为“trans”,但它实现了 tl.trans() 和 tl.permute()。(“permute”可能是一个更好的名字,但它之所以被称为“trans”是因为最初它只支持 2D 张量。)

关于编码的实现说明:

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

例如,假设张量 x 具有一种编码,使得 GPU 线程 [i,j,k] 持有张量元素 [i,j,k] 的寄存器。现在我们将 x 进行 transpose 操作,order=[2,1,0],即将其维度顺序反转。在 TritonGPU 中,我们将为 transpose(x) 的输出选择一个布局,使得 GPU 线程 [i,j,k] 持有 transpose(x) 的元素 [k,j,i]。但这与它之前持有的元素是同一个元素!我们所做的只是“重命名”线程 [i,j,k] 所持有的元素。

“真正的”transpose——即在 GPU 线程之间移动数据——发生在操作之前和/或之后出现的 convertLayout 操作中。

我们这样做是为了您可以链接多个数据移动操作(例如 transpose+reshape+concat),而无需在每次操作后都移动到共享内存。

Traits: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor, SameOperandsAndResultElementType, TensorSizeTrait, VerifyTensorLayoutsTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface), TransposeOpInterface

Effects: MemoryEffects::Effect{}

属性:

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

操作数:

操作数

描述

src

浮点数或整数或 ptr 值的 ranked tensor

结果:

结果

描述

result

浮点数或整数或 ptr 值的 ranked tensor