exp rounding mode

Signed-off-by: Boyan Li <boyanl@nvidia.com>

Author: blinxt

changelog.d/exp-rounding-mode.md

@@ -0,0 +1,4 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- Optional `rounding_mode` parameter for `ct.exp()` (supports `RoundingMode.FULL` and `RoundingMode.APPROX` for f32)

src/cuda/tile/_ir/ops.py

@@ -1522,8 +1522,7 @@ def generate_bytecode(self, ctx: BytecodeContext) -> bc.Value:
             case "neg", False: return bc.encode_NegIOp(ctx.builder, res_type_id, x,
                                                        bc.IntegerOverflow.NONE)
             case "exp", True: return bc.encode_ExpOp(ctx.builder, res_type_id, x,
-                                                     # TODO: expose rounding mode in ct.exp
-                                                     rounding_mode=bc.RoundingMode.FULL)
+                                                     rounding_mode=rounding_mode)
             case "exp2", True: return bc.encode_Exp2Op(ctx.builder, res_type_id, x,
                                                        flush_to_zero=flush_to_zero)
             case "sin", True: return bc.encode_SinOp(ctx.builder, res_type_id, x)
@@ -1653,7 +1652,6 @@ def pos_impl(x: Var):
 @impl(ct.sinh, fixed_args=["sinh", _UNARY_FLOAT])
 @impl(ct.cos, fixed_args=["cos", _UNARY_FLOAT])
 @impl(ct.cosh, fixed_args=["cosh", _UNARY_FLOAT])
-@impl(ct.exp, fixed_args=["exp", _UNARY_FLOAT])
 @impl(ct.bitwise_not, fixed_args=["bitwise_not", _UNARY_BOOL_INT])
 @impl(ct.floor, fixed_args=["floor", _UNARY_STRICT_FLOAT])
 @impl(ct.ceil, fixed_args=["ceil", _UNARY_STRICT_FLOAT])
