Fix mixed C ABI / Numba ABI internal-call lowering (#781, #789)

numba_cuda/numba/cuda/compiler.py

@@ -560,6 +560,8 @@ def compile_extra(
     locals,
     library=None,
     pipeline_class=CUDACompiler,
+    call_conv=None,
+    abi_info=None,
 ):
     """Compiler entry point
 
@@ -584,7 +586,15 @@ def compile_extra(
         compiler pipeline
     """
     pipeline = pipeline_class(
-        typingctx, targetctx, library, args, return_type, flags, locals
+        typingctx,
+        targetctx,
+        library,
+        args,
+        return_type,
+        flags,
+        locals,
+        call_conv,
+        abi_info,
     )
     return pipeline.compile_extra(func)
 
@@ -602,6 +612,8 @@ def compile_ir(
     is_lifted_loop=False,
     library=None,
     pipeline_class=CUDACompiler,
+    call_conv=None,
+    abi_info=None,
 ):
     """
     Compile a function with the given IR.
@@ -637,6 +649,8 @@ def compile_local(the_ir, the_flags):
                 return_type,
                 the_flags,
                 locals,
+                call_conv,
+                abi_info,
             )
             return pipeline.compile_ir(
                 func_ir=the_ir, lifted=lifted, lifted_from=lifted_from
@@ -675,13 +689,30 @@ def compile_local(the_ir, the_flags):
 
 
 def compile_internal(
-    typingctx, targetctx, library, func, args, return_type, flags, locals
+    typingctx,
+    targetctx,
+    library,
+    func,
+    args,
+    return_type,
+    flags,
+    locals,
+    call_conv=None,
+    abi_info=None,
 ):
     """
     For internal use only.
     """
     pipeline = CUDACompiler(
-        typingctx, targetctx, library, args, return_type, flags, locals
+        typingctx,
+        targetctx,
+        library,
+        args,
+        return_type,
+        flags,
+        locals,
+        call_conv,
+        abi_info,
     )
     return pipeline.compile_extra(func)
 
@@ -744,10 +775,12 @@ def compile_cuda(
     flags.lto = lto
 
     if abi == "c":
-        flags.call_conv = CUDACABICallConv(targetctx)
+        call_conv = CUDACABICallConv(targetctx)
+    else:
+        call_conv = CUDACallConv(targetctx)
 
-    if abi_info is not None:
-        flags.abi_info = abi_info
+    if abi_info is None:
+        abi_info = {}
 
     with utils.numba_target_override():
         cres = compile_extra(
@@ -759,6 +792,8 @@ def compile_cuda(
             flags=flags,
             locals={},
             pipeline_class=CUDACompiler,
+            call_conv=call_conv,
+            abi_info=abi_info,
         )
 
     library = cres.library

numba_cuda/numba/cuda/core/base.py

@@ -961,8 +961,9 @@ def call_internal(self, builder, fndesc, sig, args):
         status, res = self.call_internal_no_propagate(
             builder, fndesc, sig, args
         )
-        with cgutils.if_unlikely(builder, status.is_error):
-            fndesc.call_conv.return_status_propagate(builder, status)
+        if status is not None:
+            with cgutils.if_unlikely(builder, status.is_error):
+                fndesc.call_conv.return_status_propagate(builder, status)
 
         res = imputils.fix_returning_optional(self, builder, sig, status, res)
         return res

numba_cuda/numba/cuda/core/compiler.py

@@ -84,7 +84,16 @@ class CompilerBase:
     """
 
     def __init__(
-        self, typingctx, targetctx, library, args, return_type, flags, locals
+        self,
+        typingctx,
+        targetctx,
+        library,
+        args,
+        return_type,
+        flags,
+        locals,
+        call_conv=None,
+        abi_info=None,
     ):
         # Make sure the environment is reloaded
         config.reload_config()
@@ -116,6 +125,14 @@ def __init__(
         # hold this for e.g. with_lifting, null out on exit
         self.state.pipeline = self
 
+        if call_conv is None:
+            call_conv = CUDACallConv(self.state.targetctx)
+        if abi_info is None:
+            abi_info = {}
+
+        self.state.call_conv = call_conv
+        self.state.abi_info = abi_info
+
         self.state.status = _CompileStatus(
             can_fallback=self.state.flags.enable_pyobject
         )

numba_cuda/numba/cuda/core/typed_passes.py

@@ -326,7 +326,7 @@ def run_pass(self, state):
         metadata = state.metadata
         pre_stats = passmanagers.dump_refprune_stats()
 
-        call_conv = flags.call_conv
+        call_conv = state.call_conv
         if call_conv is None:
             call_conv = CUDACallConv(state.targetctx)
 
@@ -348,7 +348,7 @@ def run_pass(self, state):
                     noalias=flags.noalias,
                     abi_tags=[flags.get_mangle_string()],
                     call_conv=call_conv,
-                    abi_info=flags.abi_info,
+                    abi_info=state.abi_info,
                 )
             )
 

numba_cuda/numba/cuda/flags.py

@@ -163,24 +163,6 @@ def _optional_int_type(x):
         return x
 
 
-def _call_conv_options_type(x):
-    if x is None:
-        return None
-
-    else:
-        assert isinstance(x, BaseCallConv)
-        return x
-
-
-def _abi_info_options_type(x):
-    if x is None:
-        return {}
-
-    else:
-        assert isinstance(x, dict)
-        return x
-
-
 class CUDAFlags(Flags):
     nvvm_options = Option(
         type=_nvvm_options_type,
@@ -196,7 +178,3 @@ class CUDAFlags(Flags):
         type=_optional_int_type, default=None, doc="Max registers"
     )
     lto = Option(type=bool, default=False, doc="Enable Link-time Optimization")
