Enable multiple session runs for TensorRT (#724)

* Update BUILD.md

* Update README.md

* Update tensorrt_execution_provider.cc

remap node index to handle the case that nodes in graph may be deleted and node index is not continuous.

* Update onnxruntime_providers.cmake

Solve conflicts to onnx-tensorrt

* Update tensorrt_execution_provider.h

* Update tensorrt_execution_provider.cc

* Update tensorrt_execution_provider.h

* Update tensorrt_execution_provider.cc

* Update tensorrt_execution_provider.h

* Update tensorrt_execution_provider.cc

* Update tensorrt_execution_provider.cc

* Update tensorrt_execution_provider.cc

* Update build.py

BUILD.md

@@ -125,11 +125,12 @@ The TensorRT execution provider for ONNX Runtime is built and tested with CUDA 9
  - The path to the CUDA `bin` directory must be added to the PATH environment variable so that `nvcc` is found.
  - The path to the CUDNN installation (path to folder that contains libcudnn.so) must be provided via the CUDNN_PATH environment variable, or `--cudnn_home parameter`. 
 - The path to TensorRT installation must be provided via the `--tensorrt_home parameter`.
+- There are two kind of protobufs: protobuf and protobuf-lite. Currently TensorRT execution provider only supports protobuf by using `--use_full_protobuf` option.
 
 You can build from source on Linux by using the following `cmd` from the onnxruntime directory:
 
 ```
-./build.sh --cudnn_home <path to CUDNN e.g. /usr/lib/x86_64-linux-gnu/> --cuda_home <path to folder for CUDA e.g. /usr/local/cuda> --use_tensorrt --tensorrt_home <path to TensorRT home> (Linux)
+./build.sh --use_full_protobuf --cudnn_home <path to CUDNN e.g. /usr/lib/x86_64-linux-gnu/> --cuda_home <path to folder for CUDA e.g. /usr/local/cuda> --use_tensorrt --tensorrt_home <path to TensorRT home> (Linux)
 
 ```
 

cmake/onnxruntime_providers.cmake

@@ -114,8 +114,8 @@ if (onnxruntime_USE_TENSORRT)
   set(CUDA_INCLUDE_DIRS ${onnxruntime_CUDA_HOME}/include)
   set(TENSORRT_ROOT ${onnxruntime_TENSORRT_HOME})
   include_directories(${ONNXRUNTIME_ROOT}/../cmake/external/onnx)
+  set(OLD_CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS})  
   if (WIN32)
-    set(OLD_CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS})
     set(OLD_CMAKE_CUDA_FLAGS ${CMAKE_CUDA_FLAGS})
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /wd4996 /wd4244 /wd4267 /wd4099 /wd4551 /wd4505 /wd4515 /wd4706 /wd4456")
     if (CMAKE_BUILD_TYPE STREQUAL "Debug")
@@ -125,9 +125,13 @@ if (onnxruntime_USE_TENSORRT)
     set(PROTOBUF_LIBRARY libprotobuf)
     set(DISABLED_WARNINGS_FOR_TRT /wd4267 /wd4244 /wd4996)
   endif()
-  add_subdirectory(${ONNXRUNTIME_ROOT}/../cmake/external/onnx-tensorrt)
+  if ( CMAKE_COMPILER_IS_GNUCC )
+    set(CMAKE_CXX_FLAGS  "${CMAKE_CXX_FLAGS} -Wno-unused-parameter -Wno-missing-field-initializers")
+  endif()  
+  set(CXX_VERSION_DEFINED TRUE)
+  add_subdirectory(${ONNXRUNTIME_ROOT}/../cmake/external/onnx-tensorrt)  
+  set(CMAKE_CXX_FLAGS ${OLD_CMAKE_CXX_FLAGS})  
   if (WIN32)
-    set(CMAKE_CXX_FLAGS ${OLD_CMAKE_CXX_FLAGS})
     set(CMAKE_CUDA_FLAGS ${OLD_CMAKE_CUDA_FLAGS})
     unset(PROTOBUF_LIBRARY)
     unset(OLD_CMAKE_CXX_FLAGS)

onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc

