diff --git a/build.sh b/build.sh index e0f494f9b62..ec3d8ac530b 100755 --- a/build.sh +++ b/build.sh @@ -27,6 +27,7 @@ usage() echo " [-P on|off] [-z [on|off]] [-M on|off] [-V 9.2|10.1|310|910] [-I arm64|arm32|x86_64] [-K] \\" echo " [-B on|off] [-E] [-l on|off] [-n full|lite|off] [-T on|off] [-H on|off] \\" echo " [-A [cpp|java|object-c] [-C on|off] [-o on|off] [-S on|off] [-k on|off] [-W sse|neon|avx|off] \\" + echo " [-L Tensor-RT path] \\" echo "" echo "Options:" echo " -d Debug mode" @@ -64,6 +65,7 @@ usage() echo " -k Enable make clean, clean up compilation generated cache " echo " -W Enable x86_64 SSE or AVX instruction set, use [sse|avx|neon|off], default off" echo " -H Enable hidden" + echo " -L Link and specify Tensor-RT library path, default disable Tensor-RT lib linking" } # check value of input is 'on' or 'off' @@ -122,9 +124,11 @@ checkopts() ENABLE_NPU="off" ENABLE_HIDDEN="on" LITE_ENABLE_GPU="" + TENSORRT_HOME="" # Process the options - while getopts 'drvj:c:t:hsb:a:g:p:ie:m:l:I:RP:D:zM:V:K:B:En:T:A:C:o:S:k:W:H:' opt + while getopts 'drvj:c:t:hsb:a:g:p:ie:m:l:I:RP:D:zM:V:K:B:En:T:A:C:o:S:k:W:H:L:' opt do + CASE_SENSIVE_ARG=${OPTARG} OPTARG=$(echo ${OPTARG} | tr '[A-Z]' '[a-z]') case "${opt}" in d) @@ -333,6 +337,11 @@ checkopts() ENABLE_HIDDEN="$OPTARG" echo "${OPTARG} hidden" ;; + L) + ENABLE_TRT="on" + TENSORRT_HOME="$CASE_SENSIVE_ARG" + echo "Link Tensor-RT library. Path: ${CASE_SENSIVE_ARG}" + ;; *) echo "Unknown option ${opt}!" usage @@ -489,6 +498,9 @@ build_mindspore() if [[ "X$ENABLE_HIDDEN" = "Xoff" ]]; then CMAKE_ARGS="${CMAKE_ARGS} -DENABLE_HIDDEN=OFF" fi + if [[ "X$ENABLE_TRT" == "Xon" ]]; then + CMAKE_ARGS="${CMAKE_ARGS} -DTENSORRT_HOME=${TENSORRT_HOME}" + fi echo "${CMAKE_ARGS}" if [[ "X$INC_BUILD" = "Xoff" ]]; then cmake ${CMAKE_ARGS} ../.. diff --git a/mindspore/ccsrc/CMakeLists.txt b/mindspore/ccsrc/CMakeLists.txt index 93583b4578a..91a377bb441 100644 --- a/mindspore/ccsrc/CMakeLists.txt +++ b/mindspore/ccsrc/CMakeLists.txt @@ -117,11 +117,15 @@ if(ENABLE_GPU) "runtime/device/gpu/trt_loader.cc" ) - if(DEFINED ENV{TENSORRT_HOME} AND NOT $ENV{TENSORRT_HOME} STREQUAL "") - message("Enable GPU inference. Tensor-RT dir: $ENV{TENSORRT_HOME}") + if(NOT ${TENSORRT_HOME} STREQUAL "") + find_path(TENSORRT_HOME_INCLUDE NvInfer.h HINTS ${TENSORRT_HOME}/include) + if(TENSORRT_HOME_INCLUDE STREQUAL TENSORRT_HOME_INCLUDE-NOTFOUND) + message(FATAL_ERROR "Tensor-RT dir not exist ${TENSORRT_HOME}") + endif() + message("Enable GPU inference. Tensor-RT include dir: ${TENSORRT_HOME_INCLUDE}") set(ENABLE_GPU_INFER TRUE) add_compile_definitions(ENABLE_GPU_INFER) - include_directories($ENV{TENSORRT_HOME}/include) + include_directories(${TENSORRT_HOME_INCLUDE}) list(APPEND GPU_SRC_LIST ${CMAKE_CURRENT_SOURCE_DIR}/runtime/device/gpu/trt_loader.cc) endif()