sklam · Mar 23, 2023 · Aug 11, 2022 · Aug 11, 2022 · Aug 11, 2022 · Aug 11, 2022
diff --git a/docs/source/cuda/ufunc.rst b/docs/source/cuda/ufunc.rst
@@ -69,20 +69,6 @@ the CUDA ufunc functionality.  This may be accomplished as follows::
     def matmulcore(A, B, C):
         ...
 
-There are times when the gufunc kernel uses too many of a GPU's
-resources, which can cause the kernel launch to fail.  The user can
-explicitly control the maximum size of the thread block by setting
-the `max_blocksize` attribute on the compiled gufunc object.
-
-::
-
-    from numba import guvectorize
-
-    @guvectorize(...,>
-    def very_complex_kernel(A, B, C):
-        ...
-
-    very_complex_kernel.max_blocksize = 32  # limits to 32 threads per block
 
 .. comment
 

@@ -2,7 +2,7 @@
 import numpy.core.umath_tests as ut
 
 from collections import namedtuple
-from numba import void, float32, float64
+from numba import void, int32, float32, float64
 from numba import guvectorize
 from numba import cuda
 from numba.cuda.testing import skip_on_cudasim, CUDATestCase
@@ -12,7 +12,7 @@
 from numba.tests.support import override_config
 
 
-def _get_matmulcore_gufunc(dtype=float32, max_blocksize=None):
+def _get_matmulcore_gufunc(dtype=float32):
     @guvectorize([void(dtype[:, :], dtype[:, :], dtype[:, :])],
                  '(m,n),(n,p)->(m,p)',
 >
@@ -25,18 +25,15 @@ def matmulcore(A, B, C):
                 for k in range(n):
                     C[i, j] += A[i, k] * B[k, j]
 
-    gufunc = matmulcore
-    if max_blocksize:
-        gufunc.max_blocksize = max_blocksize
-    return gufunc
+    return matmulcore
 
 
 @skip_on_cudasim('ufunc API unsupported in the simulator')
 class TestCUDAGufunc(CUDATestCase):
 
     def test_gufunc_small(self):
 
-        gufunc = _get_matmulcore_gufunc(max_blocksize=512)
+        gufunc = _get_matmulcore_gufunc()
 
         matrix_ct = 2
         A = np.arange(matrix_ct * 2 * 4, dtype=np.float32).reshape(matrix_ct, 2,
@@ -50,7 +47,7 @@ def test_gufunc_small(self):
 
     def test_gufunc_auto_transfer(self):
 
-        gufunc = _get_matmulcore_gufunc(max_blocksize=512)
+        gufunc = _get_matmulcore_gufunc()
 
         matrix_ct = 2
         A = np.arange(matrix_ct * 2 * 4, dtype=np.float32).reshape(matrix_ct, 2,
@@ -66,7 +63,7 @@ def test_gufunc_auto_transfer(self):
 
     def test_gufunc(self):
 
-        gufunc = _get_matmulcore_gufunc(max_blocksize=512)
+        gufunc = _get_matmulcore_gufunc()
 
         matrix_ct = 1001 # an odd number to test thread/block division in CUDA
         A = np.arange(matrix_ct * 2 * 4, dtype=np.float32).reshape(matrix_ct, 2,
@@ -80,7 +77,7 @@ def test_gufunc(self):
 
     def test_gufunc_hidim(self):
 
-        gufunc = _get_matmulcore_gufunc(max_blocksize=512)
+        gufunc = _get_matmulcore_gufunc()
 
         matrix_ct = 100 # an odd number to test thread/block division in CUDA
         A = np.arange(matrix_ct * 2 * 4, dtype=np.float32).reshape(4, 25, 2, 4)
@@ -105,24 +102,9 @@ def test_gufunc_new_axis(self):
         res2 = gufunc(X, np.tile(Y, (10, 1, 1)))
         np.testing.assert_allclose(gold, res2)
 
-    def test_gufunc_adjust_blocksize(self):
-
-        gufunc = _get_matmulcore_gufunc(max_blocksize=512)
-
-        matrix_ct = 1001 # an odd number to test thread/block division in CUDA
-        A = np.arange(matrix_ct * 2 * 4, dtype=np.float32).reshape(matrix_ct, 2,
-                                                                   4)
-        B = np.arange(matrix_ct * 4 * 5, dtype=np.float32).reshape(matrix_ct, 4,
-                                                                   5)
-
-        gufunc.max_blocksize = 32
-        C = gufunc(A, B)
-        Gold = ut.matrix_multiply(A, B)
-        self.assertTrue(np.allclose(C, Gold))
-
     def test_gufunc_stream(self):
 
-        gufunc = _get_matmulcore_gufunc(max_blocksize=512)
+        gufunc = _get_matmulcore_gufunc()
 
         #cuda.driver.flush_pending_free()
         matrix_ct = 1001 # an odd number to test thread/block division in CUDA
@@ -153,6 +135,21 @@ def copy(A, B):
             for i in range(B.size):
                 B[i] = A[i]
 
+        A = np.arange(10, dtype=np.float32) + 1
+        B = np.zeros_like(A)
+        copy(A, out=B)
+        np.testing.assert_allclose(A, B)
+
+    def test_copy_unspecified_return(self):
+        # Ensure that behaviour is correct when the return type is not
+        # specified in the signature.
+        @guvectorize([(float32[:], float32[:])],
+                     '(x)->(x)',
+>
+        def copy(A, B):
+            for i in range(B.size):
+                B[i] = A[i]
+
         A = np.arange(10, dtype=np.float32) + 1
         B = np.zeros_like(A)
         copy(A, out=B)
@@ -266,7 +263,7 @@ def foo(inp, out):
         with self.assertRaises(ValueError) as raises:
             foo(inp, out, out=out)
 
-        msg = "cannot specify 'out' as both a positional and keyword argument"
+        msg = "cannot specify argument 'out' as both positional and keyword"
         self.assertEqual(str(raises.exception), msg)
 
     def check_tuple_arg(self, a, b):
@@ -305,9 +302,138 @@ def test_tuple_of_array_arg(self):
         self.check_tuple_arg(a, b)
 
     def test_gufunc_name(self):
-        gufunc = _get_matmulcore_gufunc(max_blocksize=512)
+        gufunc = _get_matmulcore_gufunc()
         self.assertEqual(gufunc.__name__, 'matmulcore')
 
+    def test_bad_return_type(self):
+        with self.assertRaises(TypeError) as te:
+            @guvectorize([int32(int32[:], int32[:])], '(m)->(m)',>
+            def f(x, y):
+                pass
+
+        msg = str(te.exception)
+        self.assertIn('guvectorized functions cannot return values', msg)
+        self.assertIn('specifies int32 return type', msg)
+
+    def test_incorrect_number_of_pos_args(self):
+        @guvectorize([(int32[:], int32[:], int32[:])],
+                     '(m),(m)->(m)',>
+        def f(x, y, z):
+            pass
+
+        arr = np.arange(5)
+
+        # Inputs only, too few
+        with self.assertRaises(TypeError) as te:
+            f(arr)
+
+        msg = str(te.exception)
+        self.assertIn('gufunc accepts 2 positional arguments', msg)
+        self.assertIn('or 3 positional arguments', msg)
+        self.assertIn('Got 1 positional argument.', msg)
+
+        # Inputs and outputs, too many
+        with self.assertRaises(TypeError) as te:
+            f(arr, arr, arr, arr)
+
+        msg = str(te.exception)
+        self.assertIn('gufunc accepts 2 positional arguments', msg)
+        self.assertIn('or 3 positional arguments', msg)
+        self.assertIn('Got 4 positional arguments.', msg)
+
+
+@skip_on_cudasim('ufunc API unsupported in the simulator')
+class TestMultipleOutputs(CUDATestCase):
+    def test_multiple_outputs_same_type_passed_in(self):
+        @guvectorize([void(float32[:], float32[:], float32[:])],
+                     '(x)->(x),(x)',
+>
+        def copy(A, B, C):
+            for i in range(B.size):
+                B[i] = A[i]
+                C[i] = A[i]
+
+        A = np.arange(10, dtype=np.float32) + 1
+        B = np.zeros_like(A)
+        C = np.zeros_like(A)
+        copy(A, B, C)
+        np.testing.assert_allclose(A, B)
+        np.testing.assert_allclose(A, C)
+
+    def test_multiple_outputs_distinct_values(self):
+
+        @guvectorize([void(float32[:], float32[:], float32[:])],
+                     '(x)->(x),(x)',
+>
+        def copy_and_double(A, B, C):
+            for i in range(B.size):
+                B[i] = A[i]
+                C[i] = A[i] * 2
+
+        A = np.arange(10, dtype=np.float32) + 1
+        B = np.zeros_like(A)
+        C = np.zeros_like(A)
+        copy_and_double(A, B, C)
+        np.testing.assert_allclose(A, B)
+        np.testing.assert_allclose(A * 2, C)
+
+    def test_multiple_output_allocation(self):
+        @guvectorize([void(float32[:], float32[:], float32[:])],
+                     '(x)->(x),(x)',
+>
+        def copy_and_double(A, B, C):
+            for i in range(B.size):
+                B[i] = A[i]
+                C[i] = A[i] * 2
+
+        A = np.arange(10, dtype=np.float32) + 1
+        B, C = copy_and_double(A)
+        np.testing.assert_allclose(A, B)
+        np.testing.assert_allclose(A * 2, C)
+
+    def test_multiple_output_dtypes(self):
+
+        @guvectorize([void(int32[:], int32[:], float64[:])],
+                     '(x)->(x),(x)',
+>
+        def copy_and_multiply(A, B, C):
+            for i in range(B.size):
+                B[i] = A[i]
+                C[i] = A[i] * 1.5
+
+        A = np.arange(10, dtype=np.int32) + 1
+        B = np.zeros_like(A)
+        C = np.zeros_like(A, dtype=np.float64)
+        copy_and_multiply(A, B, C)
+        np.testing.assert_allclose(A, B)
+        np.testing.assert_allclose(A * np.float64(1.5), C)
+
+    def test_incorrect_number_of_pos_args(self):
+        @guvectorize([(int32[:], int32[:], int32[:], int32[:])],
+                     '(m),(m)->(m),(m)',>
+        def f(x, y, z, w):
+            pass
+
+        arr = np.arange(5)
+
+        # Inputs only, too few
+        with self.assertRaises(TypeError) as te:
+            f(arr)
+
+        msg = str(te.exception)
+        self.assertIn('gufunc accepts 2 positional arguments', msg)
+        self.assertIn('or 4 positional arguments', msg)
+        self.assertIn('Got 1 positional argument.', msg)
+
+        # Inputs and outputs, too many
+        with self.assertRaises(TypeError) as te:
+            f(arr, arr, arr, arr, arr)
+
+        msg = str(te.exception)
+        self.assertIn('gufunc accepts 2 positional arguments', msg)
+        self.assertIn('or 4 positional arguments', msg)
+        self.assertIn('Got 5 positional arguments.', msg)
+
 
 if __name__ == '__main__':
     unittest.main()
@@ -12,17 +12,8 @@ class CUDAUFuncDispatcher(object):
 
     def __init__(self, types_to_retty_kernels, pyfunc):
         self.functions = types_to_retty_kernels
-        self._maxblocksize = 0  # ignored
         self.__name__ = pyfunc.__name__
 
-    @property
-    def max_blocksize(self):
-        return self._maxblocksize
-
-    @max_blocksize.setter
-    def max_blocksize(self, blksz):
-        self._max_blocksize = blksz
-
     def __call__(self, *args, **kws):
         """
         *args: numpy arrays or DeviceArrayBase (created by cuda.to_device).
@@ -94,6 +85,10 @@ class _CUDAGUFuncCallSteps(GUFuncCallSteps):
         '_stream',
     ]
 
+    def __init__(self, nin, nout, args, kwargs):
+        super().__init__(nin, nout, args, kwargs)
+        self._stream = kwargs.get('stream', 0)
+
     def is_device_array(self, obj):
         return cuda.is_cuda_array(obj)
 
@@ -114,12 +109,9 @@ def to_host(self, devary, hostary):
         out = devary.copy_to_host(hostary, stream=self._stream)
         return out
 
-    def device_array(self, shape, dtype):
+    def allocate_device_array(self, shape, dtype):
         return cuda.device_array(shape=shape, dtype=dtype, stream=self._stream)
 
-    def prepare_inputs(self):
-        self._stream = self.kwargs.get('stream', 0)
-
     def launch_kernel(self, kernel, nelem, args):
         kernel.forall(nelem, stream=self._stream)(*args)
 
@@ -177,7 +169,7 @@ def to_device(self, hostary, stream):
     def to_host(self, devary, stream):
         return devary.copy_to_host(stream=stream)
 
-    def device_array(self, shape, dtype, stream):
+    def allocate_device_array(self, shape, dtype, stream):
         return cuda.device_array(shape=shape, dtype=dtype, stream=stream)
 
     def broadcast_device(self, ary, shape):