Skip to content

Commit d951ebc

Browse files
sklamsklam
authored
Merge pull request #8341 from gmarkall/issue-8303-2
CUDA: Support multiple outputs for Generalized Ufuncs
2 parents 2829f1f + 8c166f8 commit d951ebc

4 files changed

Lines changed: 355 additions & 192 deletions

File tree

‎docs/source/cuda/ufunc.rst‎

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -69,20 +69,6 @@ the CUDA ufunc functionality. This may be accomplished as follows::
6969
def matmulcore(A, B, C):
7070
...
7171

72-
There are times when the gufunc kernel uses too many of a GPU's
73-
resources, which can cause the kernel launch to fail. The user can
74-
explicitly control the maximum size of the thread block by setting
75-
the `max_blocksize` attribute on the compiled gufunc object.
76-
77-
::
78-
79-
from numba import guvectorize
80-
81-
@guvectorize(...,>
82-
def very_complex_kernel(A, B, C):
83-
...
84-
85-
very_complex_kernel.max_blocksize = 32 # limits to 32 threads per block
8672

8773
.. comment
8874

‎numba/cuda/tests/cudapy/test_gufunc.py‎

Lines changed: 154 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
import numpy.core.umath_tests as ut
33

44
from collections import namedtuple
5-
from numba import void, float32, float64
5+
from numba import void, int32, float32, float64
66
from numba import guvectorize
77
from numba import cuda
88
from numba.cuda.testing import skip_on_cudasim, CUDATestCase
@@ -12,7 +12,7 @@
1212
from numba.tests.support import override_config
1313

1414

15-
def _get_matmulcore_gufunc(dtype=float32, max_blocksize=None):
15+
def _get_matmulcore_gufunc(dtype=float32):
1616
@guvectorize([void(dtype[:, :], dtype[:, :], dtype[:, :])],
1717
'(m,n),(n,p)->(m,p)',
1818
target='cuda')
@@ -25,18 +25,15 @@ def matmulcore(A, B, C):
2525
for k in range(n):
2626
C[i, j] += A[i, k] * B[k, j]
2727

28-
gufunc = matmulcore
29-
if max_blocksize:
30-
gufunc.max_blocksize = max_blocksize
31-
return gufunc
28+
return matmulcore
3229

3330

3431
@skip_on_cudasim('ufunc API unsupported in the simulator')
3532
class TestCUDAGufunc(CUDATestCase):
3633

3734
def test_gufunc_small(self):
3835

39-
gufunc = _get_matmulcore_gufunc(max_blocksize=512)
36+
gufunc = _get_matmulcore_gufunc()
4037

4138
matrix_ct = 2
4239
A = np.arange(matrix_ct * 2 * 4, dtype=np.float32).reshape(matrix_ct, 2,
@@ -50,7 +47,7 @@ def test_gufunc_small(self):
5047

5148
def test_gufunc_auto_transfer(self):
5249

53-
gufunc = _get_matmulcore_gufunc(max_blocksize=512)
50+
gufunc = _get_matmulcore_gufunc()
5451

5552
matrix_ct = 2
5653
A = np.arange(matrix_ct * 2 * 4, dtype=np.float32).reshape(matrix_ct, 2,
@@ -66,7 +63,7 @@ def test_gufunc_auto_transfer(self):
6663

6764
def test_gufunc(self):
6865

69-
gufunc = _get_matmulcore_gufunc(max_blocksize=512)
66+
gufunc = _get_matmulcore_gufunc()
7067

7168
matrix_ct = 1001 # an odd number to test thread/block division in CUDA
7269
A = np.arange(matrix_ct * 2 * 4, dtype=np.float32).reshape(matrix_ct, 2,
@@ -80,7 +77,7 @@ def test_gufunc(self):
8077

8178
def test_gufunc_hidim(self):
8279

83-
gufunc = _get_matmulcore_gufunc(max_blocksize=512)
80+
gufunc = _get_matmulcore_gufunc()
8481

8582
matrix_ct = 100 # an odd number to test thread/block division in CUDA
8683
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):
105102
res2 = gufunc(X, np.tile(Y, (10, 1, 1)))
106103
np.testing.assert_allclose(gold, res2)
107104

108-
def test_gufunc_adjust_blocksize(self):
109-
110-
gufunc = _get_matmulcore_gufunc(max_blocksize=512)
111-
112-
matrix_ct = 1001 # an odd number to test thread/block division in CUDA
113-
A = np.arange(matrix_ct * 2 * 4, dtype=np.float32).reshape(matrix_ct, 2,
114-
4)
115-
B = np.arange(matrix_ct * 4 * 5, dtype=np.float32).reshape(matrix_ct, 4,
116-
5)
117-
118-
gufunc.max_blocksize = 32
119-
C = gufunc(A, B)
120-
Gold = ut.matrix_multiply(A, B)
121-
self.assertTrue(np.allclose(C, Gold))
122-
123105
def test_gufunc_stream(self):
124106

125-
gufunc = _get_matmulcore_gufunc(max_blocksize=512)
107+
gufunc = _get_matmulcore_gufunc()
126108

127109
#cuda.driver.flush_pending_free()
128110
matrix_ct = 1001 # an odd number to test thread/block division in CUDA
@@ -153,6 +135,21 @@ def copy(A, B):
153135
for i in range(B.size):
154136
B[i] = A[i]
155137

138+
A = np.arange(10, dtype=np.float32) + 1
139+
B = np.zeros_like(A)
140+
copy(A, out=B)
141+
np.testing.assert_allclose(A, B)
142+
143+
def test_copy_unspecified_return(self):
144+
# Ensure that behaviour is correct when the return type is not
145+
# specified in the signature.
146+
@guvectorize([(float32[:], float32[:])],
147+
'(x)->(x)',
148+
target='cuda')
149+
def copy(A, B):
150+
for i in range(B.size):
151+
B[i] = A[i]
152+
156153
A = np.arange(10, dtype=np.float32) + 1
157154
B = np.zeros_like(A)
158155
copy(A, out=B)
@@ -266,7 +263,7 @@ def foo(inp, out):
266263
with self.assertRaises(ValueError) as raises:
267264
foo(inp, out, out=out)
268265

269-
msg = "cannot specify 'out' as both a positional and keyword argument"
266+
msg = "cannot specify argument 'out' as both positional and keyword"
270267
self.assertEqual(str(raises.exception), msg)
271268

272269
def check_tuple_arg(self, a, b):
@@ -305,9 +302,138 @@ def test_tuple_of_array_arg(self):
305302
self.check_tuple_arg(a, b)
306303

307304
def test_gufunc_name(self):
308-
gufunc = _get_matmulcore_gufunc(max_blocksize=512)
305+
gufunc = _get_matmulcore_gufunc()
309306
self.assertEqual(gufunc.__name__, 'matmulcore')
310307

308+
def test_bad_return_type(self):
309+
with self.assertRaises(TypeError) as te:
310+
@guvectorize([int32(int32[:], int32[:])], '(m)->(m)', target='cuda')
311+
def f(x, y):
312+
pass
313+
314+
msg = str(te.exception)
315+
self.assertIn('guvectorized functions cannot return values', msg)
316+
self.assertIn('specifies int32 return type', msg)
317+
318+
def test_incorrect_number_of_pos_args(self):
319+
@guvectorize([(int32[:], int32[:], int32[:])],
320+
'(m),(m)->(m)', target='cuda')
321+
def f(x, y, z):
322+
pass
323+
324+
arr = np.arange(5)
325+
326+
# Inputs only, too few
327+
with self.assertRaises(TypeError) as te:
328+
f(arr)
329+
330+
msg = str(te.exception)
331+
self.assertIn('gufunc accepts 2 positional arguments', msg)
332+
self.assertIn('or 3 positional arguments', msg)
333+
self.assertIn('Got 1 positional argument.', msg)
334+
335+
# Inputs and outputs, too many
336+
with self.assertRaises(TypeError) as te:
337+
f(arr, arr, arr, arr)
338+
339+
msg = str(te.exception)
340+
self.assertIn('gufunc accepts 2 positional arguments', msg)
341+
self.assertIn('or 3 positional arguments', msg)
342+
self.assertIn('Got 4 positional arguments.', msg)
343+
344+
345+
@skip_on_cudasim('ufunc API unsupported in the simulator')
346+
class TestMultipleOutputs(CUDATestCase):
347+
def test_multiple_outputs_same_type_passed_in(self):
348+
@guvectorize([void(float32[:], float32[:], float32[:])],
349+
'(x)->(x),(x)',
350+
target='cuda')
351+
def copy(A, B, C):
352+
for i in range(B.size):
353+
B[i] = A[i]
354+
C[i] = A[i]
355+
356+
A = np.arange(10, dtype=np.float32) + 1
357+
B = np.zeros_like(A)
358+
C = np.zeros_like(A)
359+
copy(A, B, C)
360+
np.testing.assert_allclose(A, B)
361+
np.testing.assert_allclose(A, C)
362+
363+
def test_multiple_outputs_distinct_values(self):
364+
365+
@guvectorize([void(float32[:], float32[:], float32[:])],
366+
'(x)->(x),(x)',
367+
target='cuda')
368+
def copy_and_double(A, B, C):
369+
for i in range(B.size):
370+
B[i] = A[i]
371+
C[i] = A[i] * 2
372+
373+
A = np.arange(10, dtype=np.float32) + 1
374+
B = np.zeros_like(A)
375+
C = np.zeros_like(A)
376+
copy_and_double(A, B, C)
377+
np.testing.assert_allclose(A, B)
378+
np.testing.assert_allclose(A * 2, C)
379+
380+
def test_multiple_output_allocation(self):
381+
@guvectorize([void(float32[:], float32[:], float32[:])],
382+
'(x)->(x),(x)',
383+
target='cuda')
384+
def copy_and_double(A, B, C):
385+
for i in range(B.size):
386+
B[i] = A[i]
387+
C[i] = A[i] * 2
388+
389+
A = np.arange(10, dtype=np.float32) + 1
390+
B, C = copy_and_double(A)
391+
np.testing.assert_allclose(A, B)
392+
np.testing.assert_allclose(A * 2, C)
393+
394+
def test_multiple_output_dtypes(self):
395+
396+
@guvectorize([void(int32[:], int32[:], float64[:])],
397+
'(x)->(x),(x)',
398+
target='cuda')
399+
def copy_and_multiply(A, B, C):
400+
for i in range(B.size):
401+
B[i] = A[i]
402+
C[i] = A[i] * 1.5
403+
404+
A = np.arange(10, dtype=np.int32) + 1
405+
B = np.zeros_like(A)
406+
C = np.zeros_like(A, dtype=np.float64)
407+
copy_and_multiply(A, B, C)
408+
np.testing.assert_allclose(A, B)
409+
np.testing.assert_allclose(A * np.float64(1.5), C)
410+
411+
def test_incorrect_number_of_pos_args(self):
412+
@guvectorize([(int32[:], int32[:], int32[:], int32[:])],
413+
'(m),(m)->(m),(m)', target='cuda')
414+
def f(x, y, z, w):
415+
pass
416+
417+
arr = np.arange(5)
418+
419+
# Inputs only, too few
420+
with self.assertRaises(TypeError) as te:
421+
f(arr)
422+
423+
msg = str(te.exception)
424+
self.assertIn('gufunc accepts 2 positional arguments', msg)
425+
self.assertIn('or 4 positional arguments', msg)
426+
self.assertIn('Got 1 positional argument.', msg)
427+
428+
# Inputs and outputs, too many
429+
with self.assertRaises(TypeError) as te:
430+
f(arr, arr, arr, arr, arr)
431+
432+
msg = str(te.exception)
433+
self.assertIn('gufunc accepts 2 positional arguments', msg)
434+
self.assertIn('or 4 positional arguments', msg)
435+
self.assertIn('Got 5 positional arguments.', msg)
436+
311437

312438
if __name__ == '__main__':
313439
unittest.main()

‎numba/cuda/vectorizers.py‎

Lines changed: 6 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,17 +12,8 @@ class CUDAUFuncDispatcher(object):
1212

1313
def __init__(self, types_to_retty_kernels, pyfunc):
1414
self.functions = types_to_retty_kernels
15-
self._maxblocksize = 0 # ignored
1615
self.__name__ = pyfunc.__name__
1716

18-
@property
19-
def max_blocksize(self):
20-
return self._maxblocksize
21-
22-
@max_blocksize.setter
23-
def max_blocksize(self, blksz):
24-
self._max_blocksize = blksz
25-
2617
def __call__(self, *args, **kws):
2718
"""
2819
*args: numpy arrays or DeviceArrayBase (created by cuda.to_device).
@@ -94,6 +85,10 @@ class _CUDAGUFuncCallSteps(GUFuncCallSteps):
9485
'_stream',
9586
]
9687

88+
def __init__(self, nin, nout, args, kwargs):
89+
super().__init__(nin, nout, args, kwargs)
90+
self._stream = kwargs.get('stream', 0)
91+
9792
def is_device_array(self, obj):
9893
return cuda.is_cuda_array(obj)
9994

@@ -114,12 +109,9 @@ def to_host(self, devary, hostary):
114109
out = devary.copy_to_host(hostary, stream=self._stream)
115110
return out
116111

117-
def device_array(self, shape, dtype):
112+
def allocate_device_array(self, shape, dtype):
118113
return cuda.device_array(shape=shape, dtype=dtype, stream=self._stream)
119114

120-
def prepare_inputs(self):
121-
self._stream = self.kwargs.get('stream', 0)
122-
123115
def launch_kernel(self, kernel, nelem, args):
124116
kernel.forall(nelem, stream=self._stream)(*args)
125117

@@ -177,7 +169,7 @@ def to_device(self, hostary, stream):
177169
def to_host(self, devary, stream):
178170
return devary.copy_to_host(stream=stream)
179171

180-
def device_array(self, shape, dtype, stream):
172+
def allocate_device_array(self, shape, dtype, stream):
181173
return cuda.device_array(shape=shape, dtype=dtype, stream=stream)
182174

183175
def broadcast_device(self, ary, shape):

0 commit comments

Comments
 (0)