Skip to content

CUDA: Support multiple outputs for Generalized Ufuncs #8341

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sklam merged 33 commits into from
Mar 23, 2023
Merged

CUDA: Support multiple outputs for Generalized Ufuncs #8341

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
33 commits
Select commit Hold shift + click to select a range
b973ce2
Device GUFuncs: Enforce void or unspecified return type
gmarkall Aug 11, 2022
4a7b0eb
Device ufuncs: remove long-deprecated argtypes / restype support
gmarkall Aug 11, 2022
507c352
Device ufuncs: Make the kernelmap hold multiple output dtypes
gmarkall Aug 11, 2022
a43cdfc
Device ufuncs: return multiple output dtypes from schedule
gmarkall Aug 11, 2022
7a0a7bc
Device ufuncs: support callsteps holding multiple outputs
gmarkall Aug 11, 2022
9ba159c
Device ufuncs: More generality for multiple outputs in callsteps
gmarkall Aug 11, 2022
e927fc7
Device ufuncs: allow multiple kernel return values in callsteps
gmarkall Aug 11, 2022
e8b6c97
Device ufuncs: correct parameter passing for multiple outputs
gmarkall Aug 11, 2022
08e1835
CUDA ufuncs: another multiple output test
gmarkall Aug 11, 2022
b1e2dd6
Device ufuncs: fix allocation of multiple outputs
gmarkall Aug 11, 2022
6acf970
Merge remote-tracking branch 'numba/main' into issue-8303-2
gmarkall Nov 17, 2022
a650d54
Merge remote-tracking branch 'numba/main' into issue-8303-2
gmarkall Dec 2, 2022
24592e4
Skip ufunc multiple output tests on simulator
gmarkall Dec 2, 2022
d3e16c4
Merge remote-tracking branch 'numba/main' into issue-8303-2
gmarkall Feb 22, 2023
575592f
Device ufuncs: some refactors for clarity
gmarkall Feb 22, 2023
00da70f
Device ufuncs: rename `norm_inputs` to `inputs`
gmarkall Feb 22, 2023
f63ac9e
Device ufuncs: corrext description of kernelmap
gmarkall Feb 22, 2023
dac4659
Device gufunc call steps: remove some state
gmarkall Feb 22, 2023
c82575c
Device ufuncs: remove kernel_returnvalues from state
gmarkall Feb 22, 2023
b30def6
Device gufuncs: Implement call steps as an ABC
gmarkall Feb 22, 2023
8ade7fc
Device gufuncs: refactor input and output preparation
gmarkall Feb 22, 2023
e4d9e42
Device GUFuncs: Add docs and refactor
gmarkall Feb 22, 2023
583762a
Merge remote-tracking branch 'numba/main' into issue-8303-2
gmarkall Mar 7, 2023
4879b05
Remove `max_blocksize` from CUDA GUFuncs
gmarkall Mar 7, 2023
56a98d2
Merge remote-tracking branch 'numba/main' into issue-8303-2
gmarkall Mar 12, 2023
d73b7ec
Some stylistic fixes from PR #8341 feedback
gmarkall Mar 15, 2023
1a5f2f5
Rename [G]UFuncMechanism.device_array() to allocate_device_array()
gmarkall Mar 15, 2023
00eee0c
deviceufunc: clarify generality of implementation (PR #8341 feedback)
gmarkall Mar 15, 2023
efca2c5
deviceufunc.py: refer to normalized args
gmarkall Mar 15, 2023
f2a4ae3
deviceufunc.py: Clarify some docstrings, ensure `post_process_outputs()`
gmarkall Mar 15, 2023
3fe2a03
CUDA GUFuncs: Ensure the correct number of args is passed
gmarkall Mar 16, 2023
1b13d1d
CUDA GUfuncs: edit pos and kwarg argument exception message
gmarkall Mar 16, 2023
8c166f8
Device ufuncs: simplify formatting of "positional argument[s]"
gmarkall Mar 17, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Device ufuncs: correct parameter passing for multiple outputs 
Also add a test for a CUDA multiple output gufunc now that it works.
  • Loading branch information
@gmarkall
gmarkall committed Aug 11, 2022
commit e8b6c97fb2b682378b3ace9d9d084d9088ffbe94
17 changes: 17 additions & 0 deletions numba/cuda/tests/cudapy/test_gufunc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,23 @@ def copy(A, B):
copy(A, out=B)
self.assertTrue(np.allclose(A, B))

def test_copy_multiple_outputs(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_copy_odd(self):

@guvectorize([void(float32[:], float32[:])],
Expand Down
16 changes: 9 additions & 7 deletions numba/np/ufunc/deviceufunc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -637,10 +637,10 @@ def __call__(self, *args, **kws):
callsteps.adjust_input_types(indtypes)
callsteps.allocate_outputs(schedule, outdtypes)
callsteps.prepare_kernel_parameters()
newparams, newretval = self._broadcast(schedule,
callsteps.kernel_parameters,
callsteps.kernel_returnvalues)
callsteps.launch_kernel(kernel, schedule.loopn, newparams + [newretval])
newparams = self._broadcast(schedule,
callsteps.kernel_parameters,
callsteps.kernel_returnvalues)
callsteps.launch_kernel(kernel, schedule.loopn, newparams)
return callsteps.post_process_result()

def _schedule(self, inputs, outs):
Expand Down Expand Up @@ -682,7 +682,6 @@ def _search_matching_signature(self, idtypes):
raise TypeError("no matching signature")

def _broadcast(self, schedule, params, retvals):
retval = retvals[0]
assert schedule.loopn > 0, "zero looping dimension"

odim = 1 if not schedule.loopdims else schedule.loopn
Expand All @@ -695,8 +694,11 @@ def _broadcast(self, schedule, params, retvals):
else:
# Broadcast vector input
newparams.append(self._broadcast_array(p, odim, cs))
newretval = retval.reshape(odim, *schedule.oshapes[0])
return newparams, newretval

newretvals = []
for retval, oshape in zip(retvals, schedule.oshapes):
newretvals.append(retval.reshape(odim, *oshape))
return tuple(newparams) + tuple(newretvals)

def _broadcast_array(self, ary, newdim, innerdim):
newshape = (newdim,) + innerdim
Expand Down