numba-cuda 0.22.0__cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/cudadrv/libs.py

@@ -0,0 +1,220 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""CUDA Toolkit libraries lookup utilities.
+
+CUDA Toolkit libraries can be available via either:
+
+- the `cuda-nvcc` and `cuda-nvrtc` conda packages,
+- a user supplied location from CUDA_HOME,
+- a system wide location,
+- package-specific locations (e.g. the Debian NVIDIA packages),
+- or can be discovered by the system loader.
+"""
+
+import os
+import sys
+import ctypes
+
+from numba.cuda.misc.findlib import find_lib
+from numba.cuda.cuda_paths import get_cuda_paths
+from numba.cuda.cudadrv.driver import locate_driver_and_loader, load_driver
+from numba.cuda.cudadrv.error import CudaSupportError
+from numba.cuda.core import config
+
+
+if sys.platform == "win32":
+    _dllnamepattern = "%s.dll"
+    _staticnamepattern = "%s.lib"
+elif sys.platform == "darwin":
+    _dllnamepattern = "lib%s.dylib"
+    _staticnamepattern = "lib%s.a"
+else:
+    _dllnamepattern = "lib%s.so"
+    _staticnamepattern = "lib%s.a"
+
+
+def get_libdevice():
+    d = get_cuda_paths()
+    paths = d["libdevice"].info
+    return paths
+
+
+def open_libdevice():
+    with open(get_libdevice(), "rb") as bcfile:
+        return bcfile.read()
+
+
+def get_cudalib(lib, static=False):
+    """
+    Find the path of a CUDA library based on a search of known locations. If
+    the search fails, return a generic filename for the library (e.g.
+    'libnvvm.so' for 'nvvm') so that we may attempt to load it using the system
+    loader's search mechanism.
+    """
+    if lib in {"nvrtc", "nvvm"}:
+        return get_cuda_paths()[lib].info or _dllnamepattern % lib
+
+    dir_type = "static_cudalib_dir" if static else "cudalib_dir"
+    libdir = get_cuda_paths()[dir_type].info
+
+    candidates = find_lib(lib, libdir, static=static)
+    namepattern = _staticnamepattern if static else _dllnamepattern
+    return max(candidates) if candidates else namepattern % lib
+
+
+def get_cuda_include_dir():
+    """
+    Find the path to cuda include dir based on a list of default locations.
+    Note that this does not list the `CUDA_INCLUDE_PATH` entry in user
+    configuration.
+    """
+
+    return get_cuda_paths()["include_dir"].info
+
+
+def check_cuda_include_dir(path):
+    if path is None or not os.path.exists(path):
+        raise FileNotFoundError(f"{path} not found")
+
+    if not os.path.exists(os.path.join(path, "cuda_runtime.h")):
+        raise FileNotFoundError(f"Unable to find cuda_runtime.h from {path}")
+
+
+def open_cudalib(lib):
+    path = get_cudalib(lib)
+    return ctypes.CDLL(path)
+
+
+def check_static_lib(path):
+    if not os.path.isfile(path):
+        raise FileNotFoundError(f"{path} not found")
+
+
+def _get_source_variable(lib, static=False):
+    if lib == "nvvm":
+        return get_cuda_paths()["nvvm"].by
+    elif lib == "nvrtc":
+        return get_cuda_paths()["nvrtc"].by
+    elif lib == "libdevice":
+        return get_cuda_paths()["libdevice"].by
+    elif lib == "include_dir":
+        return get_cuda_paths()["include_dir"].by
+    else:
+        dir_type = "static_cudalib_dir" if static else "cudalib_dir"
+        return get_cuda_paths()[dir_type].by
+
+
+def test():
+    """Test library lookup.  Path info is printed to stdout."""
+    failed = False
+
+    # Check for the driver
+    try:
+        dlloader, candidates = locate_driver_and_loader()
+        print("Finding driver from candidates:")
+        for location in candidates:
+            print(f"\t{location}")
+        print(f"Using loader {dlloader}")
+        print("\tTrying to load driver", end="...")
+        dll, path = load_driver(dlloader, candidates)
+        print("\tok")
+        print(f"\t\tLoaded from {path}")
+    except CudaSupportError as e:
+        print(f"\tERROR: failed to open driver: {e}")
+        failed = True
+
+    # Find the absolute location of the driver on Linux. Various driver-related
+    # issues have been reported by WSL2 users, and it is almost always due to a
+    # Linux (i.e. not- WSL2) driver being installed in a WSL2 system.
+    # Providing the absolute location of the driver indicates its version
+    # number in the soname (e.g. "libcuda.so.530.30.02"), which can be used to
+    # look up whether the driver was intended for "native" Linux.
+    if sys.platform == "linux" and not failed:
+        pid = os.getpid()
+        mapsfile = os.path.join(os.path.sep, "proc", f"{pid}", "maps")
+        try:
+            with open(mapsfile) as f:
+                maps = f.read()
+        # It's difficult to predict all that might go wrong reading the maps
+        # file - in case various error conditions ensue (the file is not found,
+        # not readable, etc.) we use OSError to hopefully catch any of them.
+        except OSError:
+            # It's helpful to report that this went wrong to the user, but we
+            # don't set failed to True because this doesn't have any connection
+            # to actual CUDA functionality.
+            print(
+                f"\tERROR: Could not open {mapsfile} to determine absolute "
+                "path to libcuda.so"
+            )
+        else:
+            # In this case we could read the maps, so we can report the
+            # relevant ones to the user
+            locations = set(s for s in maps.split() if "libcuda.so" in s)
+            print("\tMapped libcuda.so paths:")
+            for location in locations:
+                print(f"\t\t{location}")
+
+    # Checks for dynamic libraries
+    libs = "nvvm nvrtc".split()
+    for lib in libs:
+        path = get_cudalib(lib)
+        print("Finding {} from {}".format(lib, _get_source_variable(lib)))
+        print("\tLocated at", path)
+
+        try:
+            print("\tTrying to open library", end="...")
+            open_cudalib(lib)
+            print("\tok")
+        except OSError as e:
+            print("\tERROR: failed to open %s:\n%s" % (lib, e))
+            failed = True
+
+    # Check for cudadevrt (the only static library)
+    lib = "cudadevrt"
+    path = get_cudalib(lib, static=True)
+    print(
+        "Finding {} from {}".format(lib, _get_source_variable(lib, static=True))
+    )
+    print("\tLocated at", path)
+
+    try:
+        print("\tChecking library", end="...")
+        check_static_lib(path)
+        print("\tok")
+    except FileNotFoundError as e:
+        print("\tERROR: failed to find %s:\n%s" % (lib, e))
+        failed = True
+
+    # Check for libdevice
+    where = _get_source_variable("libdevice")
+    print(f"Finding libdevice from {where}")
+    path = get_libdevice()
+    print("\tLocated at", path)
+
+    try:
+        print("\tChecking library", end="...")
+        check_static_lib(path)
+        print("\tok")
+    except FileNotFoundError as e:
+        print("\tERROR: failed to find %s:\n%s" % (lib, e))
+        failed = True
+
+    # Check cuda include paths
+
+    print("Include directory configuration variable:")
+    print(f"\tCUDA_INCLUDE_PATH={config.CUDA_INCLUDE_PATH}")
+
+    where = _get_source_variable("include_dir")
+    print(f"Finding include directory from {where}")
+    include = get_cuda_include_dir()
+    print("\tLocated at", include)
+    try:
+        print("\tChecking include directory", end="...")
+        check_cuda_include_dir(include)
+        print("\tok")
+    except FileNotFoundError as e:
+        print("\tERROR: failed to find cuda include directory:\n%s" % e)
+        failed = True
+
+    return not failed

numba_cuda/numba/cuda/cudadrv/linkable_code.py

@@ -0,0 +1,184 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import io
+import os
+from typing import Union, Type
+
+from .mappings import FILE_EXTENSION_MAP
+
+
+class LinkableCode:
+    """An object that holds code to be linked from memory.
+
+    :param data: A buffer, StringIO or BytesIO containing the data to link.
+                 If a file object is passed, the content in the object is
+                 read when `data` property is accessed.
+    :param name: The name of the file to be referenced in any compilation or
+                 linking errors that may be produced.
+    :param setup_callback: A function called prior to the launch of a kernel
+                           contained within a module that has this code object
+                           linked into it.
+    :param teardown_callback: A function called just prior to the unloading of
+                              a module that has this code object linked into
+                              it.
+    :param nrt: If True, assume this object contains NRT function calls and
+                add NRT source code to the final link.
+    """
+
+    def __init__(
+        self,
+        data,
+        name=None,
+        setup_callback=None,
+        teardown_callback=None,
+        nrt=False,
+    ):
+        if setup_callback and not callable(setup_callback):
+            raise TypeError("setup_callback must be callable")
+        if teardown_callback and not callable(teardown_callback):
+            raise TypeError("teardown_callback must be callable")
+
+        self.nrt = nrt
+        self._name = name
+        self._data = data
+        self.setup_callback = setup_callback
+        self.teardown_callback = teardown_callback
+
+    @property
+    def name(self):
+        return self._name or self.default_name
+
+    @property
+    def data(self):
+        if isinstance(self._data, (io.StringIO, io.BytesIO)):
+            return self._data.getvalue()
+        return self._data
+
+    @staticmethod
+    def from_path(path: str):
+        """
+        Load a linkable code object from a file.
+
+        Parameters
+        ----------
+        path : str
+            The path to the file to load.
+
+        Returns
+        -------
+        LinkableCode
+            The linkable code object.
+
+        Raises
+        ------
+        ValueError
+            If the file extension is not supported.
+        """
+        root, extension = os.path.splitext(path)
+        basename = os.path.basename(root)
+        if extension in (".cu", ".ptx"):
+            mode = "r"
+        else:
+            mode = "rb"
+
+        with open(path, mode) as f:
+            data = f.read()
+
+        cls = _extension_to_linkable_code_kind(extension)
+        return cls(data, name=basename)
+
+    @classmethod
+    def from_path_or_obj(cls, path_or_obj: Union[str, "LinkableCode"]):
+        """
+        Load a linkable code object from a file or a LinkableCode object.
+
+        If a path is provided, the file is loaded and the LinkableCode object
+        is returned. If a LinkableCode object is provided, it is returned as is.
+
+        Parameters
+        ----------
+        path_or_obj : str or LinkableCode
+            The path to the file or the LinkableCode object to load.
+
+        Returns
+        -------
+        LinkableCode
+            The linkable code object.
+
+        Raises
+        ------
+        ValueError
+            If the file extension is not supported.
+        """
+        if isinstance(path_or_obj, str):
+            return cls.from_path(path_or_obj)
+        return path_or_obj
+
+
+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):
+    """An ELF Fatbin in memory."""
+
+    kind = FILE_EXTENSION_MAP["fatbin"]
+    default_name = "<unnamed-fatbin>"
+
+
+class Cubin(LinkableCode):
+    """An ELF Cubin 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>"
+
+
+class LTOIR(LinkableCode):
+    """An LTOIR file in memory."""
+
+    kind = FILE_EXTENSION_MAP["ltoir"]
+    default_name = "<unnamed-ltoir>"
+
+
+def _extension_to_linkable_code_kind(extension: str) -> Type[LinkableCode]:
+    if extension == ".cu":
+        return CUSource
+    elif extension == ".ptx":
+        return PTXSource
+    elif extension == ".fatbin":
+        return Fatbin
+    elif extension == ".cubin":
+        return Cubin
+    elif extension == ".a":
+        return Archive
+    elif extension == ".o":
+        return Object
+    elif extension == ".ltoir":
+        return LTOIR
+    else:
+        raise ValueError(f"Unknown extension: {extension}")

