CUDA: Support multiple outputs for Generalized Ufuncs by gmarkall · Pull Request #8341 · numba/numba

gmarkall · 2022-08-11T15:12:16Z

Support for multiple outputs from CUDA gufuncs - previously only one output was supported, so it was possible to write a signature like '(x)->(x)', but not '(x)->(x),(x)'.

This PR consists of:

Changes to numba.np.ufunc.deviceufunc:
- A lot of small changes to pluralise things returning to the output (e.g. outdtype -> outdtypes)
- Some simplifications of the code in DeviceGUFuncVectorize.add() where I could. The logic is a bit contrived, but I've tried to make it clearer and better-explained in the comments.
- A rewrite of DeviceGUFuncVectorize.__call__() and the GUFuncCallSteps class. The code in these was so far off looking like it does what it actually does that I ended up evolving them to something that looks completely different. For review purposes, it probably makes the most sense to check out the code locally and read these implementations, and ignore the diff, because a side-by-side comparison is (IMHO) impossible to make sense of.
- This also forced a small change in _CUDAGufuncCallSteps, to remove the prepare_inputs() method, which shadowed the new method of the same name in GUFuncCallSteps, and didn't actually prepare any input - its functionality, which set self._stream, has been moved to the __init__() method, which still has an equivalent effect from that location - i.e. the _stream is set by the time it is required, and also not too soon.
Removal of the max_blocksize property of CUDA GUFuncs. This has always (literally, always) been ignored, so it makes no sense to keep it.
Some tidy-ups in the test_gufunc.py file in addition to the new tests for multiple outputs.

Usually I'd squash this down into a tidy individually-reviewable patch series, but because it was so hard to understand how things worked I have left the evolutionary changes I made in, because if there's a regression that will make it much easier to bisect - I generally kept things passing tests in each commit as I went.

CC @s-m-e in support of your Poliastro use cases.

Fixes #8303.

This only adds a check to prevent users providing an invalid signature - it does not affect functionality.

Also add a test for a CUDA multiple output gufunc now that it works.

In this tests the outputs are expected to have different values

Also add some more tests

gmarkall · 2022-08-11T15:16:25Z

gpuci run tests

gmarkall · 2022-11-17T16:51:36Z

gpuci run tests

gmarkall · 2022-12-02T11:24:52Z

gpuci run tests

gmarkall · 2022-12-02T11:56:14Z

A note to myself: the following patch was supposed to be a start at refactoring or simplifying the logic to more explicitly spell out what it does:

diff --git a/numba/np/ufunc/deviceufunc.py b/numba/np/ufunc/deviceufunc.py
index bccc2b39c..058c54a7f 100644
--- a/numba/np/ufunc/deviceufunc.py
+++ b/numba/np/ufunc/deviceufunc.py
@@ -731,7 +731,6 @@ class GUFuncCallSteps(object):
         'norm_inputs',
         'kernel_returnvalues',
         'kernel_parameters',
-        '_is_device_array',
         '_need_device_conversion',
     ]
 
@@ -756,17 +755,16 @@ class GUFuncCallSteps(object):
                 self.outputs.append(output)
             user_outputs_are_device.append(user_output_is_device)
 
-        self._is_device_array = [self.is_device_array(a) for a in self.args]
-        self._need_device_conversion = (not any(self._is_device_array) and
-                                        not any(user_outputs_are_device))
+        _is_device_array = [self.is_device_array(a) for a in self.args]
+        # - If any of the arguments are device arrays, we leave the output on
+        #   the device.
+        self._need_device_conversion = (not any(_is_device_array) and
+                                         not any(user_outputs_are_device))
+        print(f"Need device conversion: {self._need_device_conversion}")
 
         # Normalize inputs
-        inputs = []
-        for a, isdev in zip(self.args, self._is_device_array):
-            if isdev:
-                inputs.append(self.as_device_array(a))
-            else:
-                inputs.append(np.asarray(a))
+        inputs = [a if self.is_device_array(a) else np.asarray(a)
+                  for a in self.args]
         self.norm_inputs = inputs[:nin]
 
         # Check if there are extra arguments for outputs.
@@ -804,12 +802,8 @@ class GUFuncCallSteps(object):
         self.kernel_returnvalues = retvals
 
     def prepare_kernel_parameters(self):
-        params = []
-        for inp, isdev in zip(self.norm_inputs, self._is_device_array):
-            if isdev:
-                params.append(inp)
-            else:
-                params.append(self.to_device(inp))
+        params = [p if self.is_device_array(p) else self.to_device(p)
+                  for p in self.norm_inputs]
         assert all(self.is_device_array(a) for a in params)
         self.kernel_parameters = params

but it introduces a failure:

