diff --git a/.gitignore b/.gitignore index 2b2944940462e5..e2f20396e62190 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,7 @@ torch.egg-info/ */**/__pycache__ torch/csrc/generic/TensorMethods.cpp torch/lib/*.so* +torch/lib/*.dylib* torch/lib/*.h torch/lib/build torch/lib/tmp_install @@ -12,3 +13,5 @@ torch/csrc/nn/THNN.cpp torch/csrc/nn/THCUNN.cwrap torch/csrc/nn/THCUNN.cpp */**/*.pyc +*/**/*.so* +*/**/*.dylib* diff --git a/setup.py b/setup.py index a3597e18016d9d..d13fe3c61ed4a1 100644 --- a/setup.py +++ b/setup.py @@ -151,9 +151,16 @@ def run(self): if WITH_CUDA: if platform.system() == 'Darwin': - include_dirs += ['/Developer/NVIDIA/CUDA-7.5/include'] + cuda_path = '/Developer/NVIDIA/CUDA-7.5' + cuda_include_path = cuda_path + '/include' + cuda_lib_path = cuda_path + '/lib' else: - include_dirs += ['/usr/local/cuda/include'] + cuda_path = '/usr/local/cuda' + cuda_include_path = cuda_path + '/include' + cuda_lib_path = cuda_path + '/lib64' + 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'] main_libraries += ['THC'] main_sources += [ @@ -167,6 +174,13 @@ def run(self): extra_compile_args += ['-O0', '-g'] extra_link_args += ['-O0', '-g'] + +def make_relative_rpath(path): + if platform.system() == 'Darwin': + return '-Wl,-rpath,@loader_path/' + path + else: + return '-Wl,-rpath,$ORIGIN/' + path + ################################################################################ # Declare extensions and package ################################################################################ @@ -180,7 +194,7 @@ def run(self): language='c++', extra_compile_args=extra_compile_args, include_dirs=include_dirs, - extra_link_args=extra_link_args + ['-Wl,-rpath,$ORIGIN/lib'], + extra_link_args=extra_link_args + [make_relative_rpath('lib')] ) extensions.append(C) @@ -190,7 +204,7 @@ def run(self): language='c++', extra_compile_args=extra_compile_args, include_dirs=include_dirs, - extra_link_args=extra_link_args + ['-Wl,-rpath,$ORIGIN/../lib'], + extra_link_args=extra_link_args + [make_relative_rpath('../lib')] ) extensions.append(THNN) @@ -201,7 +215,7 @@ def run(self): language='c++', extra_compile_args=extra_compile_args, include_dirs=include_dirs, - extra_link_args=extra_link_args + ['-Wl,-rpath,$ORIGIN/../lib'], + extra_link_args=extra_link_args + [make_relative_rpath('../lib')] ) extensions.append(THCUNN) packages += ['torch.cuda', 'torch.legacy.cunn'] diff --git a/torch/csrc/cuda/Module.cpp b/torch/csrc/cuda/Module.cpp index bdb45fd5a12e29..36da32c07f34bc 100644 --- a/torch/csrc/cuda/Module.cpp +++ b/torch/csrc/cuda/Module.cpp @@ -177,7 +177,7 @@ PyObject * THCPModule_initExtension(PyObject *self) PyObject *torch_module = PyImport_ImportModule("torch.cuda"); if (!torch_module) { THPUtils_setError("class loader couldn't access torch module"); - return false; + return NULL; } PyObject* module_dict = PyModule_GetDict(torch_module); return PyBool_FromLong(THCPModule_initCuda(module_dict)); diff --git a/torch/lib/build_all.sh b/torch/lib/build_all.sh index 36e5ea2d64c774..eb3aa1b96951d9 100755 --- a/torch/lib/build_all.sh +++ b/torch/lib/build_all.sh @@ -4,8 +4,14 @@ set -e BASE_DIR=$(pwd) cd torch/lib INSTALL_DIR=$(pwd)/tmp_install -BASIC_FLAGS=" -DTH_INDEX_BASE=0 -I$INSTALL_DIR/include -I$INSTALL_DIR/include/TH -I$INSTALL_DIR/include/THC -L$INSTALL_DIR/lib " -FLAGS="$BASIC_FLAGS -Wl,-rpath,\$ORIGIN" +BASIC_FLAGS=" -DTH_INDEX_BASE=0 -I$INSTALL_DIR/include -I$INSTALL_DIR/include/TH -I$INSTALL_DIR/include/THC " +LDFLAGS="-L$INSTALL_DIR/lib " +if [[ $(uname) == 'Darwin' ]]; then + LDFLAGS="$LDFLAGS -Wl,-rpath,@loader_path" +else + LDFLAGS="$LDFLAGS -Wl,-rpath,\$ORIGIN" +fi +FLAGS="$BASIC_FLAGS $LDFLAGS" function build() { mkdir -p build/$1 cd build/$1 @@ -18,6 +24,15 @@ function build() { -DTH_INCLUDE_PATH="$INSTALL_DIR/include" make install -j$(getconf _NPROCESSORS_ONLN) cd ../.. + + if [[ $(uname) == 'Darwin' ]]; then + cd tmp_install/lib + for lib in *.dylib; do + echo "Updating install_name for $lib" + install_name_tool -id @rpath/$lib $lib + done + cd ../.. + fi } mkdir -p tmp_install