Allow setting CUDA capability.

kaldi-asr · Mar 18, 2021 · ea33e7b · ea33e7b
1 parent 1b2399b
commit ea33e7b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -124,8 +124,10 @@ endif()
 
 find_package(CUDA)
 if(CUDA_FOUND)
+    set(CUDA_CAPABILITY 50 CACHE STRING "CUDA capability without '.'")
     set(CUDA_PROPAGATE_HOST_FLAGS ON)
     set(KALDI_CUDA_NVCC_FLAGS "--default-stream=per-thread;-std=c++${CMAKE_CXX_STANDARD}")
+    list(APPEND KALDI_CUDA_NVCC_FLAGS "-gencode arch=compute_${CUDA_CAPABILITY},code=sm_${CUDA_CAPABILITY}")
     if(MSVC)
         list(APPEND KALDI_CUDA_NVCC_FLAGS "-Xcompiler /permissive-,/FS,/wd4819,/EHsc,/bigobj")
         list(APPEND KALDI_CUDA_NVCC_FLAGS "-Xcompiler /wd4018,/wd4244,/wd4267,/wd4291,/wd4305")