From b39257a5e6ed6768d4872691ce3be5d6072928e6 Mon Sep 17 00:00:00 2001 From: Chi Lo <54722500+chilo-ms@users.noreply.github.com> Date: Mon, 1 Aug 2022 23:57:30 -0700 Subject: [PATCH] Enable support of multi-level nested control flow ops model for TRT EP (#12147) * Make multiple-level nested control flow op model work * find correct input index * find correct input index (cont.) * enable nested layer unit tests for TRT EP * add comment * add Scan op to current workaround support of control flow op --- .../tensorrt/tensorrt_execution_provider.cc | 29 +++++++++++++++++-- .../test/framework/inference_session_test.cc | 14 +++++++-- 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index a827bc5279..ece582fd29 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -616,7 +616,14 @@ std::unique_ptr TensorrtExecutionProvider::GetSubGraph(SubGraph if (node->GetOutputEdgesCount() > node->OutputDefs().size()) { for (auto it = node->OutputEdgesBegin(), end = node->OutputEdgesEnd(); it != end; ++it) { const auto& node_idx = it->GetNode().Index(); - const auto& output = (it->GetNode()).InputDefs()[it->GetDstArgIndex()]; + const onnxruntime::NodeArg* output; + // The dst_arg_index from GetDstArgIndex() could be the index for explicit/implicit input defs of the node. + // We need to get the correct input index accordingly. (See Graph::BuildConnections() in graph.cc for more details) + if (it->GetDstArgIndex() < static_cast(it->GetNode().InputDefs().size())) { + output = (it->GetNode()).InputDefs()[it->GetDstArgIndex()]; + } else { + output = (it->GetNode()).ImplicitInputDefs()[it->GetDstArgIndex() - static_cast(it->GetNode().InputDefs().size())]; + } if (node_set.find(node_idx) != node_set.end()) { const auto& iter = fused_inputs.find(output); if (iter != fused_inputs.end()) { @@ -897,6 +904,10 @@ bool TensorrtExecutionProvider::DetectTensorRTGraphCycles(SubGraphCollection_t& input_to_nodes_map[input->Name()].insert(node_name); } + for (const auto& input : node->ImplicitInputDefs()) { + input_to_nodes_map[input->Name()].insert(node_name); + } + for (const auto& output : node->OutputDefs()) { node_to_outputs_map[node_name].insert(output->Name()); } @@ -970,7 +981,21 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph, const int number_of_ort_nodes = graph.NumberOfNodes(); std::vector nodes_vector(number_of_ort_nodes); std::iota(std::begin(nodes_vector), std::end(nodes_vector), 0); - SubGraphCollection_t supported_nodes_vector, parser_nodes_vector = {{nodes_vector, false}}; + std::vector filtered_nodes_vector; + const std::vector& node_index = graph.GetNodesInTopologicalOrder(); + + // We currently exclude "If" and "Loop" control flow ops from original node vector before calling TensorRT parser. + // The reason is, these control flow ops have subgraph which might contain TRT fused node after ORT partition. + // If this is the case, TensorRT parser will complain the non-recognized TRT fused node and fail. + for (const auto& index : nodes_vector) { + const auto& node = graph.GetNode(node_index[index]); + if (node->OpType() == "If" || node->OpType() == "Loop" || node->OpType() == "Scan") { + continue; + } + filtered_nodes_vector.push_back(index); + } + + SubGraphCollection_t supported_nodes_vector, parser_nodes_vector = {{filtered_nodes_vector, false}}; bool early_termination = false; supported_nodes_vector = GetSupportedList(parser_nodes_vector, 0, max_partition_iterations_, graph, &early_termination); if (early_termination) { diff --git a/onnxruntime/test/framework/inference_session_test.cc b/onnxruntime/test/framework/inference_session_test.cc index 041286bedc..4ccf5298eb 100644 --- a/onnxruntime/test/framework/inference_session_test.cc +++ b/onnxruntime/test/framework/inference_session_test.cc @@ -35,6 +35,9 @@ #include "core/providers/cuda/cuda_provider_factory.h" #include "core/providers/cuda/gpu_data_transfer.h" #endif +#ifdef USE_TENSORRT +#include "core/providers/tensorrt/tensorrt_provider_options.h" +#endif #ifdef USE_ROCM #include "core/providers/rocm/rocm_provider_factory.h" #include "core/providers/rocm/gpu_data_transfer.h" @@ -1447,6 +1450,7 @@ TEST(InferenceSessionTests, Test3LayerNestedSubgraph) { float_tensor.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); ONNX_NAMESPACE::TypeProto bool_tensor; bool_tensor.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_BOOL); + bool_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); auto& if_cond_input = graph.GetOrCreateNodeArg("if_cond_input", &bool_tensor); auto& graph_if_input = graph.GetOrCreateNodeArg("graph_if_input", nullptr); @@ -1485,7 +1489,9 @@ TEST(InferenceSessionTests, Test3LayerNestedSubgraph) { so.session_logid = "InferenceSessionTests.Test3LayerNestedSubgraph"; InferenceSession session_object{so, GetEnvironment()}; -#if USE_CUDA +#if USE_TENSORRT + ASSERT_STATUS_OK(session_object.RegisterExecutionProvider(DefaultTensorrtExecutionProvider())); +#elif USE_CUDA ASSERT_STATUS_OK(session_object.RegisterExecutionProvider(DefaultCudaExecutionProvider())); #elif USE_ROCM ASSERT_STATUS_OK(session_object.RegisterExecutionProvider(DefaultRocmExecutionProvider())); @@ -1562,11 +1568,13 @@ TEST(InferenceSessionTests, Test2LayerNestedSubgraph) { ONNX_NAMESPACE::TypeProto float_tensor_input; float_tensor_input.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + float_tensor_input.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); ONNX_NAMESPACE::TypeProto float_tensor; float_tensor.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); float_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_param("__graph_0__float_unknown"); ONNX_NAMESPACE::TypeProto bool_tensor; bool_tensor.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_BOOL); + bool_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); // graph inputs auto& input_0 = graph.GetOrCreateNodeArg("input_0", &float_tensor_input); @@ -1617,7 +1625,9 @@ TEST(InferenceSessionTests, Test2LayerNestedSubgraph) { so.session_logid = "InferenceSessionTests.Test2LayerNestedSubgraph"; InferenceSession session_object{so, GetEnvironment()}; -#if USE_CUDA +#if USE_TENSORRT + ASSERT_STATUS_OK(session_object.RegisterExecutionProvider(DefaultTensorrtExecutionProvider())); +#elif USE_CUDA ASSERT_STATUS_OK(session_object.RegisterExecutionProvider(DefaultCudaExecutionProvider())); #elif USE_ROCM ASSERT_STATUS_OK(session_object.RegisterExecutionProvider(DefaultRocmExecutionProvider()));