mirror of
https://github.com/NVIDIA/cutlass.git
synced 2026-07-12 02:05:34 +00:00
2214 lines
68 KiB
Python
2214 lines
68 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
|
|
|
"""
|
|
Unified NVIDIA GPU math API.
|
|
|
|
This module exposes the full set of NVIDIA GPU math capabilities as a
|
|
single, stable Python surface. Each call is dispatched to the cleanest
|
|
layer that currently provides the requested semantics:
|
|
|
|
1. MLIR ``math`` / ``arith`` dialect ops — target-agnostic,
|
|
vectorizable, pattern-match friendly. The default and
|
|
``fastmath`` paths use this layer.
|
|
2. ``nvvm.*`` dialect ops — first-class MLIR ops for a subset of
|
|
NVIDIA-specific capabilities (``nvvm.fabs``, ``nvvm.fmin`` /
|
|
``fmax``, ``nvvm.rcp_approx_ftz_f``). Used where available.
|
|
3. ``llvm.call_intrinsic "llvm.nvvm.*"`` — direct LLVM NVVM
|
|
intrinsics for PTX modifier combinations (approx, ftz,
|
|
specific rounding modes, ``div.full``) that don't yet have
|
|
first-class MLIR ops.
|
|
|
|
Why this module exists
|
|
----------------------
|
|
Upstream MLIR's ``math`` / ``arith`` dialects are chartered
|
|
target-agnostic and intentionally don't carry NVIDIA-specific attributes
|
|
(approx, ftz, rounding mode, ``div.full``). The NVVM dialect
|
|
promotes a selected subset to first-class ops but deliberately
|
|
doesn't enumerate every PTX combination. That leaves a layer — the
|
|
full NVIDIA GPU math capability surface, presented to end users —
|
|
that neither upstream owner wants to fill. This module fills it.
|
|
|
|
As upstream ``math`` / ``arith`` gain rounding-mode attributes
|
|
(in progress) and the NVVM dialect promotes more PTX variants,
|
|
intrinsic-call paths here shrink. The user-facing API stays stable.
|
|
|
|
Supported operand types
|
|
-----------------------
|
|
- Numeric scalars (Float32, Float64, Int32, etc.)
|
|
- ArithValue / Vector (MLIR value types; subclasses like
|
|
``TensorSSA`` go through the Vector dispatch path automatically,
|
|
and :meth:`Vector._wrap_like` preserves subclass-specific metadata
|
|
on per-element results)
|
|
- Python ``float`` / ``int`` / ``bool`` literals, promoted to
|
|
``Float32`` / ``Int32`` via :func:`_coerce_operand`.
|
|
|
|
Consumers that need to layer their own framework-specific types (e.g.,
|
|
memory-tracked array wrappers) can add a thin adapter that unwraps to a
|
|
supported type, delegates here, and rewraps the result.
|
|
"""
|
|
|
|
from enum import Enum
|
|
from typing import Any, Optional, Union
|
|
|
|
from .._mlir import ir
|
|
from .._mlir.extras import types as T
|
|
from .._mlir.dialects import arith, math as math_dialect, llvm, nvgpu, nvvm
|
|
from .arith import ArithValue, is_float_type, element_type
|
|
from .op import dsl_user_op
|
|
from .vector import Vector
|
|
|
|
from ..base_dsl.typing import Numeric, Float32, Int32
|
|
|
|
# =============================================================================
|
|
# Type alias
|
|
# =============================================================================
|
|
|
|
# Accepted operand types. ``Vector`` covers CuTeDSL's ``TensorSSA`` via
|
|
# inheritance (e.g., ``TensorSSA``). Additional framework-specific wrappers
|
|
# should be layered as thin adapters in their own package and delegate here.
|
|
# Python ``float``/``int``/``bool`` literals are promoted to ``Float32`` /
|
|
# ``Int32`` automatically so callers can write ``mul(x, 2.0)`` or
|
|
# ``add(x, 1)`` without manual wrapping.
|
|
MathOperand = Union[Numeric, Vector, float, int, bool]
|
|
|
|
|
|
# =============================================================================
|
|
# Enums
|
|
# =============================================================================
|
|
|
|
|
|
class RoundingMode(str, Enum):
|
|
"""IEEE 754 rounding modes for floating-point operations."""
|
|
|
|
NEAREST_EVEN = "rn"
|
|
ZERO = "rz"
|
|
UP = "rp"
|
|
DOWN = "rm"
|
|
|
|
|
|
# =============================================================================
|
|
# Internal helpers
|
|
# =============================================================================
|
|
|
|
|
|
def _coerce_operand(x: MathOperand) -> MathOperand:
|
|
"""Promote Python numeric literals to Numeric scalars.
|
|
|
|
Python ``bool``/``int`` -> :class:`Int32`, ``float`` -> :class:`Float32`.
|
|
``bool`` is checked before ``int`` because ``bool`` subclasses ``int``.
|
|
Non-literal operands (Numeric, ArithValue, Vector) pass through
|
|
unchanged.
|
|
"""
|
|
if isinstance(x, bool):
|
|
return Int32(int(x))
|
|
if isinstance(x, float):
|
|
return Float32(x)
|
|
if isinstance(x, int):
|
|
return Int32(x)
|
|
return x
|
|
|
|
|
|
def _numeric_type_name(x: MathOperand) -> str:
|
|
"""Return the scalar element type name of a coerced operand.
|
|
|
|
Returns the class name of the scalar type (e.g. ``"Float32"``,
|
|
``"Int32"``). For ``Vector`` the dtype's class name is returned
|
|
(``type(x.dtype).__name__``).
|
|
"""
|
|
if isinstance(x, Vector):
|
|
return type(x.dtype).__name__
|
|
return type(x).__name__
|
|
|
|
|
|
def _validate_same_numeric_type(op_name: str, *operands: MathOperand) -> None:
|
|
"""Raise :class:`TypeError` if operands don't share one numeric type.
|
|
|
|
Python literals are promoted via :func:`_coerce_operand` before the
|
|
check, so the comparison is purely on Numeric / Vector element types
|
|
after promotion.
|
|
|
|
Mixing typed values (e.g. ``fma(Float64_val, Float32_val, Int32_val)``)
|
|
would otherwise produce cryptic MLIR type-mismatch errors deep in the
|
|
compiler. This helper surfaces the problem immediately.
|
|
|
|
We promote Python literals (for ergonomics) but do *not* auto-promote
|
|
heterogeneous typed values — if a user mixes Float32 and Float16, they
|
|
need to convert explicitly.
|
|
"""
|
|
coerced = [_coerce_operand(op) for op in operands]
|
|
# Raw ArithValue operands (e.g. reduction results threaded through an
|
|
# threaded through a task loop) don't carry a Numeric class, so comparing
|
|
# ``type(op).__name__`` would produce false positives. Vector subclasses
|
|
# ArithValue here, so we must exclude it explicitly.
|
|
if any(isinstance(op, ArithValue) and not isinstance(op, Vector) for op in coerced):
|
|
return
|
|
type_names = [_numeric_type_name(op) for op in coerced]
|
|
if len(set(type_names)) > 1:
|
|
raise TypeError(
|
|
f"All operands to {op_name} must have the same numeric type, "
|
|
f"got {', '.join(type_names)}"
|
|
)
|
|
|
|
|
|
def _check_vector_consistency(op_name: str, *operands: MathOperand) -> None:
|
|
"""Validate that vector-like and scalar operands aren't mixed.
|
|
|
|
MLIR math/arith ops require all operands to share a shape. We forbid
|
|
mixing a Vector with a Numeric scalar early so the user sees a clear
|
|
TypeError instead of a cryptic MLIR error later.
|
|
"""
|
|
coerced = [_coerce_operand(op) for op in operands]
|
|
is_vector_like = [isinstance(op, Vector) for op in coerced]
|
|
if any(is_vector_like) and not all(is_vector_like):
|
|
vector_type = next(
|
|
type(op).__name__ for op, is_v in zip(coerced, is_vector_like) if is_v
|
|
)
|
|
scalar_type = next(
|
|
type(op).__name__ for op, is_v in zip(coerced, is_vector_like) if not is_v
|
|
)
|
|
raise TypeError(
|
|
f"{op_name}: Expected all inputs to be of type {vector_type}, "
|
|
f"but got {scalar_type} for another operand"
|
|
)
|
|
|
|
|
|
def _get_ir_value(
|
|
x: MathOperand, *, ip: Optional[ir.InsertionPoint] = None
|
|
) -> ir.Value:
|
|
"""Extract MLIR IR value from any supported operand type."""
|
|
x = _coerce_operand(x)
|
|
if isinstance(x, ir.Value):
|
|
# Vector (and its subclasses) are already ir.Values — use directly.
|
|
return x
|
|
if not isinstance(x, Numeric):
|
|
raise TypeError(f"Expected a Vector or Numeric, but got {type(x).__name__}")
|
|
# Numeric (Float32, etc.) — has ir_value() method
|
|
return x.ir_value(ip=ip)
|
|
|
|
|
|
def _get_element_type(x: MathOperand) -> ir.Type:
|
|
"""Get the scalar element type of an operand."""
|
|
v = _get_ir_value(x)
|
|
return element_type(v.type)
|
|
|
|
|
|
def _is_float(x: MathOperand) -> bool:
|
|
"""Check if operand is a floating-point type."""
|
|
return is_float_type(_get_element_type(x))
|
|
|
|
|
|
def _is_f32(x: MathOperand) -> bool:
|
|
return _get_element_type(x) == T.f32()
|
|
|
|
|
|
def _is_f64(x: MathOperand) -> bool:
|
|
return _get_element_type(x) == T.f64()
|
|
|
|
|
|
def _is_unsigned_int(x: MathOperand) -> bool:
|
|
"""Check if operand is an unsigned integer type."""
|
|
return not _is_float(x) and getattr(x, "signed", True) is False
|
|
|
|
|
|
def _wrap_result(original: MathOperand, result_ir: ir.Value) -> MathOperand:
|
|
"""Wrap an MLIR result back into the original operand's type."""
|
|
if isinstance(original, Vector):
|
|
# Polymorphic wrap: Vector subclasses (e.g., cute.TensorSSA) use
|
|
# ``_wrap_like`` to preserve their own metadata (CuTe nested shape,
|
|
# layout) when a math op returns a new SSA value. Fall back to a
|
|
# plain Vector construction when the base-class hook isn't present
|
|
# (compat with older Vector that doesn't yet define ``_wrap_like``).
|
|
if hasattr(original, "_wrap_like"):
|
|
return original._wrap_like(result_ir)
|
|
return Vector(result_ir, dtype=original._dtype)
|
|
if isinstance(original, ArithValue):
|
|
return result_ir # ArithValue IS an ir.Value
|
|
# Coerce Python literals so we return a Numeric subclass instance instead
|
|
# of attempting to construct e.g. ``float(ir.Value)``.
|
|
original = _coerce_operand(original)
|
|
# Numeric — wrap in the same Numeric subclass
|
|
return type(original)(result_ir)
|
|
|
|
|
|
def _get_fastmath_flag(fastmath: bool) -> arith.FastMathFlags:
|
|
return arith.FastMathFlags.fast if fastmath else arith.FastMathFlags.none
|
|
|
|
|
|
def _get_type_suffix(x: MathOperand) -> str:
|
|
"""Get NVVM type suffix: 'f' for f32, 'd' for f64."""
|
|
if _is_f32(x):
|
|
return "f"
|
|
elif _is_f64(x):
|
|
return "d"
|
|
else:
|
|
raise TypeError("NVVM math intrinsics require f32 or f64 scalar operands")
|
|
|
|
|
|
def _get_llvm_type(x: MathOperand) -> ir.Type:
|
|
"""Get LLVM result type for NVVM intrinsic."""
|
|
if _is_f32(x):
|
|
return T.f32()
|
|
elif _is_f64(x):
|
|
return T.f64()
|
|
else:
|
|
raise TypeError("NVVM math intrinsics require f32 or f64 scalar operands")
|
|
|
|
|
|
def _const_one(
|
|
ty: ir.Type,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> ir.Value:
|
|
"""Constant 1.0 of a scalar-float or vector-of-float type."""
|
|
if isinstance(ty, ir.VectorType):
|
|
attr = ir.DenseElementsAttr.get_splat(
|
|
ty, ir.FloatAttr.get(ty.element_type, 1.0)
|
|
)
|
|
return arith.constant(ty, attr, loc=loc, ip=ip)
|
|
return arith.constant(ty, 1.0, loc=loc, ip=ip)
|
|
|
|
|
|
# =============================================================================
|
|
# Validation
|
|
# =============================================================================
|
|
|
|
|
|
def _validate_fastmath_exclusive(
|
|
op_name: str,
|
|
fastmath: bool,
|
|
*,
|
|
approx: bool = False,
|
|
ftz: bool = False,
|
|
rounding: Optional[RoundingMode] = None,
|
|
full: bool = False,
|
|
) -> None:
|
|
"""Raise if fastmath (math/arith dialect) is combined with NVVM attributes.
|
|
|
|
Three valid modes:
|
|
- No flags: strict MLIR math/arith op
|
|
- fastmath=True: MLIR op with fastmath<fast> attribute
|
|
- approx/ftz/rounding: NVVM intrinsic
|
|
|
|
fastmath and approx/ftz/rounding are mutually exclusive because
|
|
they target different MLIR dialects.
|
|
"""
|
|
nvvm_attrs = []
|
|
if approx:
|
|
nvvm_attrs.append("approx")
|
|
if ftz:
|
|
nvvm_attrs.append("ftz")
|
|
if rounding is not None:
|
|
nvvm_attrs.append("rounding")
|
|
if full:
|
|
nvvm_attrs.append("full")
|
|
if fastmath and nvvm_attrs:
|
|
attr_str = "/".join(nvvm_attrs)
|
|
raise ValueError(
|
|
f"{op_name}: fastmath and {attr_str} are mutually exclusive. "
|
|
f"fastmath uses math/arith dialect (compiler decides lowering); "
|
|
f"{attr_str} emit NVVM intrinsics (exact instruction control)."
|
|
)
|
|
|
|
|
|
def _validate_ftz_requires_approx(op_name: str, approx: bool, ftz: bool) -> None:
|
|
"""Raise if ftz is set without approx on transcendentals."""
|
|
if ftz and not approx:
|
|
raise ValueError(
|
|
f"{op_name}: ftz requires approx=True for transcendental "
|
|
f"functions. There is no non-approximate ftz variant."
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# NVVM intrinsic helpers
|
|
#
|
|
# Callers land here only when neither arith/math nor the NVVM dialect has a
|
|
# matching op for the requested (op, attributes, type) combination. Dialect-op
|
|
# coverage audit at time of writing:
|
|
# - Scalar float add/sub/mul/fma with rounding+ftz: NVVM has only
|
|
# *Packed*F32x2 variants; no scalar ops (nvvm.MulOp is integer).
|
|
# - Scalar transcendental .approx (sin/cos/ex2/lg2/tanh/rsqrt/sqrt): no
|
|
# NVVM dialect ops.
|
|
# - Scalar div.approx / div.full / div.<rnd>: no NVVM dialect ops.
|
|
# - Scalar sqrt.<rnd>: no NVVM dialect ops.
|
|
# - Scalar rcp.<rnd> / rcp.approx (no ftz): no NVVM dialect ops.
|
|
# Ops that *do* have dialect coverage are used directly: nvvm.fabs (abs.ftz),
|
|
# nvvm.fmin/fmax (min/max.ftz), nvvm.rcp_approx_ftz_f (rcp scalar f32
|
|
# approx+ftz), nvgpu.rcp (vector rcp approx[.ftz]).
|
|
# =============================================================================
|
|
|
|
# NVVM intrinsics that use overloaded types (no .f/.d suffix).
|
|
_NVVM_OVERLOADED_OPS = frozenset({"ex2", "tanh"})
|
|
|
|
|
|
def _call_nvvm_unary(
|
|
x: MathOperand,
|
|
op_name: str,
|
|
rounding: Optional[RoundingMode] = None,
|
|
ftz: bool = False,
|
|
approx: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Call a unary NVVM intrinsic."""
|
|
if isinstance(x, Vector):
|
|
# NVVM intrinsics are scalar-only. For vectors (and Arrays that wrap
|
|
# Vectors), we can't emit them without per-element unrolling (which
|
|
# defeats vectorization). Direct users to the math/arith path (via
|
|
# fastmath=True) so the MLIR pipeline can vectorize properly.
|
|
attr = "approx" if approx else ("rounding" if rounding is not None else "ftz")
|
|
raise TypeError(
|
|
f"{op_name}: {attr} is scalar-only (NVVM intrinsics don't "
|
|
f"vectorize). For vector inputs, use fastmath=True — the "
|
|
f"math/arith dialect lowering will be vectorized by the compiler."
|
|
)
|
|
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
result_type = _get_llvm_type(x)
|
|
type_suffix = _get_type_suffix(x)
|
|
|
|
intrinsic_name = f"llvm.nvvm.{op_name}"
|
|
if approx:
|
|
intrinsic_name += ".approx"
|
|
elif rounding is not None:
|
|
intrinsic_name += f".{rounding.value}"
|
|
elif ftz:
|
|
# ftz without explicit rounding: default to rn
|
|
intrinsic_name += ".rn"
|
|
if ftz:
|
|
intrinsic_name += ".ftz"
|
|
if op_name not in _NVVM_OVERLOADED_OPS:
|
|
intrinsic_name += f".{type_suffix}"
|
|
|
|
result = llvm.call_intrinsic(
|
|
result_type,
|
|
intrinsic_name,
|
|
[x_ir],
|
|
[],
|
|
[],
|
|
loc=loc,
|
|
ip=ip,
|
|
)
|
|
return _wrap_result(x, result)
|
|
|
|
|
|
def _call_nvvm_binary(
|
|
x: MathOperand,
|
|
y: MathOperand,
|
|
op_name: str,
|
|
rounding: Optional[RoundingMode] = None,
|
|
ftz: bool = False,
|
|
approx: bool = False,
|
|
full: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Call a binary NVVM intrinsic."""
|
|
if isinstance(x, Vector) or isinstance(y, Vector):
|
|
# See _call_nvvm_unary: scalar-only; direct users to fastmath for
|
|
# vectorizable math/arith lowering.
|
|
attr = (
|
|
"approx"
|
|
if approx
|
|
else ("full" if full else ("rounding" if rounding is not None else "ftz"))
|
|
)
|
|
raise TypeError(
|
|
f"{op_name}: {attr} is scalar-only (NVVM intrinsics don't "
|
|
f"vectorize). For vector inputs, use fastmath=True — the "
|
|
f"math/arith dialect lowering will be vectorized by the compiler."
|
|
)
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
y_ir = _get_ir_value(y, ip=ip)
|
|
result_type = _get_llvm_type(x)
|
|
type_suffix = _get_type_suffix(x)
|
|
|
|
intrinsic_name = f"llvm.nvvm.{op_name}"
|
|
if approx:
|
|
intrinsic_name += ".approx"
|
|
elif full:
|
|
intrinsic_name += ".full"
|
|
elif rounding is not None:
|
|
intrinsic_name += f".{rounding.value}"
|
|
elif ftz:
|
|
# ftz without explicit rounding: default to rn
|
|
intrinsic_name += ".rn"
|
|
if ftz:
|
|
intrinsic_name += ".ftz"
|
|
if not full:
|
|
intrinsic_name += f".{type_suffix}"
|
|
|
|
result = llvm.call_intrinsic(
|
|
result_type,
|
|
intrinsic_name,
|
|
[x_ir, y_ir],
|
|
[],
|
|
[],
|
|
loc=loc,
|
|
ip=ip,
|
|
)
|
|
return _wrap_result(x, result)
|
|
|
|
|
|
def _call_nvvm_ternary(
|
|
a: MathOperand,
|
|
b: MathOperand,
|
|
c: MathOperand,
|
|
op_name: str,
|
|
rounding: Optional[RoundingMode] = None,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Call a ternary NVVM intrinsic (e.g., fma)."""
|
|
if any(isinstance(v, Vector) for v in (a, b, c)):
|
|
# See _call_nvvm_unary: scalar-only; direct users to fastmath for
|
|
# vectorizable math/arith lowering.
|
|
attr = "rounding" if rounding is not None else "ftz"
|
|
raise TypeError(
|
|
f"{op_name}: {attr} is scalar-only (NVVM intrinsics don't "
|
|
f"vectorize). For vector inputs, use fastmath=True — the "
|
|
f"math/arith dialect lowering will be vectorized by the compiler."
|
|
)
|
|
a_ir = _get_ir_value(a, ip=ip)
|
|
b_ir = _get_ir_value(b, ip=ip)
|
|
c_ir = _get_ir_value(c, ip=ip)
|
|
result_type = _get_llvm_type(a)
|
|
type_suffix = _get_type_suffix(a)
|
|
|
|
intrinsic_name = f"llvm.nvvm.{op_name}"
|
|
if rounding is not None:
|
|
intrinsic_name += f".{rounding.value}"
|
|
else:
|
|
intrinsic_name += ".rn"
|
|
if ftz:
|
|
intrinsic_name += ".ftz"
|
|
intrinsic_name += f".{type_suffix}"
|
|
|
|
result = llvm.call_intrinsic(
|
|
result_type,
|
|
intrinsic_name,
|
|
[a_ir, b_ir, c_ir],
|
|
[],
|
|
[],
|
|
loc=loc,
|
|
ip=ip,
|
|
)
|
|
return _wrap_result(a, result)
|
|
|
|
|
|
def _needs_nvvm_intrinsic(rounding: Optional[RoundingMode], ftz: bool) -> bool:
|
|
"""Check if we need NVVM intrinsics (any explicit rounding or ftz)."""
|
|
if ftz:
|
|
return True
|
|
if rounding is not None:
|
|
return True
|
|
return False
|
|
|
|
|
|
# =============================================================================
|
|
# Unary math helper
|
|
# =============================================================================
|
|
|
|
|
|
def _unary_math_op(
|
|
x: MathOperand,
|
|
float_op: Any,
|
|
int_op: Any = None,
|
|
fastmath: bool = False,
|
|
op_name: str = "math_op",
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Apply a unary MLIR math operation."""
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
|
|
if _is_float(x):
|
|
result = float_op(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
else:
|
|
if int_op is None:
|
|
raise TypeError(f"{op_name} is only supported for floating-point types")
|
|
result = int_op(x_ir, loc=loc, ip=ip)
|
|
|
|
return _wrap_result(x, result)
|
|
|
|
|
|
# =============================================================================
|
|
# Transcendentals with NVVM approx intrinsics
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def sin(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
approx: bool = False,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise sine of the input operand.
|
|
|
|
:param a: Input operand (in radians)
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the sine of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = sin(a) # Compute sine
|
|
"""
|
|
_validate_fastmath_exclusive("sin", fastmath, approx=approx, ftz=ftz)
|
|
_validate_ftz_requires_approx("sin", approx, ftz)
|
|
if approx:
|
|
return _call_nvvm_unary(x, "sin", approx=True, ftz=ftz, loc=loc, ip=ip)
|
|
return _unary_math_op(x, math_dialect.sin, None, fastmath, "sin", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def cos(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
approx: bool = False,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise cosine of the input operand.
|
|
|
|
:param a: Input operand (in radians)
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the cosine of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = cos(a) # Compute cosine
|
|
"""
|
|
_validate_fastmath_exclusive("cos", fastmath, approx=approx, ftz=ftz)
|
|
_validate_ftz_requires_approx("cos", approx, ftz)
|
|
if approx:
|
|
return _call_nvvm_unary(x, "cos", approx=True, ftz=ftz, loc=loc, ip=ip)
|
|
return _unary_math_op(x, math_dialect.cos, None, fastmath, "cos", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def exp(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise exponential of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the exponential of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = exp(a) # Compute exponential
|
|
"""
|
|
return _unary_math_op(x, math_dialect.exp, None, fastmath, "exp", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def exp2(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
approx: bool = False,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise base-2 exponential of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing 2 raised to the power of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = exp2(a) # Compute 2^x
|
|
"""
|
|
_validate_fastmath_exclusive("exp2", fastmath, approx=approx, ftz=ftz)
|
|
_validate_ftz_requires_approx("exp2", approx, ftz)
|
|
if approx:
|
|
return _call_nvvm_unary(x, "ex2", approx=True, ftz=ftz, loc=loc, ip=ip)
|
|
return _unary_math_op(x, math_dialect.exp2, None, fastmath, "exp2", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def log(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise natural logarithm of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the natural logarithm of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = log(a) # Compute natural logarithm
|
|
"""
|
|
return _unary_math_op(x, math_dialect.log, None, fastmath, "log", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def log2(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
approx: bool = False,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise base-2 logarithm of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the base-2 logarithm of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = log2(a) # Compute log base 2
|
|
"""
|
|
_validate_fastmath_exclusive("log2", fastmath, approx=approx, ftz=ftz)
|
|
_validate_ftz_requires_approx("log2", approx, ftz)
|
|
if approx:
|
|
return _call_nvvm_unary(x, "lg2", approx=True, ftz=ftz, loc=loc, ip=ip)
|
|
return _unary_math_op(x, math_dialect.log2, None, fastmath, "log2", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def log10(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise base-10 logarithm of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the base-10 logarithm of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = log10(a) # Compute log base 10
|
|
"""
|
|
return _unary_math_op(
|
|
x, math_dialect.log10, None, fastmath, "log10", loc=loc, ip=ip
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def tanh(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
approx: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise hyperbolic tangent of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the hyperbolic tangent of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = tanh(a) # Compute hyperbolic tangent
|
|
"""
|
|
_validate_fastmath_exclusive("tanh", fastmath, approx=approx)
|
|
if approx:
|
|
return _call_nvvm_unary(x, "tanh", approx=True, loc=loc, ip=ip)
|
|
return _unary_math_op(x, math_dialect.tanh, None, fastmath, "tanh", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def rsqrt(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
approx: bool = False,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise reciprocal square root of the input operand.
|
|
|
|
Computes 1/√x element-wise.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the reciprocal square root of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = rsqrt(a) # Compute 1/√x
|
|
"""
|
|
_validate_fastmath_exclusive("rsqrt", fastmath, approx=approx, ftz=ftz)
|
|
_validate_ftz_requires_approx("rsqrt", approx, ftz)
|
|
if approx:
|
|
return _call_nvvm_unary(x, "rsqrt", approx=True, ftz=ftz, loc=loc, ip=ip)
|
|
return _unary_math_op(
|
|
x, math_dialect.rsqrt, None, fastmath, "rsqrt", loc=loc, ip=ip
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def sqrt(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
approx: bool = False,
|
|
ftz: bool = False,
|
|
rounding: Optional[RoundingMode] = None,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise square root of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the square root of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = sqrt(a) # Compute square root
|
|
"""
|
|
_validate_fastmath_exclusive(
|
|
"sqrt", fastmath, approx=approx, ftz=ftz, rounding=rounding
|
|
)
|
|
if approx:
|
|
return _call_nvvm_unary(x, "sqrt", approx=True, ftz=ftz, loc=loc, ip=ip)
|
|
if _needs_nvvm_intrinsic(rounding, ftz):
|
|
return _call_nvvm_unary(x, "sqrt", rounding=rounding, ftz=ftz, loc=loc, ip=ip)
|
|
return _unary_math_op(x, math_dialect.sqrt, None, fastmath, "sqrt", loc=loc, ip=ip)
|
|
|
|
|
|
# =============================================================================
|
|
# MLIR-only transcendentals (no NVVM approx intrinsic)
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def erf(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise error function of the input operand.
|
|
|
|
The error function is defined as:
|
|
erf(x) = 2/√π ∫[0 to x] exp(-t²) dt
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the error function value for each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = erf(a) # Compute error function
|
|
"""
|
|
return _unary_math_op(x, math_dialect.erf, None, fastmath, "erf", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def tan(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise tangent of the input operand.
|
|
|
|
:param a: Input operand (in radians)
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the tangent of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = tan(a) # Compute tangent
|
|
"""
|
|
return _unary_math_op(x, math_dialect.tan, None, fastmath, "tan", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def acos(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise arc cosine of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the arc cosine of each element in input operand
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = acos(a) # Compute arc cosine
|
|
"""
|
|
return _unary_math_op(x, math_dialect.acos, None, fastmath, "acos", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def asin(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise arc sine of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the arc sine of each element in input operand
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = asin(a) # Compute arc sine
|
|
"""
|
|
return _unary_math_op(x, math_dialect.asin, None, fastmath, "asin", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def atan(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise arc tangent of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the arc tangent of each element in input operand
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = atan(a) # Compute arc tangent
|
|
"""
|
|
return _unary_math_op(x, math_dialect.atan, None, fastmath, "atan", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def atan2(
|
|
y: MathOperand,
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise arc tangent of two tensors.
|
|
|
|
Computes atan2(a, b) element-wise. The function atan2(a, b) is the angle in radians
|
|
between the positive x-axis and the point given by the coordinates (b, a).
|
|
|
|
:param a: First input operand (y-coordinates)
|
|
:type a: MathOperand
|
|
:param b: Second input operand (x-coordinates)
|
|
:type b: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the arc tangent of a/b element-wise
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
theta = atan2(y, x) # Compute angles
|
|
"""
|
|
_check_vector_consistency("atan2", y, x)
|
|
_validate_same_numeric_type("atan2", y, x)
|
|
y_ir = _get_ir_value(y, ip=ip)
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
result = math_dialect.atan2(y_ir, x_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
return _wrap_result(y, result)
|
|
|
|
|
|
@dsl_user_op
|
|
def sinh(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise hyperbolic sine of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the hyperbolic sine of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = sinh(a) # Compute hyperbolic sine
|
|
"""
|
|
return _unary_math_op(x, math_dialect.sinh, None, fastmath, "sinh", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def cosh(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise hyperbolic cosine of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the hyperbolic cosine of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = cosh(a) # Compute hyperbolic cosine
|
|
"""
|
|
return _unary_math_op(x, math_dialect.cosh, None, fastmath, "cosh", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def pow(
|
|
base: MathOperand,
|
|
exponent: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise power of the input tensors.
|
|
|
|
:param a: Input operand (base)
|
|
:type a: MathOperand
|
|
:param b: Input operand (exponent)
|
|
:type b: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing a raised to the power b for each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
z = pow(a, b) # Compute a^b
|
|
"""
|
|
_check_vector_consistency("pow", base, exponent)
|
|
_validate_same_numeric_type("pow", base, exponent)
|
|
base_ir = _get_ir_value(base, ip=ip)
|
|
exp_ir = _get_ir_value(exponent, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
result = math_dialect.powf(base_ir, exp_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
return _wrap_result(base, result)
|
|
|
|
|
|
@dsl_user_op
|
|
def cbrt(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise cube root of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the cube root of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = cbrt(a) # Compute cube root
|
|
"""
|
|
return _unary_math_op(x, math_dialect.cbrt, None, fastmath, "cbrt", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def expm1(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise exp(x) - 1 of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing exp(x) - 1 of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = expm1(a) # Compute exp(x) - 1
|
|
"""
|
|
return _unary_math_op(
|
|
x, math_dialect.expm1, None, fastmath, "expm1", loc=loc, ip=ip
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def log1p(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise log(1 + x) of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing log(1 + x) of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = log1p(a) # Compute log(1 + x)
|
|
"""
|
|
return _unary_math_op(
|
|
x, math_dialect.log1p, None, fastmath, "log1p", loc=loc, ip=ip
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Absolute value, sign, rounding
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def abs(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Absolute value (float or integer).
|
|
|
|
:param ftz: Flush denormals to zero (float only, uses nvvm.fabs)
|
|
"""
|
|
_validate_fastmath_exclusive("abs", fastmath, ftz=ftz)
|
|
if ftz and _is_float(x):
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
result = nvvm.fabs(x_ir, ftz=True, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
return _unary_math_op(
|
|
x,
|
|
math_dialect.absf,
|
|
math_dialect.absi,
|
|
fastmath,
|
|
"abs",
|
|
loc=loc,
|
|
ip=ip,
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def copysign(
|
|
mag: MathOperand,
|
|
sign: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise copysign, combining magnitude of a with sign of b.
|
|
|
|
:param a: Input operand (magnitude source)
|
|
:type a: MathOperand
|
|
:param b: Input operand (sign source)
|
|
:type b: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the copysign of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
z = copysign(a, b) # Copy sign of b to magnitude of a
|
|
"""
|
|
_check_vector_consistency("copysign", mag, sign)
|
|
_validate_same_numeric_type("copysign", mag, sign)
|
|
mag_ir = _get_ir_value(mag, ip=ip)
|
|
sign_ir = _get_ir_value(sign, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
result = math_dialect.copysign(
|
|
mag_ir, sign_ir, fastmath=fastmath_flag, loc=loc, ip=ip
|
|
)
|
|
return _wrap_result(mag, result)
|
|
|
|
|
|
@dsl_user_op
|
|
def ceil(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise ceiling of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the ceiling of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = ceil(a) # Compute ceiling
|
|
"""
|
|
return _unary_math_op(x, math_dialect.ceil, None, fastmath, "ceil", loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def floor(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise floor of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the largest integer less than or equal to each element in input operand
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = floor(a) # Compute floor
|
|
"""
|
|
return _unary_math_op(
|
|
x, math_dialect.floor, None, fastmath, "floor", loc=loc, ip=ip
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def round(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise round to nearest integer (ties away from zero) of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the rounded value of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = round(a) # Round to nearest integer
|
|
"""
|
|
return _unary_math_op(
|
|
x, math_dialect.round, None, fastmath, "round", loc=loc, ip=ip
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def roundeven(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise round to nearest integer (ties to even) of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the rounded value of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = roundeven(a) # Round to nearest integer (ties to even)
|
|
"""
|
|
return _unary_math_op(
|
|
x, math_dialect.roundeven, None, fastmath, "roundeven", loc=loc, ip=ip
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def trunc(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise truncation toward zero of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the truncated value of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = trunc(a) # Truncate toward zero
|
|
"""
|
|
return _unary_math_op(
|
|
x, math_dialect.trunc, None, fastmath, "trunc", loc=loc, ip=ip
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Min / Max
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def min(
|
|
x: MathOperand,
|
|
y: MathOperand,
|
|
propagate_nan: bool = False,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Element-wise minimum.
|
|
|
|
:param propagate_nan: If True, NaN propagates (IEEE 754 minimum).
|
|
If False, NaN is ignored (IEEE 754 minimumNumber).
|
|
:param ftz: Flush denormals to zero (float only, uses nvvm.fmin).
|
|
"""
|
|
_check_vector_consistency("min", x, y)
|
|
_validate_same_numeric_type("min", x, y)
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
y_ir = _get_ir_value(y, ip=ip)
|
|
if _is_float(x):
|
|
if ftz:
|
|
result = nvvm.fmin(x_ir, y_ir, ftz=True, nan=propagate_nan, loc=loc, ip=ip)
|
|
elif propagate_nan:
|
|
result = arith.minimumf(x_ir, y_ir, loc=loc, ip=ip)
|
|
else:
|
|
result = arith.minnumf(x_ir, y_ir, loc=loc, ip=ip)
|
|
elif _is_unsigned_int(x):
|
|
result = arith.minui(x_ir, y_ir, loc=loc, ip=ip)
|
|
else:
|
|
result = arith.minsi(x_ir, y_ir, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
|
|
|
|
@dsl_user_op
|
|
def max(
|
|
x: MathOperand,
|
|
y: MathOperand,
|
|
propagate_nan: bool = False,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Element-wise maximum.
|
|
|
|
:param propagate_nan: If True, NaN propagates (IEEE 754 maximum).
|
|
If False, NaN is ignored (IEEE 754 maximumNumber).
|
|
:param ftz: Flush denormals to zero (float only, uses nvvm.fmax).
|
|
"""
|
|
_check_vector_consistency("max", x, y)
|
|
_validate_same_numeric_type("max", x, y)
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
y_ir = _get_ir_value(y, ip=ip)
|
|
if _is_float(x):
|
|
if ftz:
|
|
result = nvvm.fmax(x_ir, y_ir, ftz=True, nan=propagate_nan, loc=loc, ip=ip)
|
|
elif propagate_nan:
|
|
result = arith.maximumf(x_ir, y_ir, loc=loc, ip=ip)
|
|
else:
|
|
result = arith.maxnumf(x_ir, y_ir, loc=loc, ip=ip)
|
|
elif _is_unsigned_int(x):
|
|
result = arith.maxui(x_ir, y_ir, loc=loc, ip=ip)
|
|
else:
|
|
result = arith.maxsi(x_ir, y_ir, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
|
|
|
|
# =============================================================================
|
|
# FMA
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def fma(
|
|
a: MathOperand,
|
|
b: MathOperand,
|
|
c: MathOperand,
|
|
fastmath: bool = False,
|
|
rounding: Optional[RoundingMode] = None,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Fused multiply-add: a * b + c."""
|
|
_check_vector_consistency("fma", a, b, c)
|
|
_validate_same_numeric_type("fma", a, b, c)
|
|
_validate_fastmath_exclusive("fma", fastmath, ftz=ftz, rounding=rounding)
|
|
if _needs_nvvm_intrinsic(rounding, ftz):
|
|
return _call_nvvm_ternary(a, b, c, "fma", rounding, ftz, loc=loc, ip=ip)
|
|
a_ir = _get_ir_value(a, ip=ip)
|
|
b_ir = _get_ir_value(b, ip=ip)
|
|
c_ir = _get_ir_value(c, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
result = math_dialect.fma(a_ir, b_ir, c_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
return _wrap_result(a, result)
|
|
|
|
|
|
# =============================================================================
|
|
# Arithmetic with fastmath / rounding / ftz
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def add(
|
|
x: MathOperand,
|
|
y: MathOperand,
|
|
fastmath: bool = False,
|
|
rounding: Optional[RoundingMode] = None,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Floating-point addition."""
|
|
_check_vector_consistency("add", x, y)
|
|
_validate_same_numeric_type("add", x, y)
|
|
_validate_fastmath_exclusive("add", fastmath, ftz=ftz, rounding=rounding)
|
|
if _needs_nvvm_intrinsic(rounding, ftz):
|
|
return _call_nvvm_binary(x, y, "add", rounding, ftz, loc=loc, ip=ip)
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
y_ir = _get_ir_value(y, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
result = arith.addf(x_ir, y_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
|
|
|
|
@dsl_user_op
|
|
def sub(
|
|
x: MathOperand,
|
|
y: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Floating-point subtraction.
|
|
|
|
Note: NVVM sub uses a different intrinsic convention (rounding as an
|
|
integer arg), so rounding/ftz are not supported here. Callers that
|
|
need explicit rounding control on subtraction should emit the
|
|
PTX-level ``sub.<rn|rz|rm|rp>[.ftz].f32`` intrinsic directly.
|
|
"""
|
|
_check_vector_consistency("sub", x, y)
|
|
_validate_same_numeric_type("sub", x, y)
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
y_ir = _get_ir_value(y, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
result = arith.subf(x_ir, y_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
|
|
|
|
@dsl_user_op
|
|
def mul(
|
|
x: MathOperand,
|
|
y: MathOperand,
|
|
fastmath: bool = False,
|
|
rounding: Optional[RoundingMode] = None,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Floating-point multiplication."""
|
|
_check_vector_consistency("mul", x, y)
|
|
_validate_same_numeric_type("mul", x, y)
|
|
_validate_fastmath_exclusive("mul", fastmath, ftz=ftz, rounding=rounding)
|
|
if _needs_nvvm_intrinsic(rounding, ftz):
|
|
return _call_nvvm_binary(x, y, "mul", rounding, ftz, loc=loc, ip=ip)
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
y_ir = _get_ir_value(y, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
result = arith.mulf(x_ir, y_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
|
|
|
|
@dsl_user_op
|
|
def div(
|
|
x: MathOperand,
|
|
y: MathOperand,
|
|
fastmath: bool = False,
|
|
approx: bool = False,
|
|
ftz: bool = False,
|
|
full: bool = False,
|
|
rounding: Optional[RoundingMode] = None,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Floating-point division."""
|
|
_check_vector_consistency("div", x, y)
|
|
_validate_same_numeric_type("div", x, y)
|
|
_validate_fastmath_exclusive(
|
|
"div", fastmath, approx=approx, ftz=ftz, rounding=rounding, full=full
|
|
)
|
|
if fastmath:
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
y_ir = _get_ir_value(y, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(True)
|
|
result = arith.divf(x_ir, y_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
if approx:
|
|
return _call_nvvm_binary(x, y, "div", approx=True, ftz=ftz, loc=loc, ip=ip)
|
|
if full:
|
|
return _call_nvvm_binary(x, y, "div", full=True, ftz=ftz, loc=loc, ip=ip)
|
|
if _needs_nvvm_intrinsic(rounding, ftz):
|
|
return _call_nvvm_binary(x, y, "div", rounding, ftz, loc=loc, ip=ip)
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
y_ir = _get_ir_value(y, ip=ip)
|
|
result = arith.divf(x_ir, y_ir, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
|
|
|
|
@dsl_user_op
|
|
def rcp(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
approx: bool = False,
|
|
rounding: Optional[RoundingMode] = None,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Reciprocal (1/x)."""
|
|
_validate_fastmath_exclusive(
|
|
"rcp", fastmath, approx=approx, ftz=ftz, rounding=rounding
|
|
)
|
|
if fastmath:
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
one = _const_one(x_ir.type, loc=loc, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(True)
|
|
result = arith.divf(one, x_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
if approx:
|
|
if isinstance(x, Vector):
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
result = nvgpu.rcp(
|
|
x_ir,
|
|
rounding=nvgpu.RcpRoundingMode.APPROX,
|
|
ftz=ftz or None,
|
|
loc=loc,
|
|
ip=ip,
|
|
)
|
|
return _wrap_result(x, result)
|
|
# Scalar approx+ftz on f32 has a first-class NVVM dialect op.
|
|
# Prefer it over llvm.call_intrinsic for better pattern-match/
|
|
# canonicalization visibility.
|
|
if ftz and _is_f32(x):
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
result = nvvm.rcp_approx_ftz_f(x_ir, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
return _call_nvvm_unary(x, "rcp", approx=True, ftz=ftz, loc=loc, ip=ip)
|
|
if _needs_nvvm_intrinsic(rounding, ftz):
|
|
return _call_nvvm_unary(x, "rcp", rounding=rounding, ftz=ftz, loc=loc, ip=ip)
|
|
# Default semantics: IEEE round-to-nearest reciprocal.
|
|
# Scalar f32/f64: emit nvvm.rcp.rn.{f32,f64} (one PTX rcp instruction).
|
|
# Vector or scalar element type without an NVVM scalar-rcp intrinsic
|
|
# (e.g., f16): lower to arith.divf(1, x) — same IEEE semantics,
|
|
# vectorizable by the compiler.
|
|
if isinstance(x, Vector) or not (_is_f32(x) or _is_f64(x)):
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
one = _const_one(x_ir.type, loc=loc, ip=ip)
|
|
result = arith.divf(one, x_ir, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
return _call_nvvm_unary(
|
|
x, "rcp", rounding=RoundingMode.NEAREST_EVEN, loc=loc, ip=ip
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Inverse hyperbolic
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def acosh(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise inverse hyperbolic cosine of the input operand.
|
|
|
|
:param a: Input operand (must be >= 1.0)
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the inverse hyperbolic cosine of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = acosh(a) # Compute inverse hyperbolic cosine
|
|
"""
|
|
return _unary_math_op(
|
|
x, math_dialect.acosh, None, fastmath, "acosh", loc=loc, ip=ip
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def asinh(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise inverse hyperbolic sine of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the inverse hyperbolic sine of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = asinh(a) # Compute inverse hyperbolic sine
|
|
"""
|
|
return _unary_math_op(
|
|
x, math_dialect.asinh, None, fastmath, "asinh", loc=loc, ip=ip
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def atanh(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise inverse hyperbolic tangent of the input operand.
|
|
|
|
:param a: Input operand (must be in (-1, 1))
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the inverse hyperbolic tangent of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = atanh(a) # Compute inverse hyperbolic tangent
|
|
"""
|
|
return _unary_math_op(
|
|
x, math_dialect.atanh, None, fastmath, "atanh", loc=loc, ip=ip
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Complementary error function
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def erfc(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Compute element-wise complementary error function of the input operand.
|
|
|
|
:param a: Input operand
|
|
:type a: MathOperand
|
|
:param fastmath: Enable fast math optimizations, defaults to False
|
|
:type fastmath: bool, optional
|
|
:param loc: Source location information, defaults to None
|
|
:type loc: Optional[Location]
|
|
:param ip: Insertion point for IR generation, defaults to None
|
|
:type ip: Optional[InsertionPoint]
|
|
:return: Result containing the complementary error function of each element
|
|
:rtype: MathOperand
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
y = erfc(a) # Compute complementary error function
|
|
"""
|
|
return _unary_math_op(x, math_dialect.erfc, None, fastmath, "erfc", loc=loc, ip=ip)
|
|
|
|
|
|
# =============================================================================
|
|
# Combined sin+cos
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def sincos(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> tuple:
|
|
"""Combined sine and cosine: returns (sin(x), cos(x)).
|
|
|
|
More efficient than separate sin() and cos() calls.
|
|
"""
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
op = math_dialect.SincosOp(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
return (_wrap_result(x, op.sin), _wrap_result(x, op.cos))
|
|
|
|
|
|
# =============================================================================
|
|
# Power variants
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def fpowi(
|
|
base: MathOperand,
|
|
exponent: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Float raised to integer power: base^exponent (exponent is integer)."""
|
|
base_ir = _get_ir_value(base, ip=ip)
|
|
exp_ir = _get_ir_value(exponent, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
result = math_dialect.fpowi(base_ir, exp_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
return _wrap_result(base, result)
|
|
|
|
|
|
@dsl_user_op
|
|
def ipowi(
|
|
base: MathOperand,
|
|
exponent: MathOperand,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Integer raised to integer power: base^exponent."""
|
|
base_ir = _get_ir_value(base, ip=ip)
|
|
exp_ir = _get_ir_value(exponent, ip=ip)
|
|
result = math_dialect.ipowi(base_ir, exp_ir, loc=loc, ip=ip)
|
|
return _wrap_result(base, result)
|
|
|
|
|
|
# =============================================================================
|
|
# Clamp (native math dialect op)
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def clamp(
|
|
x: MathOperand,
|
|
lo: MathOperand,
|
|
hi: MathOperand,
|
|
fastmath: bool = False,
|
|
propagate_nan: bool = False,
|
|
ftz: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Clamp value to range [lo, hi].
|
|
|
|
Lowered unconditionally as ``max(lo, min(x, hi))``. The math dialect's
|
|
``clampf`` op currently has no LLVM translation interface registered, so
|
|
a direct lowering fails at JIT time on scalar f16 inputs (and any other
|
|
type). Composition via min/max picks up the right per-type lowering
|
|
(``arith.minnumf`` / ``arith.maximumf`` / ``nvvm.fmin`` / ``nvvm.fmax``).
|
|
"""
|
|
_validate_fastmath_exclusive("clamp", fastmath, ftz=ftz)
|
|
inner = min(x, hi, propagate_nan=propagate_nan, ftz=ftz, loc=loc, ip=ip)
|
|
return max(lo, inner, propagate_nan=propagate_nan, ftz=ftz, loc=loc, ip=ip)
|
|
|
|
|
|
# =============================================================================
|
|
# FP classification predicates
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def isnan(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Test if value is NaN. Returns i1."""
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
return math_dialect.isnan(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def isinf(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Test if value is infinite. Returns i1."""
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
return math_dialect.isinf(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def isfinite(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Test if value is finite. Returns i1."""
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
return math_dialect.isfinite(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def isnormal(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Test if value is normal (not zero, subnormal, inf, or NaN). Returns i1."""
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
return math_dialect.isnormal(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
|
|
|
|
# =============================================================================
|
|
# Integer math
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def absi(
|
|
x: MathOperand,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Integer absolute value."""
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
result = math_dialect.absi(x_ir, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
|
|
|
|
# =============================================================================
|
|
# Negation and remainder (arith dialect)
|
|
# =============================================================================
|
|
|
|
|
|
@dsl_user_op
|
|
def neg(
|
|
x: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Floating-point negation."""
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
result = arith.negf(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
|
|
|
|
@dsl_user_op
|
|
def rem(
|
|
x: MathOperand,
|
|
y: MathOperand,
|
|
fastmath: bool = False,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> MathOperand:
|
|
"""Floating-point remainder."""
|
|
x_ir = _get_ir_value(x, ip=ip)
|
|
y_ir = _get_ir_value(y, ip=ip)
|
|
fastmath_flag = _get_fastmath_flag(fastmath)
|
|
result = arith.remf(x_ir, y_ir, fastmath=fastmath_flag, loc=loc, ip=ip)
|
|
return _wrap_result(x, result)
|
|
|
|
|
|
# =============================================================================
|
|
# __all__
|
|
# =============================================================================
|
|
|
|
__all__ = [
|
|
"RoundingMode",
|
|
"MathOperand",
|
|
# Transcendentals with NVVM approx (hardware instructions)
|
|
"sin",
|
|
"cos",
|
|
"exp2",
|
|
"log2",
|
|
"tanh",
|
|
"rsqrt",
|
|
"sqrt",
|
|
# Transcendentals without hardware approx
|
|
"exp",
|
|
"log",
|
|
"log10",
|
|
# Trigonometric
|
|
"tan",
|
|
"acos",
|
|
"asin",
|
|
"atan",
|
|
"atan2",
|
|
# Hyperbolic
|
|
"sinh",
|
|
"cosh",
|
|
"acosh",
|
|
"asinh",
|
|
"atanh",
|
|
# Error functions
|
|
"erf",
|
|
"erfc",
|
|
# Power / root / exponential
|
|
"pow",
|
|
"fpowi",
|
|
"ipowi",
|
|
"cbrt",
|
|
"expm1",
|
|
"log1p",
|
|
# Combined
|
|
"sincos",
|
|
# Absolute value, sign, rounding
|
|
"abs",
|
|
"absi",
|
|
"copysign",
|
|
"neg",
|
|
"ceil",
|
|
"floor",
|
|
"round",
|
|
"roundeven",
|
|
"trunc",
|
|
# Clamp
|
|
"clamp",
|
|
# Min / Max
|
|
"min",
|
|
"max",
|
|
# FP predicates
|
|
"isnan",
|
|
"isinf",
|
|
"isfinite",
|
|
"isnormal",
|
|
# FMA
|
|
"fma",
|
|
# Arithmetic
|
|
"add",
|
|
"sub",
|
|
"mul",
|
|
"div",
|
|
"rem",
|
|
"rcp",
|
|
]
|