Fix cmake scripts for CUDA and MSVC (#3713)

* Fix wrong CUDA generators and allow for new ones * Fix CUDA detection for other generators * Simplify the changed code * Remove useless flags for MSVC
pytorch · Nov 15, 2017 · 737aba3 · 737aba3
1 parent 3bb2308
commit 737aba3
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 16 deletions.
diff --git a/aten/CMakeLists.txt b/aten/CMakeLists.txt
@@ -25,11 +25,13 @@ if("${isSystemDir}" STREQUAL "-1")
   set(CMAKE_INSTALL_RPATH "${CMAKE_INSTALL_PREFIX}/lib")
 endif()
 
-# C++11
-set(CMAKE_CXX_FLAGS "--std=c++11 ${CMAKE_CXX_FLAGS}")
+IF(NOT MSVC)
+  # C++11
+  set(CMAKE_CXX_FLAGS "--std=c++11 ${CMAKE_CXX_FLAGS}")
 
-set(CMAKE_C_FLAGS "-fexceptions ${CMAKE_C_FLAGS}")
-set(CMAKE_CXX_FLAGS "-fexceptions ${CMAKE_CXX_FLAGS}")
+  set(CMAKE_C_FLAGS "-fexceptions ${CMAKE_C_FLAGS}")
+  set(CMAKE_CXX_FLAGS "-fexceptions ${CMAKE_CXX_FLAGS}")
+ENDIF(NOT MSVC)
 
 
 # Top-level build config
@@ -53,11 +55,13 @@ IF(MSVC)
   ENDIF()
 ENDIF(MSVC)
 
-IF (CMAKE_VERSION VERSION_LESS "3.1")
-  SET(CMAKE_C_FLAGS "-std=c11 ${CMAKE_C_FLAGS}")
-ELSE ()
-  SET(CMAKE_C_STANDARD 11)
-ENDIF ()
+IF (NOT MSVC)
+  IF (CMAKE_VERSION VERSION_LESS "3.1")
+    SET(CMAKE_C_FLAGS "-std=c11 ${CMAKE_C_FLAGS}")
+  ELSE ()
+    SET(CMAKE_C_STANDARD 11)
+  ENDIF ()
+ENDIF(NOT MSVC)
 
 if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
   if(CMAKE_CXX_COMPILER_VERSION VERSION_GREATER "4.9")
@@ -80,7 +84,7 @@ ENDIF()
 LIST(APPEND CUDA_NVCC_FLAGS $ENV{TORCH_NVCC_FLAGS})
 CUDA_SELECT_NVCC_ARCH_FLAGS(NVCC_FLAGS_EXTRA $ENV{TORCH_CUDA_ARCH_LIST})
 LIST(APPEND CUDA_NVCC_FLAGS ${NVCC_FLAGS_EXTRA})
-IF(CMAKE_POSITION_INDEPENDENT_CODE)
+IF(CMAKE_POSITION_INDEPENDENT_CODE AND NOT MSVC)
   LIST(APPEND CUDA_NVCC_FLAGS "-Xcompiler -fPIC")
 ENDIF()
 

diff --git a/aten/cmake/FindCUDA/FindCUDA.cmake b/aten/cmake/FindCUDA/FindCUDA.cmake
@@ -487,8 +487,9 @@ option(CUDA_HOST_COMPILATION_CPP "Generated file extension" ON)
 # Extra user settable flags
 set(CUDA_NVCC_FLAGS "" CACHE STRING "Semi-colon delimit multiple arguments.")
 
-if(CMAKE_GENERATOR MATCHES "Visual Studio")
-  set(CUDA_HOST_COMPILER "$(VCInstallDir)bin" CACHE FILEPATH "Host side compiler used by NVCC")
+if(MSVC AND NOT "${CMAKE_C_COMPILER}" MATCHES "/cl.exe")
+  find_program(REAL_MSVC_COMPILER cl.exe)
+  set(CUDA_HOST_COMPILER "${REAL_MSVC_COMPILER}" CACHE FILEPATH "Host side compiler used by NVCC")
 else()
   if(APPLE
       AND "${CMAKE_C_COMPILER_ID}" MATCHES "Clang"

diff --git a/aten/cmake/FindCUDA/FindCUDA/select_compute_arch.cmake b/aten/cmake/FindCUDA/FindCUDA/select_compute_arch.cmake
@@ -62,8 +62,15 @@ function(CUDA_DETECT_INSTALLED_GPUS OUT_VARIABLE)
       "  return 0;\n"
       "}\n")
 
+    if(MSVC AND NOT "${CMAKE_C_COMPILER}" MATCHES "/cl.exe")
+      find_program(REAL_MSVC_COMPILER cl.exe)
+      set(CCBIN "${REAL_MSVC_COMPILER}")
+    else()
+      set(CCBIN "${CMAKE_C_COMPILER}")
+    endif()
+
     execute_process(COMMAND "${CUDA_NVCC_EXECUTABLE}" "--run" "${cufile}"
-                    "-ccbin" ${CMAKE_CXX_COMPILER}
+                    "-ccbin" ${CCBIN}
                     WORKING_DIRECTORY "${PROJECT_BINARY_DIR}/CMakeFiles/"
                     RESULT_VARIABLE nvcc_res OUTPUT_VARIABLE nvcc_out
                     ERROR_QUIET OUTPUT_STRIP_TRAILING_WHITESPACE)

diff --git a/cmake/FindCUDA/FindCUDA.cmake b/cmake/FindCUDA/FindCUDA.cmake
@@ -487,8 +487,9 @@ option(CUDA_HOST_COMPILATION_CPP "Generated file extension" ON)
 # Extra user settable flags
 set(CUDA_NVCC_FLAGS "" CACHE STRING "Semi-colon delimit multiple arguments.")
 
-if(CMAKE_GENERATOR MATCHES "Visual Studio")
-  set(CUDA_HOST_COMPILER "${CMAKE_C_COMPILER}" CACHE FILEPATH "Host side compiler used by NVCC")
+if(MSVC AND NOT "${CMAKE_C_COMPILER}" MATCHES "/cl.exe")
+  find_program(REAL_MSVC_COMPILER cl.exe)
+  set(CUDA_HOST_COMPILER "${REAL_MSVC_COMPILER}" CACHE FILEPATH "Host side compiler used by NVCC")
 else()
   if(APPLE
       AND "${CMAKE_C_COMPILER_ID}" MATCHES "Clang"

diff --git a/cmake/FindCUDA/FindCUDA/select_compute_arch.cmake b/cmake/FindCUDA/FindCUDA/select_compute_arch.cmake
@@ -62,8 +62,15 @@ function(CUDA_DETECT_INSTALLED_GPUS OUT_VARIABLE)
       "  return 0;\n"
       "}\n")
 
+    if(MSVC AND NOT "${CMAKE_C_COMPILER}" MATCHES "/cl.exe")
+      find_program(REAL_MSVC_COMPILER cl.exe)
+      set(CCBIN "${REAL_MSVC_COMPILER}")
+    else()
+      set(CCBIN "${CMAKE_C_COMPILER}")
+    endif()
+
     execute_process(COMMAND "${CUDA_NVCC_EXECUTABLE}" "--run" "${cufile}"
-                    "-ccbin" ${CMAKE_CXX_COMPILER}
+                    "-ccbin" ${CCBIN}
                     WORKING_DIRECTORY "${PROJECT_BINARY_DIR}/CMakeFiles/"
                     RESULT_VARIABLE nvcc_res OUTPUT_VARIABLE nvcc_out
                     ERROR_QUIET OUTPUT_STRIP_TRAILING_WHITESPACE)