From ab350fa4c7f3c98a5792a48042673eaaeba99a88 Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Wed, 19 Dec 2018 18:56:35 +1000 Subject: [PATCH] Re-structure the inference session initialization to (#217) - apply any transforms to the main graph and any subgraphs first - call Graph::Resolve() once on the main graph, which will recurse into the subgraphs - previously it was called after the transform on each subgraph, which results in it traversing up to the main graph to call resolve, and that resolve call recursing into all subgraphs every time. This avoids lots of unnecessary Graph::Resolve calls, and prevents subgraphs from being broken by SessionStateInitializer::InitializeAndSave calling graph_.CleanAllInitializedTensors() prior to final Graph::Resolve call. If a subgraph has optional inputs the backing initializers were removed by CleanAllInitializedTensors causing the next Resolve to incorrectly turn them into required inputs. --- .../framework/session_state_initializer.cc | 56 +------------- .../framework/session_state_initializer.h | 4 +- onnxruntime/core/graph/graph_utils.cc | 76 ++++++++++++++++--- onnxruntime/core/graph/graph_utils.h | 15 ++-- onnxruntime/core/session/inference_session.cc | 61 +++++++++++++-- 5 files changed, 131 insertions(+), 81 deletions(-) diff --git a/onnxruntime/core/framework/session_state_initializer.cc b/onnxruntime/core/framework/session_state_initializer.cc index 6d8df8500e..4c2126c48c 100644 --- a/onnxruntime/core/framework/session_state_initializer.cc +++ b/onnxruntime/core/framework/session_state_initializer.cc @@ -26,12 +26,6 @@ namespace onnxruntime { -static common::Status TransformGraph(onnxruntime::Graph& graph, - const onnxruntime::GraphTransformerManager& graph_transformer_mgr, - const ExecutionProviders& exec_providers, - KernelRegistryManager& kernel_registry_manager, - const InsertCastTransformer& insert_cast_transformer); - static common::Status SaveMLValueNameIndexMapping(const onnxruntime::Graph& graph, MLValueNameIdxMap& mlvalue_name_idx_map, const logging::Logger& logger); @@ -68,15 +62,9 @@ SessionStateInitializer::SessionStateInitializer(onnxruntime::Graph& graph, logger_{logger} { } -common::Status SessionStateInitializer::CreatePlan(const onnxruntime::GraphTransformerManager& graph_transformation_manager, - const InsertCastTransformer& insert_cast_transformer, - const std::vector& outer_scope_node_args, +common::Status SessionStateInitializer::CreatePlan(const std::vector& outer_scope_node_args, bool enable_sequential_execution) { - ORT_RETURN_IF_ERROR(TransformGraph(graph_, graph_transformation_manager, - execution_providers_, kernel_registry_manager_, - insert_cast_transformer)); - - // After transformation/partitioning, the graph now is fixed and graph viewer is created and set for execution. + // the graph now is fixed and graph viewer is created and set for execution. session_state_.SetGraphViewer(std::make_unique(graph_)); auto& mlvalue_name_idx_map = session_state_.GetMLValueNameIdxMap(); @@ -142,46 +130,6 @@ common::Status SessionStateInitializer::InitializeAndSave(bool enable_memory_pat return Status::OK(); } -common::Status TransformGraph(onnxruntime::Graph& graph, - const onnxruntime::GraphTransformerManager& graph_transformer_mgr, - const ExecutionProviders& providers, - KernelRegistryManager& kernel_registry_manager, - const InsertCastTransformer& insert_cast_transformer) { - // The transformer order: - // 1. built-in graph rewriter - // 2. each execution provider's transformer - // 3. do node placement according to kernel definition - // 4. insert copy nodes - // 5. insert cast nodes. - - // first apply the default/system/basic graph to graph optimizations. - ORT_RETURN_IF_ERROR(graph_transformer_mgr.ApplyAll(graph)); - - auto kernels{kernel_registry_manager.GetAllKernelRegistries()}; - - // Do partitioning based on execution providers' capability. - GraphPartitioner partitioner(kernel_registry_manager, providers); - ORT_RETURN_IF_ERROR(partitioner.Partition(graph)); - - // Insert copy nodes. - for (auto& provider : providers) { - if (provider->Type() != onnxruntime::kCpuExecutionProvider && - provider->Type() != onnxruntime::kMklDnnExecutionProvider && - provider->Type() != onnxruntime::kNupharExecutionProvider) { - TransformerMemcpyImpl copy_impl(graph, provider->Type()); - copy_impl.ModifyGraph(kernel_registry_manager); - } - } - - // Insert cast node/s. - bool modified = false; - ORT_RETURN_IF_ERROR(insert_cast_transformer.Apply(graph, modified)); - - ORT_RETURN_IF_ERROR(graph.Resolve()); - - return common::Status::OK(); -} - // Build the MLValue name->idx mapping common::Status SaveMLValueNameIndexMapping(const onnxruntime::Graph& graph, MLValueNameIdxMap& mlvalue_name_idx_map, diff --git a/onnxruntime/core/framework/session_state_initializer.h b/onnxruntime/core/framework/session_state_initializer.h index 95fc97d65b..fa8e4b5f98 100644 --- a/onnxruntime/core/framework/session_state_initializer.h +++ b/onnxruntime/core/framework/session_state_initializer.h @@ -29,9 +29,7 @@ class SessionStateInitializer { const logging::Logger& logger); // First perform any transformations and create the execution plan - common::Status CreatePlan(const onnxruntime::GraphTransformerManager& graph_transformation_manager, - const InsertCastTransformer& insert_cast_transformer, - const std::vector& outer_scope_node_args, + common::Status CreatePlan(const std::vector& outer_scope_node_args, bool enable_sequential_execution); // initialize tensors, and save. save kernels and input/output node mappings diff --git a/onnxruntime/core/graph/graph_utils.cc b/onnxruntime/core/graph/graph_utils.cc index 7daeba8388..3dcbc71358 100644 --- a/onnxruntime/core/graph/graph_utils.cc +++ b/onnxruntime/core/graph/graph_utils.cc @@ -4,18 +4,70 @@ namespace onnxruntime { namespace utils { - // fusion is only done for ONNX domain ops - bool IsSupportedOptypeVersionAndDomain(const Node& node, - const std::string& op_type, - ONNX_NAMESPACE::OperatorSetVersion version, - const std::string& domain) { - if (node.OpType() != op_type || - node.Op()->Deprecated() || node.Op()->SinceVersion() != version || - (!node.Domain().empty() && node.Domain() != domain)) { - return false; - } - return true; +// fusion is only done for ONNX domain ops +bool IsSupportedOptypeVersionAndDomain(const Node& node, + const std::string& op_type, + ONNX_NAMESPACE::OperatorSetVersion version, + const std::string& domain) { + if (node.OpType() != op_type || + node.Op()->Deprecated() || node.Op()->SinceVersion() != version || + (!node.Domain().empty() && node.Domain() != domain)) { + return false; } + return true; } -} // namespace onnxruntime \ No newline at end of file +Status ForAllMutableSubgraphs(Graph& graph, std::function func) { + Status status = Status::OK(); + + for (auto& node : graph.Nodes()) { + for (auto& attribute : node.GetAttributes()) { + auto& name = attribute.first; + auto& proto = attribute.second; + + // check if it has a subgraph + if (proto.has_g()) { + Graph* subgraph = node.GetMutableGraphAttribute(name); + ORT_ENFORCE(subgraph, "Main Graph instance should have populated all subgraphs when being resolved."); + + status = func(*subgraph); + ORT_RETURN_IF_ERROR(status); + + // recurse + status = ForAllMutableSubgraphs(*subgraph, func); + ORT_RETURN_IF_ERROR(status); + } + } + } + + return status; +} + +Status ForAllSubgraphs(const Graph& graph, std::function func) { + Status status = Status::OK(); + + for (auto& node : graph.Nodes()) { + for (auto& attribute : node.GetAttributes()) { + auto& name = attribute.first; + auto& proto = attribute.second; + + // check if it has a subgraph + if (proto.has_g()) { + const Graph* subgraph = node.GetGraphAttribute(name); + ORT_ENFORCE(subgraph, "Main Graph instance should have populated all subgraphs when being resolved."); + + status = func(*subgraph); + ORT_RETURN_IF_ERROR(status); + + // recurse + status = ForAllSubgraphs(*subgraph, func); + ORT_RETURN_IF_ERROR(status); + } + } + } + + return status; +} + +} // namespace utils +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/graph_utils.h b/onnxruntime/core/graph/graph_utils.h index 7ff5e93d3e..993a535bab 100644 --- a/onnxruntime/core/graph/graph_utils.h +++ b/onnxruntime/core/graph/graph_utils.h @@ -9,10 +9,13 @@ namespace onnxruntime { namespace utils { - bool IsSupportedOptypeVersionAndDomain(const Node& node, - const std::string& op_type, - ONNX_NAMESPACE::OperatorSetVersion version, - const std::string& domain = kOnnxDomainAlias); -} +bool IsSupportedOptypeVersionAndDomain(const Node& node, + const std::string& op_type, + ONNX_NAMESPACE::OperatorSetVersion version, + const std::string& domain = kOnnxDomainAlias); -} \ No newline at end of file +Status ForAllMutableSubgraphs(Graph& main_graph, std::function func); +Status ForAllSubgraphs(Graph& main_graph, std::function func); + +} // namespace utils +} // namespace onnxruntime diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 52703a8b22..f61cd61122 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -14,6 +14,7 @@ #include "core/graph/graph_viewer.h" #include "core/graph/graph_transformer.h" #include "core/graph/graph_transformer_mgr.h" +#include "core/graph/graph_utils.h" #include "core/graph/model.h" #include "core/framework/allocatormgr.h" #include "core/framework/customregistry.h" @@ -247,6 +248,44 @@ class InferenceSession::Impl { return common::Status::OK(); } + static common::Status TransformGraph(onnxruntime::Graph& graph, + const onnxruntime::GraphTransformerManager& graph_transformer_mgr, + const ExecutionProviders& providers, + KernelRegistryManager& kernel_registry_manager, + const InsertCastTransformer& insert_cast_transformer) { + // The transformer order: + // 1. built-in graph rewriter + // 2. each execution provider's transformer + // 3. do node placement according to kernel definition + // 4. insert copy nodes + // 5. insert cast nodes. + + // first apply the default/system/basic graph to graph optimizations. + ORT_RETURN_IF_ERROR(graph_transformer_mgr.ApplyAll(graph)); + + auto kernels{kernel_registry_manager.GetAllKernelRegistries()}; + + // Do partitioning based on execution providers' capability. + GraphPartitioner partitioner(kernel_registry_manager, providers); + ORT_RETURN_IF_ERROR(partitioner.Partition(graph)); + + // Insert copy nodes. + for (auto& provider : providers) { + if (provider->Type() != onnxruntime::kCpuExecutionProvider && + provider->Type() != onnxruntime::kMklDnnExecutionProvider && + provider->Type() != onnxruntime::kNupharExecutionProvider) { + TransformerMemcpyImpl copy_impl(graph, provider->Type()); + copy_impl.ModifyGraph(kernel_registry_manager); + } + } + + // Insert cast node/s. + bool modified = false; + ORT_RETURN_IF_ERROR(insert_cast_transformer.Apply(graph, modified)); + + return common::Status::OK(); + } + // memory allocations for a subgraph that are owned by InferenceSession struct SubgraphMemory { std::unique_ptr session_state; @@ -277,10 +316,8 @@ class InferenceSession::Impl { SessionStateInitializer initializer{*subgraph, *subgraph_info.session_state, execution_providers_, kernel_registry_manager_, *session_logger_}; - ORT_RETURN_IF_ERROR( - initializer.CreatePlan(graph_transformation_mgr_, insert_cast_transformer_, - node.ImplicitInputDefs(), - session_options_.enable_sequential_execution)); + ORT_RETURN_IF_ERROR(initializer.CreatePlan(node.ImplicitInputDefs(), + session_options_.enable_sequential_execution)); ORT_RETURN_IF_ERROR(initializer.InitializeAndSave(session_state_.GetEnableMemoryPattern(), subgraph_info.weights_buffers)); @@ -347,9 +384,21 @@ class InferenceSession::Impl { SessionStateInitializer session_initializer{graph, session_state_, execution_providers_, kernel_registry_manager_, *session_logger_}; - ORT_RETURN_IF_ERROR(session_initializer.CreatePlan(graph_transformation_mgr_, insert_cast_transformer_, - {}, session_options_.enable_sequential_execution)); + // apply any transformations to the main graph and any subgraphs + ORT_RETURN_IF_ERROR(TransformGraph(graph, graph_transformation_mgr_, + execution_providers_, kernel_registry_manager_, + insert_cast_transformer_)); + ORT_RETURN_IF_ERROR(utils::ForAllMutableSubgraphs(graph, [this](Graph& subgraph) { + return TransformGraph(subgraph, graph_transformation_mgr_, + execution_providers_, kernel_registry_manager_, + insert_cast_transformer_); + })); + + // now that all the transforms are done, call Resolve on the main graph. this will recurse into the subgraphs. + ORT_RETURN_IF_ERROR(graph.Resolve()); + + ORT_RETURN_IF_ERROR(session_initializer.CreatePlan({}, session_options_.enable_sequential_execution)); ORT_RETURN_IF_ERROR(session_initializer.InitializeAndSave(session_state_.GetEnableMemoryPattern(), weights_buffers_));