tilelang.language.math_intrinsics

Common math intrinsics exposed on the TileLang language surface.

Functions

ieee_add(x, y[, rounding_mode])	IEEE-compliant addition with specified rounding mode
ieee_sub(x, y[, rounding_mode])	IEEE-compliant subtraction with specified rounding mode
ieee_mul(x, y[, rounding_mode])	IEEE-compliant multiplication with specified rounding mode
ieee_fmaf(x, y, z[, rounding_mode])	IEEE-compliant fused multiply-add with specified rounding mode
ieee_frcp(x[, rounding_mode])	IEEE-compliant reciprocal with specified rounding mode
ieee_fsqrt(x[, rounding_mode])	IEEE-compliant square root with specified rounding mode
ieee_frsqrt(x)	IEEE-compliant reciprocal square root (round to nearest only)
ieee_fdiv(x, y[, rounding_mode])	IEEE-compliant division with specified rounding mode
add2(x, y)	Packed element-wise add (x + y).
sub2(x, y)	Packed element-wise subtract (x - y).
mul2(x, y)	Packed element-wise multiply (x * y).
fma2(x, y, z)	Packed fused multiply-add (x * y + z).
max2(x, y)	Packed element-wise maximum.
min2(x, y)	Packed element-wise minimum.
abs2(x)	Packed element-wise absolute value.

Module Contents

tilelang.language.math_intrinsics.ieee_add(x, y, rounding_mode='rn')

IEEE-compliant addition with specified rounding mode

Parameters:

• 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’.

Returns:

result – The result.

Return type:

PrimExpr

tilelang.language.math_intrinsics.ieee_sub(x, y, rounding_mode='rn')

IEEE-compliant subtraction with specified rounding mode

Parameters:

• x (PrimExpr) – First operand.

• y (PrimExpr) – Second operand.

• rounding_mode (str, optional) – Rounding mode: ‘rn’, ‘rz’, ‘ru’, ‘rd’. Default is ‘rn’.

Returns:

result – The result.

Return type:

PrimExpr

tilelang.language.math_intrinsics.ieee_mul(x, y, rounding_mode='rn')

IEEE-compliant multiplication with specified rounding mode

Parameters:

• x (PrimExpr) – First operand.

• y (PrimExpr) – Second operand.

• rounding_mode (str, optional) – Rounding mode: ‘rn’, ‘rz’, ‘ru’, ‘rd’. Default is ‘rn’.

Returns:

result – The result.

Return type:

PrimExpr

tilelang.language.math_intrinsics.ieee_fmaf(x, y, z, rounding_mode='rn')

IEEE-compliant fused multiply-add with specified rounding mode

Parameters:

• 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’.

Returns:

result – The result of x * y + z.

Return type:

PrimExpr

tilelang.language.math_intrinsics.ieee_frcp(x, rounding_mode='rn')

IEEE-compliant reciprocal with specified rounding mode

Parameters:

• x (PrimExpr) – Input operand.

• rounding_mode (str, optional) – Rounding mode: ‘rn’, ‘rz’, ‘ru’, ‘rd’. Default is ‘rn’.

Returns:

result – The result of 1/x.

Return type:

PrimExpr

tilelang.language.math_intrinsics.ieee_fsqrt(x, rounding_mode='rn')

IEEE-compliant square root with specified rounding mode

Parameters:

• x (PrimExpr) – Input operand.

• rounding_mode (str, optional) – Rounding mode: ‘rn’, ‘rz’, ‘ru’, ‘rd’. Default is ‘rn’.

Returns:

result – The result of sqrt(x).

Return type:

PrimExpr

tilelang.language.math_intrinsics.ieee_frsqrt(x)

IEEE-compliant reciprocal square root (round to nearest only)

Parameters:

x (PrimExpr) – Input operand.

Returns:

result – The result of 1/sqrt(x).

Return type:

PrimExpr

tilelang.language.math_intrinsics.ieee_fdiv(x, y, rounding_mode='rn')

IEEE-compliant division with specified rounding mode

Parameters:

• x (PrimExpr) – Dividend.

• y (PrimExpr) – Divisor.

• rounding_mode (str, optional) – Rounding mode: ‘rn’, ‘rz’, ‘ru’, ‘rd’. Default is ‘rn’.

Returns:

result – The result of x/y.

Return type:

PrimExpr

tilelang.language.math_intrinsics.add2(x, y)

Packed element-wise add (x + y).

Parameters:

• x (tvm.tir.PrimExpr)

• y (tvm.tir.PrimExpr)

Return type:

tvm.tir.PrimExpr

tilelang.language.math_intrinsics.sub2(x, y)

Packed element-wise subtract (x - y).

Parameters:

• x (tvm.tir.PrimExpr)

• y (tvm.tir.PrimExpr)

Return type:

tvm.tir.PrimExpr

tilelang.language.math_intrinsics.mul2(x, y)

Packed element-wise multiply (x * y).

Parameters:

• x (tvm.tir.PrimExpr)

• y (tvm.tir.PrimExpr)

Return type:

tvm.tir.PrimExpr

tilelang.language.math_intrinsics.fma2(x, y, z)

Packed fused multiply-add (x * y + z).

Parameters:

• x (tvm.tir.PrimExpr)

• y (tvm.tir.PrimExpr)

• z (tvm.tir.PrimExpr)

Return type:

tvm.tir.PrimExpr

tilelang.language.math_intrinsics.max2(x, y)

Packed element-wise maximum.

Parameters:

• x (tvm.tir.PrimExpr)

• y (tvm.tir.PrimExpr)

Return type:

tvm.tir.PrimExpr

tilelang.language.math_intrinsics.min2(x, y)

Packed element-wise minimum.

Parameters:

• x (tvm.tir.PrimExpr)

• y (tvm.tir.PrimExpr)

Return type:

tvm.tir.PrimExpr

tilelang.language.math_intrinsics.abs2(x)

Packed element-wise absolute value.

Parameters:

x (tvm.tir.PrimExpr)

Return type:

tvm.tir.PrimExpr