Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
tensorrt.cmake 1.31 KB
Newer Older
limm's avatar
add config module  
limm committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37 
# Copyright (c) OpenMMLab. All rights reserved.

include(${CMAKE_SOURCE_DIR}/cmake/modules/FindTENSORRT.cmake)
include(${CMAKE_SOURCE_DIR}/cmake/modules/FindCUDNN.cmake)
find_path(
        TENSORRT_INCLUDE_DIR NvInfer.h
        HINTS ${TENSORRT_DIR} ${CUDA_TOOLKIT_ROOT_DIR}
        PATH_SUFFIXES include)
if (TENSORRT_INCLUDE_DIR)
    message(STATUS "Found TensorRT headers at ${TENSORRT_INCLUDE_DIR}")
else ()
    message(ERROR "Cannot find TensorRT headers")
endif ()

find_library(
        TENSORRT_LIBRARY_INFER nvinfer
        HINTS ${TENSORRT_DIR} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR}
        PATH_SUFFIXES lib lib64 lib/x64)
find_library(
        TENSORRT_LIBRARY_INFER_PLUGIN nvinfer_plugin
        HINTS ${TENSORRT_DIR} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR}
        PATH_SUFFIXES lib lib64 lib/x64)
set(TENSORRT_LIBRARY ${TENSORRT_LIBRARY_INFER}
        ${TENSORRT_LIBRARY_INFER_PLUGIN})
if (TENSORRT_LIBRARY_INFER
        AND TENSORRT_LIBRARY_INFER_PLUGIN)
    message(STATUS "Found TensorRT libs at ${TENSORRT_LIBRARY}")
else ()
    message(FATAL_ERROR "Cannot find TensorRT libs")
endif ()

include(FindPackageHandleStandardArgs)
find_package_handle_standard_args(TENSORRT DEFAULT_MSG TENSORRT_INCLUDE_DIR
        TENSORRT_LIBRARY)
if (NOT TENSORRT_FOUND)
    message(ERROR "Cannot find TensorRT library.")
endif ()