• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# The CUDA toolkit 11 and higher versions include cub library,
2# but the lower version (CUDA 10) doesn't support.
3if(USE_CUDA)
4    find_package(CUDA REQUIRED)
5    set(cub_CFLAGS "-D_FORTIFY_SOURCE=2 -O2")
6    find_path(CUB_INCLUDE_DIRS
7            HINTS "${CUDA_INCLUDE_DIRS}"
8            NAMES cub/cub.cuh
9            DOC "The directory where cub library reside.")
10    include(FindPackageHandleStandardArgs)
11    find_package_handle_standard_args(cub
12            FOUND_VAR CUB_FOUND
13            REQUIRED_VARS CUB_INCLUDE_DIRS)
14    if(CUB_FOUND)
15        include_directories(${CUB_INCLUDE_DIRS})
16    else()
17        set(REQ_URL "https://github.com/NVlabs/cub/archive/1.8.0.zip")
18        set(SHA256 "6bfa06ab52a650ae7ee6963143a0bbc667d6504822cbd9670369b598f18c58c3")
19        set(INCLUDE "cub")
20        mindspore_add_pkg(cub
21                VER 1.8.0
22                HEAD_ONLY ${INCLUDE}
23                URL ${REQ_URL}
24                SHA256 ${SHA256})
25        include_directories(${cub_INC}/../)
26    endif()
27endif()
28