Merge remote-tracking branch 'NVIDIA/main' into fix-dump-llvm

Author: gmarkall

.github/workflows/pr.yaml

@@ -25,10 +25,10 @@ jobs:
       - test-wheels-pynvjitlink
       - build-docs
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-24.08
+    uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-25.04
   checks:
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@branch-24.08
+    uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@branch-25.04
     with:
       enable_check_generated_files: false
   compute-matrix:

.github/workflows/publish.yaml

@@ -41,7 +41,26 @@ jobs:
       build_type: release
       script: "ci/build_wheel.sh"
       matrix_filter: ${{ needs.compute-matrix.outputs.BUILD_MATRIX }}
-      upload_to_pypi: true
+  publish-wheels:
+    needs: build-wheels
+    runs-on: ubuntu-latest
+    steps:
+      - name: Download sdist
+        if: "!cancelled()"
+        uses: actions/download-artifact@v4
+        with:
+          name: sdist
+          path: dist
+      - name: Download wheel
+        if: "!cancelled()"
+        uses: actions/download-artifact@v4
+        with:
+          name: wheel
+          path: dist
+      - name: Publish wheel
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.RAPIDSAI_PYPI_TOKEN }}
   build-docs:
     needs:
       - build-conda

.github/workflows/wheels-build.yaml

@@ -45,10 +45,6 @@ on:
         type: string
         default: ''
 
-      upload_to_pypi:
-        type: boolean
-        required: false
-
 defaults:
   run:
     shell: bash
@@ -173,9 +169,3 @@ jobs:
         with:
           name: wheel
           path: ${{ env.wheel_path }}
-
-      - name: Publish wheel
-        if: inputs.upload_to_pypi
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          password: ${{ secrets.RAPIDSAI_PYPI_TOKEN }}

ci/build_conda.sh

@@ -13,7 +13,7 @@ rapids-print-env
 
 rapids-logger "Begin py build"
 
-rapids-conda-retry mambabuild conda/recipes/numba-cuda
+rapids-conda-retry build conda/recipes/numba-cuda
 
 package_path=(/tmp/conda-bld-output/noarch/numba-cuda-*.tar.bz2)
 echo "package_path=$package_path" >> $GITHUB_ENV

docs/source/user/cuda_array_interface.rst

@@ -513,6 +513,7 @@ The following Python libraries have adopted the CUDA Array Interface:
 - `ArrayViews <https://github.com/xnd-project/arrayviews>`_
 - `JAX <https://jax.readthedocs.io/en/latest/index.html>`_
 - `PyCUDA <https://documen.tician.de/pycuda/tutorial.html#interoperability-with-other-libraries-using-the-cuda-array-interface>`_