numba_cuda/numba/cuda/cudadrv/mappings.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from cuda.bindings.driver import CUjitInputType
+
+FILE_EXTENSION_MAP = {
+    "o": CUjitInputType.CU_JIT_INPUT_OBJECT,
+    "ptx": CUjitInputType.CU_JIT_INPUT_PTX,
+    "a": CUjitInputType.CU_JIT_INPUT_LIBRARY,
+    "lib": CUjitInputType.CU_JIT_INPUT_LIBRARY,
+    "cubin": CUjitInputType.CU_JIT_INPUT_CUBIN,
+    "fatbin": CUjitInputType.CU_JIT_INPUT_FATBINARY,
+    "ltoir": CUjitInputType.CU_JIT_INPUT_NVVM,
+}

numba_cuda/numba/cuda/cudadrv/ndarray.py

@@ -0,0 +1,26 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.cudadrv import devices, driver
+
+try:
+    from numba.core.registry import cpu_target
+
+    def _calc_array_sizeof(ndim):
+        """
+        Use the ABI size in the CPU target
+        """
+        ctx = cpu_target.target_context
+        return ctx.calc_array_sizeof(ndim)
+except ImportError:
+    pass
+
+
+def ndarray_device_allocate_data(ary):
+    """
+    Allocate gpu data buffer
+    """
+    datasize = driver.host_memory_size(ary)
+    # allocate
+    gpu_data = devices.get_context().memalloc(datasize)
+    return gpu_data

