triton.language

昇腾拓展API列表

api	简要说明
extract_slice	从输入张量中按照操作指定的偏移量、大小和步幅参数提取一个张量。
insert_slice	将一个张量（子张量）插入到另一个张量的指定位置，按照操作指定的偏移量、大小和步幅参数插入到另一个张量中。
sync_block	显式的核心间同步指令，用于协调 Cube-Vector 架构中不同核心间的执行顺序和数据一致性。
compile_hint	一个编译器提示（hint）机制，允许用户为张量附加元数据信息，这些信息会被传递到编译器后端，用于指导优化和代码生成。
multibuffer	为张量设置多缓冲，允许编译器对同一张量创建多个副本。
parallel	parallel 是一个专门用于多核心并行执行的迭代器,提供显式的多核心并行语义。
get_element	根据给定的索引，从输入张量中读取单个元素。
index_select 高性能接口	在非尾轴维度上并行 gather 多个索引，并以 tile 为单位将数据零拷贝地从全局内存（GM）直接搬运到统一缓冲区（UB）的正确位置。该操作等效于 torch.index_select 的高性能实现，适用于嵌入层查找、稀疏索引访问等场景。

原子操作

api	简要说明
atomic_add	在由 pointer 指定的内存位置执行原子加法
atomic_and	在由 pointer 指定的内存位置执行原子逻辑和操作
atomic_cas	在由 pointer 指定的内存位置执行 1 个原子比较并交换操作
atomic_max	在由 pointer 指定的内存位置执行 1 个原子最大值操作
atomic_min	在由 pointer 指定的内存位置执行 1 个原子最小值操作
atomic_or	在由 pointer 指定的内存位置执行 1 个原子逻辑或操作
atomic_xchg	在由 pointer 指定的内存位置执行 1 个原子交换操作
atomic_xor	在由 pointer 指定的内存位置执行原子逻辑异或操作

比较操作

api	简要说明
eq	用于比较两个张量的元素，与==等价
le	用于比较两个张量的元素，与<=等价。
ge	用于比较两个张量的元素，与>=等价。
lt	用于比较两个张量的元素，与<等价。
gt	用于比较两个张量的元素，与>等价。
ne	用于比较两个张量的元素，与!=等价。

编译器提示操作

api	简要说明
debug_barrier	插入 1 个屏障以同步 1 个块中的所有线程
max_constancy	告知编译器 input 中的第 1 个值是常量
max_contiguous	告知编译器 input 中的第 1 个值是连续
multiple_of	告知编译器 input 中的所有值都是 value 的倍数
assume	用于向编译器提供条件假设信息，允许编译器基于已知为真的条件进行优化。
compile_hint	一个编译器提示（hint）机制，允许用户为张量附加元数据信息，这些信息会被传递到编译器后端，用于指导优化和代码生成。
multibuffer	为张量设置多缓冲，允许编译器对同一张量创建多个副本。
parallel	parallel 是一个专门用于多核心并行执行的迭代器,提供显式的多核心并行语义。
sync_block指令	显式的核心间同步指令，用于协调 Cube-Vector 架构中不同核心间的执行顺序和数据一致性。

创建操作

api	简要说明
arange	返回半开区间 [start, end) 内的连续值
cat	连接给定的块
full	返回一个张量，该张量填充了指定 shape 和 dtype 的标量值
zeros	返回一个张量，该张量用指定 shape 和 dtype 填充了标量值 0
zeros_like	返回一个 shape 和 dtype 与给定张量相同的全零张量
cast	将张量转换为指定的 dtype

调试操作

api	简要说明
static_print	在编译时打印数值
static_assert	在编译时断言条件
device_print	在运行时从设备打印数值
device_assert	在运行时从设备上断言条件

索引与元素操作

api	简要说明
flip	沿着维度 dim 翻转张量 x
where	根据 condition 返回来自 x 或 y 的元素组成的张量
swizzle2d	将行主序排列为 size_i * size_j 的矩阵的索引，转换为每组 size_g 行的列主序矩阵的索引
get_element	根据给定的索引，从输入张量中读取单个元素。
index_select 高性能接口	在非尾轴维度上并行 gather 多个索引，并以 tile 为单位将数据零拷贝地从全局内存（GM）直接搬运到统一缓冲区（UB）的正确位置。该操作等效于 torch.index_select 的高性能实现，适用于嵌入层查找、稀疏索引访问等场景。
gather	对srctensor沿axis维度按照index执行gather操作

内联汇编

api	简要说明
inline_asm_elementwise	在张量上执行内联汇编

迭代器

api	简要说明
range	永远向上计数的迭代器
static_range	永远向上计数的迭代器

线性代数操作

api	简要说明
dot	返回两个块的矩阵乘积
dot_scaled	计算以缩放格式表示两个矩阵块的矩阵乘积

逻辑操作

