diff --git a/cmake/modules/ROCM.cmake b/cmake/modules/ROCM.cmake index b22e9f2532b7..ca8682392191 100644 --- a/cmake/modules/ROCM.cmake +++ b/cmake/modules/ROCM.cmake @@ -61,7 +61,7 @@ if(USE_ROCM) find_package(rocprim REQUIRED) find_package(rocthrust REQUIRED) set_source_files_properties(src/runtime/contrib/thrust/thrust.cu PROPERTIES LANGUAGE CXX) - file(GLOB CONTRIB_THRUST_SRC src/runtime/contrib/thrust/thrust.cu) + file(GLOB CONTRIB_THRUST_SRC src/runtime/contrib/thrust/*.cu) list(APPEND RUNTIME_SRCS ${CONTRIB_THRUST_SRC}) list(APPEND TVM_RUNTIME_LINKER_LIBS roc::rocthrust) endif(USE_THRUST) diff --git a/src/runtime/contrib/thrust/thrust.cu b/src/runtime/contrib/thrust/thrust.cu index 7295d4c47c3f..df83b57847a0 100644 --- a/src/runtime/contrib/thrust/thrust.cu +++ b/src/runtime/contrib/thrust/thrust.cu @@ -26,6 +26,7 @@ #include #include #include +#include #include #include