From 0e3e098cd0c009920ec1ef04b7d78c17a1c7af9c Mon Sep 17 00:00:00 2001
From: BenWeber42 <benjamin.weber@inf.ethz.ch>
Date: Wed, 12 Jun 2024 16:38:05 +0200
Subject: [PATCH] Made `dace::math::pow` forward to `std::pow` more generic
 (#1583)

This overload solution will forward the argument types separately and
use the same return type.

---------

Co-authored-by: Berke Ates <berke@ates.ch>
Co-authored-by: Berke Ates <beates@student.ethz.ch>
Co-authored-by: Philipp Schaad <schaad.phil@gmail.com>
Co-authored-by: Philip Mueller <philip.mueller@cscs.ch>
---
 dace/runtime/include/dace/math.h             | 26 +++---------------
 tests/python_frontend/power_operator_test.py | 29 ++++++++++++++++++++
 2 files changed, 33 insertions(+), 22 deletions(-)

diff --git a/dace/runtime/include/dace/math.h b/dace/runtime/include/dace/math.h
index afc08a64d3..0a9d153767 100644
--- a/dace/runtime/include/dace/math.h
+++ b/dace/runtime/include/dace/math.h
@@ -473,36 +473,29 @@ namespace dace
             return (thrust::complex<T>)thrust::pow(a, b);
         }
 #endif
-        template<typename T>
-        DACE_CONSTEXPR DACE_HDFI T pow(const T& a, const T& b)
+        template<typename T, typename U>
+        DACE_CONSTEXPR DACE_HDFI auto pow(const T& a, const U& b)
         {
-            return (T)std::pow(a, b);
+            return std::pow(a, b);
         }
 
 #ifndef DACE_XILINX
         static DACE_CONSTEXPR DACE_HDFI int pow(const int& a, const int& b)
         {
-/*#ifndef __CUDA_ARCH__
-            return std::pow(a, b);
-#else*/
             if (b < 0) return 0;
             int result = 1;
             for (int i = 0; i < b; ++i)
                 result *= a;
             return result;
-//#endif
         }
+
         static DACE_CONSTEXPR DACE_HDFI unsigned int pow(const unsigned int& a,
                                        const unsigned int& b)
         {
-/*#ifndef __CUDA_ARCH__
-            return std::pow(a, b);
-#else*/
             unsigned int result = 1;
             for (unsigned int i = 0; i < b; ++i)
                 result *= a;
             return result;
-//#endif
         }
 #endif
 
@@ -514,17 +507,6 @@ namespace dace
             return result;
         }
 
-        template<typename T>
-        DACE_CONSTEXPR DACE_HDFI T pow(const T& a, const int& b)
-        {
-            return (T)std::pow(a, (T)b);
-        }
-        template<typename T>
-        DACE_CONSTEXPR DACE_HDFI T pow(const T& a, const unsigned int& b)
-        {
-            return (T)std::pow(a, (T)b);
-        }
-
         template<typename T, typename std::enable_if<std::is_integral<T>::value>::type* = nullptr>
         DACE_CONSTEXPR DACE_HDFI T ifloor(const T& a)
         {
diff --git a/tests/python_frontend/power_operator_test.py b/tests/python_frontend/power_operator_test.py
index f1e9638f06..3657da6d7e 100644
--- a/tests/python_frontend/power_operator_test.py
+++ b/tests/python_frontend/power_operator_test.py
@@ -1,4 +1,6 @@
 # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved.
+
+import pytest
 import dace
 import numpy as np
 
@@ -35,8 +37,35 @@ def test_pow_neg_exp():
     pow_neg_exp(a=res)
     assert (res[0] == 0.01)
 
+in_types = [dace.float32, dace.float64, dace.int8, dace.int16, dace.int32, dace.int64]
+
+@pytest.mark.parametrize("a_type", in_types)
+@pytest.mark.parametrize("b_type", in_types)
+def test_pow_types(a_type, b_type):
+    @dace.program
+    def pow_types(A: a_type[1], B: b_type[1], R: dace.float64[1]):
+        with dace.tasklet(dace.Language.Python):
+            scalar_a << A[0]
+            scalar_b << B[0]
+            scalar_r >> R[0]
+            scalar_r = scalar_a ** scalar_b
+
+    # a ** b needs to fit into the smallest type (int8)
+    a = np.random.rand(1) * 4
+    b = np.random.rand(1) * 4
+    r = np.random.rand(1).astype(np.float64)
+
+    a = a.astype(a_type.as_numpy_dtype())
+    b = b.astype(b_type.as_numpy_dtype())
+
+    pow_types(A=a, B=b, R=r)
+    assert np.allclose(r, a ** b)
 
 if __name__ == "__main__":
     test_pow_num_literals()
     test_pow_op_preced()
     test_pow_neg_exp()
+
+    for a_type in in_types:
+        for b_type in in_types:
+          test_pow_types(a_type, b_type)