From 450c73d3a705e41e9cc470c88ffde868bc87e16f Mon Sep 17 00:00:00 2001
From: Sergei Bastrakov <sergey.bastrakov@gmail.com>
Date: Tue, 24 Sep 2019 11:41:08 +0200
Subject: [PATCH] Choose the value for the -arch nvcc flag depending on CUDA
 version

The original code has hardcoded 2.0, not supported since CUDA 9.
Closes #162.
---
 CMakeLists.txt | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 4376bd13..1713b8ac 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -22,7 +22,12 @@ endif()
 # CUDA
 ###############################################################################
 find_package(CUDA REQUIRED)
-set(CUDA_NVCC_FLAGS "-arch=sm_20;-use_fast_math;")
+if(CUDA_VERSION VERSION_LESS 9.0)
+    set(COMPUTE_CAPABILITY "20")
+else()
+    set(COMPUTE_CAPABILITY "30")
+endif()
+set(CUDA_NVCC_FLAGS "-arch=sm_${COMPUTE_CAPABILITY};-use_fast_math;")
 set(CUDA_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR})
 include_directories(${CUDA_INCLUDE_DIRS})
 cuda_include_directories(${CUDA_INCLUDE_DIRS})