Try to fix TORCH_CUDA_ARCH_LIST for PyTorch again (#7936)

ssnl · soumith · commit c5b895ac5071 · 2018-05-30T16:30:21.000-04:00
* try again

* use DEFINED

* use a loop

* Minor fixes
diff --git a/cmake/public/cuda.cmake b/cmake/public/cuda.cmake
@@ -299,28 +299,62 @@ function(caffe2_select_nvcc_arch_flags out_variable)
     unset(CUDA_ARCH_PTX CACHE)
   endif()
 
-  if($ENV{TORCH_CUDA_ARCH_LIST})
-    # Pass CUDA architecture directly
-    set(__cuda_arch_bin $ENV{TORCH_CUDA_ARCH_LIST})
-    message(STATUS "Set CUDA arch from TORCH_CUDA_ARCH_LIST: ${__cuda_arch_bin}")
-  elseif(${CUDA_ARCH_NAME} STREQUAL "Kepler")
-    set(__cuda_arch_bin "30 35")
-  elseif(${CUDA_ARCH_NAME} STREQUAL "Maxwell")
-    set(__cuda_arch_bin "50")
-  elseif(${CUDA_ARCH_NAME} STREQUAL "Pascal")
-    set(__cuda_arch_bin "60 61")
-  elseif(${CUDA_ARCH_NAME} STREQUAL "Volta")
-    set(__cuda_arch_bin "70")
-  elseif(${CUDA_ARCH_NAME} STREQUAL "All")
-    set(__cuda_arch_bin ${Caffe2_known_gpu_archs})
-  elseif(${CUDA_ARCH_NAME} STREQUAL "Manual")
-    set(__cuda_arch_bin ${CUDA_ARCH_BIN})
-    set(__cuda_arch_ptx ${CUDA_ARCH_PTX})
-  elseif(${CUDA_ARCH_NAME} STREQUAL "Auto")
-    caffe2_detect_installed_gpus(__cuda_arch_bin)
+  set(CUDA_ARCH_LIST)
+  if(DEFINED ENV{TORCH_CUDA_ARCH_LIST})
+    set(TORCH_CUDA_ARCH_LIST $ENV{TORCH_CUDA_ARCH_LIST})
+    string(REGEX REPLACE "[ \t]+" ";" TORCH_CUDA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}")
+    list(APPEND CUDA_ARCH_LIST ${TORCH_CUDA_ARCH_LIST})
+    message(STATUS "Set CUDA arch from TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}")
   else()
-    message(FATAL_ERROR "Invalid CUDA_ARCH_NAME")
+    list(APPEND CUDA_ARCH_LIST ${CUDA_ARCH_NAME})
+    message(STATUS "Set CUDA arch from CUDA_ARCH_NAME: ${CUDA_ARCH_NAME}")
   endif()
+  list(REMOVE_DUPLICATES CUDA_ARCH_LIST)
+
+  set(__cuda_arch_bin)
+  set(__cuda_arch_ptx)
+  foreach(arch_name ${CUDA_ARCH_LIST})
+    set(arch_bin)
+    set(arch_ptx)
+    set(add_ptx FALSE)
+    # Check to see if we are compiling PTX
+    if(arch_name MATCHES "(.*)\\+PTX$")
+      set(add_ptx TRUE)
+      set(arch_name ${CMAKE_MATCH_1})
+    endif()
+    if(arch_name MATCHES "(^[0-9]\\.[0-9](\\([0-9]\\.[0-9]\\))?)$")
+      set(arch_bin ${CMAKE_MATCH_1})
+      set(arch_ptx ${arch_bin})
+    else()
+      # Look for it in our list of known architectures
+     if(${arch_name} STREQUAL "Kepler")
+        set(arch_bin "30 35")
+      elseif(${arch_name} STREQUAL "Maxwell")
+        set(arch_bin "50")
+      elseif(${arch_name} STREQUAL "Pascal")
+        set(arch_bin "60 61")
+      elseif(${arch_name} STREQUAL "Volta")
+        set(arch_bin "70")
+      elseif(${arch_name} STREQUAL "All")
+        set(arch_bin ${Caffe2_known_gpu_archs})
+      elseif(${arch_name} STREQUAL "Manual")
+        set(arch_bin ${CUDA_ARCH_BIN})
+        set(arch_ptx ${CUDA_ARCH_PTX})
+        set(add_ptx TRUE)
+      elseif(${arch_name} STREQUAL "Auto")
+        caffe2_detect_installed_gpus(arch_bin)
+      else()
+        message(FATAL_ERROR "Unknown CUDA architecture name ${arch_name}")
+      endif()
+    endif()
+    list(APPEND __cuda_arch_bin ${arch_bin})
+    if(add_ptx)
+      if (NOT arch_ptx)
+        set(arch_ptx ${arch_bin})
+      endif()
+      list(APPEND __cuda_arch_ptx ${arch_ptx})
+    endif()
+  endforeach()
 
   # Remove dots and convert to lists
   string(REGEX REPLACE "\\." "" __cuda_arch_bin "${__cuda_arch_bin}")