api	简要说明
and	逻辑与操作
or	逻辑或操作
not	逻辑非操作
logical_and	用于对两个张量进行逐元素逻辑与运算
logical_or	用于对两个张量进行逐元素逻辑或运算
not	将tensor的值按位取反。
invert	将tensor每个值按比特位进行翻转。
lshift	根据给定的位移数将tensor张量进行左移位。
rshift	根据给定的位移数将tensor张量进行右移位。
xor	计算两个元素的异或值。

数学操作

api	简要说明
add	四则运算加法 ‘+’
sub	四则运算减法 ‘-’
mul	四则运算乘法 ‘*’
div	四则运算除法 ‘/’
abs	计算 x 的逐元素绝对值
neg	将tensor的值取负。
cdiv	计算 x 除以 div 的向上取整除法
ceil	计算 x 的逐元素向上取整值
clamp	将输入张量 x 的值限制在 [min, max] 范围内
cos	计算 x 的逐元素余弦值
div_rn	计算 x 和 y 的逐元素精确除法（根据 IEEE 标准四舍五入到最近的值）
erf	计算 x 的逐元素误差函数
exp	计算 x 的逐元素指数
exp2	计算 x 的逐元素指数（以 2 为底）
fdiv	计算 x 和 y 的逐元素快速除法
floor	计算 x 的逐元素向下取整
fma	计算 x、y 和 z 的逐元素融合乘加运算
log	计算 x 的逐元素自然对数
log2	计算 x 的逐元素对数（以 2 为底）
mod	取模运算
maximum	计算 x 和 y 的逐元素最大值
minimum	计算 x 和 y 的逐元素最小值
rsqrt	计算 x 的逐元素的平方根倒数
sigmoid	计算 x 的逐元素 sigmoid 函数值
sin	Computes the element-wise sine of x. 计算 x 的逐元素正弦值
softmax	计算 x 的逐元素 softmax 值
sqrt	计算 x 的逐元素快速平方根
sqrt_rn	计算 x 的逐元素精确平方根（根据 IEEE 标准四舍五入到最近的值）
umulhi	计算 x 和 y 的 2N 位乘积的逐元素最高有效 N 位

内存/指针操作

api	简要说明
load	返回一个张量，其值从由指针定义的内存位置加载
store	将数据张量存储到由指针定义的内存位置
make_block_ptr	返回指向父张量中某个块的指针
advance	推进一个块指针
load_tensor_descriptor	从张量描述符加载数据块
make_tensor_descriptor	创建张量描述符对象

编程模型

api	简要说明
tensor	表示一个值或指针的 N 维数组
program_id	沿指定轴返回当前程序实例的 id
num_programs	沿指定轴返回当前程序实例的数量

随机数生成

api	简要说明
randint4x	给定 1 个seed 标量和 1 个offset 块，返回 4 个 int32 类型的随机块
randint	给定 1 个 seed 标量和 1 个 offset 块，返回 1 个 int32 类型的随机块
rand	给定 1 个 seed 标量和 1 个 offset 块，返回 1 个在 U(0,1) 中的 float32 类型的随机块
randn	给定 1 个 seed 标量和 1 个 offset 块，返回 1 个在 N(0,1) 中的 float32 类型的随机块

归约操作

api	简要说明
argmax	返回沿指定 axis 轴上 input 张量中所有元素的最大索引
argmin	返回沿指定 axis 轴上 input 张量中所有元素的最小索引
max	返回沿指定 axis 轴上 input 张量中所有元素的最大值
min	返回沿指定 axis 轴上 input 张量中所有元素的最小值
reduce	将 combine_fn 应用于沿指定 axis 的 input 张量中的所有元素
sum	返回 input 张量中，沿指定 axis 的所有元素的总和
xor_sum	返回 input 张量中，沿指定 axis 的所有元素的异或和

扫描/排序操作

api	简要说明
associative_scan	沿指定 axis 将 combine_fn 应用于 input 张量的每个元素和携带的值，并更新携带的值
cumprod	返回沿指定 axis 的 input 张量中所有元素的累积乘积
cumsum	返回沿指定 axis 的 input 张量中所有元素的累积和
histogram	基于 input 张量计算 1 个具有 num_bins 个 bin 的直方图，每个 bin 宽度为 1，起始值为0
sort	沿着指定维度对张量进行排序

Shape操作

api	简要说明
broadcast	尝试将两个给定的块广播到一个共同兼容的 shape
broadcast_to	尝试将给定的张量广播到新的 shape
expand_dims	通过插入新的长度为 1 的维度来扩展张量的形状
interleave	沿着最后一个维度交错两个张量的值
join	在一个新的次要维度中连接给定的张量
permute	排列张量的维度
ravel	返回 x 的连续扁平视图
reshape	返回一个具有与输入相同元素数但具有提供的形状的张量
split	将张量沿其最后一个维度分成两部分，该维度大小必须为 2
trans	排列张量的维度。
view	返回具有与输入相同元素但形状不同的张量
extract_slice	从输入张量中按照操作指定的偏移量、大小和步幅参数提取一个张量。
insert_slice	将一个张量（子张量）插入到另一个张量的指定位置，即将一个张量按照操作指定的偏移量、大小和步幅参数插入到另一个张量中。