+- `Paddle <https://www.paddlepaddle.org.cn/>`_
 - `DALI: the NVIDIA Data Loading Library <https://github.com/NVIDIA/DALI>`_ :
 
     - `TensorGPU objects

numba_cuda/VERSION

@@ -1 +1 @@
-0.5.0
+0.6.0

numba_cuda/numba/cuda/dispatcher.py

@@ -4,8 +4,9 @@
 import sys
 import ctypes
 import functools
+from collections import defaultdict
 
-from numba.core import config, serialize, sigutils, types, typing, utils
+from numba.core import config, ir, serialize, sigutils, types, typing, utils
 from numba.core.caching import Cache, CacheImpl
 from numba.core.compiler_lock import global_compiler_lock
 from numba.core.dispatcher import Dispatcher
@@ -42,6 +43,55 @@
 reshape_funcs = ['nocopy_empty_reshape', 'numba_attempt_nocopy_reshape']
 
 
+def get_cres_link_objects(cres):
+    """Given a compile result, return a set of all linkable code objects that
+    are required for it to be fully linked."""
+
+    link_objects = set()
+
+    # List of calls into declared device functions
+    device_func_calls = [
+        (name, v) for name, v in cres.fndesc.typemap.items() if (
+            isinstance(v, cuda_types.CUDADispatcher)
+        )
+    ]
+
+    # List of tuples with SSA name of calls and corresponding signature
+    call_signatures = [
+        (call.func.name, sig)
+        for call, sig in cres.fndesc.calltypes.items() if (
+            isinstance(call, ir.Expr) and call.op == 'call'
+        )
+    ]
+
+    # Map SSA names to all invoked signatures
+    call_signature_d = defaultdict(list)
+    for name, sig in call_signatures:
+        call_signature_d[name].append(sig)
+
+    # Add the link objects from the current function's callees
+    for name, v in device_func_calls:
+        for sig in call_signature_d.get(name, []):
+            called_cres = v.dispatcher.overloads[sig.args]
+            called_link_objects = get_cres_link_objects(called_cres)
+            link_objects.update(called_link_objects)
+
+    # From this point onwards, we are only interested in ExternFunction
+    # declarations - these are the calls made directly in this function to
+    # them.
+    for name, v in cres.fndesc.typemap.items():
+        if not isinstance(v, Function):
+            continue
+
+        if not isinstance(v.typing_key, ExternFunction):
+            continue
+
+        for obj in v.typing_key.link:
+            link_objects.add(obj)
+
+    return link_objects
+
+
 class _Kernel(serialize.ReduceMixin):
     '''
     CUDA Kernel specialized for a given set of argument types. When called, this
@@ -159,15 +209,8 @@ def link_to_library_functions(library_functions, library_path,
 
         self.maybe_link_nrt(link, tgt_ctx, asm)
 
-        for k, v in cres.fndesc.typemap.items():
-            if not isinstance(v, Function):
-                continue
-
-            if not isinstance(v.typing_key, ExternFunction):
-                continue
-
-            for obj in v.typing_key.link:
-                lib.add_linking_file(obj)
+        for obj in get_cres_link_objects(cres):
+            lib.add_linking_file(obj)
 
         for filepath in link:
             lib.add_linking_file(filepath)

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

@@ -205,6 +205,14 @@ def rgba_caller(x, channels):
 }
 """)
 
+times3_cu = cuda.CUSource("""
+extern "C" __device__
+int times3(int *out, int a)
+{
+  *out = a * 3;
+  return 0;
+}
+""")
 
 times4_cu = cuda.CUSource("""
 extern "C" __device__
@@ -351,6 +359,123 @@ def kernel(x, seed):
         kernel[1, 1](x, 1)
         np.testing.assert_equal(x[0], 323845807)
 
+    def test_declared_in_called_function(self):
+        times2 = cuda.declare_device('times2', 'int32(int32)', link=times2_cu)
+
+        @cuda.jit
+        def device_func(x):
+            return times2(x)
+
+        @cuda.jit
+        def kernel(r, x):
+            i = cuda.grid(1)
+            if i < len(r):
+                r[i] = device_func(x[i])
+
+        x = np.arange(10, dtype=np.int32)
+        r = np.empty_like(x)
+
+        kernel[1, 32](r, x)
+
+        np.testing.assert_equal(r, x * 2)
+
+    def test_declared_in_called_function_twice(self):
+        times2 = cuda.declare_device('times2', 'int32(int32)', link=times2_cu)
+
+        @cuda.jit
+        def device_func_1(x):
+            return times2(x)
+
+        @cuda.jit
+        def device_func_2(x):
+            return device_func_1(x)
+
+        @cuda.jit
+        def kernel(r, x):
+            i = cuda.grid(1)
+            if i < len(r):
+                r[i] = device_func_2(x[i])
+
+        x = np.arange(10, dtype=np.int32)
+        r = np.empty_like(x)
+
+        kernel[1, 32](r, x)
+
+        np.testing.assert_equal(r, x * 2)
+
+    def test_declared_in_called_function_two_calls(self):
+        times2 = cuda.declare_device('times2', 'int32(int32)', link=times2_cu)
+
+        @cuda.jit
+        def device_func(x):
+            return times2(x)
+
+        @cuda.jit
+        def kernel(r, x):
+            i = cuda.grid(1)
+            if i < len(r):
+                r[i] = device_func(x[i]) + device_func(x[i] + i)
+
+        x = np.arange(10, dtype=np.int32)
+        r = np.empty_like(x)
+
+        kernel[1, 32](r, x)
+
+        np.testing.assert_equal(r, x * 6)
+
+    def test_call_declared_function_twice(self):
+        times2 = cuda.declare_device('times2', 'int32(int32)', link=times2_cu)
+
+        @cuda.jit
+        def kernel(r, x):
+            i = cuda.grid(1)
+            if i < len(r):
+                r[i] = times2(x[i]) + times2(x[i] + i)
+
+        x = np.arange(10, dtype=np.int32)
+        r = np.empty_like(x)
+
+        kernel[1, 32](r, x)
+
+        np.testing.assert_equal(r, x * 6)
+
+    def test_declared_in_called_function_and_parent(self):
+        times2 = cuda.declare_device('times2', 'int32(int32)', link=times2_cu)
+
+        @cuda.jit
+        def device_func(x):
+            return times2(x)
+
+        @cuda.jit
+        def kernel(r, x):
+            i = cuda.grid(1)
+            if i < len(r):
+                r[i] = device_func(x[i]) + times2(x[i])
+
+        x = np.arange(10, dtype=np.int32)
+        r = np.empty_like(x)
+
+        kernel[1, 32](r, x)
+
+        np.testing.assert_equal(r, x * 4)
+
+    def test_call_two_different_declared_functions(self):
+        times2 = cuda.declare_device('times2', 'int32(int32)', link=times2_cu)
+        times3 = cuda.declare_device('times3', 'int32(int32)', link=times3_cu)
+
+        @cuda.jit
+        def kernel(r, x):
+            i = cuda.grid(1)
+            if i < len(r):
+                r[i] = times2(x[i]) + times3(x[i])
+
+        x = np.arange(10, dtype=np.int32)
+        r = np.empty_like(x)
+
+        kernel[1, 32](r, x)
+
+        np.testing.assert_equal(r, x * 5)
+
 
 if __name__ == '__main__':
     unittest.main()