From 7abcd36c57ebbca20ba4400b4fad8f44208164de Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach aka wash Date: Fri, 5 Jul 2024 12:53:35 -0700 Subject: [PATCH] Search for nvdisasm in the conda environment / `CUDA_HOME` / default CUDA install location in addition to the path. Fixes #9. --- numba_cuda/numba/cuda/codegen.py | 4 +- numba_cuda/numba/cuda/cuda_paths.py | 66 ++++++++++++++++++++++++++++- numba_cuda/numba/cuda/dispatcher.py | 8 ++-- numba_cuda/numba/cuda/testing.py | 6 +-- 4 files changed, 75 insertions(+), 9 deletions(-) diff --git a/numba_cuda/numba/cuda/codegen.py b/numba_cuda/numba/cuda/codegen.py index 6009dcb..72ff0f2 100644 --- a/numba_cuda/numba/cuda/codegen.py +++ b/numba_cuda/numba/cuda/codegen.py @@ -4,6 +4,7 @@ from numba.core.codegen import Codegen, CodeLibrary from .cudadrv import devices, driver, nvvm, runtime from numba.cuda.cudadrv.libs import get_cudalib +from numba.cuda.cuda_paths import get_cuda_paths import os import subprocess @@ -24,7 +25,8 @@ def run_nvdisasm(cubin, flags): f.write(cubin) try: - cp = subprocess.run(['nvdisasm', *flags, fname], check=True, + nvdisasm = get_cuda_paths()['nvdisasm'].info + cp = subprocess.run([nvdisasm, *flags, fname], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) except FileNotFoundError as e: diff --git a/numba_cuda/numba/cuda/cuda_paths.py b/numba_cuda/numba/cuda/cuda_paths.py index 8ff9664..557f3fc 100644 --- a/numba_cuda/numba/cuda/cuda_paths.py +++ b/numba_cuda/numba/cuda/cuda_paths.py @@ -2,6 +2,7 @@ import re import os from collections import namedtuple +from shutil import which from numba.core.config import IS_WIN32 from numba.misc.findlib import find_lib, find_file @@ -10,6 +11,26 @@ _env_path_tuple = namedtuple('_env_path_tuple', ['by', 'info']) +def find_executable(name, bindirs=None): + # This should probably go into numba.misc.findlib + if bindirs is None: + return which(name) # Check the path if we've been given no directories. + + if isinstance(bindirs, str): + bindirs = [bindirs,] + else: + bindirs = list(bindirs) + files = [] + for bdir in bindirs: + try: + entries = os.listdir(bdir) + except FileNotFoundError: + continue + candidates = [os.path.join(bdir, ent) for ent in entries if name == ent] + files.extend([c for c in candidates if os.path.isfile(c)]) + return files + + def _find_valid_path(options): """Find valid path from *options*, which is a list of 2-tuple of (name, path). Return first pair where *path* is not None. @@ -52,6 +73,18 @@ def _get_nvvm_path_decision(): return by, path +def _get_nvdisasm_path_decision(): + options = [ + ('Conda environment', get_conda_ctk()), + ('Conda environment (NVIDIA package)', get_nvidia_nvdisasm_ctk()), + ('CUDA_HOME', get_cuda_home('bin')), + ('System', get_system_ctk('bin')), + ('Path', None), + ] + by, path = _find_valid_path(options) + return by, path + + def _get_libdevice_paths(): by, libdir = _get_libdevice_path_decision() # Search for pattern @@ -161,6 +194,29 @@ def get_nvidia_nvvm_ctk(): return os.path.dirname(max(paths)) +def get_nvidia_nvdisasm_ctk(): + """Return path to directory containing the nvdisasm executable. + """ + is_conda_env = os.path.exists(os.path.join(sys.prefix, 'conda-meta')) + if not is_conda_env: + return + + # Assume the existence of nvdisasm in the conda env implies that a CUDA + # toolkit conda package is installed. + + # Try the location used on Linux and the Windows 11.x packages + bindir = os.path.join(sys.prefix, 'bin') + if not os.path.exists(bindir) or not os.path.isdir(bindir): + return + + paths = find_executable('nvdisasm', bindir) + if not paths: + return + + # Use the directory name of the max path + return os.path.dirname(max(paths)) + + def get_nvidia_libdevice_ctk(): """Return path to directory containing the libdevice library. """ @@ -220,6 +276,13 @@ def _get_nvvm_path(): return _env_path_tuple(by, path) +def _get_nvdisasm_path(): + by, path = _get_nvdisasm_path_decision() + candidates = find_executable('nvdisasm', path) + path = max(candidates) if candidates else None + return _env_path_tuple(by, path) + + def get_cuda_paths(): """Returns a dictionary mapping component names to a 2-tuple of (source_variable, info). @@ -238,9 +301,10 @@ def get_cuda_paths(): # Not in cache d = { 'nvvm': _get_nvvm_path(), + 'nvdisasm': _get_nvdisasm_path(), 'libdevice': _get_libdevice_paths(), 'cudalib_dir': _get_cudalib_dir(), - 'static_cudalib_dir': _get_static_cudalib_dir(), + 'static_cudalib_dir': _get_static_cudalib_dir() } # Cache result get_cuda_paths._cached_result = d diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py index 16680a2..af354ef 100644 --- a/numba_cuda/numba/cuda/dispatcher.py +++ b/numba_cuda/numba/cuda/dispatcher.py @@ -247,7 +247,7 @@ def inspect_sass_cfg(self): ''' Returns the CFG of the SASS for this kernel. - Requires nvdisasm to be available on the PATH. + Requires nvdisasm to be available. ''' return self._codelibrary.get_sass_cfg() @@ -255,7 +255,7 @@ def inspect_sass(self): ''' Returns the SASS code for this kernel. - Requires nvdisasm to be available on the PATH. + Requires nvdisasm to be available. ''' return self._codelibrary.get_sass() @@ -995,7 +995,7 @@ def inspect_sass_cfg(self, signature=None): The CFG for the device in the current context is returned. - Requires nvdisasm to be available on the PATH. + Requires nvdisasm to be available. ''' if self.targetoptions.get('device'): raise RuntimeError('Cannot get the CFG of a device function') @@ -1017,7 +1017,7 @@ def inspect_sass(self, signature=None): SASS for the device in the current context is returned. - Requires nvdisasm to be available on the PATH. + Requires nvdisasm to be available. ''' if self.targetoptions.get('device'): raise RuntimeError('Cannot inspect SASS of a device function') diff --git a/numba_cuda/numba/cuda/testing.py b/numba_cuda/numba/cuda/testing.py index ae62d33..7a78053 100644 --- a/numba_cuda/numba/cuda/testing.py +++ b/numba_cuda/numba/cuda/testing.py @@ -3,7 +3,7 @@ import shutil from numba.tests.support import SerialMixin -from numba.cuda.cuda_paths import get_conda_ctk +from numba.cuda.cuda_paths import get_cuda_paths, get_conda_ctk from numba.cuda.cudadrv import driver, devices, libs from numba.core import config from numba.tests.support import TestCase @@ -97,12 +97,12 @@ def skip_under_cuda_memcheck(reason): def skip_without_nvdisasm(reason): - nvdisasm_path = shutil.which('nvdisasm') + nvdisasm_path = get_cuda_paths()['nvdisasm'].info return unittest.skipIf(nvdisasm_path is None, reason) def skip_with_nvdisasm(reason): - nvdisasm_path = shutil.which('nvdisasm') + nvdisasm_path = get_cuda_paths()['nvdisasm'].info return unittest.skipIf(nvdisasm_path is not None, reason)