@@ -1682,6 +1680,7 @@ def unary_impl_with_rd_and_ftz(fn: str, behavior: _UnaryBehavior,
 
 
 @impl(ct.tanh, fixed_args=["tanh", _UNARY_FLOAT])
+@impl(ct.exp, fixed_args=["exp", _UNARY_FLOAT])
 def unary_impl_with_rd(fn: str, behavior: _UnaryBehavior, x: Var, rounding_mode: Var) -> Var:
     rounding_mode = require_optional_constant_enum(rounding_mode, RoundingMode)
     return unary(fn, behavior, x, rounding_mode=rounding_mode)

src/cuda/tile/_ir/ops_utils.py

@@ -45,6 +45,7 @@ class MathOpDef:
 _RD_TRUEDIV = {**_RD_BASIC, RoundingMode.FULL: None, RoundingMode.APPROX: None}
 _RD_SQRT = {**_RD_BASIC, RoundingMode.APPROX: None}
 _RD_TANH = {RoundingMode.FULL: None, RoundingMode.APPROX: BytecodeVersion.V_13_2}
+_RD_EXP = {RoundingMode.FULL: None, RoundingMode.APPROX: BytecodeVersion.V_13_3}
 
 BINOP_REGISTRY = {
     "add": MathOpDef(lambda x, y: x + y, _RD_BASIC, support_flush_to_zero=True),
@@ -91,7 +92,7 @@ def _invert(x: int | bool, bool_action: Literal['raise'] | Literal['not']):
 UNARYOP_REGISTRY = {
     "abs": MathOpDef(abs),
     "neg": MathOpDef(lambda x: -x),
-    "exp": MathOpDef(math.exp),
+    "exp": MathOpDef(math.exp, _RD_EXP),
     "exp2": MathOpDef(lambda x: 2 ** x, support_flush_to_zero=True),
     "sin": MathOpDef(math.sin),
     "sinh": MathOpDef(math.sinh),
@@ -113,7 +114,7 @@ def _invert(x: int | bool, bool_action: Literal['raise'] | Literal['not']):
 
 
 def get_default_rounding_mode(opname: Optional[str] = None):
-    return RoundingMode.FULL if opname == 'tanh' else RoundingMode.RN
+    return RoundingMode.FULL if opname in ('tanh', 'exp') else RoundingMode.RN
 
 
 rounding_mode_to_bytecode = {

src/cuda/tile/_stub.py

@@ -3255,10 +3255,23 @@ def wrapped(*args, **kwargs):
     return wrapped
 
 
-@_doc_unary_op
 @stub
-def exp(x, /) -> TileOrScalar:
+def exp(x, /, *, rounding_mode: Optional[RoundingMode] = None) -> TileOrScalar:
     """
+    Perform `exp` on a tile.
+
+    Args:
+        x (Tile):
+        rounding_mode (RoundingMode): Supported values:
+
+            - ``RoundingMode.FULL`` (f32 only)
+            - ``RoundingMode.APPROX`` (f32 only)
+
+            (since CTK 13.3)
+
+    Returns:
+        Tile:
+
     Examples:
 
         .. testcode::
@@ -3474,8 +3487,10 @@ def tanh(x, /, *, rounding_mode: Optional[RoundingMode] = None) -> TileOrScalar:
         x (Tile):
         rounding_mode (RoundingMode): Supported values:
 
-            - ``RoundingMode.FULL``
-            - ``RoundingMode.APPROX`` (since CTK 13.2)
+            - ``RoundingMode.FULL`` (f32 only)
+            - ``RoundingMode.APPROX`` (f32 only)
+
+            (since CTK 13.2)
 
     Returns:
         Tile:

test/test_bytecode_version_compat.py

@@ -53,3 +53,22 @@ def kernel(x, y):
 
     # Should not raise version error
     compile_with_version(kernel, (tensor(), tensor()), "13.1")
+
+
+def test_exp_rounding_mode_requires_13_3():
+    def kernel(x, y):
+        tx = ct.load(x, 0, shape=64)
+        ct.store(y, 0, tile=ct.exp(tx, rounding_mode=RoundingMode.APPROX))
+
+    with pytest.raises(TileUnsupportedFeatureError,
+                       match=r"exp rounding_mode=approx requires tileiras 13\.3"):
+        compile_with_version(kernel, (tensor(), tensor()), "13.2")
+
+
+def test_exp_without_rounding_mode_works_on_13_1():
+    def kernel(x, y):
+        tx = ct.load(x, 0, shape=64)
+        ct.store(y, 0, tile=ct.exp(tx))
+
+    # Should not raise version error
+    compile_with_version(kernel, (tensor(), tensor()), "13.1")

test/test_unary_elementwise.py

@@ -336,7 +336,7 @@ def test_scalar_rounding(shape, tile, is_constant, dtype, op, tmp_path):
 @requires_tileiras(BytecodeVersion.V_13_2)
 @pytest.mark.use_mlir
 @pytest.mark.parametrize("dtype", float_dtypes, ids=dtype_id)
-@pytest.mark.parametrize("rounding_mode", [RMd.FULL, RMd.APPROX])
+@pytest.mark.parametrize("rounding_mode", [RMd.FULL, RMd.APPROX, None])
 def test_array_tanh_rounding_mode(shape, tile, dtype, rounding_mode, tmp_path):
     should_raise_dtype = rounding_mode in [RMd.FULL, RMd.APPROX] and dtype != torch.float32
     x = make_tensor(shape, dtype=dtype, device='cuda')
@@ -353,7 +353,7 @@ def test_array_tanh_rounding_mode(shape, tile, dtype, rounding_mode, tmp_path):
             launch_unary(kernel, x, y, tile)
     else:
         bytecode = get_bytecode(kernel, (x, y, tile))
-        if rounding_mode is RMd.FULL:
+        if rounding_mode in (RMd.FULL, None):
             # FULL is the default, not included in mlir text
             check_directive = "// CHECK: %[[RES:.*]] = tanh %[[A:.*]]{{[[:space:]]*}}:"
         else:
@@ -363,3 +363,35 @@ def test_array_tanh_rounding_mode(shape, tile, dtype, rounding_mode, tmp_path):
         filecheck(bytecode, check_directive)
         launch_unary(kernel, x, y, tile)
         assert_close(y, y_ref)
+
+
+@requires_tileiras(BytecodeVersion.V_13_3)
+@pytest.mark.use_mlir
+@pytest.mark.parametrize("dtype", float_dtypes, ids=dtype_id)
+@pytest.mark.parametrize("rounding_mode", [RMd.FULL, RMd.APPROX, None])
+def test_array_exp_rounding_mode(shape, tile, dtype, rounding_mode, tmp_path):
+    should_raise_dtype = rounding_mode in [RMd.FULL, RMd.APPROX] and dtype != torch.float32
+    x = make_tensor(shape, dtype=dtype, device='cuda')
+    y_ref = torch.exp(x)
+    y = torch.zeros_like(y_ref, device="cuda")
+    kernel = array_kernel("exp_rounding_mode",
+                          f"ty = ct.exp(tx, rounding_mode={rounding_mode})",
+                          tmp_path,
+                          globals={"RoundingMode": RMd})
+    if should_raise_dtype:
+        with pytest.raises(TileTypeError,
+                           match=fr"Rounding mode {rounding_mode.value} can only be used for "
+                           "float32 type"):
+            launch_unary(kernel, x, y, tile)
+    else:
+        bytecode = get_bytecode(kernel, (x, y, tile))
+        if rounding_mode in (RMd.FULL, None):
+            # FULL is the default, not included in mlir text
+            check_directive = "// CHECK: %[[RES:.*]] = exp %[[A:.*]]{{[[:space:]]*}}:"
+        else:
+            check_directive = (
+                f"// CHECK: %[[RES:.*]] = exp %[[A:.*]] rounding<{rounding_mode.value}>"
+            )
+        filecheck(bytecode, check_directive)
+        launch_unary(kernel, x, y, tile)
+        assert_close(y, y_ref)