From 05d20343ee04990a7082ccb162faab9dd9a8305c Mon Sep 17 00:00:00 2001 From: stevenlix <38092805+stevenlix@users.noreply.github.com> Date: Wed, 22 Dec 2021 12:19:56 -0800 Subject: [PATCH] Remove duplicated constant initializer copies for TensorRT nodes (#10105) * add new field constant_initializers in metadef and remove constant initializers from trt node inputs * remove redundancy * use GetConstantInitializer() to get constant initializers * add ORT_ENFORCE check Co-authored-by: Ubuntu --- include/onnxruntime/core/graph/indexed_sub_graph.h | 1 + onnxruntime/core/graph/function.cc | 10 ++++++++++ .../providers/shared_library/provider_interfaces.h | 1 + .../shared_library/provider_wrappedtypes.h | 2 ++ .../tensorrt/tensorrt_execution_provider.cc | 13 +++++++++++++ onnxruntime/core/session/provider_bridge_ort.cc | 1 + 6 files changed, 28 insertions(+) diff --git a/include/onnxruntime/core/graph/indexed_sub_graph.h b/include/onnxruntime/core/graph/indexed_sub_graph.h index 43275cebe1..0ec595399e 100644 --- a/include/onnxruntime/core/graph/indexed_sub_graph.h +++ b/include/onnxruntime/core/graph/indexed_sub_graph.h @@ -39,6 +39,7 @@ struct IndexedSubGraph { std::vector inputs; ///< Inputs of customized SubGraph/FunctionProto. std::vector outputs; ///< Outputs of customized SubGraph/FunctionProto. + std::vector constant_initializers; ///< Constant initializers of customized SubGraph/FunctionProto. NodeAttributes attributes; ///< Attributes of customized SubGraph/FunctionProto. std::string doc_string; ///< Doc string of customized SubGraph/FunctionProto. diff --git a/onnxruntime/core/graph/function.cc b/onnxruntime/core/graph/function.cc index 43a1638354..613915934e 100644 --- a/onnxruntime/core/graph/function.cc +++ b/onnxruntime/core/graph/function.cc @@ -457,6 +457,16 @@ FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph, } } + for (const auto& constant_initializer : meta_def->constant_initializers) { + const ONNX_NAMESPACE::TensorProto* initializer = graph.GetConstantInitializer(constant_initializer, true); + ORT_ENFORCE(initializer != nullptr, "Initializer " + constant_initializer + " is not found or is not constant initializer."); + // meta_def->constant_initializers could have duplicates so make sure we only add once + const ONNX_NAMESPACE::TensorProto* subgraph_initializer = nullptr; + if (!function_body_graph.GetInitializedTensor(constant_initializer, subgraph_initializer)) { + function_body_graph.AddInitializedTensor(*initializer); + } + } + //TODO: if we reuse the nodes in parent graph, maybe we don't need to resolve it. auto status = function_body_graph.Resolve(); ORT_ENFORCE(status.IsOK(), status.ErrorMessage()); diff --git a/onnxruntime/core/providers/shared_library/provider_interfaces.h b/onnxruntime/core/providers/shared_library/provider_interfaces.h index bd0c7a581c..37ce00a4e8 100644 --- a/onnxruntime/core/providers/shared_library/provider_interfaces.h +++ b/onnxruntime/core/providers/shared_library/provider_interfaces.h @@ -429,6 +429,7 @@ struct ProviderHost { virtual ONNX_NAMESPACE::OperatorStatus& IndexedSubGraph_MetaDef__status(IndexedSubGraph_MetaDef* p) = 0; virtual std::vector& IndexedSubGraph_MetaDef__inputs(IndexedSubGraph_MetaDef* p) = 0; virtual std::vector& IndexedSubGraph_MetaDef__outputs(IndexedSubGraph_MetaDef* p) = 0; + virtual std::vector& IndexedSubGraph_MetaDef__constant_initializers(IndexedSubGraph_MetaDef* p) = 0; virtual NodeAttributes& IndexedSubGraph_MetaDef__attributes(IndexedSubGraph_MetaDef* p) = 0; virtual std::string& IndexedSubGraph_MetaDef__doc_string(IndexedSubGraph_MetaDef* p) = 0; diff --git a/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h b/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h index bf93e2c415..572f6f7ca1 100644 --- a/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h +++ b/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h @@ -348,6 +348,8 @@ struct IndexedSubGraph_MetaDef final { const std::vector& inputs() const { return g_host->IndexedSubGraph_MetaDef__inputs(const_cast(this)); } std::vector& inputs() { return g_host->IndexedSubGraph_MetaDef__inputs(this); } const std::vector& outputs() const { return g_host->IndexedSubGraph_MetaDef__outputs(const_cast(this)); } + const std::vector& constant_initializers() const { return g_host->IndexedSubGraph_MetaDef__constant_initializers(const_cast(this)); } + std::vector& constant_initializers() { return g_host->IndexedSubGraph_MetaDef__constant_initializers(this); } std::vector& outputs() { return g_host->IndexedSubGraph_MetaDef__outputs(this); } NodeAttributes& attributes() { return g_host->IndexedSubGraph_MetaDef__attributes(this); } diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 3f0751151f..7deec5ca6d 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -714,10 +714,15 @@ std::unique_ptr TensorrtExecutionProvider::GetSubGraph(SubGraph int input_order = 0; int output_order = 0; + std::vector initializers; for (const auto& index : graph_nodes_index.first) { sub_graph->Nodes().push_back(node_index[index]); const auto& node = graph.GetNode(node_index[index]); for (const auto& input : node->InputDefs()) { + if (graph.IsConstantInitializer(input->Name(), true)) { + initializers.push_back(input->Name()); + continue; + } const auto& it = fused_outputs.find(input); if (it != fused_outputs.end()) { fused_outputs.erase(it); @@ -729,6 +734,10 @@ std::unique_ptr TensorrtExecutionProvider::GetSubGraph(SubGraph } for (const auto& input : node->ImplicitInputDefs()) { + if (graph.IsConstantInitializer(input->Name(), true)) { + initializers.push_back(input->Name()); + continue; + } const auto& it = fused_outputs.find(input); if (it != fused_outputs.end()) { fused_outputs.erase(it); @@ -810,6 +819,10 @@ std::unique_ptr TensorrtExecutionProvider::GetSubGraph(SubGraph } } + for (const auto& initializer : initializers) { + meta_def->constant_initializers().push_back(initializer); + } + for (const auto& output : outputs) { if (output.second->Exists()) { meta_def->outputs().push_back(output.second->Name()); diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc index cb76556779..bec88849df 100644 --- a/onnxruntime/core/session/provider_bridge_ort.cc +++ b/onnxruntime/core/session/provider_bridge_ort.cc @@ -509,6 +509,7 @@ struct ProviderHostImpl : ProviderHost { ONNX_NAMESPACE::OperatorStatus& IndexedSubGraph_MetaDef__status(IndexedSubGraph_MetaDef* p) override { return p->status; } std::vector& IndexedSubGraph_MetaDef__inputs(IndexedSubGraph_MetaDef* p) override { return p->inputs; } std::vector& IndexedSubGraph_MetaDef__outputs(IndexedSubGraph_MetaDef* p) override { return p->outputs; } + std::vector& IndexedSubGraph_MetaDef__constant_initializers(IndexedSubGraph_MetaDef* p) override { return p->constant_initializers; } NodeAttributes& IndexedSubGraph_MetaDef__attributes(IndexedSubGraph_MetaDef* p) override { return p->attributes; } std::string& IndexedSubGraph_MetaDef__doc_string(IndexedSubGraph_MetaDef* p) override { return p->doc_string; }