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.

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<NodeArg*>& outer_scope_node_args,
+common::Status SessionStateInitializer::CreatePlan(const std::vector<NodeArg*>& 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<onnxruntime::GraphViewer>(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,

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<NodeArg*>& outer_scope_node_args,
+  common::Status CreatePlan(const std::vector<NodeArg*>& outer_scope_node_args,
                             bool enable_sequential_execution);
 
   // initialize tensors, and save. save kernels and input/output node mappings

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
+Status ForAllMutableSubgraphs(Graph& graph, std::function<Status(Graph&)> 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<Status(const Graph&)> 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

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);
 
-}
+Status ForAllMutableSubgraphs(Graph& main_graph, std::function<Status(Graph&)> func);
+Status ForAllSubgraphs(Graph& main_graph, std::function<Status(Graph&)> func);
+
+}  // namespace utils
+}  // namespace onnxruntime

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<SessionState> 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_));