Add exact-name bitcode library lookup

cuda_pathfinder/cuda/pathfinder/__init__.py

@@ -41,9 +41,15 @@
 from cuda.pathfinder._static_libs.find_bitcode_lib import (
     find_bitcode_lib as find_bitcode_lib,
 )
+from cuda.pathfinder._static_libs.find_bitcode_lib import (
+    find_bitcode_lib_by_name as find_bitcode_lib_by_name,
+)
 from cuda.pathfinder._static_libs.find_bitcode_lib import (
     locate_bitcode_lib as locate_bitcode_lib,
 )
+from cuda.pathfinder._static_libs.find_bitcode_lib import (
+    locate_bitcode_lib_by_name as locate_bitcode_lib_by_name,
+)
 from cuda.pathfinder._static_libs.find_static_lib import (
     SUPPORTED_STATIC_LIBS as _SUPPORTED_STATIC_LIBS,
 )
@@ -77,7 +83,8 @@
 SUPPORTED_BINARY_UTILITIES = _SUPPORTED_BINARIES
 
 #: Tuple of supported bitcode library names that can be resolved
-#: via ``locate_bitcode_lib()`` and ``find_bitcode_lib()``.
+#: via ``locate_bitcode_lib()``, ``locate_bitcode_lib_by_name()``,
+#: ``find_bitcode_lib()``, and ``find_bitcode_lib_by_name()``.
 #: Example value: ``"device"``.
 SUPPORTED_BITCODE_LIBS = _SUPPORTED_BITCODE_LIBS
 

cuda_pathfinder/cuda/pathfinder/_static_libs/find_bitcode_lib.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import functools
+import ntpath
 import os
 from dataclasses import dataclass
 from typing import NoReturn, TypedDict
