diff --git a/cmake/modules/CUDA.cmake b/cmake/modules/CUDA.cmake index 1284f85bec..1ebf4e2234 100644 --- a/cmake/modules/CUDA.cmake +++ b/cmake/modules/CUDA.cmake @@ -109,6 +109,30 @@ if(USE_CUDA) # Add CUDA builtins to RelaxVM tvm_file_glob(GLOB RELAX_VM_CUDA_BUILTIN_SRC_CC src/runtime/relax_vm/cuda/*.cc) list(APPEND RUNTIME_SRCS ${RELAX_VM_CUDA_BUILTIN_SRC_CC}) + + if(USE_CUTLASS) + if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + set(CMAKE_CUDA_COMPILER ${CUDA_TOOLKIT_ROOT_DIR}/bin/nvcc) + execute_process( + COMMAND ${CMAKE_CUDA_COMPILER} --version + OUTPUT_VARIABLE NVCC_VERSION_OUTPUT + ) + if(NVCC_VERSION_OUTPUT MATCHES "release ([^,]+),") + set(CUDA_VERSION "${CMAKE_MATCH_1}") + endif(NVCC_VERSION_OUTPUT MATCHES "release ([^,]+),") + message(STATUS "CUDA_VERSION=${CUDA_VERSION}") + + if(CUDA_VERSION VERSION_GREATER_EQUAL 11.8) + set(CMAKE_CUDA_ARCHITECTURES "75;80;86;89;90") + else(CUDA_VERSION VERSION_GREATER_EQUAL 11.8) + set(CMAKE_CUDA_ARCHITECTURES "75;80;86") + endif(CUDA_VERSION VERSION_GREATER_EQUAL 11.8) + message(STATUS "Set CMAKE_CUDA_ARCHITECTURES=${CMAKE_CUDA_ARCHITECTURES}") + else(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + message(STATUS "Found CMAKE_CUDA_ARCHITECTURES=${CMAKE_CUDA_ARCHITECTURES}") + endif(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + endif(USE_CUTLASS) + else(USE_CUDA) list(APPEND COMPILER_SRCS src/target/opt/build_cuda_off.cc) endif(USE_CUDA)