Various fixups

ci/test_conda.sh

@@ -36,7 +36,6 @@ rapids-mamba-retry install \
 
 rapids-logger "Build Tests"
 pushd test_binary_generation
-export GPU_CC=`nvidia-smi --query-gpu=compute_cap --format=csv | grep -v compute_cap | sed 's/\.//'`
 make
 popd
 

ci/test_wheel.sh

@@ -26,7 +26,6 @@ python -m pip install $(echo ./dist/pynvjitlink*.whl)
 
 rapids-logger "Build Tests"
 pushd test_binary_generation
-export GPU_CC=`nvidia-smi --query-gpu=compute_cap --format=csv | grep -v compute_cap | sed 's/\.//'`
 make
 popd
 

examples/jit_link_data.py

@@ -1,3 +1,13 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+# Demonstrates the use of a CUSource object to link code to a @cuda.jit
+# function where the linked source is supplied as a string in memory, rather
+# than on-disk. The CUSource object is passed in the `link` list, just as paths
+# to files are passed when linking source files from disk.
+#
+# In addition to CUSource files, PTXSource objects can be used to link PTX from
+# memory (not shown in this example).
+
 from numba import cuda
 from pynvjitlink import patch
 

pynvjitlink/patch.py

@@ -44,6 +44,9 @@
 
 
 class LinkableCode:
+    """An object that can be passed in the `link` list argument to `@cuda.jit`
+    kernels to supply code to be linked from memory."""
+
     def __init__(self, data, name=None):
         self.data = data
         self._name = name
@@ -54,31 +57,43 @@ def name(self):
 
 
 class PTXSource(LinkableCode):
+    """PTX Source code in memory"""
+
     kind = FILE_EXTENSION_MAP["ptx"]
     default_name = "<unnamed-ptx>"
 
 
 class CUSource(LinkableCode):
+    """CUDA C/C++ Source code in memory"""
+
     kind = "cu"
     default_name = "<unnamed-cu>"
 
 
 class Fatbin(LinkableCode):
+    """A fatbin ELF in memory"""
+
     kind = FILE_EXTENSION_MAP["fatbin"]
     default_name = "<unnamed-fatbin>"
 
 
 class Cubin(LinkableCode):
+    """A cubin ELF in memory"""
+
     kind = FILE_EXTENSION_MAP["cubin"]
     default_name = "<unnamed-cubin>"
 
 
 class Archive(LinkableCode):
+    """An archive of objects in memory"""
+
     kind = FILE_EXTENSION_MAP["a"]
     default_name = "<unnamed-archive>"
 
 
 class Object(LinkableCode):
+    """An object file in memory"""
+
     kind = FILE_EXTENSION_MAP["o"]
     default_name = "<unnamed-object>"
 
@@ -144,14 +159,12 @@ def add_file_guess_ext(self, path_or_code):
 
         # Otherwise, we should have been given a LinkableCode object
         if not isinstance(path_or_code, LinkableCode):
-            raise TypeError("Exoected path to file or a LinkableCode object")
-
-        kind = path_or_code.kind
+            raise TypeError("Expected path to file or a LinkableCode object")
 
-        if kind == "cu":
-            return self.add_cu(path_or_code.data, path_or_code.name)
-
-        self.add_data(path_or_code.data, path_or_code.kind, path_or_code.name)
+        if path_or_code.kind == "cu":
+            self.add_cu(path_or_code.data, path_or_code.name)
+        else:
+            self.add_data(path_or_code.data, path_or_code.kind, path_or_code.name)
 
     def add_file(self, path, kind):
         try:
@@ -224,8 +237,11 @@ def patch_numba_linker():
         msg = f"Cannot patch Numba: {_numba_error}"
         raise RuntimeError(msg)
 
+    # Replace the built-in linker that uses the Driver API with our linker that
+    # uses nvJitLink
     Linker.new = new_patched_linker
 
+    # Add linkable code objects to Numba's top-level API
     cuda.Archive = Archive
     cuda.CUSource = CUSource
     cuda.Cubin = Cubin

pynvjitlink/tests/conftest.py

@@ -27,8 +27,7 @@ def alt_gpu_compute_capability(gpu_compute_capability):
 
 @pytest.fixture(scope="session")
 def absent_gpu_compute_capability(gpu_compute_capability, alt_gpu_compute_capability):
-    """A compute capability that does not match the current GPU"""
-    # A compute capability not used in any cubin or fatbin test binary
+    """A compute capability not used in any cubin or fatbin test binary"""
     cc_majors = {6, 7, 8}
     cc_majors.remove(gpu_compute_capability[0])
     cc_majors.remove(alt_gpu_compute_capability[0])

pynvjitlink/tests/test_patch.py

@@ -3,6 +3,7 @@
 import pytest
 import sys
 
+from numba import cuda
 from pynvjitlink import patch, NvJitLinkError
 from pynvjitlink.patch import (
     PatchedLinker,
@@ -109,6 +110,16 @@ def test_add_file_guess_ext_linkable_code(file, gpu_compute_capability, request)
     patched_linker.add_file_guess_ext(file)
 
 
+def test_add_file_guess_ext_invalid_input(
+    device_functions_cubin, gpu_compute_capability
+):
+    # Feeding raw data as bytes to add_file_guess_ext should raise, because
+    # there's no way to know what kind of file to treat it as
+    patched_linker = PatchedLinker(cc=gpu_compute_capability)
+    with pytest.raises(TypeError, match="Expected path to file or a LinkableCode"):
+        patched_linker.add_file_guess_ext(device_functions_cubin)
+
+
 @pytest.mark.skipif(
     not _numba_version_ok,
     reason=f"Requires Numba == {required_numba_ver[0]}.{required_numba_ver[1]}",
@@ -125,8 +136,6 @@ def test_add_file_guess_ext_linkable_code(file, gpu_compute_capability, request)
     ),
 )
 def test_jit_with_linkable_code(file, request):
-    from numba import cuda
-
     file = request.getfixturevalue(file)
     patch_numba_linker()
 
@@ -142,5 +151,20 @@ def kernel(result):
     assert result[0] == 3
 
 
+@pytest.mark.skipif(
+    not _numba_version_ok,
+    reason=f"Requires Numba == {required_numba_ver[0]}.{required_numba_ver[1]}",
+)
+def test_jit_with_invalid_linkable_code(device_functions_cubin):
+    # Attempting to pass raw bytes to the `link` kwarg should fail as in
+    # test_add_file_guess_ext_invalid_input - this is testing the same error
+    # checking triggered through the "public" API of the patched behaviour
+    with pytest.raises(TypeError, match="Expected path to file or a LinkableCode"):
+
+        @cuda.jit("void()", link=[device_functions_cubin])
+        def kernel():
+            pass
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main())

test_binary_generation/Makefile

@@ -1,6 +1,7 @@
 # Generates the input files used by the pynvjitlink binding test suite
 
-# Default to CC 7.5 for most testing unless explicitly set otherwise
+# Test binaries are built taking into account the CC of the GPU in the test machine
+GPU_CC := $(shell nvidia-smi --query-gpu=compute_cap --format=csv | grep -v compute_cap | sed 's/\.//')
 GPU_CC ?= 75
 
 # Use CC 7.0 as an alternative in fatbin testing, unless CC is 7.x