Vendor the imputils module for CUDA refactoring

[ashermancinelli]

No modules import this module in numba-cuda yet, but other modules that have yet to be vendored will use it. I see no standalone tests for this module. It would be nice to vendor this module so it is already in the repo in case we would like to vendor any of the following modules for CUDA-specific refactoring.

These are the places it is being imported in upstream numba:

Click to expand

numba/typed/dictobject.py
numba/typed/listobject.py
numba/typed/typeddict.py
numba/typed/typedlist.py
numba/typed/dictimpl.py
numba/stencils/stencil.py
numba/parfors/parfor.py
numba/cpython/slicing.py
numba/cuda/vector_types.py
numba/cpython/new_numbers.py
numba/cpython/new_mathimpl.py
numba/cpython/old_tupleobj.py
numba/cpython/new_tupleobj.py
numba/cuda/mathimpl.py
numba/cpython/randomimpl.py
numba/cpython/old_numbers.py
numba/cpython/printimpl.py
numba/cpython/old_mathimpl.py
numba/cpython/cmathimpl.py
numba/misc/cffiimpl.py
numba/cpython/iterators.py
numba/cpython/setobj.py
numba/cuda/libdeviceimpl.py
numba/cpython/old_builtins.py
numba/cpython/enumimpl.py
numba/cuda/cudaimpl.py
numba/cpython/new_builtins.py
numba/cpython/unicode_support.py
numba/tests/test_target_extension.py
numba/cpython/listobj.py
numba/cpython/rangeobj.py
numba/cuda/printimpl.py
numba/cpython/unicode.py
numba/experimental/structref.py
numba/experimental/function_type.py
numba/experimental/jitclass/base.py
numba/core/base.py
numba/core/optional.py
numba/core/inline_closurecall.py
numba/tests/pdlike_usecase.py
numba/core/extending.py
numba/tests/test_extending.py
numba/np/npyfuncs.py
numba/np/linalg.py
numba/np/old_arraymath.py
numba/np/unsafe/ndarray.py
numba/np/arrayobj.py
numba/np/new_arraymath.py
numba/np/npyimpl.py
numba/tests/test_practical_lowering_issues.py
numba/np/npdatetime.py
numba/np/math/mathimpl.py
numba/np/math/cmathimpl.py
numba/np/math/numbers.py
numba/tests/test_nrt.py
numba/tests/test_target_overloadselector.py

Depends on at least the following modules:

• extending module (no PR at time of writing)
• numba.np module
• Vendor the inline-closure-call module #443
• Vendor the base module for CUDA-specific refactoring #445

