67 lines
1.7 KiB
CMake
67 lines
1.7 KiB
CMake
# FindTensorRT.cmake -- Locate NVIDIA TensorRT
|
|
|
|
include(FindPackageHandleStandardArgs)
|
|
|
|
# TensorRT root
|
|
if (DEFINED TensorRT_ROOT)
|
|
list(APPEND _TensorRT_SEARCH_PATHS
|
|
${TensorRT_ROOT}
|
|
"$ENV{TensorRT_ROOT}"
|
|
)
|
|
endif()
|
|
list(APPEND _TensorRT_SEARCH_PATHS /usr /usr/local)
|
|
|
|
# Header
|
|
find_path(TensorRT_INCLUDE_DIR
|
|
NAMES NvInfer.h
|
|
PATHS ${_TensorRT_SEARCH_PATHS}
|
|
PATH_SUFFIXES include
|
|
)
|
|
|
|
# Core library
|
|
find_library(TensorRT_LIBRARY
|
|
NAMES nvinfer
|
|
PATHS ${_TensorRT_SEARCH_PATHS}
|
|
PATH_SUFFIXES lib lib64 lib/x64
|
|
)
|
|
|
|
find_package_handle_standard_args(TensorRT
|
|
REQUIRED_VARS TensorRT_INCLUDE_DIR TensorRT_LIBRARY
|
|
)
|
|
|
|
if (TensorRT_FOUND)
|
|
set(TensorRT_INCLUDE_DIRS ${TensorRT_INCLUDE_DIR})
|
|
set(TensorRT_LIBRARIES ${TensorRT_LIBRARY})
|
|
|
|
# Optional components
|
|
foreach(_comp IN ITEMS nvinfer_plugin nvonnxparser nvparsers)
|
|
find_library(TensorRT_${_comp}_LIBRARY
|
|
NAMES ${_comp}
|
|
PATHS ${_TensorRT_SEARCH_PATHS}
|
|
PATH_SUFFIXES lib lib64 lib/x64
|
|
)
|
|
if (TensorRT_${_comp}_LIBRARY)
|
|
list(APPEND TensorRT_LIBRARIES ${TensorRT_${_comp}_LIBRARY})
|
|
endif()
|
|
endforeach()
|
|
|
|
# Core target
|
|
add_library(TensorRT::TensorRT UNKNOWN IMPORTED)
|
|
set_target_properties(TensorRT::TensorRT PROPERTIES
|
|
INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}"
|
|
IMPORTED_LOCATION "${TensorRT_LIBRARY}"
|
|
)
|
|
|
|
# Component targets
|
|
foreach(_comp IN ITEMS nvinfer_plugin nvonnxparser nvparsers)
|
|
if (TensorRT_${_comp}_LIBRARY)
|
|
add_library(TensorRT::${_comp} UNKNOWN IMPORTED)
|
|
set_target_properties(TensorRT::${_comp} PROPERTIES
|
|
IMPORTED_LOCATION "${TensorRT_${_comp}_LIBRARY}"
|
|
)
|
|
endif()
|
|
endforeach()
|
|
|
|
message(STATUS "Found TensorRT at ${TensorRT_INCLUDE_DIR}")
|
|
endif()
|