microsoft · edgchen1 · Jan 4, 2022 · Nov 4, 2021 · Nov 8, 2021 · Nov 9, 2021
diff --git a/cmake/onnxruntime_optimizer.cmake b/cmake/onnxruntime_optimizer.cmake
@@ -13,8 +13,8 @@ if (onnxruntime_MINIMAL_BUILD)
 
   if (onnxruntime_ENABLE_ORT_FORMAT_RUNTIME_GRAPH_OPTIMIZATION)
     list(APPEND onnxruntime_optimizer_src_patterns
-      "${ONNXRUNTIME_ROOT}/core/optimizer/ort_format_runtime_optimization/utils.h"
-      "${ONNXRUNTIME_ROOT}/core/optimizer/ort_format_runtime_optimization/utils.cc"
+      "${ONNXRUNTIME_INCLUDE_DIR}/core/optimizer/graph_transformer_utils.h"
+      "${ONNXRUNTIME_ROOT}/core/optimizer/graph_transformer_utils.cc"
       "${ONNXRUNTIME_ROOT}/core/optimizer/qdq_transformer/qdq_util.h"
       "${ONNXRUNTIME_ROOT}/core/optimizer/qdq_transformer/qdq_util.cc"
       "${ONNXRUNTIME_ROOT}/core/optimizer/qdq_transformer/selectors_actions/*.h"
@@ -37,13 +37,6 @@ else()
     "${ONNXRUNTIME_ROOT}/core/optimizer/transpose_optimizer/*.h"
     "${ONNXRUNTIME_ROOT}/core/optimizer/transpose_optimizer/*.cc"
   )
-
-  if (onnxruntime_ENABLE_ORT_FORMAT_RUNTIME_GRAPH_OPTIMIZATION)
-    list(APPEND onnxruntime_optimizer_src_patterns
-      "${ONNXRUNTIME_ROOT}/core/optimizer/ort_format_runtime_optimization/utils.h"
-      "${ONNXRUNTIME_ROOT}/core/optimizer/ort_format_runtime_optimization/utils.cc"
-    )
-  endif()
 endif()
 
 if (onnxruntime_ENABLE_TRAINING)

diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake
@@ -230,6 +230,11 @@ else()  # minimal and/or reduced ops build
   endif()
 endif()
 
