Introduce memory efficient topological sort (microsoft#20258)

### Introduce memory efficient topo sort (for training)

~~and laze initialize Priority-Based and Memory-Efficient topo sort.
Because in most cases, they are not needed, so we free the overheads of
GraphViewer construction for most use cases.~~

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

include/onnxruntime/core/graph/constants.h

@@ -56,7 +56,4 @@ constexpr const char* kAzureExecutionProvider = "AzureExecutionProvider";
 constexpr const char* kExecutionProviderSharedLibraryPath = "shared_lib_path";
 constexpr const char* kExecutionProviderSharedLibraryEntry = "provider_factory_entry_point";
 
-// For Priority based graph topology sorting.
-constexpr const char* kBackwardNodeAttributeName = "__backwardpass";
-
 }  // namespace onnxruntime

include/onnxruntime/core/graph/graph.h

@@ -1109,6 +1109,19 @@ class Graph {  // NOLINT(clang-analyzer-optin.performance.Padding): preserve exi
 
 #endif
 
+#ifdef ENABLE_TRAINING
+  /**
+   * @brief Performs topological sort with customized Kahn's algorithm on the graph/s.
+   *  This is a specialized version for training where need memory efficient topological sort.
+   * @param yield_op The YieldOp used in ORTModule training.
+   * @param shape_size_parents The shape size parents nodes.
+   * @param node_orders The output node orders.
+   */
+  void MemoryEfficientTopologicalSort(const Node* yield_op,
+                                      const InlinedHashMap<NodeIndex, InlinedVector<NodeIndex>>& shape_size_parents,
+                                      std::vector<NodeIndex>& node_orders) const;
+#endif
+
   /** Gets the map of operator domains to their opset versions. */
   const std::unordered_map<std::string, int>& DomainToVersionMap() const noexcept {
     return domain_to_version_;

include/onnxruntime/core/graph/graph_viewer.h

@@ -207,6 +207,11 @@ class GraphViewer {
   std::vector<NodeIndex> nodes_in_topological_order_with_priority_;
 #endif
 
+#ifdef ENABLE_TRAINING
+  // The NodeIndex values of the graph nodes sorted in memory efficient topological order.
+  std::vector<NodeIndex> nodes_in_mem_efficient_topological_order_;
+#endif
+
   // Graph root nodes.
   std::vector<NodeIndex> root_nodes_;
 

onnxruntime/core/framework/session_options.h

@@ -22,8 +22,9 @@
 namespace onnxruntime {
 
 enum class ExecutionOrder {
-  DEFAULT = 0,        // default topological sort
-  PRIORITY_BASED = 1  // priority-based topological sort
+  DEFAULT = 0,           // default topological sort
+  PRIORITY_BASED = 1,    // priority-based topological sort
+  MEMORY_EFFICIENT = 2,  // memory-efficient topological sort for training purposes.
 };
 
 inline std::ostream& operator<<(std::ostream& os, const ExecutionOrder& order) {
@@ -34,6 +35,9 @@ inline std::ostream& operator<<(std::ostream& os, const ExecutionOrder& order) {
     case ExecutionOrder::PRIORITY_BASED:
       os << "PRIORITY_BASED";
       break;
+    case ExecutionOrder::MEMORY_EFFICIENT:
+      os << "MEMORY_EFFICIENT";
+      break;
     default:
       os << "UNKNOWN";
       break;