diff --git a/CMakeLists.txt b/CMakeLists.txt index 42cb46547614..f69a5016e0ec 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -598,8 +598,17 @@ if(USE_CUDA) include(${CMAKE_ROOT}/Modules/FindCUDA/select_compute_arch.cmake) CUDA_SELECT_NVCC_ARCH_FLAGS(CUDA_ARCH_FLAGS ${MXNET_CUDA_ARCH}) message("-- CUDA: Using the following NVCC architecture flags ${CUDA_ARCH_FLAGS}") + set(arch_code_list) + foreach(arch_str ${CUDA_ARCH_FLAGS}) + if((arch_str MATCHES ".*sm_[0-9]+")) + string( REGEX REPLACE ".*sm_([0-9]+)" "\\1" arch_code ${arch_str} ) + list(APPEND arch_code_list ${arch_code}) + endif() + endforeach() + MESSAGE( "${arch_code_list}") + string(REPLACE ";" " " CUDA_ARCH_FLAGS_SPACES "${CUDA_ARCH_FLAGS}") - string(APPEND CMAKE_CUDA_FLAGS " ${CUDA_ARCH_FLAGS_SPACES}") + find_package(CUDAToolkit REQUIRED cublas cufft cusolver curand OPTIONAL_COMPONENTS nvToolsExt nvrtc) @@ -672,6 +681,7 @@ add_library(sample_lib SHARED ${CMAKE_CURRENT_SOURCE_DIR}/example/extensions/lib target_include_directories(sample_lib PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include/mxnet) set(MXNET_INSTALL_TARGETS mxnet) if(UNIX) + string(APPEND CMAKE_CUDA_FLAGS "${CUDA_ARCH_FLAGS_SPACES}") # Create dummy file since we want an empty shared library before linking set(DUMMY_SOURCE ${CMAKE_BINARY_DIR}/dummy.c) file(WRITE ${DUMMY_SOURCE} "") @@ -688,14 +698,14 @@ elseif(MSVC) set_target_properties(sample_lib PROPERTIES PREFIX "lib") if(USE_CUDA) - if(FIRST_CUDA AND MSVC) + if(MSVC) if(USE_SPLIT_ARCH_DLL) add_executable(gen_warp tools/windowsbuild/gen_warp.cpp) add_library(mxnet SHARED tools/windowsbuild/warp_dll.cpp ${CMAKE_BINARY_DIR}/warp_gen_cpp.cpp ${CMAKE_BINARY_DIR}/warp_gen.asm) target_link_libraries(mxnet PRIVATE cudart Shlwapi) - list(GET cuda_arch 0 mxnet_first_arch) - foreach(arch ${cuda_arch}) + list(GET arch_code_list 0 mxnet_first_arch) + foreach(arch ${arch_code_list}) add_library(mxnet_${arch} SHARED ${SOURCE}) target_compile_options( mxnet_${arch} @@ -720,7 +730,7 @@ elseif(MSVC) COMMAND gen_warp $ WORKING_DIRECTORY ${CMAKE_BINARY_DIR}/ DEPENDS $) else(USE_SPLIT_ARCH_DLL) string(REPLACE ";" " " NVCC_FLAGS_ARCH "${NVCC_FLAGS_ARCH}") - set(CMAKE_CUDA_FLAGS "${NVCC_FLAGS_ARCH}") + set(CMAKE_CUDA_FLAGS "${CUDA_ARCH_FLAGS_SPACES}") add_library(mxnet SHARED ${SOURCE}) target_compile_options( mxnet @@ -781,7 +791,7 @@ if(USE_PLUGINS_WARPCTC) endif() if(MSVC) - if(FIRST_CUDA AND USE_SPLIT_ARCH_DLL) + if(USE_SPLIT_ARCH_DLL) foreach(arch ${cuda_arch}) target_link_libraries(mxnet_${arch} PUBLIC ${mxnet_LINKER_LIBS}) target_link_libraries(mxnet_${arch} PUBLIC dmlc)