Repository: numba/numba

Issue:
On target cuda, return type of function may not match type given in signature string
- [x] I have tried using the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/main/CHANGE_LOG).
- [x] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

While testing #8308 with some kind of "load simulator" I came across this:

```python
import numba as nb
from numba import cuda
import numpy as np

COMPLEXITY = 2 ** 11

from math import cos, sin

for target in ('cpu', 'parallel', 'cuda'):

    if target == 'cuda':
    
        @cuda.jit('f4(f4)', device = True)
        def helper(scalar: float) -> float:
            res: float = 0.0
            for idx in range(COMPLEXITY):
                if idx % 2 == round(scalar) % 2:
                    res += sin(idx)
                else:
                    res -= cos(idx)
            return res
    
    else:
        
        @nb.jit('f4(f4)', nopython = True)
        def helper(scalar: float) -> float:
            res: float = 0.0
            for idx in range(COMPLEXITY):
                if idx % 2 == round(scalar) % 2:
                    res += sin(idx)
                else:
                    res -= cos(idx)
            return res

    @nb.vectorize('f4(f4)', nopython = True, target = target)
    def demo(d: float) -> float:
        return helper(d)

    for dtype in (np.float32,):
        
        data = np.arange(2**14, dtype = dtype)
        result = demo(data)
        print(target, type(demo), data.dtype, result.dtype)
        
        assert result.dtype == dtype
    
    del demo, helper
```

On target `cuda`, the return type of `demo` does not match what is requested: I am getting a double although a single is specified:

```
cpu <class 'numba.np.ufunc.dufunc.DUFunc'> float32 float32
parallel <class 'numpy.ufunc'> float32 float32

/github.numba/numba/np/ufunc/deviceufunc.py:364: RuntimeWarning: nopython kwarg for cuda target is redundant
  warnings.warn("nopython kwarg for cuda target is redundant",
/github.numba/numba/cuda/dispatcher.py:502: NumbaPerformanceWarning: Grid size 26 will likely result in GPU under-utilization due to low occupancy.
  warn(NumbaPerformanceWarning(msg))

cuda <class 'numba.cuda.vectorizers.CUDAUFuncDispatcher'> float32 float64

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Input In [2], in <cell line: 5>()
     38     result = demo(data)
     39     print(target, type(demo), data.dtype, result.dtype)
---> 41     assert result.dtype == dtype
     43 del demo, helper

AssertionError: 
```

Interestingly, if I reduce the `helper` functions to something very simple ...

```python
        @xxx('f4(f4)', nopython = True)
        def helper(scalar: float) -> float:
            return x ** 2
```

... I am getting the desired return data type. So something within the helper when looping over the `cos` and `sin` functions is not doing what it is supposed to.

---

The bug is present in 

- `main` as of now: c1e63c032899935d142aa6ead5d52b8730af53b8
- `0.56.0` as released

CPython 3.10.5, Ubuntu 20.04 LTS



Discussion:
Hi, 
I have a similar issue with automatic type deduction when I try to compile a PTX for a device function as shown below:

```py
from numba import cuda, int32


def foo(x, y):
    return x + y


ptx = cuda.compile_ptx_for_current_device(foo, (int32, int32), device=True)
assert ptx[1] == int32, f"Result type is {ptx[1]}"
```

This piece of code returns `AssertionError: Result type is int64` and seems that the `foo` arguments are deduced to be python's `int` i.e `int64`. In the case of using `float32` instead, the deduced type is correctly set to `float32`.
The problem here is that the return type request is not wired up:

https://github.com/numba/numba/blob/c1e63c032899935d142aa6ead5d52b8730af53b8/numba/cuda/decorators.py#L91

`restype` is not passed to `disp.compile`.



RE: https://github.com/numba/numba/issues/8400#issuecomment-1229502785 @cstyl 

The behavior of `i32 + i32` promoting to `i64 + i64` is due to Numba's lack of bigint support and we opted for preventing overflow. This is related to [NBEP 1](https://github.com/numba/numba/blob/94916113d92cb519a28adb10c30d8afda8049d07/docs/source/proposals/integer-typing.rst)
Starter patch:

