wehub-resource-sync
47 lines
1.5 KiB
CMake
47 lines
1.5 KiB
CMake
function(add_thrust_patches_if_necessary)
|
|
# ROCm 7.0+ has rocThrust shuffle support built-in, so no patches needed.
|
|
if(WITH_ROCM
|
|
AND DEFINED PADDLE_ROCM_VERSION
|
|
AND PADDLE_ROCM_VERSION GREATER_EQUAL 70000000)
|
|
message(STATUS "ROCm 7.0+ detected, skipping thrust patches")
|
|
return()
|
|
endif()
|
|
|
|
# ROCm < 7.0 still needs thrust patches.
|
|
if(WITH_ROCM)
|
|
set(thrust_patches "${PADDLE_SOURCE_DIR}/patches/thrust")
|
|
message(STATUS "ROCm < 7.0 detected, add thrust patches: ${thrust_patches}")
|
|
include_directories(${thrust_patches})
|
|
return()
|
|
endif()
|
|
|
|
# For CUDA, check if thrust has shuffle support
|
|
if(WITH_GPU)
|
|
set(thrust_detect_file ${PROJECT_BINARY_DIR}/detect_thrust.cu)
|
|
file(
|
|
WRITE ${thrust_detect_file}
|
|
""
|
|
"#include \"thrust/version.h\"\n"
|
|
"#include \"thrust/shuffle.h\"\n"
|
|
"#include \"stdio.h\"\n"
|
|
"int main() {\n"
|
|
" int version = THRUST_VERSION;\n"
|
|
" printf(\"%d\", version);\n"
|
|
" return 0;\n"
|
|
"}\n")
|
|
|
|
execute_process(
|
|
COMMAND "${CUDA_NVCC_EXECUTABLE}" "--run" "${thrust_detect_file}"
|
|
WORKING_DIRECTORY "${PROJECT_BINARY_DIR}/CMakeFiles/"
|
|
RESULT_VARIABLE nvcc_res
|
|
ERROR_QUIET)
|
|
if(NOT nvcc_res EQUAL 0)
|
|
set(thrust_patches "${PADDLE_SOURCE_DIR}/patches/thrust")
|
|
message(STATUS "Add thrust patches: ${thrust_patches}")
|
|
include_directories(${thrust_patches})
|
|
endif()
|
|
endif()
|
|
endfunction()
|
|
|
|
add_thrust_patches_if_necessary()
|