mmdeploy/cmake/modules/FindTENSORRT.cmake

52 lines
1.8 KiB
CMake

# Copyright (c) OpenMMLab. All rights reserved.
if (NOT DEFINED TENSORRT_DIR)
set(TENSORRT_DIR $ENV{TENSORRT_DIR})
endif ()
if (NOT TENSORRT_DIR)
message(FATAL_ERROR "Please set TENSORRT_DIR with cmake -D option.")
endif()
find_path(
TENSORRT_INCLUDE_DIR NvInfer.h
HINTS ${TENSORRT_DIR} ${CUDA_TOOLKIT_ROOT_DIR}
PATH_SUFFIXES include)
if (NOT TENSORRT_INCLUDE_DIR)
message(FATAL_ERROR "Cannot find TensorRT header NvInfer.h "
"in TENSORRT_DIR: ${TENSORRT_DIR} or in CUDA_TOOLKIT_ROOT_DIR: "
"${CUDA_TOOLKIT_ROOT_DIR}, please check if the path is correct.")
endif ()
set(__TENSORRT_LIB_COMPONENTS nvinfer;nvinfer_plugin)
foreach(__component ${__TENSORRT_LIB_COMPONENTS})
find_library(
__component_path ${__component}
HINTS ${TENSORRT_DIR} ${CUDA_TOOLKIT_ROOT_DIR}
PATH_SUFFIXES lib lib64 lib/x64)
if (NOT __component_path)
message(FATAL_ERROR "Cannot find TensorRT lib ${__component} in "
"TENSORRT_DIR: ${TENSORRT_DIR} or CUDA_TOOLKIT_ROOT_DIR: ${CUDA_TOOLKIT_ROOT_DIR}, "
"please check if the path is correct")
endif()
add_library(${__component} SHARED IMPORTED)
set_property(TARGET ${__component} APPEND PROPERTY IMPORTED_CONFIGURATIONS RELEASE)
if (MSVC)
set_target_properties(
${__component} PROPERTIES
IMPORTED_IMPLIB_RELEASE ${__component_path}
INTERFACE_INCLUDE_DIRECTORIES ${TENSORRT_INCLUDE_DIR}
)
else()
set_target_properties(
${__component} PROPERTIES
IMPORTED_LOCATION_RELEASE ${__component_path}
INTERFACE_INCLUDE_DIRECTORIES ${TENSORRT_INCLUDE_DIR}
)
endif()
unset(__component_path CACHE)
endforeach()
set(TENSORRT_LIBS ${__TENSORRT_LIB_COMPONENTS})