Add paths of cl and CUDA before loading anything

Author: woct0rdho

python/triton/__init__.py

@@ -1,6 +1,18 @@
 """isort:skip_file"""
 __version__ = '3.0.0'
 
+# Users may not know how to add cl and CUDA to PATH. Let's do it before loading anything
+import os
+if os.name == "nt":
+    from .windows_utils import find_cuda, find_msvc_winsdk
+    msvc_winsdk_inc_dirs, _ = find_msvc_winsdk()
+    if msvc_winsdk_inc_dirs:
+        cl_path = msvc_winsdk_inc_dirs[0].replace(r"\include", r"\bin\Hostx64\x64")
+        os.environ["PATH"] = cl_path + os.pathsep + os.environ["PATH"]
+    cuda_bin_path, _, _ = find_cuda()
+    if cuda_bin_path:
+        os.environ["PATH"] = cuda_bin_path + os.pathsep + os.environ["PATH"]
+
 # ---------------------------------------
 # Note: import order is significant here.
 

python/triton/runtime/build.py

@@ -8,7 +8,7 @@
 import setuptools
 
 if os.name == "nt":
-    from .windows import find_msvc_winsdk, find_python
+    from triton.windows_utils import find_msvc_winsdk, find_python
 
 
 @contextlib.contextmanager
@@ -44,12 +44,6 @@ def _build(name, src, srcdir, library_dirs, include_dirs, libraries):
     # try to avoid setuptools if possible
     cc = os.environ.get("CC")
     if cc is None:
-        # Users may not know how to add cl to PATH. Let's do it for them
-        if os.name == "nt":
-            msvc_winsdk_inc_dirs, _ = find_msvc_winsdk()
-            if msvc_winsdk_inc_dirs:
-                cl_path = msvc_winsdk_inc_dirs[0].replace(r"\include", r"\bin\Hostx64\x64")
-            os.environ["PATH"] = cl_path + os.pathsep + os.environ["PATH"]
         # TODO: support more things here.
         cl = shutil.which("cl")
         gcc = shutil.which("gcc")

python/triton/runtime/windows.py renamed to python/triton/windows_utils.py

@@ -21,10 +21,12 @@ def unparse_version(t, prefix=""):
     return prefix + ".".join([str(x) for x in t])
 
 
-def max_version(versions):
-    versions = [parse_version(x) for x in versions]
+def max_version(versions, prefix=""):
+    versions = [parse_version(x, prefix) for x in versions]
     versions = [x for x in versions if x is not None]
-    version = unparse_version(max(versions))
+    if not versions:
+        return None
+    version = unparse_version(max(versions), prefix)
     return version
 
 
@@ -105,6 +107,10 @@ def find_msvc():
         return [], []
 
     version = max_version(os.listdir(msvc_base_path))
+    if version is None:
+        print("WARNING: Failed to find MSVC.")
+        return [], []
+
     return (
         [str(msvc_base_path / version / "include")],
         [str(msvc_base_path / version / "lib" / "x64")],
@@ -144,6 +150,10 @@ def find_winsdk():
         return [], []
 
     version = max_version(os.listdir(winsdk_base_path / "Include"))
+    if version is None:
+        print("WARNING: Failed to find Windows SDK.")
+        return [], []
+
     return (
         [
             str(winsdk_base_path / "Include" / version / "shared"),
@@ -174,3 +184,25 @@ def find_python():
         rf"C:\Python{version}\libs",
     ]
     return python_lib_dirs
+
+
+@functools.cache
+def find_cuda():
+    cuda_base_path = os.environ.get("CUDA_PATH")
+    if cuda_base_path is not None:
+        cuda_base_path = Path(cuda_base_path)
+        if not cuda_base_path.exists():
+            cuda_base_path = None
+
+    if cuda_base_path is None:
+        paths = glob(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12*")
+        if not paths:
+            return None, [], []
+        # Heuristic to find the highest version
+        cuda_base_path = Path(sorted(paths)[-1])
+
+    return (
+        str(cuda_base_path / "bin"),
+        [str(cuda_base_path / "include")],
+        [str(cuda_base_path / "lib" / "x64")],
+    )

third_party/nvidia/backend/driver.py

@@ -12,16 +12,18 @@
 dirname = os.path.dirname(os.path.realpath(__file__))
 include_dir = [os.path.join(dirname, "include")]
 if os.name == "nt":
-    cuda_path = os.environ.get("CUDA_PATH")
-    include_dir += [f"{cuda_path}\\include"]
+    from triton.windows_utils import find_cuda
+    _, cuda_inc_dirs, _ = find_cuda()
+    include_dir += cuda_inc_dirs
 libdevice_dir = os.path.join(dirname, "lib")
 libraries = ['cuda']
 
 
 @functools.lru_cache()
 def libcuda_dirs():
     if os.name == "nt":
-        return [os.path.join(cuda_path, "lib", "x64")]
+        _, _, cuda_lib_dirs = find_cuda()
+        return cuda_lib_dirs
 
     env_libcuda_path = os.getenv("TRITON_LIBCUDA_PATH")
     if env_libcuda_path: