NVIDIA/cutlass
1
0
Fork 0
mirror of https://github.com/NVIDIA/cutlass.git synced 2026-05-12 01:10:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add absf and floor to cute.math (#3156)

Browse Source 
The ops are already exposed by the underlying dialect.
This commit is contained in:
Nandor Licker
2026-04-17 03:54:24 +03:00
committed by GitHub
parent 3f3db08a0a
commit ea46e277d2
2 changed files with 203 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

94
python/CuTeDSL/cutlass/cute/math.py
View File

@@ -49,6 +49,33 @@ def _math_op(func: Callable, fastmath: bool, *args, **kwargs):
@dsl_user_op
def absf(
a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise absolute value of the input tensor.
:param a: Input tensor
:type a: Union[TensorSSA, Numeric]
: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: Tensor containing the absolute value of each element in input tensor
:rtype: Union[TensorSSA, Numeric]
Example:
.. code-block::
x = cute.make_rmem_tensor(layout) # Create tensor
y = x.load() # Load values
z = absf(y) # Compute absolute value
"""
return _math_op(math.absf, fastmath, a, loc=loc, ip=ip)
def acos(
a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
) -> Union[TensorSSA, Numeric]:
@@ -166,6 +193,43 @@ def atan2(
return _math_op(math.atan2, fastmath, a, b, loc=loc, ip=ip)
@dsl_user_op
def copysign(
a: Union[TensorSSA, Numeric],
b: Union[TensorSSA, Numeric],
fastmath: bool = False,
*,
loc=None,
ip=None,
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise copysign of two tensors.
Returns a value with the magnitude of ``a`` and the sign of ``b``.
:param a: Input tensor providing magnitude
:type a: Union[TensorSSA, Numeric]
:param b: Input tensor providing sign
:type b: Union[TensorSSA, Numeric]
: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: Tensor where each element has the magnitude of ``a`` and the sign of ``b``
:rtype: Union[TensorSSA, Numeric]
Example:
.. code-block::
mag = cute.make_rmem_tensor(ptr1, layout).load() # magnitudes
sgn = cute.make_rmem_tensor(ptr2, layout).load() # signs
result = copysign(mag, sgn) # Combine magnitude and sign
"""
return _math_op(math.copysign, fastmath, a, b, loc=loc, ip=ip)
@dsl_user_op
def cos(
a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
@@ -282,6 +346,33 @@ def exp2(
@dsl_user_op
def floor(
a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise floor of the input tensor.
:param a: Input tensor
:type a: Union[TensorSSA, Numeric]
: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: Tensor containing the largest integer less than or equal to each element in input tensor
:rtype: Union[TensorSSA, Numeric]
Example:
.. code-block::
x = cute.make_rmem_tensor(layout) # Create tensor
y = x.load() # Load values
z = floor(y) # Compute floor
"""
return _math_op(math.floor, fastmath, a, loc=loc, ip=ip)
def log(
a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
) -> Union[TensorSSA, Numeric]:
@@ -508,14 +599,17 @@ def tanh(
__all__ = [
"absf",
"acos",
"asin",
"atan",
"atan2",
"copysign",
"cos",
"erf",
"exp",
"exp2",
"floor",
"log",
"log10",
"log2",