【Triton 教程】math-ops

31 阅读3分钟

Triton 是一种用于并行编程的语言和编译器。它旨在提供一个基于 Python 的编程环境,以高效编写自定义 DNN 计算内核,并能够在现代 GPU 硬件上以最大吞吐量运行。

*在线运行 Triton 学习教程 → go.hyper.ai/wS9x1

triton_language.abs

triton.language.abs(x)

计算 x 的逐元素绝对值。

参数**:**

  • x (Block) - 输入值。

triton_language.cdiv

triton.language.cdiv(x,div)

计算 x 除以 div 的向上取整结果。

参数:

  • x (Block) - 输入数字。
  • div (Block) - 除数。

这个函数也可作为 tensor 的成员函数调用,例如 x.div(...) 而不是 div(x, ...)

triton_language.ceil

triton.language.ceil(x)

计算 x 的逐元素向上取整结果。

参数**:**

  • x (Block) - 输入值。

triton_language.clamp

triton.language.clamp(x, min, max, propagate_nan: ~triton.language.core.constexpr = <PROPAGATE_NAN.NONE: 0>)

将输入张量 x 的值限制在 [min, max] 范围内。

参数:

  • x (Block) - 输入值。
  • min (Block) – 限制操作的下界值。
  • max (Block) – 限制操作的上界值。
  • propagate_nan (tl.PropagateNan) – 是否将 NaN 值传播出去。此设置仅对张量x有效。如果minmax中任一值为NaN,则最终结果将无法确定。

另请参阅
tl.PropagateNan

triton_language.cos

triton.language.cos(x)

计算 x 的逐元素余弦值。

参数**:**

  • x (Block) - 输入值。

triton_language.div_rn

triton.language.div_rn(x,y)

计算 x 和 y 的逐元素精确除法(根据 IEEE 标准四舍五入到最近的值)。

参数**:**

  • x (Block) - 输入值。
  • y (Block) - 输入值。

triton_language.erf

triton.language.erf(x)

计算 x 的逐元素误差函数。

参数**:**

  • x (Block) - 输入值。

triton_language.exp

triton.language.exp(x)

计算 x 的逐元素指数函数。

参数**:**

  • x (Block) - 输入值。

triton_language.exp2

triton.language.exp2(x)

计算 x 的逐元素指数函数(以 2 为底)。

参数**:**

  • x (Block) - 输入值。

triton_language.fdiv

triton.language.fdiv(x,y)

计算 x 和 y 的逐元素快速除法。

参数**:**

  • x (Block) - 输入值。
  • y (Block) - 输入值。

triton_language.floor

triton.language.floor(x)

计算 x 的逐元素向下取整。

参数**:**

  • x (Block) - 输入值。

triton_language.fma

triton.language.fma(x,y,z)

计算 x 、 y 和 z 的逐元素融合乘加运算。

参数**:**

  • x (Block) - 输入值。
  • y (Block) - 输入值。
  • z (Block) - 输入值。

triton_language.log

triton.language.log(x)

计算 x 的逐元素自然对数。

参数**:**

  • x (Block) - 输入值。

triton_language.log2

triton.language.log2(x)

计算 x 的逐元素对数(以 2 为底)。

参数**:**

  • x (Block) - 输入值。

triton_language.maximum

triton.language.maximum(x, y, propagate_nan: ~triton.language.core.constexpr = <PROPAGATE_NAN.NONE: 0>

计算 x 和 y 的逐元素最大值。

参数**:**

  • x (Block) - 第 1 个输入的张量。
  • y (Block) - 第 2 个输入的张量。
  • propagate_nan (tl.PropagateNan) - 是否传播 NaN 值。

另请参阅
tl.PropagateNan

triton_language.minimum

triton.language.minimum(x,y)

计算 x 和 y 的逐元素最小值。

参数**:**

  • x (Block) - 第 1 个输入的张量。
  • y (Block) - 第 2 个输入的张量。
  • propagate_nan (tl.PropagateNan) - 是否传播 NaN 值。

另请参阅
tl.PropagateNan

triton_language.rsqrt

triton.language.rsqrt(x)

计算 x 的逐元素倒数平方根。

参数**:**

  • x (Block) - 输入值。

triton_language.sigmoid

triton.language.sigmoid(x)

计算 x 的逐元素 sigmoid 函数值。

参数**:**

  • x (Block) - 输入值。

这个函数也可作为 tensor 的成员函数调用,使用 x.sigmoid(...) 而不是 sigmoid(x, ...)

triton_language.sin

triton.language.sin(x)

计算 x 的逐元素正弦值。

参数**:**

  • x (Block) - 输入值。

triton_language.softmax

triton.language.softmax(x, ieee_rounding=False)

计算 x 的逐元素 softmax 函数值。

参数**:**

  • x (Block) - 输入值。

这个函数也可作为 tensor 的成员函数调用,使用 x.softmax(...) 而不是 softmax(x, ...)

triton_language.sqrt

triton.language.sqrt(x)

计算 x 的逐元素快速平方根。

参数**:**

  • x (Block) - 输入值。

triton_language.sqrt_rn

triton.language.sqrt_rn(x)

计算 x 的逐元素精确平方根(根据 IEEE 标准四舍五入到最近值)。

参数:

  • x (Block) - 输入值。

triton_language.umulhi

triton.language.umulhi(x,y)

计算 x 和 y 的 2N 位乘积中逐元素的最高有效 N 位。

参数**:**

  • x (Block) - 输入值。
  • y (Block) - 输入值。