Enable runtime resolution of CUDA header path for NVRTC

transformer_engine/common/CMakeLists.txt

@@ -371,6 +371,7 @@ target_link_libraries(transformer_engine PUBLIC
                       CUDA::cublas
                       CUDA::cudart
                       CUDNN::cudnn_all)
+target_link_libraries(transformer_engine PRIVATE ${CMAKE_DL_LIBS})
 
 target_include_directories(transformer_engine PRIVATE
                            ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})

transformer_engine/common/__init__.py

@@ -235,24 +235,49 @@ def _get_sys_extension() -> str:
     raise RuntimeError(f"Unsupported operating system ({system})")
 
 
+def _cuda_runtime_major(cuda_runtime: ctypes.CDLL) -> Optional[int]:
+    """Return the major version of the CUDA runtime loaded with Transformer Engine."""
+
+    runtime_version = ctypes.c_int()
+    get_runtime_version = cuda_runtime.cudaRuntimeGetVersion
+    get_runtime_version.argtypes = [ctypes.POINTER(ctypes.c_int)]
+    get_runtime_version.restype = ctypes.c_int
+    if get_runtime_version(ctypes.byref(runtime_version)) != 0 or runtime_version.value <= 0:
+        return None
+    return runtime_version.value // 1000
+
+
 @functools.lru_cache(maxsize=None)
-def _nvidia_cudart_include_dir() -> str:
+def _nvidia_cudart_include_dir(cuda_major_version: int) -> str:
     """Returns the include directory for cuda_runtime.h if exists in python environment."""
 
+    # This is primarily here to support editable installs. cuda_runtime.cpp handles the
+    # resolution for install via wheel or when using the shared library ABI directly.
+
     try:
         import nvidia
     except ModuleNotFoundError:
         return ""
 
-    # Installing some nvidia-* packages, like nvshmem, create nvidia name, so "import nvidia"
-    # above doesn't throw. However, they don't set "__file__" attribute.
+    # NVIDIA packages may use either a regular package or a namespace package spread
+    # across multiple package roots.
     if nvidia.__file__ is not None:
-        nvidia_root = Path(nvidia.__file__).parent
+        nvidia_roots = (Path(nvidia.__file__).parent,)
     else:
-        nvidia_root = Path(nvidia.__path__[0])  # namespace package
+        nvidia_roots = tuple(Path(path) for path in nvidia.__path__)
+
+    layouts = [f"cu{cuda_major_version}"]
+    if cuda_major_version == 12:
+        layouts.append("cuda_runtime")
 
-    include_dir = nvidia_root / "cuda_runtime"
-    return str(include_dir) if include_dir.exists() else ""
+    for layout in layouts:
+        for nvidia_root in nvidia_roots:
+            cuda_root = nvidia_root / layout
+            if (cuda_root / "cuda_runtime.h").is_file() or (
+                cuda_root / "include" / "cuda_runtime.h"
+            ).is_file():
+                return str(cuda_root)
+    return ""
 
 
 @functools.lru_cache(maxsize=None)
@@ -382,5 +407,8 @@ def _load_core_library():
     _TE_LIB_CTYPES = _load_core_library()
 
     # Needed to find the correct headers for NVRTC kernels.
-    if not os.getenv("NVTE_CUDA_INCLUDE_DIR") and _nvidia_cudart_include_dir():
-        os.environ["NVTE_CUDA_INCLUDE_DIR"] = _nvidia_cudart_include_dir()
+    _cuda_major_version = _cuda_runtime_major(_TE_LIB_CTYPES)
+    if not os.getenv("NVTE_CUDA_INCLUDE_DIR") and _cuda_major_version is not None:
+        cuda_include_dir = _nvidia_cudart_include_dir(_cuda_major_version)
+        if cuda_include_dir:
+            os.environ["NVTE_CUDA_INCLUDE_DIR"] = cuda_include_dir

transformer_engine/common/util/cuda_runtime.cpp

@@ -7,6 +7,7 @@
 #include "../util/cuda_runtime.h"
 
 #include <cublasLt.h>
+#include <dlfcn.h>
 
 #include <filesystem>
 #include <fstream>
@@ -26,6 +27,83 @@ namespace {
 // String with build-time CUDA include path
 #include "string_path_cuda_include.h"
 
+// Get the runtime directory of the shared library that contains this code
+std::filesystem::path shared_library_directory() {
+  static const char library_anchor = 0;
+  Dl_info library_info{};
+  if (dladdr(static_cast<const void *>(&library_anchor), &library_info) == 0 ||
+      library_info.dli_fname == nullptr) {
+    return {};
+  }
+
+  std::filesystem::path library_path = library_info.dli_fname;
+  if (library_path.is_relative()) {
+    std::error_code error;
+    library_path = std::filesystem::absolute(library_path, error);
+    if (error) {
+      return {};
+    }
+  }
+
+  return library_path.parent_path();
+}
+
+std::string runtime_cuda_major_version() {
+  int runtime_version = 0;
+  // Header discovery is best-effort, so do not throw if the runtime cannot
+  // report its version.
+  if (cudaRuntimeGetVersion(&runtime_version) != cudaSuccess || runtime_version <= 0) {
+    return {};
+  }
+
+  return std::to_string(runtime_version / 1000);
+}
+
+std::filesystem::path python_cuda_directory() {
+  using Path = std::filesystem::path;
+
+  // Find the Python package root from the installed Transformer Engine package. Do not
+  // assume that the root is named site-packages or dist-packages since valid installs
+  // may use an arbitrary target directory.
+  Path te_package_directory = shared_library_directory();
+  while (true) {
+    if (te_package_directory.filename() == "transformer_engine") {
+      break;
+    }
+
+    const Path parent = te_package_directory.parent_path();
+    if (parent == te_package_directory) {
+      // Root directory reached
+      return {};
+    }
+
+    te_package_directory = parent;
+  }
+
+  const Path nvidia_directory = te_package_directory.parent_path() / "nvidia";
+  const auto cuda_major_version = runtime_cuda_major_version();
+  if (cuda_major_version.empty()) {
+    return {};
+  }
+
+  std::error_code error;
+  const Path cuda_directory = nvidia_directory / ("cu" + cuda_major_version);
+  if (std::filesystem::is_directory(cuda_directory, error)) {
+    return cuda_directory;
+  }
+
+  // CUDA 12 Python wheels use the older nvidia/cuda_runtime layout.
+  if (cuda_major_version == "12") {
+    error.clear();
+    const Path legacy_cuda_directory = nvidia_directory / "cuda_runtime";
+    if (std::filesystem::is_directory(legacy_cuda_directory, error)) {
+      return legacy_cuda_directory;
+    }
+  }
+
+  return {};
+}
+
 }  // namespace
 
 int num_devices() {
@@ -152,6 +230,7 @@ const std::string &include_directory(bool required) {
     std::vector<std::pair<std::string, Path>> search_paths = {{"NVTE_CUDA_INCLUDE_DIR", ""},
                                                               {"CUDA_HOME", ""},
                                                               {"CUDA_DIR", ""},
+                                                              {"", python_cuda_directory()},
                                                               {"", string_path_cuda_include},
                                                               {"", "/usr/local/cuda"}};
     for (auto &[env, p] : search_paths) {