tensor-rt library dynamic loadg

mindspore/ccsrc/CMakeLists.txt

@@ -88,13 +88,6 @@ if(ENABLE_GPU)
         endif()
     endif()
 
-    if(DEFINED ENV{TENSORRT_HOME} AND NOT $ENV{TENSORRT_HOME} STREQUAL "")
-        message("Enable GPU inference. Tensor-RT dir: $ENV{TENSORRT_HOME}")
-        set(ENABLE_GPU_INFER TRUE)
-        add_compile_definitions(ENABLE_GPU_INFER)
-        include_directories($ENV{TENSORRT_HOME}/include)
-    endif()
-
     if(NOT CUPTI_INCLUDE_DIRS OR CUPTI_INCLUDE_DIRS STREQUAL "")
         set(CUPTI_INCLUDE_DIRS  ${CUDA_PATH}/extras/CUPTI/include)
     endif()
@@ -122,8 +115,17 @@ if(ENABLE_GPU)
                                   "runtime/device/gpu/distribution/collective_wrapper.cc"
                                   "runtime/device/gpu/distribution/mpi_wrapper.cc"
                                   "runtime/device/gpu/distribution/nccl_wrapper.cc"
+                                  "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}")
+        set(ENABLE_GPU_INFER TRUE)
+        add_compile_definitions(ENABLE_GPU_INFER)
+        include_directories($ENV{TENSORRT_HOME}/include)
+        file(GLOB_RECURSE GPU_SRC_LIST RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "runtime/device/gpu/trt_loader.cc")
+    endif()
+
     set(NVCC_TMP_CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS})
     string(REPLACE "-std=c++17" "-std=c++11" CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}")
     set_property(SOURCE ${GPU_SRC_LIST} PROPERTY COMPILE_DEFINITIONS SUBMODULE_ID=mindspore::SubModuleId::SM_DEVICE)
@@ -316,14 +318,6 @@ else()
     target_link_libraries(mindspore -Wl,--start-group proto_input mindspore::protobuf -Wl,--end-group)
 endif()
 