[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.

[ashermancinelli]

/ok to test e07ff46

[VijayKandiah]

No modules import this module in numba-cuda yet, but other modules that have yet to be vendored will use it. I see no standalone tests for this module. It would be nice to vendor this module so it is already in the repo in case we would like to vendor any of the following modules for CUDA-specific refactoring.

I do see some places this module is currently being imported:

$ rg -i imputils
numba/cuda/cudaimpl.py
9:from numba.core.imputils import Registry

numba/cuda/extending.py
25:    from numba.core.imputils import impl_ret_borrowed

numba/cuda/_internal/cuda_fp16.py
26:from numba.core.imputils import Registry as TargetRegistry, lower_cast

numba/cuda/libdeviceimpl.py
4:from numba.core.imputils import Registry

numba/cuda/mathimpl.py
6:from numba.core.imputils import Registry

numba/cuda/printimpl.py
6:from numba.core.imputils import Registry

numba/cuda/vector_types.py
9:from numba.core.imputils import Registry as ImplRegistry

[ashermancinelli]

Thanks! Not sure how I missed those.

[ashermancinelli]

/ok to test 38a84c0

[gmarkall]

The two PRs highlighted as dependencies of this PR have now been merged - does this unblock it, or are there new dependencies that will need to be covered?

[ashermancinelli]

After updating, I found additional dependencies so this will have to remain blocked for now. Thanks Graham!

[ashermancinelli]

/ok to test f40b4bd

[ashermancinelli]

The failing test is from thirdparty cudf tests:

NotImplementedError: No definition for lowering <class 'cudf.core.udf.api.Masked'>(int64, bool) -> Masked(int64)

details

____________ test_masked_udf_scalar_args_binops_multiple[ge-data0] _____________
Traceback (most recent call last):
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/cudf/core/indexed_frame.py", line 3515, in _apply
    kernel, retty = kr.get_kernel()
                    ^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/cudf/core/udf/udf_kernel_base.py", line 167, in get_kernel
    return self._compile_or_get_kernel()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/cudf/core/udf/udf_kernel_base.py", line 184, in _compile_or_get_kernel
    kernel, scalar_return_type = self.compile_kernel()
                                 ^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/cudf/core/udf/udf_kernel_base.py", line 145, in compile_kernel
    kernel = self.compile_kernel_string(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/cudf/core/udf/udf_kernel_base.py", line 159, in compile_kernel_string
    kernel = cuda.jit(
             ^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/decorators.py", line 207, in _jit
    disp.compile(argtypes)
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba/core/compiler_lock.py", line 35, in _acquire_compile_lock
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/dispatcher.py", line 1894, in compile
    kernel = _Kernel(self.py_func, argtypes, **self.targetoptions)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba/core/compiler_lock.py", line 35, in _acquire_compile_lock
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/dispatcher.py", line 148, in __init__
    cres = compile_cuda(
           ^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba/core/compiler_lock.py", line 35, in _acquire_compile_lock
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/compiler.py", line 750, in compile_cuda
    cres = compile_extra(
           ^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/compiler.py", line 591, in compile_extra
    return pipeline.compile_extra(func)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/core/compiler.py", line 124, in compile_extra
    return self._compile_bytecode()
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/core/compiler.py", line 192, in _compile_bytecode
    return self._compile_core()
           ^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/core/compiler.py", line 168, in _compile_core
    raise e
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/core/compiler.py", line 160, in _compile_core
    pm.run(self.state)
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/core/compiler_machinery.py", line 393, in run
    raise e
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/core/compiler_machinery.py", line 386, in run
    self._runPass(idx, pass_inst, state)
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba/core/compiler_lock.py", line 35, in _acquire_compile_lock
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/core/compiler_machinery.py", line 338, in _runPass
    mutated |= check(pss.run_pass, internal_state)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/core/compiler_machinery.py", line 292, in check
    mangled = func(compiler_state)
              ^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/core/typed_passes.py", line 355, in run_pass
    lower.lower()
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/lowering.py", line 230, in lower
    self.lower_normal_function(self.fndesc)
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/lowering.py", line 269, in lower_normal_function
    entry_block_tail = self.lower_function_body()
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/lowering.py", line 300, in lower_function_body
    self.lower_block(block)
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/lowering.py", line 318, in lower_block
    self.lower_inst(inst)
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/lowering.py", line 518, in lower_inst
    val = self.lower_assign(ty, inst)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/lowering.py", line 746, in lower_assign
    return self.lower_expr(ty, value)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/lowering.py", line 1285, in lower_expr
    res = self.lower_call(resty, expr)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/lowering.py", line 1029, in lower_call
    res = self._lower_call_normal(fnty, expr, signature)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/lowering.py", line 1247, in _lower_call_normal
    impl = self.context.get_function(fnty, signature)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/core/base.py", line 575, in get_function
    return self.get_function(fn, sig, _firstcall=False)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/3.12.11/lib/python3.12/site-packages/numba_cuda/numba/cuda/core/base.py", line 577, in get_function
    raise NotImplementedError(
NotImplementedError: No definition for lowering <class 'cudf.core.udf.api.Masked'>(int64, bool) -> Masked(int64)

I'm investigating.

[brandon-b-miller]

As mentioned offline I believe this may be the source of the test failures. I'll try and confirm my suspicion here and think about what to to do beyond eventually updating cuDF.

[ashermancinelli]

I'm wondering if we need to add the registry here:

# context.py
class Context(BaseContext):
    # This list will be extended to include all the registries
    # that are needed for CUDA
    def load_additional_registries(self):

Once I can repro the cudf tests locally I'll try some things out. Thanks for taking a look!

[ashermancinelli]

Reproducing locally, and adding numba's registries to CUDATargetContext::load_additional_registries() is allowing these tests to pass. I might be missing other registries, but the cudf tests that were failing that I have retested are all passing now.

[ashermancinelli]

/ok to test e6b6bd2

[ashermancinelli]

/ok to test 6253827

[gmarkall]

This looks good.

I think the workaround installing the upstream registry is probably the best possible workaround; I feel a little nervous that it might somehow cause issues, but as I can't predict the nature of them, I think we'll just have to address them if / when they occur.