diff --git a/cmake/Modules/FindTENSORRT.cmake b/cmake/Modules/FindTENSORRT.cmake index 2b381be..be53de0 100644 --- a/cmake/Modules/FindTENSORRT.cmake +++ b/cmake/Modules/FindTENSORRT.cmake @@ -34,12 +34,27 @@ if(TENSORRT_INCLUDE_DIR) file(READ "${TENSORRT_INCLUDE_DIR}/NvInfer.h" TENSORRT_H_CONTENTS) endif() - string(REGEX MATCH "define NV_TENSORRT_MAJOR ([0-9]+)" _ "${TENSORRT_H_CONTENTS}") - set(TENSORRT_VERSION_MAJOR ${CMAKE_MATCH_1} CACHE INTERNAL "") - string(REGEX MATCH "define NV_TENSORRT_MINOR ([0-9]+)" _ "${TENSORRT_H_CONTENTS}") - set(TENSORRT_VERSION_MINOR ${CMAKE_MATCH_1} CACHE INTERNAL "") - string(REGEX MATCH "define NV_TENSORRT_PATCH ([0-9]+)" _ "${TENSORRT_H_CONTENTS}") - set(TENSORRT_VERSION_PATCH ${CMAKE_MATCH_1} CACHE INTERNAL "") + string(REGEX MATCH "define[ \t]+NV_TENSORRT_MAJOR[ \t]+([0-9]+)" _ "${TENSORRT_H_CONTENTS}") + set(TENSORRT_VERSION_MAJOR "${CMAKE_MATCH_1}") + + string(REGEX MATCH "define[ \t]+NV_TENSORRT_MINOR[ \t]+([0-9]+)" _ "${TENSORRT_H_CONTENTS}") + set(TENSORRT_VERSION_MINOR "${CMAKE_MATCH_1}") + + string(REGEX MATCH "define[ \t]+NV_TENSORRT_PATCH[ \t]+([0-9]+)" _ "${TENSORRT_H_CONTENTS}") + set(TENSORRT_VERSION_PATCH "${CMAKE_MATCH_1}") + + # TRT 10.11+ aliases NV_TENSORRT_* to TRT_*_ENTERPRISE, regex can't resolve preprocessor + # indirection, so match the underlying macros directly + if(NOT TENSORRT_VERSION_MAJOR) + string(REGEX MATCH "define[ \t]+TRT_MAJOR_ENTERPRISE[ \t]+([0-9]+)" _ "${TENSORRT_H_CONTENTS}") + set(TENSORRT_VERSION_MAJOR "${CMAKE_MATCH_1}") + + string(REGEX MATCH "define[ \t]+TRT_MINOR_ENTERPRISE[ \t]+([0-9]+)" _ "${TENSORRT_H_CONTENTS}") + set(TENSORRT_VERSION_MINOR "${CMAKE_MATCH_1}") + + string(REGEX MATCH "define[ \t]+TRT_PATCH_ENTERPRISE[ \t]+([0-9]+)" _ "${TENSORRT_H_CONTENTS}") + set(TENSORRT_VERSION_PATCH "${CMAKE_MATCH_1}") + endif() set(TENSORRT_VERSION "${TENSORRT_VERSION_MAJOR}.${TENSORRT_VERSION_MINOR}.${TENSORRT_VERSION_PATCH}"