Repository: pytorch/pytorch

Issue:
figure out a way to load libcudart.so from it's compile-time path
right now if libcudart.so is not in LD_LIBRARY_PATH, though it was found at compile-time, these lines will fail: https://github.com/pytorch/pytorch/blob/master/torch/cuda/__init__.py#L39-L43

Avoid this, by idk doing something...

I think i have a few good ideas.



Discussion:
We need to add /usr/local/cuda/lib64 (or that's mac dir) to the rpath of _C

the problem i dont think is to add it to rpath, it's because we ctypes.load

Right. We could try with the default cuda installation path if it's not found from LD_LIBRARY_PATH.

Actually, do we really want to use the compile path? We should try to detect it at runtime, right? Or are we going to ship CUDA libs in our binaries?

we first load using runtime paths, and if it fails, fallback to path known at compile time.

I fear this is going to get brittle, especially if you are shipping CUDA libs.  e.g. a user swaps HW and picks up a new driver and/or toolkit.  As @soumith says, I'm pretty sure you want to use the runtime paths.  Similarly, CUDA 7.5 won't run on Pascal and we have a cudnn built for 7.5 that will not do the right thing on CUDA 8, so those need to go together.  If a user installs the CUDA bits via deb/rpm the right things are supposed to happen.  If you push bits, we need to figure out how to make sure things line up.  But this is perhaps a much larger packaging conversation.

@thatguymike Yeah, I've implemented it to first try loading it without any path and then try `/usr/local/cuda/lib64` and `/usr/local/cuda/lib`. I'll push the commit tomorrow.


Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -9,8 +9,8 @@
 import sys
 import os
 
-# TODO: make this more robust
-WITH_CUDA = os.path.exists('/Developer/NVIDIA/CUDA-7.5/include') or os.path.exists('/usr/local/cuda/include')
+CUDA_HOME = os.getenv('CUDA_HOME', '/usr/local/cuda')
+WITH_CUDA = os.path.exists(CUDA_HOME)
 WITH_CUDNN = WITH_CUDA
 DEBUG = False
 
@@ -176,18 +176,17 @@ def run(self):
     pass
 
 if WITH_CUDA:
-    if platform.system() == 'Darwin':
-        cuda_path = '/Developer/NVIDIA/CUDA-7.5'
-        cuda_include_path = cuda_path + '/include'
-        cuda_lib_path = cuda_path + '/lib'
-    else:
-        cuda_path = '/usr/local/cuda'
-        cuda_include_path = cuda_path + '/include'
-        cuda_lib_path = cuda_path + '/lib64'
+    cuda_lib_dirs = ['lib64', 'lib']
+    cuda_include_path = os.path.join(CUDA_HOME, 'include')
+    for lib_dir in cuda_lib_dirs:
+        cuda_lib_path = os.path.join(CUDA_HOME, lib_dir)
+        if os.path.exists(cuda_lib_path):
+            break
     include_dirs.append(cuda_include_path)
     extra_link_args.append('-L' + cuda_lib_path)
     extra_link_args.append('-Wl,-rpath,' + cuda_lib_path)
     extra_compile_args += ['-DWITH_CUDA']
+    extra_compile_args += ['-DCUDA_LIB_PATH=' + cuda_lib_path]
     main_libraries += ['THC']
     main_sources += [
         "torch/csrc/cuda/Module.cpp",
diff --git a/torch/cuda/__init__.py b/torch/cuda/__init__.py
--- a/torch/cuda/__init__.py
+++ b/torch/cuda/__init__.py
@@ -2,6 +2,7 @@
 import contextlib
 import platform
 import ctypes
+import os
 import torch
 
 _initialized = False
@@ -11,10 +12,28 @@ def is_available():
     return (hasattr(torch._C, '_cuda_isDriverSufficient') and
             torch._C._cuda_isDriverSufficient())
 
-def _lazy_init():
-    global _initialized, _cudart
-    if _initialized:
-        return
+
+def _load_cudart():
+    system = platform.system()
+    lib_name = 'libcudart.' + ('dylib' if system == 'Darwin' else 'so')
+    lib_paths = [
+        lib_name,
+        os.path.join(torch._C._cuda_getLibPath(), lib_name),
+        os.path.join('/usr/local/cuda/lib64', lib_name),
+        os.path.join('/usr/local/cuda/lib', lib_name),
+    ]
+    for path in lib_paths:
+        try:
+            return ctypes.cdll.LoadLibrary(path)
+        except OSError:
+            pass
+    raise RuntimeError("couldn't find libcudart. Make sure CUDA libraries "
+        "are installed in a default location, or that they're in " +
+        ("DYLD_LIBRARY_PATH" if system == 'Darwin' else "LD_LIBRARY_PATH") +
+        ".")
+
+
+def _check_driver():
     if not hasattr(torch._C, '_cuda_isDriverSufficient'):
         raise AssertionError("Torch not compiled with CUDA enabled")
     if not torch._C._cuda_isDriverSufficient():
@@ -33,14 +52,18 @@ def _lazy_init():
 Alternatively, go to: https://pytorch.org/binaries to install
 a PyTorch version that has been compiled with your version
 of the CUDA driver.""".format(str(torch._C._cuda_getDriverVersion())))
+
+
+def _lazy_init():
+    global _initialized, _cudart
+    if _initialized:
+        return
+    _check_driver()
     assert torch._C._cuda_init()
-    _initialized = True
-    if platform.system() == 'Darwin':
-        _cudart = ctypes.cdll.LoadLibrary('libcudart.dylib')
-    else:
-        _cudart = ctypes.cdll.LoadLibrary('libcudart.so')
+    _cudart = _load_cudart()
     _cudart.cudaGetErrorName.restype = ctypes.c_char_p
     _cudart.cudaGetErrorString.restype = ctypes.c_char_p
+    _initialized = True
 
 
 def cudart():


Test patch: