add Pascal architecture (#192)

dmlc · Jan 7, 2017 · 9252a54 · 9252a54
1 parent bc4889b
commit 9252a54
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/cmake/Cuda.cmake b/cmake/Cuda.cmake
@@ -8,7 +8,7 @@ check_cxx_compiler_flag("-std=c++11"   SUPPORT_CXX11)
 
 # Known NVIDIA GPU achitectures mshadow can be compiled for.
 # This list will be used for CUDA_ARCH_NAME = All option
-set(mshadow_known_gpu_archs "20 21(20) 30 35 50")
+set(mshadow_known_gpu_archs "20 21(20) 30 35 50 60 61")
 
 ################################################################################################
 # A function for automatic detection of GPUs installed  (if autodetection is enabled)
@@ -71,7 +71,7 @@ endfunction()
 #   mshadow_select_nvcc_arch_flags(out_variable)
 function(mshadow_select_nvcc_arch_flags out_variable)
   # List of arch names
-  set(__archs_names "Fermi" "Kepler" "Maxwell" "All" "Manual")
+  set(__archs_names "Fermi" "Kepler" "Maxwell" "Pascal" "All" "Manual")
   set(__archs_name_default "All")
   if(NOT CMAKE_CROSSCOMPILING)
     list(APPEND __archs_names "Auto")
@@ -104,6 +104,8 @@ function(mshadow_select_nvcc_arch_flags out_variable)
     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 "All")
     set(__cuda_arch_bin ${mshadow_known_gpu_archs})
   elseif(${CUDA_ARCH_NAME} STREQUAL "Auto")