+if(NOT onnxruntime_MINIMAL_BUILD OR onnxruntime_EXTENDED_MINIMAL_BUILD)
+  list(APPEND onnxruntime_test_optimizer_src
+       "${TEST_SRC_DIR}/optimizer/runtime_optimization/graph_runtime_optimization_test.cc")
+endif()
+
 file(GLOB onnxruntime_test_training_src
   "${ORTTRAINING_SOURCE_DIR}/test/model/*.cc"
   "${ORTTRAINING_SOURCE_DIR}/test/gradient/*.cc"

diff --git a/include/onnxruntime/core/graph/basic_types.h b/include/onnxruntime/core/graph/basic_types.h
@@ -30,6 +30,8 @@ using NodeArgInfo = ONNX_NAMESPACE::ValueInfoProto;
 using InitializedTensorSet = std::unordered_map<std::string, const ONNX_NAMESPACE::TensorProto*>;
 using ArgNameToTypeMap = std::unordered_map<std::string, ONNX_NAMESPACE::TypeProto>;
 using ProviderType = const std::string&;
+using HashValue = uint64_t;  // TODO remove after https://github.com/microsoft/onnxruntime/pull/9710 is merged
+
 // TODO - Evaluate switching the types below to support transparent comparators and enable
 // lookups based on gsl::cstring_span<> and std::string_view.  This would reduces allocations
 // converting to std::string, but requires conversion to std::map<std::string, foo, std::less<>>

diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h
@@ -10,6 +10,12 @@
 #include <unordered_map>
 #include <unordered_set>
 
+#ifdef _WIN32
+#pragma warning(push)
+// disable some warnings from protobuf to pass Windows build
+#pragma warning(disable : 4244)
+#endif
+
 #if !defined(ORT_MINIMAL_BUILD)
 #include "onnx/defs/schema.h"
 #else
@@ -18,6 +24,10 @@
 #include "onnx/onnx_pb.h"
 #include "onnx/onnx-operators_pb.h"
 
+#ifdef _WIN32
+#pragma warning(pop)
+#endif
+
 #include "gsl/gsl"
 
 #include "core/common/common.h"
@@ -779,7 +789,9 @@ class Graph {
     return ConstGraphNodes(nodes_, std::move(filter_func));
   }
 
-  /** Gets the maximum NodeIndex value used in the Graph. */
+  /** Gets the maximum NodeIndex value used in the Graph.
+  WARNING: This actually returns the max index value used + 1.
+  */
   int MaxNodeIndex() const noexcept { return static_cast<int>(nodes_.size()); }  //assume the casting won't overflow
 
   /** Gets the number of valid Nodes in the Graph.
@@ -978,7 +990,7 @@ class Graph {
   @remarks As a new Graph instance for the fused nodes is not created, a GraphViewer can be constructed with the
            IndexedSubGraph information to provide a view of the subgraph. The original nodes are left in place
            while this is in use.
-		   Call FinalizeFuseSubGraph to remove them once the fused replacement node is fully created.
+           Call FinalizeFuseSubGraph to remove them once the fused replacement node is fully created.
   */
   Node& BeginFuseSubGraph(const IndexedSubGraph& sub_graph, const std::string& fused_node_name);
 
@@ -1198,7 +1210,20 @@ class Graph {
   RuntimeOptimizationRecordContainer& MutableRuntimeOptimizations() {
     return runtime_optimizations_;
   }
-#endif
+
+  // Stores information collected during the replay of loaded runtime optimizations
+  struct RuntimeOptimizationReplayContext {
+    std::unordered_map<NodeIndex, HashValue> produced_node_index_to_kernel_def_hash;
+  };
+
+  const RuntimeOptimizationReplayContext& RuntimeOptimizationReplayCtx() const {
+    return runtime_optimization_replay_context_;
+  }
+
+  RuntimeOptimizationReplayContext& MutableRuntimeOptimizationReplayCtx() {
+    return runtime_optimization_replay_context_;
+  }
+#endif  // defined(ORT_ENABLE_ORT_FORMAT_RUNTIME_GRAPH_OPTIMIZATION)
 
  private:
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(Graph);
@@ -1421,6 +1446,8 @@ class Graph {
   // Note: runtime_optimizations_ == *runtime_optimizations_ptr_ and must be initialized
   std::unique_ptr<RuntimeOptimizationRecordContainer> runtime_optimizations_ptr_;
   RuntimeOptimizationRecordContainer& runtime_optimizations_;
+
+  RuntimeOptimizationReplayContext runtime_optimization_replay_context_;
 #endif
 
 #if !defined(ORT_MINIMAL_BUILD)

diff --git a/include/onnxruntime/core/optimizer/graph_transformer_utils.h b/include/onnxruntime/core/optimizer/graph_transformer_utils.h
@@ -3,19 +3,29 @@
 
 #pragma once
 
-#include <gsl/gsl>
+#include <memory>
+#include <unordered_set>
+#include <vector>
 
+#include "core/framework/session_options.h"
 #include "core/optimizer/graph_transformer.h"
+
+#if !defined(ORT_MINIMAL_BUILD)
 #include "core/optimizer/rule_based_graph_transformer.h"
 #include "core/optimizer/rewrite_rule.h"
+#endif
+
+#if defined(ORT_ENABLE_ORT_FORMAT_RUNTIME_GRAPH_OPTIMIZATION)
+#include "core/optimizer/selectors_actions/selector_action_transformer_apply_contexts.h"
+#endif
 
 namespace onnxruntime {
-struct FreeDimensionOverride;
 class IExecutionProvider;
-struct RuntimeOptimizationSaveContext;
 
 namespace optimizer_utils {
 
+#if !defined(ORT_MINIMAL_BUILD)
+
 /** Generates all predefined rules for this level.
    If rules_to_enable is not empty, it returns the intersection of predefined rules and rules_to_enable.
    TODO: This is visible for testing at the moment, but we should rather make it private. */
@@ -40,6 +50,10 @@ std::vector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
     const IExecutionProvider& execution_provider /*required by constant folding*/,
     const std::unordered_set<std::string>& rules_and_transformers_to_disable = {});
 
+#endif  // !defined(ORT_MINIMAL_BUILD)
+
+#if defined(ORT_ENABLE_ORT_FORMAT_RUNTIME_GRAPH_OPTIMIZATION)
+
 /** Generates all predefined transformers which support runtime optimizations for this level.
     Any transformers or rewrite rules named in rules_and_transformers_to_disable will be excluded.
 
@@ -48,8 +62,11 @@ std::vector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
     - The set of transformers which support runtime optimizations is different. */
 std::vector<std::unique_ptr<GraphTransformer>> GenerateTransformersForRuntimeOptimizations(
     TransformerLevel level,
-    const RuntimeOptimizationSaveContext& runtime_optimization_save_context,
+    const SessionOptions& session_options,
+    const SatApplyContextVariant& apply_context,
     const std::unordered_set<std::string>& rules_and_transformers_to_disable = {});
 
+#endif  // defined(ORT_ENABLE_ORT_FORMAT_RUNTIME_GRAPH_OPTIMIZATION)
+
 }  // namespace optimizer_utils
 }  // namespace onnxruntime
diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h
@@ -72,8 +72,8 @@ static const char* const kOrtSessionOptionsConfigUseORTModelBytesDirectly = "ses
 // Save information for replaying graph optimizations later instead of applying them directly.
 //
 // When an ONNX model is loaded, ORT can perform various optimizations on the graph.
-// However, when an ORT format model is loaded, these optimizations are typically not available - this scenario must
-// be supported by minimal builds.
+// However, when an ORT format model is loaded, these optimizations may not be available because this scenario must be
+// supported by minimal builds.
 // When loading an ONNX model, ORT can optionally save the effects of some optimizations for later replay in an ORT
 // format model. These are known as "runtime optimizations" - in an ORT format model, they happen at runtime.
 //

diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc
@@ -982,21 +982,41 @@ Status SessionState::LoadFromOrtFormat(const fbs::SessionState& fbs_session_stat
   // kernel hashes for model are in top level SessionState
   const auto& compiled_kernel_hashes = GetCompiledKernelHashes();
 
+  const bool original_nodes_should_exist =
+      compiled_kernel_hashes.empty()
+#if defined(ORT_ENABLE_ORT_FORMAT_RUNTIME_GRAPH_OPTIMIZATION)
+      && graph_.RuntimeOptimizationReplayCtx().produced_node_index_to_kernel_def_hash.empty()
+#endif  // defined(ORT_ENABLE_ORT_FORMAT_RUNTIME_GRAPH_OPTIMIZATION)
+      ;
+
   // process the nodes that existed when the model was created
   for (FbsSessionStateViewer::Index i = 0, end = fbs_session_state_viewer.GetNumNodeKernelInfos(); i < end; ++i) {
     const auto node_kernel_info = fbs_session_state_viewer.GetNodeKernelInfo(i);
 
     Node* const node = graph_.GetNode(node_kernel_info.node_index);
     if (node == nullptr) {
-      // this is OK if we have compiled kernels and the original node was replaced. if not the model is invalid.
-      ORT_RETURN_IF(compiled_kernel_hashes.empty(),
+      // this is OK if we have compiled kernels/replayed runtime optimizations and the original node was replaced.
+      // if not the model is invalid.
+      ORT_RETURN_IF(original_nodes_should_exist,
                     "Can't find node with index ", node_kernel_info.node_index, ". Invalid ORT format model.");
       continue;
     }
 
     ORT_RETURN_IF_ERROR(add_kernel_by_hash(*node, node_kernel_info.kernel_def_hash));
   }
 
+#if defined(ORT_ENABLE_ORT_FORMAT_RUNTIME_GRAPH_OPTIMIZATION)
+  // process the nodes that were added by replaying any loaded runtime optimizations
+  for (const auto& [node_index, kernel_def_hash] :
+       graph_.RuntimeOptimizationReplayCtx().produced_node_index_to_kernel_def_hash) {
+    const auto* node = graph_.GetNode(node_index);
+    ORT_RETURN_IF(node == nullptr,
+                  "Can't find runtime optimization produced node with index ", node_index);
+
+    ORT_RETURN_IF_ERROR(add_kernel_by_hash(*node, kernel_def_hash));
+  }
+#endif  // defined(ORT_ENABLE_ORT_FORMAT_RUNTIME_GRAPH_OPTIMIZATION)
+
   // lookup the hashes for any nodes we compiled. the nodes indexes for compiled nodes are not in node_indices
   // as they were created at runtime.
   if (!compiled_kernel_hashes.empty()) {
@@ -1053,7 +1073,7 @@ static void ComputeConstantInitializerUseCount(const Graph& graph, std::unordere
 }
 
 Status SessionState::FinalizeSessionState(const std::basic_string<PATH_CHAR_TYPE>& graph_location,
-                                          KernelRegistryManager& kernel_registry_manager,
+                                          const KernelRegistryManager& kernel_registry_manager,
                                           const SessionOptions& session_options,
                                           const onnxruntime::fbs::SessionState* serialized_session_state,
                                           bool remove_initializers,
@@ -1210,7 +1230,7 @@ static void AccumulateAllNestedSubgraphsInfo(
 }
 
 Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_TYPE>& graph_location,
-                                              KernelRegistryManager& kernel_registry_manager,
+                                              const KernelRegistryManager& kernel_registry_manager,
                                               _In_opt_ const Node* parent_node,
                                               const SessionOptions& session_options,
                                               bool remove_initializers,
@@ -1221,6 +1241,27 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
     CreateGraphInfo();
   }
 
+#if defined(ORT_ENABLE_ORT_FORMAT_RUNTIME_GRAPH_OPTIMIZATION) && defined(ORT_MINIMAL_BUILD)
+  // remove any unused initializers
+  // not needed in a full build because unused initializers should have been removed earlier by Graph::Resolve()
+  // not needed in a minimal build with runtime optimizations disabled because only runtime optimizations are expected
+  //   to possibly result in unused initializers
+  {
+    std::vector<std::string> unused_initializer_names;
+    for (const auto& [name, tensor_proto] : graph_.GetAllInitializedTensors()) {
+      ORT_UNUSED_PARAMETER(tensor_proto);
+      int idx;
+      if (!ort_value_name_idx_map_.GetIdx(name, idx).IsOK()) {
+        unused_initializer_names.push_back(name);
+      }
+    }
+
+    for (const auto& name : unused_initializer_names) {
+      graph_.RemoveInitializedTensor(name);
+    }
+  }
+#endif  // defined(ORT_ENABLE_ORT_FORMAT_RUNTIME_GRAPH_OPTIMIZATION) && defined(ORT_MINIMAL_BUILD)
+
   // ignore any outer scope args we don't know about. this can happen if a node contains multiple subgraphs.
   std::vector<const NodeArg*> valid_outer_scope_node_args;
   if (parent_node) {

diff --git a/onnxruntime/core/framework/session_state.h b/onnxruntime/core/framework/session_state.h
@@ -303,7 +303,7 @@ class SessionState {
                            const KernelRegistryManager& kernel_registry_manager);
 
   Status FinalizeSessionState(const std::basic_string<PATH_CHAR_TYPE>& graph_loc,
-                              KernelRegistryManager& kernel_registry_manager,
+                              const KernelRegistryManager& kernel_registry_manager,
                               const SessionOptions& session_options = {},
                               const onnxruntime::fbs::SessionState* serialized_session_state = nullptr,
                               bool remove_initializers = true,
@@ -373,7 +373,7 @@ class SessionState {
 #endif
 
   Status FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_TYPE>& graph_loc,
-                                  KernelRegistryManager& kernel_registry_manager,
+                                  const KernelRegistryManager& kernel_registry_manager,
                                   _In_opt_ const Node* parent_node,
                                   const SessionOptions& session_options,
                                   bool remove_initializers,

diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc
@@ -1,10 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#ifdef _WIN32
-// disable some warnings from protobuf to pass Windows build
-#pragma warning(disable : 4244)
-#endif
+#include "core/graph/graph.h"
 
 #include <cassert>
 #include <fstream>
@@ -2858,7 +2855,7 @@ static void RemoveRepeatedFieldEntry(T& repeated_field, const TIter& entry_to_re
     // we do this so we don't have to move all the entries past the one being deleted down one.
     auto slot = entry_to_remove - repeated_field.begin();
     auto last_entry = repeated_field.end() - 1;
-    repeated_field.SwapElements(slot, num_entries - 1);
+    repeated_field.SwapElements(gsl::narrow<int>(slot), gsl::narrow<int>(num_entries - 1));
     repeated_field.erase(last_entry);
   } else {
     repeated_field.erase(entry_to_remove);

diff --git a/onnxruntime/core/graph/graph_viewer.cc b/onnxruntime/core/graph/graph_viewer.cc
@@ -1,11 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#ifdef _WIN32
-// disable some warnings from protobuf to pass Windows build
-#pragma warning(disable : 4244)
-#endif
-
 #include "core/graph/graph_viewer.h"
 #include "core/graph/indexed_sub_graph.h"
 

diff --git a/onnxruntime/core/graph/runtime_optimization_record.h b/onnxruntime/core/graph/runtime_optimization_record.h
@@ -6,6 +6,7 @@
 #include <limits>
 #include <vector>
 #include <string>
+#include <tuple>  // for std::tie
 
 #include "core/graph/basic_types.h"
 
@@ -60,6 +61,18 @@ an ORT format model. This also means that non-empty node indices here must be in
   int num_variadic_inputs;
   /** The number of variadic output values of the target node. */
   int num_variadic_outputs;
+
+  friend bool operator==(const NodesToOptimizeIndices& a, const NodesToOptimizeIndices& b) {
+    const auto tied = [](const NodesToOptimizeIndices& n) {
+      return std::tie(n.nodes, n.num_inputs, n.num_outputs, n.variadic_input, n.variadic_output,
+                      n.num_variadic_inputs, n.num_variadic_outputs);
+    };
+    return tied(a) == tied(b);
+  }
+
+  friend bool operator!=(const NodesToOptimizeIndices& a, const NodesToOptimizeIndices& b) {
+    return !(a == b);
+  }
 };
 
 struct NodeIndexAndKernelDefHash {