@@ -74,13 +75,21 @@ def _no_such_file_in_dir(dir_path: str, filename: str, error_messages: list[str]
         attachments.append(f'  Directory does not exist: "{dir_path}"')
 
 
+def _validate_filename(filename: str) -> str:
+    if not isinstance(filename, str):
+        raise TypeError(f"filename must be a string, got {type(filename).__name__}")
+    if filename in ("", ".", "..") or "\0" in filename or ntpath.basename(filename) != filename:
+        raise ValueError(f"filename must be a file name without a directory: {filename!r}")
+    return filename
+
+
 class _FindBitcodeLib:
-    def __init__(self, name: str) -> None:
+    def __init__(self, name: str, *, filename: str | None = None) -> None:
         if name not in _SUPPORTED_BITCODE_LIBS_INFO:  # Updated reference
             raise ValueError(f"Unknown bitcode library: '{name}'. Supported: {', '.join(SUPPORTED_BITCODE_LIBS)}")
         self.name: str = name
         self.config: _BitcodeLibInfo = _SUPPORTED_BITCODE_LIBS_INFO[name]  # Updated reference
-        self.filename: str = self.config["filename"]
+        self.filename: str = self.config["filename"] if filename is None else _validate_filename(filename)
         self.rel_path: str = self.config["rel_path"]
         self.site_packages_dirs: tuple[str, ...] = self.config["site_packages_dirs"]
         self.error_messages: list[str] = []
@@ -130,14 +139,8 @@ def raise_not_found_error(self) -> NoReturn:
         raise BitcodeLibNotFoundError(f'Failure finding "{self.filename}": {err}\n{att}')
 
 
-def locate_bitcode_lib(name: str) -> LocatedBitcodeLib:
-    """Locate a bitcode library by name.
-
-    Raises:
-        ValueError: If ``name`` is not a supported bitcode library.
-        BitcodeLibNotFoundError: If the bitcode library cannot be found.
-    """
-    finder = _FindBitcodeLib(name)
+def _locate_bitcode_lib(name: str, *, filename: str | None = None) -> LocatedBitcodeLib:
+    finder = _FindBitcodeLib(name, filename=filename)
 
     abs_path = finder.try_site_packages()
     if abs_path is not None:
@@ -169,6 +172,35 @@ def locate_bitcode_lib(name: str) -> LocatedBitcodeLib:
     finder.raise_not_found_error()
 
 
+def locate_bitcode_lib(name: str) -> LocatedBitcodeLib:
+    """Locate a bitcode library by name.
+
+    Raises:
+        ValueError: If ``name`` is not a supported bitcode library.
+        BitcodeLibNotFoundError: If the bitcode library cannot be found.
+    """
+    return _locate_bitcode_lib(name)
+
+
+def locate_bitcode_lib_by_name(name: str, filename: str) -> LocatedBitcodeLib:
+    """Locate an exact bitcode filename using a supported library's search paths.
+
+    This lets a library define its own filename convention, including any
+    architecture or other attributes encoded in the filename.
+
+    Args:
+        name: Supported bitcode library whose configured directories to search.
+        filename: Exact file name to locate. Directory components are not allowed.
+
+    Raises:
+        TypeError: If ``filename`` is not a string.
+        ValueError: If ``name`` is unsupported or ``filename`` is not a file name.
+        BitcodeLibNotFoundError: If ``filename`` cannot be found.
+    """
+    _validate_filename(filename)
+    return _locate_bitcode_lib(name, filename=filename)
+
+
 @functools.cache
 def find_bitcode_lib(name: str) -> str:
     """Find the absolute path to a bitcode library.
@@ -178,3 +210,22 @@ def find_bitcode_lib(name: str) -> str:
         BitcodeLibNotFoundError: If the bitcode library cannot be found.
     """
     return locate_bitcode_lib(name).abs_path
+
+
+@functools.cache
+def find_bitcode_lib_by_name(name: str, filename: str) -> str:
+    """Find an exact bitcode filename using a supported library's search paths.
+
+    This lets a library define its own filename convention, including any
+    architecture or other attributes encoded in the filename.
+
+    Args:
+        name: Supported bitcode library whose configured directories to search.
+        filename: Exact file name to find. Directory components are not allowed.
+
+    Raises:
+        TypeError: If ``filename`` is not a string.
+        ValueError: If ``name`` is unsupported or ``filename`` is not a file name.
+        BitcodeLibNotFoundError: If ``filename`` cannot be found.
+    """
+    return locate_bitcode_lib_by_name(name, filename).abs_path

cuda_pathfinder/docs/source/api.rst

@@ -35,7 +35,9 @@ CUDA bitcode and static libraries.
 
    SUPPORTED_BITCODE_LIBS
    find_bitcode_lib
+   find_bitcode_lib_by_name
    locate_bitcode_lib
+   locate_bitcode_lib_by_name
    LocatedBitcodeLib
    BitcodeLibNotFoundError
 

cuda_pathfinder/pyproject.toml

@@ -29,7 +29,7 @@ cu12 = [
     "nvidia-cusparselt-cu12",
     "nvidia-libmathdx-cu12",
     "nvidia-nccl-cu12; sys_platform != 'win32'",
-    "nvidia-nvshmem-cu12<3.7; sys_platform != 'win32'",
+    "nvidia-nvshmem-cu12; sys_platform != 'win32'",
 ]
 cu13 = [
     "cuda-toolkit[nvcc,cublas,nvrtc,cudart,cufft,curand,cusolver,cusparse,npp,nvfatbin,nvjitlink,nvjpeg,cccl,cupti,profiler,nvvm]==13.*",
@@ -43,7 +43,7 @@ cu13 = [
     "nvidia-cusparselt-cu13",
     "nvidia-libmathdx-cu13",
     "nvidia-nccl-cu13; sys_platform != 'win32'",
-    "nvidia-nvshmem-cu13<3.7; sys_platform != 'win32'",
+    "nvidia-nvshmem-cu13; sys_platform != 'win32'",
 ]
 host = [
     # TODO: remove the Python 3.15 guard once 3.15 is officially supported