From 2127a229d7f9640f171d1df005844680a2e92ed1 Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Thu, 5 Nov 2020 17:21:56 +1000 Subject: [PATCH] The IndexedSubGraph is used to create the Function body, but after that is invalid as the nodes it referred to have been removed from the main Graph. As such there's no need to store it in the FunctionImpl instance. (#5669) --- include/onnxruntime/core/graph/function.h | 5 +-- include/onnxruntime/core/graph/graph.h | 2 +- .../core/framework/graph_partitioner.cc | 33 ++++++++++--------- onnxruntime/core/graph/function.cc | 19 +++++------ onnxruntime/core/graph/function_impl.h | 5 +-- onnxruntime/core/graph/graph.cc | 24 +++++++------- 6 files changed, 40 insertions(+), 48 deletions(-) diff --git a/include/onnxruntime/core/graph/function.h b/include/onnxruntime/core/graph/function.h index 4c763f9840..6328713705 100644 --- a/include/onnxruntime/core/graph/function.h +++ b/include/onnxruntime/core/graph/function.h @@ -26,9 +26,6 @@ class Function { /** Gets the Graph instance for the Function body subgraph. */ virtual const onnxruntime::Graph& Body() const = 0; - - /** Gets the IndexedSubGraph for the Function. */ - virtual const IndexedSubGraph& GetIndexedSubGraph() const = 0; }; /** @@ -37,6 +34,6 @@ Create a new Function instance. @param customized_func the IndexedSubGraph to use for the Function. */ std::unique_ptr MakeFunction(const onnxruntime::Graph& graph, - std::unique_ptr customized_func, + const IndexedSubGraph& nodes_to_fuse, const logging::Logger& logger); } // namespace onnxruntime diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index c07a1f634f..6c3df7f469 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -906,7 +906,7 @@ class Graph { @param fused_node_name The name for the new Node. @returns Node with fused subgraph. */ - Node& FuseSubGraph(std::unique_ptr sub_graph, const std::string& fused_node_name); + Node& FuseSubGraph(const IndexedSubGraph& sub_graph, const std::string& fused_node_name); /** Directly insert the nodes in the function Node provided into this Graph. diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 1a604048ba..43f431ef5c 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -60,28 +60,25 @@ KernelDefBuilder& BuildFusedKernelDef(KernelDefBuilder& builder, const onnxrunti * \param count A counter for generating fused node names. Should be unique within this subgraph * \return Fused node. Return nullptr if there is no fuse */ -static Node* PlaceNode(Graph& graph, std::unique_ptr capability, - const KernelRegistryManager& kernel_registry_mgr, const std::string& provider_type, int& count) { - if (nullptr == capability) { - return nullptr; - } - - if (nullptr == capability->GetMetaDef()) { +static Node* PlaceNode(Graph& graph, const IndexedSubGraph& capability, + const KernelRegistryManager& kernel_registry_mgr, const std::string& provider_type, + int& count) { + if (nullptr == capability.GetMetaDef()) { // The can run a single node in the if not using meta-defs. // A fused kernel is not supported in this case. - ORT_ENFORCE(1 == capability->nodes.size()); + ORT_ENFORCE(1 == capability.nodes.size()); - auto* node = graph.GetNode(capability->nodes[0]); + auto* node = graph.GetNode(capability.nodes[0]); if (nullptr != node && node->GetExecutionProviderType().empty()) { // The node was not fused or assigned. Assign it to this . node->SetExecutionProviderType(provider_type); } } else { // The can run a fused in the . - ORT_ENFORCE(nullptr != capability->GetMetaDef()); + ORT_ENFORCE(nullptr != capability.GetMetaDef()); // Check whether any node in the was already assigned. bool sub_graph_available_for_assignment = true; - for (auto node_index : capability->nodes) { + for (auto node_index : capability.nodes) { auto node = graph.GetNode(node_index); if (nullptr == node || !node->GetExecutionProviderType().empty()) { // The node was fused or assigned, so that the whole sub-graph will not be assigned to this @@ -92,9 +89,9 @@ static Node* PlaceNode(Graph& graph, std::unique_ptr capability } if (sub_graph_available_for_assignment) { std::ostringstream oss; - oss << provider_type << "_" << capability->GetMetaDef()->name << "_" << count++; + oss << provider_type << "_" << capability.GetMetaDef()->name << "_" << count++; std::string node_name = oss.str(); - auto& fused_node = graph.FuseSubGraph(std::move(capability), node_name); + auto& fused_node = graph.FuseSubGraph(capability, node_name); fused_node.SetExecutionProviderType(provider_type); // searching in kernel registries, if no kernel registered for the fused_node, use compile approach if (!KernelRegistryManager::HasImplementationOf(kernel_registry_mgr, fused_node, provider_type)) { @@ -154,7 +151,11 @@ Status GraphPartitioner::Partition(Graph& graph, bool export_dll, FuncManager& f std::vector> capabilities = provider->GetCapability(graph_viewer, kernel_registry_mgr_.GetKernelRegistriesByProviderType(provider->Type())); for (auto& capability : capabilities) { - Node* n = PlaceNode(graph, std::move(capability->sub_graph), kernel_registry_mgr_, provider->Type(), count); + if (!capability || !capability->sub_graph) { // in theory an EP could return an empty value... + continue; + } + + Node* n = PlaceNode(graph, *capability->sub_graph, kernel_registry_mgr_, provider->Type(), count); if (n != nullptr) { nodes_need_compile.push_back(n); } @@ -197,9 +198,9 @@ Status GraphPartitioner::Partition(Graph& graph, bool export_dll, FuncManager& f if (nullptr == node_func) { continue; } - nodes_need_inline.push_back(&node); + nodes_need_inline.push_back(&node); } - } + } for (auto* node : nodes_need_inline) { // If the node has a functionbody with no kernel and cannot be inlined diff --git a/onnxruntime/core/graph/function.cc b/onnxruntime/core/graph/function.cc index 9e647ab759..439073ffd5 100644 --- a/onnxruntime/core/graph/function.cc +++ b/onnxruntime/core/graph/function.cc @@ -145,22 +145,22 @@ static void update_subgraphs_within_function_body(ONNX_NAMESPACE::GraphProto& su } FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph, - std::unique_ptr customized_func, + const IndexedSubGraph& nodes_to_fuse, const logging::Logger& logger) : parent_graph_(&graph), body_("fused_function_subgraph", false, onnxruntime::ModelMetaData(), graph.ModelPath().ToPathString(), IOnnxRuntimeOpSchemaRegistryList({graph.GetSchemaRegistry()}), graph.DomainToVersionMap(), {}, logger) { - customized_func_body_ = std::move(customized_func); auto& function_body_graph = body_.MainGraph(); - auto meta_def = customized_func_body_->GetMetaDef(); + auto meta_def = nodes_to_fuse.GetMetaDef(); op_schema_ = onnxruntime::make_unique(); op_schema_->SetName(meta_def->name); op_schema_->SetDomain(meta_def->domain); op_schema_->SetDoc(meta_def->doc_string); op_schema_->SinceVersion(meta_def->since_version); + int i = 0; std::vector function_body_graph_inputs; function_body_graph_inputs.resize(meta_def->inputs.size()); @@ -172,6 +172,7 @@ FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph, op_schema_->Input(i, input, "", *input_arg->Type()); ++i; } + i = 0; std::vector function_body_graph_outputs; function_body_graph_outputs.resize(meta_def->outputs.size()); @@ -182,14 +183,16 @@ FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph, op_schema_->Output(i, output, "", *output_arg->Type()); ++i; } + op_schema_->Finalize(); function_body_graph.SetInputs(function_body_graph_inputs); function_body_graph.SetOutputs(function_body_graph_outputs); + //Add node and node args //TODO: for better performance, we could try to transfer the nodes in parent graph to sub-graph directly, //instead of create new nodes. - for (auto& node_index : customized_func_body_->nodes) { + for (auto& node_index : nodes_to_fuse.nodes) { auto node = parent_graph_->GetNode(node_index); std::vector inputs; std::vector outputs; @@ -394,17 +397,13 @@ const onnxruntime::Graph& FunctionImpl::Body() const { return body_.MainGraph(); } -const IndexedSubGraph& FunctionImpl::GetIndexedSubGraph() const { - return *customized_func_body_; -} - const ONNX_NAMESPACE::FunctionProto* FunctionImpl::GetFuncProto() const { return &onnx_func_proto_; } std::unique_ptr MakeFunction(const onnxruntime::Graph& graph, - std::unique_ptr customized_func, + const IndexedSubGraph& nodes_to_fuse, const logging::Logger& logger) { - return onnxruntime::make_unique(graph, std::move(customized_func), logger); + return onnxruntime::make_unique(graph, nodes_to_fuse, logger); } } // namespace onnxruntime diff --git a/onnxruntime/core/graph/function_impl.h b/onnxruntime/core/graph/function_impl.h index c0c5b7d9db..e477b35f22 100644 --- a/onnxruntime/core/graph/function_impl.h +++ b/onnxruntime/core/graph/function_impl.h @@ -17,7 +17,7 @@ namespace onnxruntime { class FunctionImpl final : public Function { public: FunctionImpl(const onnxruntime::Graph& graph, - std::unique_ptr customized_func, + const IndexedSubGraph& nodes_to_fuse, const logging::Logger& logger); FunctionImpl(const onnxruntime::Graph& graph, @@ -31,13 +31,10 @@ class FunctionImpl final : public Function { const onnxruntime::Graph& Body() const override; - const IndexedSubGraph& GetIndexedSubGraph() const override; - const ONNX_NAMESPACE::FunctionProto* GetFuncProto() const; private: const onnxruntime::Graph* const parent_graph_; - std::unique_ptr customized_func_body_; std::unique_ptr op_schema_; onnxruntime::Model body_; ONNX_NAMESPACE::FunctionProto onnx_func_proto_; diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index ec6d117bfc..f020326a04 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -991,7 +991,7 @@ Graph::Graph(const Model& owning_model, // For now we convert sparse_intializer to dense tensors // since there are currently no supported ops that consume sparse - // initializers directly. We remove them from graph_proto. We will reconstitute them + // initializers directly. We remove them from graph_proto. We will reconstitute them // when saving to ORT format to save space on disk. if (graph_proto_->sparse_initializer_size() > 0) { for (const auto& sparse_tensor : graph_proto_->sparse_initializer()) { @@ -1037,8 +1037,8 @@ Graph::Graph(const Model& owning_model, auto p = name_to_initial_tensor_.emplace(tensor.name(), &tensor); if (!p.second) { LOGS(logger_, WARNING) << "Duplicate initializer (dense, sparse or ConstantNode): '" << tensor.name() - << "' the model will use the latest encountered initializer" - << ". Please, fix your model."; + << "' the model will use the latest encountered initializer" + << ". Please, fix your model."; p.first->second = &tensor; } @@ -3317,12 +3317,10 @@ IOnnxRuntimeOpSchemaCollectionPtr Graph::GetSchemaRegistry() const { return schema_registry_; } -Node& Graph::FuseSubGraph(std::unique_ptr<::onnxruntime::IndexedSubGraph> sub_graph, - const std::string& fused_node_name) { - ORT_ENFORCE(nullptr != sub_graph && nullptr != sub_graph->GetMetaDef()); - - auto func_meta_def = sub_graph->GetMetaDef(); +Node& Graph::FuseSubGraph(const IndexedSubGraph& sub_graph, const std::string& fused_node_name) { + auto* func_meta_def = sub_graph.GetMetaDef(); ORT_ENFORCE(nullptr != func_meta_def); + std::vector input_args; std::vector output_args; for (auto& arg_name : func_meta_def->inputs) { @@ -3341,12 +3339,11 @@ Node& Graph::FuseSubGraph(std::unique_ptr<::onnxruntime::IndexedSubGraph> sub_gr func_meta_def->domain); fused_node.SetNodeType(Node::Type::Fused); - function_container_.emplace_back(MakeFunction(*this, std::move(sub_graph), logger_)); + function_container_.emplace_back(MakeFunction(*this, sub_graph, logger_)); fused_node.SetFunctionBody(*function_container_.back()); // Remove nodes fused above. - auto& sub_graph_ref = function_container_.back()->GetIndexedSubGraph(); - for (auto node_index : sub_graph_ref.nodes) { + for (auto node_index : sub_graph.nodes) { auto node = GetNode(node_index); if (nullptr == node) { continue; @@ -3357,6 +3354,7 @@ Node& Graph::FuseSubGraph(std::unique_ptr<::onnxruntime::IndexedSubGraph> sub_gr } RemoveNode(node_index); } + return fused_node; } @@ -3578,8 +3576,8 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::experimental::fbs::Gr auto p = name_to_initial_tensor_.emplace(initializer->name(), initializer); if (!p.second) { LOGS(logger_, WARNING) << "Duplicate initializer (dense or ConstantNode): '" << initializer->name() - << "' the model will use the latest encountered initializer" - << ". Please, fix your model."; + << "' the model will use the latest encountered initializer" + << ". Please, fix your model."; p.first->second = initializer; } }