numba_cuda/numba/cuda/cudadrv/nvrtc.py

@@ -0,0 +1,193 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.cudadrv.error import (
+    CCSupportError,
+)
+from numba.cuda import config
+from numba.cuda.cuda_paths import get_cuda_paths
+from numba.cuda.utils import _readenv
+
+import os
+import warnings
+import functools
+
+from cuda.core.experimental import Program, ProgramOptions
+from cuda.bindings import nvrtc as bindings_nvrtc
+
+NVRTC_EXTRA_SEARCH_PATHS = _readenv(
+    "NUMBA_CUDA_NVRTC_EXTRA_SEARCH_PATHS", str, ""
+) or getattr(config, "CUDA_NVRTC_EXTRA_SEARCH_PATHS", "")
+if not hasattr(config, "CUDA_NVRTC_EXTRA_SEARCH_PATHS"):
+    config.CUDA_NVRTC_EXTRA_SEARCH_PATHS = NVRTC_EXTRA_SEARCH_PATHS
+
+
+@functools.cache
+def _get_nvrtc_version():
+    retcode, major, minor = bindings_nvrtc.nvrtcVersion()
+    if retcode != bindings_nvrtc.nvrtcResult.NVRTC_SUCCESS:
+        raise RuntimeError(f"{retcode.name} when calling nvrtcVersion()")
+    return (major, minor)
+
+
+def compile(src, name, cc, ltoir=False, lineinfo=False, debug=False):
+    """
+    Compile a CUDA C/C++ source to PTX or LTOIR for a given compute capability.
+
+    :param src: The source code to compile
+    :type src: str
+    :param name: The filename of the source (for information only)
+    :type name: str
+    :param cc: A tuple ``(major, minor)`` of the compute capability
+    :type cc: tuple
+    :param ltoir: Compile into LTOIR if True, otherwise into PTX
+    :type ltoir: bool
+    :param lineinfo: Whether to include line information in the compiled code
+    :type lineinfo: bool
+    :param debug: Whether to include debug information in the compiled code
+    :type debug: bool
+    :return: The compiled PTX or LTOIR and compilation log
+    :rtype: tuple
+    """
+    version = _get_nvrtc_version()
+
+    ver_str = lambda version: ".".join(str(v) for v in version)
+    supported_ccs = get_supported_ccs()
+    try:
+        found = max(filter(lambda v: v <= cc, [v for v in supported_ccs]))
+    except ValueError:
+        raise RuntimeError(
+            f"Device compute capability {ver_str(cc)} is less than the "
+            f"minimum supported by NVRTC {ver_str(version)}. Supported "
+            "compute capabilities are "
+            f"{', '.join([ver_str(v) for v in supported_ccs])}."
+        )
+
+    if found != cc:
+        warnings.warn(
+            f"Device compute capability {ver_str(cc)} is not supported by "
+            f"NVRTC {ver_str(version)}. Using {ver_str(found)} instead."
+        )
+
+    # Compilation options:
+    # - Compile for the current device's compute capability.
+    # - The CUDA include path is added.
+    # - Relocatable Device Code (rdc) is needed to prevent device functions
+    #   being optimized away.
+    major, minor = found
+
+    arch = f"sm_{major}{minor}"
+
+    cuda_include_dir = get_cuda_paths()["include_dir"].info
+    cuda_includes = [f"{cuda_include_dir}"]
+
+    cudadrv_path = os.path.dirname(os.path.abspath(__file__))
+    numba_cuda_path = os.path.dirname(cudadrv_path)
+
+    nvrtc_ver_major = version[0]
+    if nvrtc_ver_major == 12:
+        numba_include = f"{os.path.join(numba_cuda_path, 'include', '12')}"
+        # For CUDA 12 wheels, `cuda_include_dir` is `site-packages/nvidia/cuda_runtime/include`
+        # We need to find CCCL at `site-packages/nvidia/cuda_cccl/include`
+        # For CUDA 12 conda / system install, CCCL is just in the `include` directory
+        cuda_includes.append(
+            f"{os.path.join(cuda_include_dir, '..', '..', 'cuda_cccl', 'include')}"
+        )
+    elif nvrtc_ver_major == 13:
+        numba_include = f"{os.path.join(numba_cuda_path, 'include', '13')}"
+        # For CUDA 13 wheels, `cuda_include_dir` is `site-packages/nvidia/cu13/include`
+        # We need to find CCCL at `site-packages/nvidia/cu13/include/cccl`
+        # For CUDA 13 conda / system install, CCCL is in the `include/cccl` directory
+        cuda_includes.append(f"{os.path.join(cuda_include_dir, 'cccl')}")
+
+    if config.CUDA_NVRTC_EXTRA_SEARCH_PATHS:
+        extra_includes = config.CUDA_NVRTC_EXTRA_SEARCH_PATHS.split(":")
+    else:
+        extra_includes = []
+
+    nrt_include = os.path.join(numba_cuda_path, "memory_management")
+
+    includes = [numba_include, *cuda_includes, nrt_include, *extra_includes]
+
+    options = ProgramOptions(
+        arch=arch,
+        include_path=includes,
+        relocatable_device_code=True,
+        link_time_optimization=ltoir,
+        name=name,
+        debug=debug,
+        lineinfo=lineinfo,
+    )
+
+    class Logger:
+        def __init__(self):
+            self.log = []
+
+        def write(self, msg):
+            self.log.append(msg)
+
+    logger = Logger()
+    if isinstance(src, bytes):
+        src = src.decode("utf8")
+
+    prog = Program(src, "c++", options=options)
+    result = prog.compile("ltoir" if ltoir else "ptx", logs=logger)
+    log = ""
+    if logger.log:
+        log = logger.log
+        joined_logs = "\n".join(log)
+        warnings.warn(f"NVRTC log messages: {joined_logs}")
+    return result, log
+
+
+def find_closest_arch(mycc):
+    """
+    Given a compute capability, return the closest compute capability supported
+    by the CUDA toolkit.
+
+    :param mycc: Compute capability as a tuple ``(MAJOR, MINOR)``
+    :return: Closest supported CC as a tuple ``(MAJOR, MINOR)``
+    """
+    supported_ccs = get_supported_ccs()
+
+    for i, cc in enumerate(supported_ccs):
+        if cc == mycc:
+            # Matches
+            return cc
+        elif cc > mycc:
+            # Exceeded
+            if i == 0:
+                # CC lower than supported
+                msg = (
+                    "GPU compute capability %d.%d is not supported"
+                    "(requires >=%d.%d)" % (mycc + cc)
+                )
+                raise CCSupportError(msg)
+            else:
+                # return the previous CC
+                return supported_ccs[i - 1]
+
+    # CC higher than supported
+    return supported_ccs[-1]  # Choose the highest
+
+
+def get_arch_option(major, minor):
+    """Matches with the closest architecture option"""
+    if config.FORCE_CUDA_CC:
+        arch = config.FORCE_CUDA_CC
+    else:
+        arch = find_closest_arch((major, minor))
+    return "compute_%d%d" % arch
+
+
+def get_lowest_supported_cc():
+    return min(get_supported_ccs())
+
+
+def get_supported_ccs():
+    retcode, archs = bindings_nvrtc.nvrtcGetSupportedArchs()
+    if retcode != bindings_nvrtc.nvrtcResult.NVRTC_SUCCESS:
+        raise RuntimeError(
+            f"{retcode.name} when calling nvrtcGetSupportedArchs()"
+        )
+    return [(arch // 10, arch % 10) for arch in archs]