tilelang.language.math_intrinsics¶
Common math intrinsics exposed on the TileLang language surface.
Functions¶
|
IEEE-compliant addition with specified rounding mode |
|
IEEE-compliant subtraction with specified rounding mode |
|
IEEE-compliant multiplication with specified rounding mode |
|
IEEE-compliant fused multiply-add with specified rounding mode |
|
IEEE-compliant reciprocal with specified rounding mode |
|
IEEE-compliant square root with specified rounding mode |
|
IEEE-compliant reciprocal square root (round to nearest only) |
|
IEEE-compliant division with specified rounding mode |
|
Packed FP32x2 add. |
|
Packed FP32x2 multiply. |
|
Packed FP32x2 fused multiply-add (x * y + z). |
Module Contents¶
- tilelang.language.math_intrinsics.ieee_add(x, y, rounding_mode='rn')¶
IEEE-compliant addition with specified rounding mode
- 参数:
x (PrimExpr) -- First operand.
y (PrimExpr) -- Second operand.
rounding_mode (str, optional) -- Rounding mode: 'rn' (round to nearest), 'rz' (round toward zero), 'ru' (round toward positive infinity), 'rd' (round toward negative infinity). Default is 'rn'.
- 返回:
result -- The result.
- 返回类型:
PrimExpr
- tilelang.language.math_intrinsics.ieee_sub(x, y, rounding_mode='rn')¶
IEEE-compliant subtraction with specified rounding mode
- 参数:
x (PrimExpr) -- First operand.
y (PrimExpr) -- Second operand.
rounding_mode (str, optional) -- Rounding mode: 'rn', 'rz', 'ru', 'rd'. Default is 'rn'.
- 返回:
result -- The result.
- 返回类型:
PrimExpr
- tilelang.language.math_intrinsics.ieee_mul(x, y, rounding_mode='rn')¶
IEEE-compliant multiplication with specified rounding mode
- 参数:
x (PrimExpr) -- First operand.
y (PrimExpr) -- Second operand.
rounding_mode (str, optional) -- Rounding mode: 'rn', 'rz', 'ru', 'rd'. Default is 'rn'.
- 返回:
result -- The result.
- 返回类型:
PrimExpr
- tilelang.language.math_intrinsics.ieee_fmaf(x, y, z, rounding_mode='rn')¶
IEEE-compliant fused multiply-add with specified rounding mode
- 参数:
x (PrimExpr) -- First operand.
y (PrimExpr) -- Second operand.
z (PrimExpr) -- Third operand (addend).
rounding_mode (str, optional) -- Rounding mode: 'rn', 'rz', 'ru', 'rd'. Default is 'rn'.
- 返回:
result -- The result of x * y + z.
- 返回类型:
PrimExpr
- tilelang.language.math_intrinsics.ieee_frcp(x, rounding_mode='rn')¶
IEEE-compliant reciprocal with specified rounding mode
- 参数:
x (PrimExpr) -- Input operand.
rounding_mode (str, optional) -- Rounding mode: 'rn', 'rz', 'ru', 'rd'. Default is 'rn'.
- 返回:
result -- The result of 1/x.
- 返回类型:
PrimExpr
- tilelang.language.math_intrinsics.ieee_fsqrt(x, rounding_mode='rn')¶
IEEE-compliant square root with specified rounding mode
- 参数:
x (PrimExpr) -- Input operand.
rounding_mode (str, optional) -- Rounding mode: 'rn', 'rz', 'ru', 'rd'. Default is 'rn'.
- 返回:
result -- The result of sqrt(x).
- 返回类型:
PrimExpr
- tilelang.language.math_intrinsics.ieee_frsqrt(x)¶
IEEE-compliant reciprocal square root (round to nearest only)
- 参数:
x (PrimExpr) -- Input operand.
- 返回:
result -- The result of 1/sqrt(x).
- 返回类型:
PrimExpr
- tilelang.language.math_intrinsics.ieee_fdiv(x, y, rounding_mode='rn')¶
IEEE-compliant division with specified rounding mode
- 参数:
x (PrimExpr) -- Dividend.
y (PrimExpr) -- Divisor.
rounding_mode (str, optional) -- Rounding mode: 'rn', 'rz', 'ru', 'rd'. Default is 'rn'.
- 返回:
result -- The result of x/y.
- 返回类型:
PrimExpr
- tilelang.language.math_intrinsics.fadd2(x, y)¶
Packed FP32x2 add.
Lowers to PTX add.rn.f32x2 on supported NVIDIA architectures/toolchains, and falls back to per-lane scalar operations otherwise.
- 参数:
x (PrimExpr) -- First operand. Must be dtype
float32x2.y (PrimExpr) -- Second operand. Must be dtype
float32x2.
- 返回:
result -- A
float32x2result.- 返回类型:
PrimExpr
- tilelang.language.math_intrinsics.fmul2(x, y)¶
Packed FP32x2 multiply.
Lowers to PTX mul.rn.f32x2 on supported NVIDIA architectures/toolchains, and falls back to per-lane scalar operations otherwise.
- 参数:
x (tvm.tir.PrimExpr)
y (tvm.tir.PrimExpr)
- 返回类型:
tvm.tir.PrimExpr
- tilelang.language.math_intrinsics.fma2(x, y, z)¶
Packed FP32x2 fused multiply-add (x * y + z).
Lowers to PTX fma.rn.f32x2 on supported NVIDIA architectures/toolchains, and falls back to per-lane scalar operations otherwise.
- 参数:
x (tvm.tir.PrimExpr)
y (tvm.tir.PrimExpr)
z (tvm.tir.PrimExpr)
- 返回类型:
tvm.tir.PrimExpr