# Set up CUDA environment

# Summary of CMake flags and compile definitions:
#
# CMake flags:
#   CUDA_INCLUDE_DIRS: path of CUDA headers
#   CUDA_LIBRARIES: names of cuda libraries
#
# Compile definitions:
#   WITH_CUDA: whether to enable CUDA support

# NOTES:
#
# 1. When using NVHPC SDK, you must also set CMAKE_C_COMPILER explicitly, otherwise
#    there may be compiling errors complaining no compatible NVC++ compiler available.
#
# 2. It's recommended to use GCC with CUDA, as CUDA has strict restriction on the
#    version of clang/intel compilers.
#
# 3. Perform 'module load nvhpc' in the terminal rather than in the script if you run
#    into errors about missing cublas_v2.h.
#
# References:
# CUDA toolkit: https://developer.nvidia.cn/cuda-downloads
# NVHPC toolkit: https://developer.nvidia.com/hpc-sdk-downloads
# Difference of CUDA and NVHPC:
#   https://forums.developer.nvidia.com/t/what-is-the-difference-between-cuda-toolkit-and-nvidia-hpc-sdk/289582
message(STATUS "")
enable_language(CUDA)

cmake_policy(SET CMP0146 OLD)
find_package(CUDA REQUIRED)

if(WIN32)
    set(_cuda_lib_dir ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64)
    find_library(CUSPARSE_LIBRARY NAMES cusparse PATHS ${_cuda_lib_dir})
    find_library(CUBLAS_LIBRARY NAMES cublas PATHS ${_cuda_lib_dir})
else()
    # On Linux search for libraries in $LD_LIBRARY_PATH. It works for both
    # CUDA toolkit and HPC SDK.
    find_library(CUSPARSE_LIBRARY NAMES cusparse PATHS ENV LD_LIBRARY_PATH)
    find_library(CUBLAS_LIBRARY NAMES cublas PATHS ENV LD_LIBRARY_PATH)
endif()
unset(_cuda_lib_dir)

if(NOT CUSPARSE_LIBRARY IN_LIST CUDA_LIBRARIES)
    list(APPEND CUDA_LIBRARIES ${CUSPARSE_LIBRARY})  
endif()  
if(NOT CUBLAS_LIBRARY IN_LIST CUDA_LIBRARIES)  
    list(APPEND CUDA_LIBRARIES ${CUBLAS_LIBRARY})  
endif()

# NOTE: We make it a mandatory option that the user must specify CUDA_ARCHITECTURES
# via command-line option. Otherwise, the programs will be unusable. ARCH of GPUs
# can be found at: https://developer.nvidia.com/cuda-gpus#collapseOne
#
# Some examples:
#   A800: 80
#   1050Ti: 61
if(DEFINED CUDA_ARCHITECTURES)
    set(_cuda_arch "compute_${CUDA_ARCHITECTURES}")
    set(_cuda_code "sm_${CUDA_ARCHITECTURES}")
else()
    message(FATAL_ERROR "CUDA_ARCHITECTURES not specified!")
endif()
set(CUDA_NVCC_FLAGS " ${CUDA_NVCC_FLAGS} --compile -Xcompiler=-fPIC; -std=c++17 -gencode arch=${_cuda_arch},code=${_cuda_code}")
unset(_cuda_arch)
unset(_cuda_code)

if(CUDA_FOUND)
    message(STATUS "  CUDA support enabled")
    message(STATUS "    CUDA Version :" ${CUDA_VERSION_STRING})
    message(STATUS "    CUDA library :")
    foreach(__lib ${CUDA_LIBRARIES})
        message(STATUS "      " ${__lib})
    endforeach()
    message(STATUS "    CUDA FLAGS :" ${CUDA_NVCC_FLAGS})
    message(STATUS "    CUDA Executable :" ${CUDA_NVCC_EXECUTABLE})
else()
    message(FATAL_ERROR "  CUDA support not found")
endif()
message(STATUS "")

# Compile definitions and header directories should be defined PUBLIC,
# otherwise there may be inconsistencies in headers, especially for C++
# which uses templates.
function(tbp_set_cuda arg)
    target_compile_definitions(${arg} PUBLIC WITH_CUDA)
    target_include_directories(${arg} PUBLIC ${CUDA_INCLUDE_DIRS})
    # DO NOT use PUBLIC for target_link_libraries. For more details, see
    # cmake-xxxx/share/cmake-xxxx/Modules/FindCUDA.cmake for more details.
    target_link_libraries(${arg} ${CUDA_LINK_LIBRARIES_KEYWORD} ${CUDA_LIBRARIES})
endfunction()