@@ -52,10 +52,9 @@ TensorrtExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph,
   onnxruntime::Model model(graph.Name(), true, ModelMetaData(), IOnnxRuntimeOpSchemaRegistryList(), graph.DomainToVersionMap());
   onnxruntime::Graph& graph_build = model.MainGraph();
   const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder();
-  for (const auto& index : node_index) {
-    const Node* node = graph.GetNode(index);
-    graph_build.AddNode(*node);
-  }
+  for (const auto& node : graph.Nodes()) {
+    graph_build.AddNode(node);
+  }  
   ORT_ENFORCE(graph_build.Resolve().IsOK());
   ONNX_NAMESPACE::ModelProto model_proto = model.ToProto();
   model_proto.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION);
@@ -74,22 +73,24 @@ TensorrtExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph,
 
   // Find inputs, initializers and outputs for each supported subgraph
   std::vector<std::unique_ptr<ComputeCapability>> result;
-  std::set<int> supported_nodes_set;
   int counter = 0;
   for (const auto& group : supported_nodes_vector) {
     if (!group.empty()) {
-      std::set<size_t> node_set(group.begin(), group.end());
+      std::unordered_set<size_t> node_set;
+      node_set.reserve(group.size());
+      for (const auto& index : group) {
+        node_set.insert(node_index[index]);
+      }      
       std::unique_ptr<IndexedSubGraph> sub_graph = std::make_unique<IndexedSubGraph>();
       // Find inputs and outputs of the subgraph
-      std::map<const NodeArg *, int> fused_inputs, fused_outputs, fused_outputs_to_add;
-      std::set<const NodeArg*> erased;
+      std::unordered_map<const NodeArg *, int> fused_inputs, fused_outputs, fused_outputs_to_add;
+      std::unordered_set<const NodeArg*> erased;
       int input_order = 0;
       int output_order = 0;
 
       for (const auto& index : group) {
-        supported_nodes_set.insert(index);
-        sub_graph->nodes.push_back(index);
-        const auto& node = graph.GetNode(index);
+        sub_graph->nodes.push_back(node_index[index]);
+        const auto& node = graph.GetNode(node_index[index]);
 
         for (const auto& input : node->InputDefs()) {
           const auto& it = fused_outputs.find(input);
@@ -110,10 +111,10 @@ TensorrtExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph,
         // to the output list
         if (node->GetOutputEdgesCount() > node->OutputDefs().size()) {
           for (auto it = node->OutputEdgesBegin(), end = node->OutputEdgesEnd(); it != end; ++it) {
-            const auto& node_index = it->GetNode().Index();
+            const auto& node_idx = it->GetNode().Index();
             const auto& output = (it->GetNode()).InputDefs()[it->GetDstArgIndex()];
 
-            if (node_set.find(node_index) != node_set.end()) {
+            if (node_set.find(node_idx) != node_set.end()) {
               const auto& iter = fused_inputs.find(output);
 
               if (iter != fused_inputs.end()) {
@@ -201,6 +202,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
     // Build map from input name to its index in input definitions
     std::unordered_map<std::string, int> input_map;
     const auto& input_defs = fused_node->InputDefs();
+    input_map.reserve(input_defs.size());
     for (int i = 0, end = input_defs.size(); i < end; ++i) {
       input_map[input_defs[i]->Name()] = i;
     }
@@ -208,6 +210,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
     // Build map from output name to its index in output definitions
     std::unordered_map<std::string, int> output_map;
     const auto& output_defs = fused_node->OutputDefs();
+    output_map.reserve(output_defs.size());
     for (int i = 0, end = output_defs.size(); i < end; ++i) {
       output_map[output_defs[i]->Name()] = i;
     }
@@ -237,7 +240,8 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
     // for the case that node's output is connected to more than one EdgeEnd nodes and some of them don't belong to the graph
     ONNX_NAMESPACE::ModelProto model_proto = model.ToProto();
     const auto& graph_output = model_proto.graph().output();
-    std::set<string> graph_outputs_set;
+    std::unordered_set<string> graph_outputs_set;
+    graph_outputs_set.reserve(graph_output.size());
     for (int i = 0, end = graph_output.size(); i < end; ++i) {
       graph_outputs_set.insert(graph_output[i].name());
     }
@@ -275,7 +279,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
     // Create TensorRT engine
     string string_buf;
     model_proto.SerializeToString(&string_buf);
-
+    
     TensorrtLogger trt_logger(nvinfer1::ILogger::Severity::kWARNING);
     auto trt_builder = unique_pointer<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(trt_logger));
     auto trt_network = unique_pointer<nvinfer1::INetworkDefinition>(trt_builder->createNetwork());
@@ -349,7 +353,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
     compute_info.create_state_func = [=](ComputeContext* context, FunctionState* state) {
       std::unique_ptr<TensorrtFuncState> p = std::make_unique<TensorrtFuncState>();
       *p = {context->allocate_func, context->release_func, context->allocator_handle, parsers_[context->node_name].get(), engines_[context->node_name].get(), contexts_[context->node_name].get(),
-            input_info_[context->node_name], output_info_[context->node_name], output_shapes_[context->node_name]};
+            input_info_[context->node_name], output_info_[context->node_name], output_shapes_[context->node_name], &tensorrt_mu_};
       *state = p.release();
       return 0;
     };
@@ -399,6 +403,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
       }
 
       // Run TRT inference
+      std::lock_guard<OrtMutex> lock(*(trt_state->tensorrt_mu_ptr));
       trt_state->context->enqueue(batch_size, &buffers[0], stream, nullptr);
 
       // Copy TRT outputs to output tensors

onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h

@@ -7,12 +7,13 @@
 #include "core/framework/op_kernel.h"
 #include "NvInfer.h"
 #include "NvOnnxParser.h"
+#include "core/platform/ort_mutex.h"
 
 namespace onnxruntime {
 
 static const int kMaxBatchSize = 1;
 static const int kMaxWorkSpaceSize = 1 << 30;
-
+    
 class TensorrtLogger : public nvinfer1::ILogger {
     nvinfer1::ILogger::Severity verbosity_;
 public:
@@ -51,6 +52,7 @@ struct TensorrtFuncState {
   std::vector<std::vector<int>> input_info;
   std::vector<std::vector<int>> output_info;
   std::vector<std::vector<int64_t>> output_shapes;
+  OrtMutex* tensorrt_mu_ptr = nullptr;
 };
 
 // Logical device representation.
@@ -87,6 +89,7 @@ class TensorrtExecutionProvider : public IExecutionProvider {
   template <typename T>
   using unique_pointer = std::unique_ptr<T, InferDeleter>;
 
+  OrtMutex tensorrt_mu_;
   int device_id_;
   std::unordered_map<std::string, unique_pointer<nvonnxparser::IParser>> parsers_;
   std::unordered_map<std::string, unique_pointer<nvinfer1::ICudaEngine>> engines_;

onnxruntime/test/perftest/README.md

@@ -4,7 +4,7 @@ onnxruntime_perf_test [options...] model_path result_file
 Options:
         -m [test_mode]: Specifies the test mode. Value coulde be 'duration' or 'times'.
                 Provide 'duration' to run the test for a fix duration, and 'times' to repeated for a certain times. Default:'duration'.
-        -e [cpu|cuda|mkldnn]: Specifies the provider 'cpu','cuda','mkldnn'. Default:'cpu'.
+        -e [cpu|cuda|mkldnn|tensorrt]: Specifies the provider 'cpu','cuda','mkldnn' or 'tensorrt'. Default:'cpu'.
         -r [repeated_times]: Specifies the repeated times if running in 'times' test mode.Default:1000.
         -t [seconds_to_run]: Specifies the seconds to run for 'duration' mode. Default:600.
         -p [profile_file]: Specifies the profile name to enable profiling and dump the profile data to the file.

tools/ci_build/build.py

@@ -548,9 +548,6 @@ def run_onnx_tests(build_dir, configs, onnx_test_data_dir, provider, enable_para
           else:
             run_subprocess([exe,'-x'] + cmd, cwd=cwd)
         else:
-          if provider == 'tensorrt':
-            run_subprocess([exe, '-c', '1'] + cmd, cwd=cwd)
-          else:
             run_subprocess([exe] + cmd, cwd=cwd)
 
 def build_python_wheel(source_dir, build_dir, configs, use_cuda, use_tensorrt):