(tvm) branch main updated: Support integer types in exp TIR expression operator (#18390)

[cbalint13]

This is an automated email from the ASF dual-hosted git repository.

cbalint13 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 30fcca2016 Support integer types in exp TIR expression operator 
(#18390)
30fcca2016 is described below

commit 30fcca201612b75815998ec84bcd4f41fa6fe564
Author: Qingchao Shen <qingchaos...@outlook.com>
AuthorDate: Mon Oct 27 17:25:20 2025 +0800

    Support integer types in exp TIR expression operator (#18390)
    
    This PR addresses the issue where tvm.tir.exp does not support integer 
types (e.g., int32, int64), causing an InternalError during LLVM code 
generation with the message.
    
    The issue arises because the llvm.exp intrinsic expects floating-point 
inputs, but no type conversion is performed for integer inputs.
    
    This change aligns the behavior of tir.exp with libraries like PyTorch and 
NumPy, which implicitly convert integer inputs to floating-point types for 
their exponential functions.
    
    Fix #18381
---
 python/tvm/tir/op.py                     |  2 ++
 tests/python/tir-base/test_tir_intrin.py | 11 +++++++++++
 2 files changed, 13 insertions(+)

diff --git a/python/tvm/tir/op.py b/python/tvm/tir/op.py
index 9a912bbb6b..7f3badcfeb 100644
--- a/python/tvm/tir/op.py
+++ b/python/tvm/tir/op.py
@@ -2116,6 +2116,8 @@ def exp(x):
         The result.
     """
     x = tir.convert(x)
+    if "int" in x.dtype:
+        x = tir.Cast("float32", x)
     return call_intrin(x.dtype, "tir.exp", x)
 
 
diff --git a/tests/python/tir-base/test_tir_intrin.py 
b/tests/python/tir-base/test_tir_intrin.py
index 1492816429..afeefba2a3 100644
--- a/tests/python/tir-base/test_tir_intrin.py
+++ b/tests/python/tir-base/test_tir_intrin.py
@@ -66,6 +66,7 @@ def test_round_intrinsics_on_int():
 
 def test_unary_intrin():
     test_funcs = [
+        (tvm.tir.exp, lambda x: np.exp(x)),
         (tvm.tir.exp10, lambda x: np.power(10, x)),
         (tvm.tir.log2, lambda x: np.log2(x)),
         (tvm.tir.log10, lambda x: np.log10(x)),
@@ -118,6 +119,16 @@ def test_unary_intrin():
             func(a2, b2)
             # all outputs should be NaN
             assert np.all(np.isnan(b2.numpy()))
+        if name == "exp":
+            n = 8
+            out_np = np.random.randint(-20, 20, size=n).astype(A.dtype)
+            a2 = tvm.runtime.tensor(out_np, dev)
+            b2 = tvm.runtime.tensor(np.empty_like(out_np), dev)
+            func(a2, b2)
+            assert b2.numpy().dtype == np.float32
+            # Verify correctness against NumPy exp
+            expected = np.exp(out_np.astype(np.float32))
+            np.testing.assert_allclose(b2.numpy(), expected, rtol=1e-5, 
atol=1e-5)
 
     for func in test_funcs:
         atol = rtol = 1e-3 if func[0].__name__ in ["asin", "acos", "atan"] 
else 1e-5