-
-    call_conv = Option(type=_call_conv_options_type, default=None, doc="")
-
-    abi_info = Option(type=_abi_info_options_type, default=None, doc="ABI info")

numba_cuda/numba/cuda/simulator/compiler.py

@@ -30,7 +30,16 @@ def define_typed_pipeline(state, name="typed"):
 
 class CompilerBase:
     def __init__(
-        self, typingctx, targetctx, library, args, return_type, flags, locals
+        self,
+        typingctx,
+        targetctx,
+        library,
+        args,
+        return_type,
+        flags,
+        locals,
+        call_conv,
+        abi_info,
     ):
         pass
 

numba_cuda/numba/cuda/target.py

@@ -11,6 +11,7 @@
 
 from numba.cuda import types
 from numba.cuda import HAS_NUMBA
+from numba.cuda.core.callconv import CUDACallConv
 from numba.cuda.core.compiler_lock import global_compiler_lock
 from numba.cuda.core.errors import NumbaWarning
 from numba.cuda.core.base import BaseContext
@@ -411,6 +412,10 @@ def _compile_subroutine_no_cache(
             flags.no_cpython_wrapper = True
             flags.no_cfunc_wrapper = True
 
+            # compile_subroutine always uses CUDACallConv
+            call_conv = CUDACallConv(self)
+            abi_info = {}
+
             cres = compiler.compile_internal(
                 self.typing_context,
                 self,
@@ -420,6 +425,8 @@ def _compile_subroutine_no_cache(
                 sig.return_type,
                 flags,
                 locals=locals,
+                call_conv=call_conv,
+                abi_info=abi_info,
             )
 
             # Allow inlining the function inside callers

numba_cuda/numba/cuda/tests/cudapy/test_compiler.py

@@ -4,6 +4,8 @@
 import os
 from math import sqrt
 from numba import cuda
+from numba.core.extending import intrinsic
+
 from numba.cuda import float32, int16, int32, int64, types, uint32, void
 from numba.cuda import (
     compile,
@@ -16,6 +18,8 @@
 from numba.cuda.cudadrv import nvrtc
 from numba.cuda.testing import skip_on_cudasim, unittest, CUDATestCase
 
+from numba.cuda.core.callconv import CUDACallConv
+
 TEST_BIN_DIR = os.getenv("NUMBA_CUDA_TEST_BIN_DIR")
 if TEST_BIN_DIR:
     test_device_functions_a = os.path.join(
@@ -701,6 +705,69 @@ def f(z, x, y):
             str(code_list[1].code.decode()), r"\.section\s+\.debug_info"
         )
 
+    def test_compile_jitted_subroutine(self):
+        # Reproducer from gh-781
+        # https://github.com/NVIDIA/numba-cuda/issues/781
+        def foo(x):
+            return 2 * x
+
+        # Create a wrapper that takes void* arguments
+        def create_void_ptr_wrapper():
+            """Create a wrapper that takes void* input and output pointers."""
+
+            # Make foo a device function
+            foo_device = cuda.jit(device=True)(foo)
+
+            # The inner signature: int32 -> int32
+            inner_sig = types.int32(types.int32)
+
+            # The wrapper signature: void(void*, void*) - input ptr, output ptr
+            wrapper_sig = types.void(types.voidptr, types.voidptr)
+
+            @intrinsic
+            def wrapper_impl(typingctx, arg0, arg1):
+                def codegen(context, builder, sig, args):
+                    input_ptr, output_ptr = args
+
+                    # Cast input void* to int32*, load value
+                    int32_llvm_type = context.get_value_type(types.int32)
+                    typed_input_ptr = builder.bitcast(
+                        input_ptr, int32_llvm_type.as_pointer()
+                    )
+                    input_val = builder.load(typed_input_ptr)
+
+                    # Call the inner function
+                    cres = context.compile_subroutine(
+                        builder, foo_device, inner_sig, caching=False
+                    )
+
+                    # Wrapper function is compiled with cabi, but inner function
+                    # is compiled with numba-abi. So cres should have CUDACallConv.
+                    assert isinstance(cres.fndesc.call_conv, CUDACallConv)
+
+                    _, result = context.call_internal_no_propagate(
+                        builder, cres.fndesc, inner_sig, [input_val]
+                    )
+
+                    # Cast output void* to int32*, store result
+                    typed_output_ptr = builder.bitcast(
+                        output_ptr, int32_llvm_type.as_pointer()
+                    )
+                    builder.store(result, typed_output_ptr)
+
+                    return context.get_dummy_value()
+
+                return wrapper_sig, codegen
+
+            def wrapper_func(input_ptr, output_ptr):
+                return wrapper_impl(input_ptr, output_ptr)
+
+            return wrapper_func, wrapper_sig
+
+        wrapper, wrapper_sig = create_void_ptr_wrapper()
+
+        cuda.compile(wrapper, wrapper_sig.args, output="ltoir")
+
 
 @skip_on_cudasim("Compilation unsupported in the simulator")
 class TestCompileForCurrentDevice(CUDATestCase):

numba_cuda/numba/cuda/tests/cudapy/test_ir_utils.py

@@ -57,6 +57,7 @@ def mk_pipeline(
                     from numba.cuda.descriptor import cuda_target
 
                     target_context = cuda_target.target_context
+
                 return cls(
                     typing_context,
                     target_context,

numba_cuda/numba/cuda/typing/templates.py

@@ -689,6 +689,7 @@ def generic(self, args, kws):
             # spoof a compiler pipline like the one that will be in use
             tyctx = fcomp.targetdescr.typing_context
             tgctx = fcomp.targetdescr.target_context
+
             compiler_inst = fcomp.pipeline_class(
                 tyctx,
                 tgctx,