Merge pull request #287 from iblis17/issue-286

build: improve cuda detection
dmlc · Oct 25, 2017 · a79e33c · a79e33c
2 parents 81b134f + 077a8cb
commit a79e33c
Showing 1 changed file with 16 additions and 6 deletions.
diff --git a/deps/build.jl b/deps/build.jl
@@ -37,7 +37,6 @@ end
 
 HAS_CUDA = false
 HAS_CUDNN = false
-cudnnver = -1
 let cudalib = Libdl.find_library(["libcuda", "nvcuda.dll"], CUDAPATHS)
   HAS_CUDA = !isempty(cudalib) && Libdl.dlopen_e(cudalib) != C_NULL
 end
@@ -53,18 +52,18 @@ end
 if HAS_CUDA  # then check cudnn
   let cudnnlib = Libdl.find_library("libcudnn", CUDAPATHS)
     HAS_CUDNN = !isempty(cudnnlib) && Libdl.dlopen_e(cudnnlib) != C_NULL
-    if HAS_CUDNN
-      # TODO: do more version check?
-      cudnnver = dec(ccall((:cudnnGetVersion, cudnnlib), Csize_t, ()))
+    if HAS_CUDNN && !haskey(ENV, "CUDA_HOME")  # inference `CUDA_HOME`
+      ENV["CUDA_HOME"] = dirname(dirname(cudnnlib))
     end
   end
 end
 
 if HAS_CUDA
   info("Found a CUDA installation.")
   if HAS_CUDNN
-    info("Found a CuDNN installation (version -> $cudnnver).")
+    info("Found a CuDNN installation.")
   end
+  info("CUDA_HOME -> $(ENV["CUDA_HOME"])")
 else
   info("Did not find a CUDA installation, using CPU-only version of MXNet.")
 end
@@ -170,7 +169,11 @@ if !libmxnet_detected
           ChangeDirectory(_mxdir)
           `git submodule deinit --force .`
           `git fetch`
-          `git checkout $libmxnet_curr_ver`
+          if libmxnet_curr_ver != "master"
+            `git checkout $libmxnet_curr_ver`
+          else
+            `git merge --ff origin/$libmxnet_curr_ver`
+          end
           `git submodule update --init --recursive`
           `git -C mshadow checkout -- make/mshadow.mk`
           `make clean`
@@ -192,9 +195,16 @@ if !libmxnet_detected
           if HAS_CUDA
             @build_steps begin
               `sed -i -s 's/USE_CUDA = 0/USE_CUDA = 1/' config.mk`
+              # address https://github.com/apache/incubator-mxnet/pull/7856
+              `sed -i -s "s/ADD_LDFLAGS =\(.*\)/ADD_LDFLAGS =\1 -lcublas -lcusolver -lcurand -lcudart/" config.mk`
               if haskey(ENV, "CUDA_HOME")
                 `sed -i -s "s@USE_CUDA_PATH = NONE@USE_CUDA_PATH = $(ENV["CUDA_HOME"])@" config.mk`
               end
+              if haskey(ENV, "CUDA_HOME")
+                # address https://github.com/apache/incubator-mxnet/pull/7838
+                flag = "-L$(ENV["CUDA_HOME"])/lib64 -L$(ENV["CUDA_HOME"])/lib"
+                `sed -i -s "s@ADD_LDFLAGS =\(.*\)@ADD_LDFLAGS =\1 $flag@" config.mk`
+              end
               if HAS_CUDNN
                 `sed -i -s 's/USE_CUDNN = 0/USE_CUDNN = 1/' config.mk`
               end