chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,462 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
cmake_minimum_required(VERSION 3.31 FATAL_ERROR)
|
||||
|
||||
project(TRTPyBinds LANGUAGES CXX)
|
||||
|
||||
option(TRT_BUILD_ENABLE_NEW_PYTHON_FLOW "Use new build logic based on the main CMake build." OFF)
|
||||
|
||||
if (${TRT_BUILD_ENABLE_NEW_PYTHON_FLOW})
|
||||
|
||||
if(MSVC)
|
||||
set(DEFAULT_PY_EXT_PATH "${TOOLS_BASE}/win32")
|
||||
else()
|
||||
set(DEFAULT_PY_EXT_PATH "/externals")
|
||||
endif()
|
||||
set(TRT_BUILD_PYTHON_EXTERNALS_PATH ${DEFAULT_PY_EXT_PATH} CACHE PATH "Path to the parent folder of the many versioned python headers/libs.")
|
||||
|
||||
set(BUILD_PYTHON_PY_VERSIONS 3.8 3.9 3.10 3.11 3.12 3.13 3.14)
|
||||
|
||||
|
||||
set(TRT_BUILD_PYTHON_PY_VERSIONS ${BUILD_PYTHON_PY_VERSIONS} CACHE STRING "The list of python versions to build TensorRT bindings for.")
|
||||
|
||||
message(STATUS "TRT_BUILD_PYTHON_PY_VERSIONS: ${TRT_BUILD_PYTHON_PY_VERSIONS}")
|
||||
|
||||
if(NOT ${TRT_BUILD_WINML})
|
||||
set(TRT_PYTHON_MODULE_NAMES
|
||||
"tensorrt"
|
||||
"tensorrt_lean"
|
||||
"tensorrt_dispatch")
|
||||
else()
|
||||
set(TRT_PYTHON_MODULE_NAMES "tensorrt_rtx")
|
||||
endif()
|
||||
|
||||
# The "main" tensorrt bindings depend on the parser, so if we aren't building it, we need to skip it.
|
||||
if(NOT ${TRT_BUILD_ONNX_PARSER})
|
||||
message(STATUS "Not building the tensorrt python bindings as the ONNX Parser was disabled.")
|
||||
list(REMOVE_ITEM TRT_PYTHON_MODULE_NAMES "tensorrt")
|
||||
endif()
|
||||
|
||||
set(CMAKE_CXX_STANDARD 20 CACHE STRING "")
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED ON CACHE BOOL "")
|
||||
set(CMAKE_CXX_EXTENSIONS OFF CACHE BOOL "")
|
||||
|
||||
find_package(
|
||||
Python3
|
||||
COMPONENTS Interpreter
|
||||
REQUIRED
|
||||
)
|
||||
|
||||
# Disable automatic python detection since we need to build bindings for many python versions in one go.
|
||||
set(PYBIND11_NOPYTHON ON CACHE INTERNAL "")
|
||||
include(FetchPyBind11)
|
||||
|
||||
# Pybind11 would normally enable this by default, but does not do so under NOPYTHON mode, so we do it manually.
|
||||
if(${TRT_BUILD_PLATFORM} STREQUAL ${TRT_PLATFORM_X86})
|
||||
set(CMAKE_INTERPROCEDURAL_OPTIMIZATION ON)
|
||||
endif()
|
||||
|
||||
add_custom_target(tensorrt_python_bindings)
|
||||
|
||||
# Creates the binding library for the specified module and python version.
|
||||
#
|
||||
# \param moduleName The module name to create the bindings for. One of "tensorrt", "tensorrt_dispatch", or "tensorrt_lean".
|
||||
# \param pyVersion The python version to create bindings for, i.e. "3.12".
|
||||
function(createBindingLibrary moduleName pyVersion)
|
||||
|
||||
set(libName tensorrt_bindings_${moduleName}_${pyVersion})
|
||||
|
||||
add_library(${libName} MODULE)
|
||||
|
||||
# Set options unique to the "full" bindings.
|
||||
# The subdirs will use the value of TRT_PYTHON_IS_FULL_BINDINGS to set sources appropriately.
|
||||
if(${moduleName} STREQUAL "tensorrt" OR ${moduleName} STREQUAL "tensorrt_rtx")
|
||||
target_compile_definitions(${libName} PRIVATE
|
||||
tensorrt_EXPORTS=1
|
||||
)
|
||||
set(TRT_PYTHON_IS_FULL_BINDINGS ON)
|
||||
else()
|
||||
set(TRT_PYTHON_IS_FULL_BINDINGS OFF)
|
||||
endif()
|
||||
|
||||
function(add_${libName}_source)
|
||||
target_sources(${libName} PRIVATE ${ARGN})
|
||||
endfunction()
|
||||
|
||||
# Create an indirect refernce to the add_${libName}_source function which can be called by the subdirectories.
|
||||
# This allows each subdir to add files to the individual targets with unique binary dirs on each call.
|
||||
set(ADD_SOURCES_FUNCTION add_${libName}_source)
|
||||
set(SUBDIR_BINARY_DIR_PREFIX subbuild/${libName})
|
||||
add_subdirectory(src ${SUBDIR_BINARY_DIR_PREFIX}/src)
|
||||
|
||||
target_link_libraries(${libName} PRIVATE
|
||||
pybind11::module
|
||||
)
|
||||
|
||||
if(MSVC)
|
||||
target_link_libraries(${libName} PRIVATE
|
||||
pybind11::windows_extras
|
||||
)
|
||||
else()
|
||||
# This allows us to use TRT libs shipped with standalone wheels.
|
||||
set_target_properties(${libName} PROPERTIES SKIP_BUILD_RPATH ON)
|
||||
target_link_options(${libName} PRIVATE "LINKER:--rpath=$ORIGIN,--disable-new-dtags")
|
||||
endif()
|
||||
|
||||
# Find the main python headers in the relevant python<ver> subfolder in the externals.
|
||||
find_path(
|
||||
PYTHON_INCLUDES Python.h
|
||||
HINTS ${TRT_BUILD_PYTHON_EXTERNALS_PATH}/python${pyVersion}
|
||||
PATH_SUFFIXES include
|
||||
NO_CACHE
|
||||
REQUIRED
|
||||
NO_CMAKE_FIND_ROOT_PATH
|
||||
)
|
||||
|
||||
# Most of the headers are in that path we just found, except "pyconfig.h", which is platform-specific
|
||||
# and in a platform-specific directory with an inconsistent naming scheme.
|
||||
# So... go hunt for that. It's "mostly" located at /externals/python<ver>/include/<triple>/python<ver>/
|
||||
# Except on windows, where instead of <triple> it's just "win".
|
||||
if(${TRT_BUILD_PLATFORM} STREQUAL ${TRT_PLATFORM_X86})
|
||||
set(PYCONFIG_H_PATH "x86_64-linux-gnu/python${pyVersion}")
|
||||
elseif(${TRT_BUILD_PLATFORM} STREQUAL ${TRT_PLATFORM_AARCH64})
|
||||
set(PYCONFIG_H_PATH "aarch64-linux-gnu/python${pyVersion}")
|
||||
else()
|
||||
message(FATAL_ERROR "The current platform \"${TRT_BUILD_PLATFORM}\" cannot be used to build the TRT Python Bindings.")
|
||||
endif()
|
||||
|
||||
find_path(
|
||||
PYCONFIG_INCLUDE pyconfig.h
|
||||
HINTS ${PYTHON_INCLUDES}/${PYCONFIG_H_PATH}
|
||||
NO_CACHE
|
||||
REQUIRED
|
||||
NO_CMAKE_FIND_ROOT_PATH
|
||||
)
|
||||
|
||||
# Add the python headers as SYSTEM headers to silence warnings.
|
||||
target_include_directories(${libName} SYSTEM PRIVATE
|
||||
${PYTHON_INCLUDES}
|
||||
${PYCONFIG_INCLUDE}
|
||||
)
|
||||
|
||||
target_include_directories(${libName} PRIVATE
|
||||
"include"
|
||||
"docstrings"
|
||||
)
|
||||
|
||||
# Setup links against the TRT Libraries.
|
||||
if(${moduleName} STREQUAL "tensorrt")
|
||||
set(TRT_LIBS tensorrt nvonnxparser)
|
||||
if(${TRT_BUILD_PLUGINS})
|
||||
list(APPEND TRT_LIBS tensorrt_plugins)
|
||||
endif()
|
||||
elseif(${moduleName} STREQUAL "tensorrt_rtx")
|
||||
set(TRT_LIBS tensorrt nvonnxparser)
|
||||
elseif(${moduleName} STREQUAL "tensorrt_lean")
|
||||
set(TRT_LIBS tensorrt_lean_runtime)
|
||||
elseif(${moduleName} STREQUAL "tensorrt_dispatch")
|
||||
set(TRT_LIBS tensorrt_dispatch_runtime)
|
||||
else()
|
||||
message(FATAL_ERROR "Unknown TensorRT module " ${moduleName})
|
||||
endif()
|
||||
|
||||
target_link_libraries(${libName} PRIVATE
|
||||
${TRT_LIBS}
|
||||
$<COMPILE_ONLY:TRT::cudart> # We need the cuda headers to compile, but we don't link against them.
|
||||
$<COMPILE_ONLY:trt_global_definitions>
|
||||
)
|
||||
|
||||
# Tell the files what module they are currently building.
|
||||
target_compile_definitions(${libName} PRIVATE
|
||||
TENSORRT_MODULE=${moduleName}
|
||||
)
|
||||
# Remove the `lib` prefix from the binding .so's and correct the output name.
|
||||
set_target_properties(${libName}
|
||||
PROPERTIES PREFIX ""
|
||||
CXX_VISIBILITY_PRESET hidden
|
||||
VISIBILITY_INLINES_HIDDEN ON
|
||||
LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/${moduleName}_bindings-py${pyVersion}
|
||||
RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/${moduleName}_bindings-py${pyVersion}
|
||||
ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/${moduleName}_bindings-py${pyVersion}
|
||||
OUTPUT_NAME ${moduleName}
|
||||
)
|
||||
|
||||
add_dependencies(tensorrt_python_bindings ${libName})
|
||||
|
||||
if(MSVC)
|
||||
# Prevent pybind11 from sharing resources with other, potentially ABI incompatible modules
|
||||
# https://github.com/pybind/pybind11/issues/2898
|
||||
add_definitions(-DPYBIND11_COMPILER_TYPE="_${PROJECT_NAME}_abi")
|
||||
|
||||
# The python lib is python<maj><minor>.lib, but pyVersion is <maj>.<minor>, so we need to remove the dot.
|
||||
string(REPLACE "." "" pyVerStr ${pyVersion})
|
||||
|
||||
if(NOT TARGET python${pyVerStr})
|
||||
# Windows needs an explicit link against the python library.
|
||||
find_library(
|
||||
PYTHON${pyVerStr}_LIBRARY_PATH python${pyVerStr}.lib
|
||||
HINTS ${TRT_BUILD_PYTHON_EXTERNALS_PATH}/python${pyVersion}
|
||||
PATH_SUFFIXES lib
|
||||
REQUIRED
|
||||
NO_CMAKE_FIND_ROOT_PATH
|
||||
)
|
||||
|
||||
add_library(python${pyVerStr} STATIC IMPORTED)
|
||||
set_target_properties(python${pyVerStr} PROPERTIES IMPORTED_LOCATION "${PYTHON${pyVerStr}_LIBRARY_PATH}")
|
||||
endif()
|
||||
|
||||
target_link_libraries(${libName} PRIVATE python${pyVerStr})
|
||||
endif()
|
||||
endfunction()
|
||||
|
||||
# Enumerate all the combinations and create the per-python per-module targets.
|
||||
foreach(moduleName IN LISTS TRT_PYTHON_MODULE_NAMES)
|
||||
foreach(pyVersion IN LISTS TRT_BUILD_PYTHON_PY_VERSIONS)
|
||||
createBindingLibrary(${moduleName} ${pyVersion})
|
||||
endforeach()
|
||||
endforeach()
|
||||
|
||||
# Enter the packaging subdir to actually build the wheels.
|
||||
add_subdirectory(packaging)
|
||||
|
||||
else() # TRT_BUILD_ENABLE_NEW_PYTHON_FLOW - old flow is below this line
|
||||
|
||||
set(TRT_BUILD_WINML OFF)
|
||||
if (${TENSORRT_MODULE} STREQUAL "tensorrt_rtx")
|
||||
# The "old" flow doesn't already have TRT_BUILD_WINML set, because it's
|
||||
# coming from the legacy make build. So if we get the tensorrt_rtx module,
|
||||
# we assume we're in the TensorRT-RTX build.
|
||||
set(TRT_BUILD_WINML ON)
|
||||
endif()
|
||||
|
||||
# Ensure TRT_WINML is either 0 or 1
|
||||
if (NOT TRT_WINML)
|
||||
set(TRT_WINML 0)
|
||||
if (TRT_BUILD_WINML)
|
||||
set(TRT_WINML 1)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# Sets variable to a value if variable is unset.
|
||||
macro(set_ifndef var val)
|
||||
if(NOT DEFINED ${var})
|
||||
set(${var} ${val})
|
||||
endif()
|
||||
endmacro()
|
||||
|
||||
function(message)
|
||||
if(VERBOSE)
|
||||
_message(${ARGN})
|
||||
endif()
|
||||
endfunction()
|
||||
|
||||
# -------- CMAKE OPTIONS --------
|
||||
|
||||
set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/${TENSORRT_MODULE}/)
|
||||
set(CPP_STANDARD
|
||||
17
|
||||
CACHE STRING "CPP Standard Version")
|
||||
set(CMAKE_CXX_STANDARD ${CPP_STANDARD})
|
||||
|
||||
# -------- PATHS --------
|
||||
find_package(CUDAToolkit REQUIRED)
|
||||
message(STATUS "EXT_PATH: ${EXT_PATH}")
|
||||
message(STATUS "TENSORRT_BUILD: ${TENSORRT_BUILD}")
|
||||
message(STATUS "CMAKE_BINARY_DIR: ${CMAKE_BINARY_DIR}")
|
||||
message(STATUS "CUDAToolkit_LIBRARY_ROOT: ${CUDAToolkit_LIBRARY_ROOT}")
|
||||
message(STATUS "CUDAToolkit_INCLUDE_DIRS: ${CUDAToolkit_INCLUDE_DIRS}")
|
||||
message(STATUS "CMAKE_BUILD_TYPE: ${CMAKE_BUILD_TYPE}")
|
||||
|
||||
set_ifndef(TENSORRT_ROOT ../)
|
||||
message(STATUS "TENSORRT_ROOT: ${TENSORRT_ROOT}")
|
||||
|
||||
set_ifndef(WIN_EXTERNALS ${EXT_PATH})
|
||||
message(STATUS "WIN_EXTERNALS: ${WIN_EXTERNALS}")
|
||||
|
||||
# Convert to an absolute path.
|
||||
set_ifndef(ONNX_INC_DIR ${TENSORRT_ROOT}/parsers/onnx)
|
||||
find_path(
|
||||
PYBIND11_DIR pybind11/pybind11.h
|
||||
HINTS ${EXT_PATH} ${WIN_EXTERNALS}
|
||||
PATH_SUFFIXES pybind11/include
|
||||
REQUIRED
|
||||
)
|
||||
|
||||
message(STATUS "ONNX_INC_DIR: ${ONNX_INC_DIR}")
|
||||
message(STATUS "PYBIND11_DIR: ${PYBIND11_DIR}")
|
||||
|
||||
# Source Files
|
||||
if(${TENSORRT_MODULE} STREQUAL "tensorrt" OR ${TENSORRT_MODULE} STREQUAL "tensorrt_rtx")
|
||||
# tensorrt full dependencies
|
||||
file(GLOB_RECURSE SOURCE_FILES src/*.cpp)
|
||||
set(_NEED_EXPORTS_DEF ON)
|
||||
else()
|
||||
# tensorrt_lean and tensorrt_dispatch dependencies
|
||||
set(SOURCE_FILES src/pyTensorRT.cpp src/utils.cpp src/infer/pyCore.cpp src/infer/pyPlugin.cpp
|
||||
src/infer/pyFoundationalTypes.cpp)
|
||||
endif()
|
||||
|
||||
set(PYTHON python${PYTHON_MAJOR_VERSION}.${PYTHON_MINOR_VERSION})
|
||||
set(PYTHON_LIB_NAME python${PYTHON_MAJOR_VERSION}${PYTHON_MINOR_VERSION})
|
||||
message(STATUS "PYTHON: ${PYTHON}")
|
||||
message(STATUS "TENSORRT_MODULE: ${TENSORRT_MODULE}")
|
||||
|
||||
set(PY_MODULE_NAME ${TENSORRT_MODULE})
|
||||
|
||||
# Find headers
|
||||
set(LIBPATH_SUFFIX lib)
|
||||
if(MSVC)
|
||||
find_path(
|
||||
PY_INCLUDE Python.h
|
||||
HINTS ${WIN_EXTERNALS}/${PYTHON} ${EXT_PATH}/${PYTHON}
|
||||
PATH_SUFFIXES include
|
||||
REQUIRED
|
||||
)
|
||||
find_path(
|
||||
PY_LIB_DIR ${PYTHON_LIB_NAME}.lib
|
||||
HINTS ${WIN_EXTERNALS}/${PYTHON} ${EXT_PATH}/${PYTHON}
|
||||
PATH_SUFFIXES ${LIBPATH_SUFFIX}
|
||||
REQUIRED
|
||||
)
|
||||
message(STATUS "PY_LIB_DIR: ${PY_LIB_DIR}")
|
||||
else()
|
||||
find_path(
|
||||
PY_INCLUDE Python.h
|
||||
HINTS ${EXT_PATH}/${PYTHON} /usr/include/${PYTHON}
|
||||
PATH_SUFFIXES include
|
||||
REQUIRED
|
||||
)
|
||||
endif()
|
||||
|
||||
message(STATUS "PY_INCLUDE: ${PY_INCLUDE}")
|
||||
|
||||
if(MSVC)
|
||||
set(PY_TARGET_DIR win)
|
||||
else()
|
||||
set(PY_TARGET_DIR ${TARGET}-linux-gnu)
|
||||
if(${TARGET} STREQUAL ppc64le)
|
||||
set(PY_TARGET_DIR powerpc64le-linux-gnu)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# The per-platform pyconfig.h is located at /externals/python<ver>/<triple>/python<ver>/.
|
||||
# Not sure why it's setup that way.
|
||||
find_path(
|
||||
PY_CONFIG_INCLUDE pyconfig.h
|
||||
HINTS ${PY_INCLUDE}
|
||||
PATH_SUFFIXES ${PY_TARGET_DIR}/${PYTHON} ${PY_TARGET_DIR}/${PYTHON}m
|
||||
REQUIRED
|
||||
)
|
||||
message(STATUS "PY_CONFIG_INCLUDE: ${PY_CONFIG_INCLUDE}")
|
||||
|
||||
# -------- GLOBAL COMPILE OPTIONS --------
|
||||
|
||||
include_directories(${TENSORRT_ROOT}/include ${PROJECT_SOURCE_DIR}/include ${CUDAToolkit_INCLUDE_DIRS}
|
||||
${PROJECT_SOURCE_DIR}/docstrings ${ONNX_INC_DIR} ${PYBIND11_DIR})
|
||||
link_directories(${TENSORRT_BUILD} ${TENSORRT_LIBPATH})
|
||||
|
||||
if(MSVC)
|
||||
# Prevent pybind11 from sharing resources with other, potentially ABI incompatible modules
|
||||
# https://github.com/pybind/pybind11/issues/2898
|
||||
add_definitions(-DPYBIND11_COMPILER_TYPE="_${PROJECT_NAME}_abi")
|
||||
endif()
|
||||
|
||||
if(MSVC)
|
||||
include_directories(SYSTEM BEFORE ${ADDITIONAL_PLATFORM_INCLUDE_DIRS})
|
||||
link_libraries(${ADDITIONAL_PLATFORM_LIB_FLAGS})
|
||||
link_directories(${PY_LIB_DIR})
|
||||
|
||||
set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} /bigobj")
|
||||
if(${NV_GEN_PDB})
|
||||
# PDB is only useful in release mode.
|
||||
set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /Zi /bigobj")
|
||||
set(CMAKE_SHARED_LINKER_FLAGS_RELEASE "${CMAKE_SHARED_LINKER_FLAGS_RELEASE} /DEBUG /OPT:REF /OPT:ICF")
|
||||
endif()
|
||||
else()
|
||||
set(CMAKE_CXX_FLAGS
|
||||
"${CMAKE_CXX_FLAGS} ${GLIBCXX_USE_CXX11_ABI_FLAG} -fvisibility=hidden -std=c++${CPP_STANDARD} -Wno-deprecated-declarations"
|
||||
)
|
||||
endif()
|
||||
|
||||
|
||||
# ---------- MODULE DEPENDENCIES ----------
|
||||
add_compile_definitions(TENSORRT_MODULE=${TENSORRT_MODULE})
|
||||
|
||||
if("${TRT_VCAST}" STREQUAL "1" OR "${TRT_VCAST_SAFE}" STREQUAL "1")
|
||||
set(vfc_suffix "_static")
|
||||
else()
|
||||
set(vfc_suffix "")
|
||||
endif()
|
||||
|
||||
|
||||
if(MSVC)
|
||||
set(nvinfer_lib_name "${TRT_NVINFER_NAME}_${TENSORRT_MAJOR_VERSION}${TRT_LIB_SUFFIX}")
|
||||
set(nvinfer_plugin_lib_name "nvinfer_plugin_${TENSORRT_MAJOR_VERSION}")
|
||||
set(nvonnxparser_lib_name "${TRT_ONNXPARSER_NAME}_${TENSORRT_MAJOR_VERSION}${TRT_LIB_SUFFIX}")
|
||||
set(nvinfer_lean_lib_name "nvinfer_lean_${TENSORRT_MAJOR_VERSION}${vfc_suffix}")
|
||||
set(nvinfer_dispatch_lib_name "nvinfer_dispatch_${TENSORRT_MAJOR_VERSION}${vfc_suffix}")
|
||||
else()
|
||||
if(${TRT_BUILD_WINML})
|
||||
set(nvinfer_lib_name "-l:lib${TRT_NVINFER_NAME}.so")
|
||||
else()
|
||||
set(nvinfer_lib_name "${TRT_NVINFER_NAME}")
|
||||
endif()
|
||||
set(nvinfer_plugin_lib_name "nvinfer_plugin")
|
||||
set(nvonnxparser_lib_name "${TRT_ONNXPARSER_NAME}")
|
||||
set(nvinfer_lean_lib_name "nvinfer_lean${vfc_suffix}")
|
||||
set(nvinfer_dispatch_lib_name "nvinfer_dispatch${vfc_suffix}")
|
||||
endif()
|
||||
message(STATUS "NVInferLib = ${nvinfer_lib_name}")
|
||||
message(STATUS "NVOnnxParserLib = ${nvonnx_parser_lib_name}")
|
||||
|
||||
if(${TENSORRT_MODULE} STREQUAL "tensorrt" OR ${TENSORRT_MODULE} STREQUAL "tensorrt_rtx")
|
||||
set(TRT_LIBS ${nvinfer_lib_name} ${nvonnxparser_lib_name})
|
||||
list(APPEND TRT_LIBS ${nvinfer_plugin_lib_name})
|
||||
elseif(${TENSORRT_MODULE} STREQUAL "tensorrt_lean")
|
||||
set(TRT_LIBS ${nvinfer_lean_lib_name})
|
||||
elseif(${TENSORRT_MODULE} STREQUAL "tensorrt_dispatch")
|
||||
set(TRT_LIBS ${nvinfer_dispatch_lib_name})
|
||||
else()
|
||||
message(FATAL_ERROR "Unknown TensorRT module " ${TENSORRT_MODULE})
|
||||
endif()
|
||||
|
||||
# -------- BUILDING --------
|
||||
set(LIB_NAME ${PY_MODULE_NAME})
|
||||
|
||||
# Set up target
|
||||
add_library(${LIB_NAME} SHARED ${SOURCE_FILES})
|
||||
if(_NEED_EXPORTS_DEF)
|
||||
target_compile_definitions(${LIB_NAME} PRIVATE
|
||||
tensorrt_EXPORTS=1
|
||||
)
|
||||
endif()
|
||||
message(STATUS "Python library name ${LIB_NAME}")
|
||||
target_include_directories(${LIB_NAME} BEFORE PUBLIC ${PY_CONFIG_INCLUDE} ${PY_INCLUDE})
|
||||
if(MSVC)
|
||||
# For some reason, we must explicitly link against the Python library on Windows.
|
||||
message(STATUS "Dependencies ${TRT_LIBS} ${PYTHON_LIB_NAME}")
|
||||
target_link_libraries(${LIB_NAME} PRIVATE ${TRT_LIBS} ${PYTHON_LIB_NAME})
|
||||
else()
|
||||
message(STATUS "Dependencies ${TRT_LIBS} ")
|
||||
target_link_libraries(${LIB_NAME} PRIVATE ${TRT_LIBS})
|
||||
endif()
|
||||
|
||||
# Note that we have to remove the `lib` prefix from the binding .so's
|
||||
set_target_properties(${LIB_NAME} PROPERTIES PREFIX "")
|
||||
|
||||
endif()
|
||||
Reference in New Issue
Block a user