[Refactor][NFC][Cleanups] Update imports to upstream numba to use the numba.cuda modules

[atmnp]

This PR updates the vast majority of imports to upstream numba modules, that had already been vendored in for future CUDA-specific changes. After this PR, there should be exactly 3 (not strictly necessary) upstream numba imports remaining - one in testing vectorization for in test_math, one in random.py, and the global compiler_lock.

It seems from the existing comments that removing the CPU jit in random.py would result in considerable performance drops, so I've just added a guard for numba being available.

There are modules such as numba.core.misc.{mergesort,quicksort} that had not yet been vendored in, so this PR also vendors those modules in. They had to be modified to not use the CPU jit, since these are not necessarily supported within numba-cuda. It's still unclear to me if this is necessary or is done correctly, since we don't have tests for these yet. My impression is that these are to support jitting of mergesort/quicksort on arrays, which I am fairly confident we cannot support on the CUDA target until more development is done on the Numba-CUDA runtime. I am unsure if we can remove these uses but I would prefer if we did.

There are many tests that rely on CPU jitting, these tests are now guarded to not run in environments where the numba package is not available. This is done using a global importlib check in numba-cuda's __init__.py. This cannot be moved into numba.cuda.config since the upstream numba doesn't have this flag, and we intend on forwarding calls to upstream numba when it is available. All previously sporadic guards for this should have been updated to use this global flag now.

Support for prange, pndindex, and stencils has been removed, likewise for typed.Dict, and typed.List. There was a pass that relied on this, which was the MakeFunctionIntoJitFunction pass, which implicitly converts functions to jitted functions, which is also a use case that is not currently tested. Also, this removes all mention of the legacy type system flag since we implicitly assume the legacy type system in place.

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[atmnp]

/ok to test

[atmnp]

/ok to test

[atmnp]

/ok to test

[atmnp]

/ok to test

[gmarkall]

We can't delete the MakeFunctionToJitFunction pass - it is load-bearing, but our testsuite has insufficient coverage (as noted in previous situations).

An example that removing the pass breaks is:

from numba import cuda
import numpy as np

N = 10

def factory(consumer_func):
    @cuda.jit
    def func():
        def inner():
            return N
        return consumer_func(inner)
    return func


@cuda.jit
def consumer(func):
    return func()

jitted_func = factory(consumer)

@cuda.jit
def kernel(out):
    out[0] = jitted_func()


out = np.zeros(1, dtype=np.uint64)
kernel[1, 1](out)
assert out[0] == N

which runs fine on main, and here results in:

ailed in cuda mode pipeline (step: nopython frontend)
Invalid use of Literal[Expr](make_function(name=None, code=<code object inner at 0x79a49fae2800, file "/home/gmarkall/numbadev/issues/numba-cuda-561/example.py", line 9>, closure=None, defaults=None)) with parameters ()
No type info available for Literal[Expr](make_function(name=None, code=<code object inner at 0x79a49fae2800, file "/home/gmarkall/numbadev/issues/numba-cuda-561/example.py", line 9>, closure=None, defaults=None)) as a callable.
During: resolving callee type: Literal[Expr](make_function(name=None, code=<code object inner at 0x79a49fae2800, file "/home/gmarkall/numbadev/issues/numba-cuda-561/example.py", line 9>, closure=None, defaults=None))
During: typing of call at /home/gmarkall/numbadev/issues/numba-cuda-561/example.py (17)


File "example.py", line 17:
def consumer(func):
    return func()
    ^

During: Pass nopython_type_inference
During: resolving callee type: type(CUDADispatcher(<function consumer at 0x79a49fb03240>))
During: typing of call at /home/gmarkall/numbadev/issues/numba-cuda-561/example.py (11)


File "example.py", line 11:
        def inner():
            <source elided>
            return N
        return consumer_func(inner)
        ^

During: Pass nopython_type_inference
During: resolving callee type: type(CUDADispatcher(<function factory.<locals>.func at 0x79a3e1e44860>))
During: typing of call at /home/gmarkall/numbadev/issues/numba-cuda-561/example.py (23)


File "example.py", line 23:
def kernel(out):
    out[0] = jitted_func()
    ^

During: Pass nopython_type_inference

there is a test, test_make_function_to_jitted_function, that I have half-completed porting to Numba-CUDA, and will add in another PR.

[gmarkall]

Whilst looking at the impact of moving MakeFunctionToJitFunction, I noticed that the following:

from numba import cuda
 # ...
kernel[1, 1]((1, 2, 3))

cuda.synchronize()

results in

...
  File "/home/gmarkall/numbadev/numba/numba/core/typing/templates.py", line 749, in _get_jit_decorator
    jitter = jit_registry[target_hw]
             ~~~~~~~~~~~~^^^^^^^^^^^
  File "/home/gmarkall/numbadev/numba/numba/core/registry.py", line 71, in __getitem__
    return super(DelayedRegistry, self).__getitem__(item)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: <class 'numba.core.target_extension.CUDA'>

so there's a bug impacting some of the code I believe this PR touches that needs fixing. I'm looking into a fix for this.

Fix for the target registration is in #566

[gmarkall]

/ok to test

[atmnp]

/ok to test