-if(ENABLE_GPU_INFER)
-    find_library(trt_plugin libnvinfer_plugin.so $ENV{TENSORRT_HOME}/lib)
-    find_library(trt_nvinfo libnvinfer.so $ENV{TENSORRT_HOME}/lib)
-    find_library(trt_parser libnvparsers.so $ENV{TENSORRT_HOME}/lib)
-    target_link_libraries(mindspore ${trt_plugin} ${trt_nvinfo} ${trt_parser})
-    set(MINDSPORE_RPATH ${MINDSPORE_RPATH}:$ENV{TENSORRT_HOME}/lib)
-endif()
-
 # set c_expression building
 set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
 set_property(SOURCE "pipeline/jit/init.cc" PROPERTY

mindspore/ccsrc/backend/kernel_compiler/gpu/trt/trt_kernel.cc

@@ -22,6 +22,7 @@
 #include <functional>
 #include "backend/kernel_compiler/gpu/data/dataset_utils.h"
 #include "backend/kernel_compiler/gpu/trt/trt_utils.h"
+#include "runtime/device/gpu/trt_loader.h"
 
 namespace mindspore {
 namespace kernel {
@@ -48,7 +49,11 @@ bool TrtKernel::Init(const CNodePtr &kernel_node) {
     output_size_list_.push_back(size_in_byte);
   }
 
-  runtime_ = TrtPtr(nvinfer1::createInferRuntime(Singleton<TrtLogger>::Instance()));
+  auto trt_loader = Singleton<device::gpu::TrtLoader>::Instance();
+  if (!trt_loader.nvinfer_loaded()) {
+    MS_LOG(EXCEPTION) << "Install Tensor-RT and export LD_LIBRARY_PATH=${TENSORRT_HOME}/lib:$LD_LIBRARY_PATH."
+  }
+  runtime_ = trt_loader.CreateInferRuntime(&Singleton<TrtLogger>::Instance());
   MS_EXCEPTION_IF_NULL(runtime_);
   serialize_ = GetAttr<std::string>(kernel_node, "serialize_model");
   engine_ = TrtPtr(runtime_->deserializeCudaEngine(serialize_.c_str(), serialize_.size(), nullptr));

mindspore/ccsrc/backend/kernel_compiler/gpu/trt/trt_utils.h

@@ -28,6 +28,8 @@
 #include "utils/log_adapter.h"
 #include "utils/singleton.h"
 #include "utils/convert_utils_base.h"
+#include "utils/shape_utils.h"
+#include "ir/dtype/type.h"
 
 namespace mindspore {
 class TrtUtils {

mindspore/ccsrc/runtime/device/gpu/trt_loader.cc

@@ -0,0 +1,67 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * 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.
+ */
+
+#include "runtime/device/gpu/trt_loader.h"
+
+#include <dlfcn.h>
+#include <memory>
+#include <NvInferRuntimeCommon.h>
+#include "backend/kernel_compiler/gpu/trt/trt_utils.h"
+
+namespace mindspore {
+namespace device {
+namespace gpu {
+TrtLoader::TrtLoader()
+    : nvinfer_loaded_(false), nvinfer_handle_(nullptr), create_infer_builder_(nullptr), create_infer_runtime_(nullptr) {
+  nvinfer_handle_ = dlopen("libnvinfer.so", RTLD_NOW | RTLD_LOCAL);
+  if (nvinfer_handle_ == nullptr) {
+    MS_LOG(WARNING) << "Can not open libnvinfer.so. " << dlerror()
+                    << ". Install Tensor-RT and export LD_LIBRARY_PATH=${TENSORRT_HOME}/lib:$LD_LIBRARY_PATH.";
+    MS_LOG(WARNING) << "Inference with native backend.";
+    return;
+  }
+
+  create_infer_builder_ = (CreateInferBuilder_t)dlsym(nvinfer_handle_, "createInferBuilder_INTERNAL");
+  if (create_infer_builder_ == nullptr) {
+    MS_LOG(WARNING) << "Failed to get createInferBuilder_INTERNAL symbol. " << dlerror();
+    return;
+  }
+
+  create_infer_runtime_ = (CreateInferRuntime_t)dlsym(nvinfer_handle_, "createInferRuntime_INTERNAL");
+  if (create_infer_runtime_ == nullptr) {
+    MS_LOG(WARNING) << "Failed to get createInferRuntime_INTERNAL symbol. " << dlerror();
+    return;
+  }
+
+  nvinfer_loaded_ = true;
+}
+
+TrtLoader::~TrtLoader() {
+  if (nvinfer_handle_ != nullptr) {
+    dlclose(nvinfer_handle_);
+  }
+}
+
+std::shared_ptr<nvinfer1::IBuilder> TrtLoader::CreateInferBuilder(nvinfer1::ILogger *logger) {
+  return TrtPtr<nvinfer1::IBuilder>(create_infer_builder_(*logger, NV_TENSORRT_VERSION));
+}
+
+std::shared_ptr<nvinfer1::IRuntime> TrtLoader::CreateInferRuntime(nvinfer1::ILogger *logger) {
+  return TrtPtr<nvinfer1::IRuntime>(create_infer_runtime_(*logger, NV_TENSORRT_VERSION));
+}
+}  // namespace gpu
+}  // namespace device
+}  // namespace mindspore

mindspore/ccsrc/runtime/device/gpu/trt_loader.h

@@ -0,0 +1,45 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * 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.
+ */
+#ifndef MINDSPORE_CCSRC_BACKEND_RUNTIME_DEVICE_GPU_TRT_LOADER_H_
+#define MINDSPORE_CCSRC_BACKEND_RUNTIME_DEVICE_GPU_TRT_LOADER_H_
+
+#include <memory>
+#include <NvInfer.h>
+
+namespace mindspore {
+namespace device {
+namespace gpu {
+typedef nvinfer1::IBuilder *(*CreateInferBuilder_t)(nvinfer1::ILogger &, int);
+typedef nvinfer1::IRuntime *(*CreateInferRuntime_t)(nvinfer1::ILogger &, int);
+class TrtLoader {
+ public:
+  TrtLoader();
+  ~TrtLoader();
+
+  std::shared_ptr<nvinfer1::IBuilder> CreateInferBuilder(nvinfer1::ILogger *logger);
+  std::shared_ptr<nvinfer1::IRuntime> CreateInferRuntime(nvinfer1::ILogger *logger);
+  bool nvinfer_loaded() { return nvinfer_loaded_; }
+
+ private:
+  bool nvinfer_loaded_;
+  void *nvinfer_handle_;
+  CreateInferBuilder_t create_infer_builder_;
+  CreateInferRuntime_t create_infer_runtime_;
+};
+}  // namespace gpu
+}  // namespace device
+}  // namespace mindspore
+#endif  // MINDSPORE_CCSRC_BACKEND_RUNTIME_DEVICE_GPU_TRT_LOADER_H_