From 8f4e7df1b386990491bfe9d6594839b11a6fe433 Mon Sep 17 00:00:00 2001
From: Philip Mueller <philip.mueller@cscs.ch>
Date: Mon, 11 Nov 2024 16:50:44 +0100
Subject: [PATCH] Added a pow specialization for integer.

The new specialization works for integers.
So if all arguments are integer the return value is guaranteed to be integer as well.
This should also solve the issue with the fft test mentioned by Tal.

https://github.com/spcl/dace/pull/1742/files#diff-a81c4594b19da98f6516283a50f6958b04976ae2913f097358141b2681868f3eR158
---
 dace/runtime/include/dace/math.h | 36 +++++++++++++++++++-------------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/dace/runtime/include/dace/math.h b/dace/runtime/include/dace/math.h
index 4dae494a8a..3b9defb176 100644
--- a/dace/runtime/include/dace/math.h
+++ b/dace/runtime/include/dace/math.h
@@ -512,32 +512,40 @@ namespace dace
             return (thrust::complex<T>)thrust::pow(a, b);
         }
 #endif
-        template<typename T, typename U>
+        template<
+            typename T,
+            typename U,
+            typename = std::enable_if_t<!(std::is_integral<T>::value && std::is_integral<U>::value)>
+        >
         DACE_CONSTEXPR DACE_HDFI auto pow(const T& a, const U& b)
         {
             return std::pow(a, b);
         }
 
 #ifndef DACE_XILINX
-        static DACE_CONSTEXPR DACE_HDFI int pow(const int& a, const int& b)
-        {
-            if (b < 0) return 0;
-            int result = 1;
-            for (int i = 0; i < b; ++i)
-                result *= a;
-            return result;
-        }
+        template<typename T>
+        using IntPowReturnType_t = std::conditional_t<std::is_unsigned<T>::value, unsigned long long int, long long int>;
 
-        static DACE_CONSTEXPR DACE_HDFI unsigned int pow(const unsigned int& a,
-                                       const unsigned int& b)
+        template<typename T1, typename T2>
+        static DACE_CONSTEXPR DACE_HDFI
+        std::enable_if_t<std::is_integral<T1>::value && std::is_integral<T2>::value, IntPowReturnType_t<T1> >
+        pow(const T1& a, const T2& b)
         {
-            unsigned int result = 1;
-            for (unsigned int i = 0; i < b; ++i)
+            /* TODO: The return value is always an integer, this is different from the behaviour of `std::pow` that
+             *  always return a float. We should probably patch the code generator to generate `ipow` calls if all
+             *  arguments are integers. */
+            if(std::is_signed<T2>::value && (b < 0))
+                return 0;
+            using IterationBound_t = typename std::make_unsigned<T2>::type;
+            IntPowReturnType_t<T1> result = 1;
+            const IterationBound_t stop = b;
+            for (IterationBound_t i = 0; i < stop; ++i)
                 result *= a;
             return result;
-        }
+	};
 #endif
 
+
         template<typename T>
         DACE_HDFI T ipow(const T& a, const unsigned int& b) {
             T result = a;