======================================================================
ERROR: test_gufunc_arg (numba.cuda.tests.cudapy.test_cuda_array_interface.TestCudaArrayInterface)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/gmarkall/numbadev/numba/numba/cuda/tests/cudapy/test_cuda_array_interface.py", line 110, in test_gufunc_arg
    out = vadd(arr, val)
  File "/home/gmarkall/numbadev/numba/numba/np/ufunc/deviceufunc.py", line 635, in __call__
    indtypes, schedule, outdtypes, kernel = self._schedule(
  File "/home/gmarkall/numbadev/numba/numba/np/ufunc/deviceufunc.py", line 647, in _schedule
    input_shapes = [a.shape for a in inputs]
  File "/home/gmarkall/numbadev/numba/numba/np/ufunc/deviceufunc.py", line 647, in <listcomp>
    input_shapes = [a.shape for a in inputs]
AttributeError: 'ForeignArray' object has no attribute 'shape'

gmarkall · 2023-02-22T09:34:14Z

gpuci run tests

Some small comments and structural edits, but mainly focused on being clearer about the logic deciding on whether to copy the result to the host.

The only work that the term `norm` is doing is making things more confusing.

gmarkall · 2023-02-22T11:17:17Z

gpuci run tests

gmarkall · 2023-03-15T12:29:43Z

gpuci run tests

gmarkall · 2023-03-15T14:08:43Z

@stuartarchibald Many thanks for the review - all comments should now be addressed so this should be ready for another look.

(PR numba#8341 feedback)

gmarkall · 2023-03-16T11:32:45Z

gpuci run tests

gmarkall · 2023-03-16T11:33:14Z

@stuartarchibald Many thanks for the round of review - I've addressed both comments, so this should be ready for another look.

stuartarchibald · 2023-03-17T11:32:46Z

@stuartarchibald Many thanks for the round of review - I've addressed both comments, so this should be ready for another look.

Many thanks, one minor suggestion RE string formatting else looks good.

PR numba#8341 feedback.

gmarkall · 2023-03-17T11:39:20Z

gpuci run tests

gmarkall · 2023-03-17T11:40:14Z

@stuartarchibald Many thanks for the suggestion, which I've amended slightly. Assuming gpuci is OK this should be OK for another look.

stuartarchibald

Thanks for the patch and fixes.

stuartarchibald · 2023-03-17T14:32:43Z

Buildfarm ID: numba_smoketest_cuda_yaml_189.

stuartarchibald · 2023-03-17T15:47:46Z

Buildfarm ID: numba_smoketest_cuda_yaml_189.

This has not yet completed, but the build job for win=64, py=3.9, np=1.23, CUDASIM=1, CUDA=11.1 had a "fatal exception" from File <__array_function__ internals>, line 180 in dot, from numba/cuda/tests/cudapy/test_matmul.py, line 67 in test_func. Have restarted this job to see if it does it again.

stuartarchibald · 2023-03-22T14:22:05Z

Buildfarm ID: numba_smoketest_cuda_yaml_189.

This has not yet completed, but the build job for win=64, py=3.9, np=1.23, CUDASIM=1, CUDA=11.1 had a "fatal exception" from File <__array_function__ internals>, line 180 in dot, from numba/cuda/tests/cudapy/test_matmul.py, line 67 in test_func. Have restarted this job to see if it does it again.

It passed on restart. Am trying another independent run:

Buildfarm ID: numba_smoketest_cuda_yaml_190.

stuartarchibald · 2023-03-22T17:44:16Z

Buildfarm ID: numba_smoketest_cuda_yaml_189.

This has not yet completed, but the build job for win=64, py=3.9, np=1.23, CUDASIM=1, CUDA=11.1 had a "fatal exception" from File <__array_function__ internals>, line 180 in dot, from numba/cuda/tests/cudapy/test_matmul.py, line 67 in test_func. Have restarted this job to see if it does it again.

It passed on restart. Am trying another independent run:

Buildfarm ID: numba_smoketest_cuda_yaml_190.

This passed.

stuartarchibald · 2023-03-22T18:01:54Z

Buildfarm ID: numba_smoketest_cuda_yaml_189.

This has not yet completed, but the build job for win=64, py=3.9, np=1.23, CUDASIM=1, CUDA=11.1 had a "fatal exception" from File <__array_function__ internals>, line 180 in dot, from numba/cuda/tests/cudapy/test_matmul.py, line 67 in test_func. Have restarted this job to see if it does it again.

It passed on restart. Am trying another independent run:
Buildfarm ID: numba_smoketest_cuda_yaml_190.

This passed.

Apologies, was looking at 189 again, 190 is still running, will post when it is done.

stuartarchibald · 2023-03-23T12:38:32Z

Buildfarm ID: numba_smoketest_cuda_yaml_189.

This has not yet completed, but the build job for win=64, py=3.9, np=1.23, CUDASIM=1, CUDA=11.1 had a "fatal exception" from File <__array_function__ internals>, line 180 in dot, from numba/cuda/tests/cudapy/test_matmul.py, line 67 in test_func. Have restarted this job to see if it does it again.

It passed on restart. Am trying another independent run:
Buildfarm ID: numba_smoketest_cuda_yaml_190.

This passed.

Apologies, was looking at 189 again, 190 is still running, will post when it is done.

Build numba_smoketest_cuda_yaml_190 has passed testing.

gmarkall · 2023-03-23T12:41:08Z

I've manually tested this on Windows without issue.

stuartarchibald · 2023-03-23T12:42:58Z

I've manually tested this on Windows without issue.

Many thanks for checking @gmarkall, this gives confidence that #8341 (comment) is/was unrelated. Will approve this now and keep a watch on future CUDA builds.

gmarkall added 10 commits August 11, 2022 10:14

Device GUFuncs: Enforce void or unspecified return type

b973ce2

This only adds a check to prevent users providing an invalid signature - it does not affect functionality.

Device ufuncs: remove long-deprecated argtypes / restype support

4a7b0eb

Device ufuncs: Make the kernelmap hold multiple output dtypes

507c352

Device ufuncs: return multiple output dtypes from schedule

a43cdfc

Device ufuncs: support callsteps holding multiple outputs

7a0a7bc

Device ufuncs: More generality for multiple outputs in callsteps

9ba159c

Device ufuncs: allow multiple kernel return values in callsteps

e927fc7

Device ufuncs: correct parameter passing for multiple outputs

e8b6c97

Also add a test for a CUDA multiple output gufunc now that it works.

CUDA ufuncs: another multiple output test

08e1835

In this tests the outputs are expected to have different values

Device ufuncs: fix allocation of multiple outputs

b1e2dd6

Also add some more tests

gmarkall requested review from sklam and stuartarchibald as code owners August 11, 2022 15:12

gmarkall added CUDA CUDA related issue/PR Effort - long Long size effort needed 2 - In Progress labels Aug 11, 2022

gmarkall added highpriority Effort - medium Medium size effort needed and removed Effort - long Long size effort needed labels Nov 11, 2022

Merge remote-tracking branch 'numba/main' into issue-8303-2

6acf970

gmarkall added 2 commits December 2, 2022 11:09

Merge remote-tracking branch 'numba/main' into issue-8303-2

a650d54

Skip ufunc multiple output tests on simulator

24592e4

Merge remote-tracking branch 'numba/main' into issue-8303-2

d3e16c4

gmarkall added 2 commits February 22, 2023 11:11

Device ufuncs: some refactors for clarity

575592f

Some small comments and structural edits, but mainly focused on being clearer about the logic deciding on whether to copy the result to the host.

Device ufuncs: rename norm_inputs to inputs

00da70f

The only work that the term `norm` is doing is making things more confusing.

gmarkall added 4 - Waiting on reviewer Waiting for reviewer to respond to author and removed 4 - Waiting on author Waiting for author to respond to review labels Mar 15, 2023

gmarkall added 2 commits March 16, 2023 11:28

CUDA GUFuncs: Ensure the correct number of args is passed

3fe2a03

CUDA GUfuncs: edit pos and kwarg argument exception message

1b13d1d

(PR numba#8341 feedback)

stuartarchibald reviewed Mar 17, 2023

View reviewed changes

Comment thread numba/np/ufunc/deviceufunc.py Outdated

Device ufuncs: simplify formatting of "positional argument[s]"

8c166f8

PR numba#8341 feedback.

stuartarchibald approved these changes Mar 17, 2023

View reviewed changes

stuartarchibald added 4 - Waiting on CI Review etc done, waiting for CI to finish Pending BuildFarm For PRs that have been reviewed but pending a push through our buildfarm and removed 4 - Waiting on reviewer Waiting for reviewer to respond to author labels Mar 17, 2023

sklam merged commit d951ebc into numba:main Mar 23, 2023

Conversation

gmarkall commented Aug 11, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gmarkall commented Aug 11, 2022

Uh oh!

gmarkall commented Nov 17, 2022

Uh oh!

gmarkall commented Dec 2, 2022

Uh oh!

gmarkall commented Dec 2, 2022

Uh oh!

gmarkall commented Feb 22, 2023

Uh oh!

gmarkall commented Feb 22, 2023

Uh oh!

gmarkall commented Mar 15, 2023

Uh oh!

gmarkall commented Mar 15, 2023

Uh oh!

gmarkall commented Mar 16, 2023

Uh oh!

gmarkall commented Mar 16, 2023

Uh oh!

Uh oh!

stuartarchibald commented Mar 17, 2023

Uh oh!

gmarkall commented Mar 17, 2023

Uh oh!

gmarkall commented Mar 17, 2023

Uh oh!

stuartarchibald left a comment

Choose a reason for hiding this comment

Uh oh!

stuartarchibald commented Mar 17, 2023

Uh oh!

stuartarchibald commented Mar 17, 2023

Uh oh!

stuartarchibald commented Mar 22, 2023

Uh oh!

stuartarchibald commented Mar 22, 2023

Uh oh!

stuartarchibald commented Mar 22, 2023

Uh oh!

stuartarchibald commented Mar 23, 2023

Uh oh!

gmarkall commented Mar 23, 2023

Uh oh!

stuartarchibald commented Mar 23, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

gmarkall commented Aug 11, 2022 •

edited

Loading