Fix cuda detection script (#22527)

peterjc123 · facebook-github-bot · commit ce8c9d9bd5fa · 2019-07-08T07:06:59.000-07:00
Summary: Fixes #22507 Pull Request resolved: #22527 Differential Revision: D16126220 Pulled By: ezyang fbshipit-source-id: eb05141282b0f058324da1b3d3cb34566f222a67
diff --git a/tools/setup_helpers/__init__.py b/tools/setup_helpers/__init__.py
@@ -1,4 +1,5 @@
 import os
+import sys
 
 
 def escape_path(path):
@@ -7,3 +8,17 @@ def escape_path(path):
     if os.path.sep != '/' and path is not None:
         return path.replace(os.path.sep, '/')
     return path
+
+
+def which(thefile):
+    path = os.environ.get("PATH", os.defpath).split(os.pathsep)
+    for d in path:
+        fname = os.path.join(d, thefile)
+        fnames = [fname]
+        if sys.platform == 'win32':
+            exts = os.environ.get('PATHEXT', '').split(os.pathsep)
+            fnames += [fname + ext for ext in exts]
+        for name in fnames:
+            if os.access(name, os.F_OK | os.X_OK) and not os.path.isdir(name):
+                return name
+    return None
diff --git a/tools/setup_helpers/cmake.py b/tools/setup_helpers/cmake.py
@@ -11,7 +11,7 @@
 import distutils.sysconfig
 from distutils.version import LooseVersion
 
-from . import escape_path
+from . import escape_path, which
 from .env import (IS_64BIT, IS_DARWIN, IS_WINDOWS,
                   DEBUG, REL_WITH_DEB_INFO,
                   check_env_flag, check_negative_env_flag)
@@ -22,20 +22,6 @@
 from .numpy_ import USE_NUMPY, NUMPY_INCLUDE_DIR
 
 
-def _which(thefile):
-    path = os.environ.get("PATH", os.defpath).split(os.pathsep)
-    for d in path:
-        fname = os.path.join(d, thefile)
-        fnames = [fname]
-        if IS_WINDOWS:
-            exts = os.environ.get('PATHEXT', '').split(os.pathsep)
-            fnames += [fname + ext for ext in exts]
-        for name in fnames:
-            if os.access(name, os.F_OK | os.X_OK) and not os.path.isdir(name):
-                return name
-    return None
-
-
 def _mkdir_p(d):
     try:
         os.makedirs(d)
@@ -47,7 +33,7 @@ def _mkdir_p(d):
 # Use ninja if it is on the PATH. Previous version of PyTorch required the
 # ninja python package, but we no longer use it, so we do not have to import it
 USE_NINJA = (not check_negative_env_flag('USE_NINJA') and
-             _which('ninja') is not None)
+             which('ninja') is not None)
 
 
 class CMake:
@@ -80,9 +66,9 @@ def _get_cmake_command():
         cmake_command = 'cmake'
         if IS_WINDOWS:
             return cmake_command
-        cmake3 = _which('cmake3')
+        cmake3 = which('cmake3')
         if cmake3 is not None:
-            cmake = _which('cmake')
+            cmake = which('cmake')
             if cmake is not None:
                 bare_version = CMake._get_version(cmake)
                 if (bare_version < LooseVersion("3.5.0") and
diff --git a/tools/setup_helpers/cuda.py b/tools/setup_helpers/cuda.py
@@ -4,27 +4,18 @@
 import ctypes.util
 from subprocess import Popen, PIPE
 
+from . import escape_path, which
 from .env import IS_WINDOWS, IS_LINUX, IS_DARWIN, check_env_flag, check_negative_env_flag
 
 LINUX_HOME = '/usr/local/cuda'
 WINDOWS_HOME = glob.glob('C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v*.*')
 
 
 def find_nvcc():
-    if IS_WINDOWS:
-        proc = Popen(['where', 'nvcc.exe'], stdout=PIPE, stderr=PIPE)
-    else:
-        proc = Popen(['which', 'nvcc'], stdout=PIPE, stderr=PIPE)
-    out, err = proc.communicate()
-    out = out.decode().strip()
-    if len(out) > 0:
-        if IS_WINDOWS:
-            if out.find('\r\n') != -1:
-                out = out.split('\r\n')[0]
-            out = os.path.abspath(os.path.join(os.path.dirname(out), ".."))
-            out = out.replace('\\', '/')
-            out = str(out)
-        return os.path.dirname(out)
+    nvcc = which('nvcc')
+    if nvcc is not None:
+        nvcc = escape_path(nvcc)
+        return os.path.dirname(nvcc)
     else:
         return None
 