```diff
diff --git a/numba/cuda/decorators.py b/numba/cuda/decorators.py
index a755ea0c8..9e514f5d4 100644
--- a/numba/cuda/decorators.py
+++ b/numba/cuda/decorators.py
@@ -121,7 +121,7 @@ def jit(func_or_sig=None, device=False, inline=False, link=[], debug=None,
                 if device:
                     from numba.core import typeinfer
                     with typeinfer.register_dispatcher(disp):
-                        disp.compile_device(argtypes)
+                        disp.compile_device(argtypes, restype)
                 else:
                     disp.compile(argtypes)
 
diff --git a/numba/cuda/dispatcher.py b/numba/cuda/dispatcher.py
index a79120d85..f9515a62c 100644
--- a/numba/cuda/dispatcher.py
+++ b/numba/cuda/dispatcher.py
@@ -823,7 +823,7 @@ class CUDADispatcher(Dispatcher, serialize.ReduceMixin):
 
         return call_template, pysig, args, kws
 
-    def compile_device(self, args):
+    def compile_device(self, args, return_type=None):
         """Compile the device function for the given argument types.
 
         Each signature is compiled once by caching the compiled function inside
@@ -844,7 +844,7 @@ class CUDADispatcher(Dispatcher, serialize.ReduceMixin):
                     'fastmath': fastmath
                 }
 
-                cres = compile_cuda(self.py_func, None, args,
+                cres = compile_cuda(self.py_func, return_type, args,
                                     debug=debug,
                                     inline=inline,
                                     fastmath=fastmath,
```

which fixes it for this example from https://numba.discourse.group/t/numba-vectorize-with-target-cuda-return-type-is-not-respected/1779/2?u=gmarkall: 

```python
import numpy as np
from numba import vectorize, float32

def vec_fn(x):
    return max(x, 0.0)

cpu_vec_fn = vectorize([float32(float32)])(vec_fn)
cuda_vec_fn = vectorize([float32(float32)], target='cuda')(vec_fn)

print(cpu_vec_fn(np.array([1.0], dtype="float32")).dtype)
print(cuda_vec_fn(np.array([1.0], dtype="float32")).dtype)
```

which prints

```
float32
float64
```

prior to application of the patch, and 

```
float32
float32
```

after application.
> RE: [#8400 (comment)](https://github.com/numba/numba/issues/8400#issuecomment-1229502785) @cstyl
> 
> The behavior of `i32 + i32` promoting to `i64 + i64` is due to Numba's lack of bigint support and we opted for preventing overflow. This is related to [NBEP 1](https://github.com/numba/numba/blob/94916113d92cb519a28adb10c30d8afda8049d07/docs/source/proposals/integer-typing.rst)

I think it should be possible to give more control over the return type if it took a signature instead of a tuple of arguments, so the caller could do:

```python
ptx, resty = cuda.compile_ptx_for_current_device(foo, int32(int32, int32), device=True)
```

and `resty` will be `int32`. I will aim to incorporate this into a fix.

Resolving patch:
diff --git a/numba/cuda/decorators.py b/numba/cuda/decorators.py
--- a/numba/cuda/decorators.py
+++ b/numba/cuda/decorators.py
@@ -121,7 +121,7 @@ def _jit(func):
                 if device:
                     from numba.core import typeinfer
                     with typeinfer.register_dispatcher(disp):
-                        disp.compile_device(argtypes)
+                        disp.compile_device(argtypes, restype)
                 else:
                     disp.compile(argtypes)
 
diff --git a/numba/cuda/dispatcher.py b/numba/cuda/dispatcher.py
--- a/numba/cuda/dispatcher.py
+++ b/numba/cuda/dispatcher.py
@@ -846,7 +846,7 @@ def get_call_template(self, args, kws):
 
         return call_template, pysig, args, kws
 
-    def compile_device(self, args):
+    def compile_device(self, args, return_type=None):
         """Compile the device function for the given argument types.
 
         Each signature is compiled once by caching the compiled function inside
@@ -868,7 +868,7 @@ def compile_device(self, args):
                 }
 
                 cc = get_current_device().compute_capability
-                cres = compile_cuda(self.py_func, None, args,
+                cres = compile_cuda(self.py_func, return_type, args,
                                     debug=debug,
                                     inline=inline,
                                     fastmath=fastmath,


Test patch:
diff --git a/numba/cuda/tests/cudapy/test_vectorize.py b/numba/cuda/tests/cudapy/test_vectorize.py
--- a/numba/cuda/tests/cudapy/test_vectorize.py
+++ b/numba/cuda/tests/cudapy/test_vectorize.py
@@ -1,26 +1,48 @@
 import numpy as np
 
 from collections import namedtuple
+from itertools import product
 from numba import vectorize
 from numba import cuda, int32, float32, float64
 from numba.cuda.testing import skip_on_cudasim
 from numba.cuda.testing import CUDATestCase
 import unittest
 
-sig = [int32(int32, int32),
-       float32(float32, float32),
-       float64(float64, float64)]
 
-test_dtypes = np.float32, np.int32
+# Signatures to test with - these are all homogeneous in dtype, so the output
+# dtype should match the input dtype - the output should not have been cast
+# upwards, as reported in #8400: https://github.com/numba/numba/issues/8400
+signatures = [int32(int32, int32),
+              float32(float32, float32),
+              float64(float64, float64)]
+
+# The order here is chosen such that each subsequent dtype might have been
+# casted to a previously-used dtype. This is unlikely to be an issue for CUDA,
+# but there might be future circumstances in which it becomes relevant, perhaps
+# if it supported Dynamic UFuncs, and we want to ensure that an implementation
+# for a the given dtype is used rather than casting the input upwards.
+dtypes = (np.float64, np.float32, np.int32)
+
+# NumPy ndarray orders
+orders = ('C', 'F')
+
+# Input sizes corresponding to operations:
+# - Less than one warp,
+# - Less than one block,
+# - Greater than one block (i.e. many blocks)
+input_sizes = (8, 100, 2 ** 10 + 1)
 
 
 @skip_on_cudasim('ufunc API unsupported in the simulator')
 class TestCUDAVectorize(CUDATestCase):
+    # Presumably chosen as an odd number unlikely to coincide with the total
+    # thread count, and large enough to ensure a significant number of blocks
+    # are used.
     N = 1000001
 
     def test_scalar(self):
 
-        @vectorize(sig, target='cuda')
+        @vectorize(signatures, target='cuda')
         def vector_add(a, b):
             return a + b
 
@@ -31,157 +53,134 @@ def vector_add(a, b):
 
     def test_1d(self):
 
-        @vectorize(sig, target='cuda')
+        @vectorize(signatures, target='cuda')
         def vector_add(a, b):
             return a + b
 
-        cuda_ufunc = vector_add
-
-        # build python ufunc
-        np_ufunc = np.add
-
-        # test it out
-        def test(ty):
+        for ty in dtypes:
             data = np.array(np.random.random(self.N), dtype=ty)
-
-            result = cuda_ufunc(data, data)
-            gold = np_ufunc(data, data)
-            self.assertTrue(np.allclose(gold, result), (gold, result))
-
-        test(np.double)
-        test(np.float32)
-        test(np.int32)
+            expected = np.add(data, data)
+            actual = vector_add(data, data)
+            np.testing.assert_allclose(expected, actual)
+            self.assertEqual(actual.dtype, ty)
 
     def test_1d_async(self):
 
-        @vectorize(sig, target='cuda')
+        @vectorize(signatures, target='cuda')
         def vector_add(a, b):
             return a + b
 
-        cuda_ufunc = vector_add
-
-        # build python ufunc
-        np_ufunc = np.add
+        stream = cuda.stream()
 
-        # test it out
-        def test(ty):
+        for ty in dtypes:
             data = np.array(np.random.random(self.N), dtype=ty)
-
-            stream = cuda.stream()
             device_data = cuda.to_device(data, stream)
-            dresult = cuda_ufunc(device_data, device_data, stream=stream)
-            result = dresult.copy_to_host()
-            stream.synchronize()
 
-            gold = np_ufunc(data, data)
+            dresult = vector_add(device_data, device_data, stream=stream)
+            actual = dresult.copy_to_host()
 
-            self.assertTrue(np.allclose(gold, result), (gold, result))
+            expected = np.add(data, data)
 
-        test(np.double)
-        test(np.float32)
-        test(np.int32)
+            np.testing.assert_allclose(expected, actual)
+            self.assertEqual(actual.dtype, ty)
 
     def test_nd(self):
 
-        @vectorize(sig, target='cuda')
+        @vectorize(signatures, target='cuda')
         def vector_add(a, b):
             return a + b
 
-        cuda_ufunc = vector_add
-
-        def test(dtype, order, nd, size=4):
-            data = np.random.random((size,) * nd).astype(dtype)
-            data[data != data] = 2.4
-            data[data == float('inf')] = 3.8
-            data[data == float('-inf')] = -3.8
-            data2 = np.array(data.T, order=order)  # .copy(order=order)
-
-            result = data + data2
-            our_result = cuda_ufunc(data, data2)
-            self.assertTrue(np.allclose(result, our_result),
-                            (dtype, order, result, our_result))
-
-        for nd in range(1, 8):
-            for dtype in test_dtypes:
-                for order in ('C', 'F'):
-                    test(dtype, order, nd)
-
-    def test_ufunc_attrib(self):
-        self.reduce_test(8)
-        self.reduce_test(100)
-        self.reduce_test(2 ** 10 + 1)
-        self.reduce_test2(8)
-        self.reduce_test2(100)
-        self.reduce_test2(2 ** 10 + 1)
+        for nd, dtype, order in product(range(1, 8), dtypes, orders):
+            shape = (4,) * nd
+            data = np.random.random(shape).astype(dtype)
+            data2 = np.array(data.T, order=order)
+
+            expected = data + data2
+            actual = vector_add(data, data2)
+            np.testing.assert_allclose(expected, actual)
+            self.assertEqual(actual.dtype, dtype)
 
     def test_output_arg(self):
-        @vectorize(sig, target='cuda')
+        @vectorize(signatures, target='cuda')
         def vector_add(a, b):
             return a + b
 
         A = np.arange(10, dtype=np.float32)
         B = np.arange(10, dtype=np.float32)
-        C = np.empty_like(A)
-        vector_add(A, B, out=C)
-        self.assertTrue(np.allclose(A + B, C))
 
-    def reduce_test(self, n):
-        @vectorize(sig, target='cuda')
+        expected = A + B
+        actual = np.empty_like(A)
+        vector_add(A, B, out=actual)
+
+        np.testing.assert_allclose(expected, actual)
+        self.assertEqual(expected.dtype, actual.dtype)
+
+    def test_reduce(self):
+        @vectorize(signatures, target='cuda')
         def vector_add(a, b):
             return a + b
 
-        cuda_ufunc = vector_add
-        x = np.arange(n, dtype=np.int32)
-        gold = np.add.reduce(x)
-        result = cuda_ufunc.reduce(x)
-        self.assertEqual(result, gold)
+        dtype = np.int32
+
+        for n in input_sizes:
+            x = np.arange(n, dtype=dtype)
+            expected = np.add.reduce(x)
+            actual = vector_add.reduce(x)
+            np.testing.assert_allclose(expected, actual)
+            # np.add.reduce is special-cased to return an int64 for any int
+            # arguments, so we can't compare against its returned dtype when
+            # we're checking the general reduce machinery (which just happens
+            # to be using addition). Instead, compare against the input dtype.
+            self.assertEqual(dtype, actual.dtype)
 
-    def reduce_test2(self, n):
+    def test_reduce_async(self):
 
-        @vectorize(sig, target='cuda')
+        @vectorize(signatures, target='cuda')
         def vector_add(a, b):
             return a + b
 
-        cuda_ufunc = vector_add
-
-        x = np.arange(n, dtype=np.int32)
-        gold = np.add.reduce(x)
         stream = cuda.stream()
-        dx = cuda.to_device(x, stream)
-        result = cuda_ufunc.reduce(dx, stream=stream)
-        self.assertEqual(result, gold)
-
-    def test_auto_transfer(self):
-        @vectorize(sig, target='cuda')
+        dtype = np.int32
+
+        for n in input_sizes:
+            x = np.arange(n, dtype=dtype)
+            expected = np.add.reduce(x)
+            dx = cuda.to_device(x, stream)
+            actual = vector_add.reduce(dx, stream=stream)
+            np.testing.assert_allclose(expected, actual)
+            # Compare against the input dtype as in test_reduce().
+            self.assertEqual(dtype, actual.dtype)
+
+    def test_manual_transfer(self):
+        @vectorize(signatures, target='cuda')
         def vector_add(a, b):
             return a + b
 
-        cuda_ufunc = vector_add
-
         n = 10
         x = np.arange(n, dtype=np.int32)
         dx = cuda.to_device(x)
-        y = cuda_ufunc(x, dx).copy_to_host()
-        np.testing.assert_equal(y, x + x)
+        expected = x + x
+        actual = vector_add(x, dx).copy_to_host()
+        np.testing.assert_equal(expected, actual)
+        self.assertEqual(expected.dtype, actual.dtype)
 
-    def test_ufunc_output_ravel(self):
-        @vectorize(sig, target='cuda')
+    def test_ufunc_output_2d(self):
+        @vectorize(signatures, target='cuda')
         def vector_add(a, b):
             return a + b
 
-        cuda_ufunc = vector_add
-
         n = 10
         x = np.arange(n, dtype=np.int32).reshape(2, 5)
         dx = cuda.to_device(x)
-        cuda_ufunc(dx, dx, out=dx)
+        vector_add(dx, dx, out=dx)
 
-        got = dx.copy_to_host()
-        expect = x + x
-        np.testing.assert_equal(got, expect)
+        expected = x + x
+        actual = dx.copy_to_host()
+        np.testing.assert_equal(expected, actual)
+        self.assertEqual(expected.dtype, actual.dtype)
 
     def check_tuple_arg(self, a, b):
-        @vectorize(sig, target='cuda')
+        @vectorize(signatures, target='cuda')
         def vector_add(a, b):
             return a + b