Sp-ization Algorithm

lib/utils/include/utils/containers.h

@@ -15,6 +15,7 @@
 #include <iostream>
 #include <map>
 #include <numeric>
+#include <set>
 #include <sstream>
 #include <string>
 #include <unordered_map>
@@ -292,6 +293,17 @@ bidict<K, V> merge_maps(bidict<K, V> const &lhs, bidict<K, V> const &rhs) {
   return result;
 }
 
+template <typename C>
+auto invert_map(C const &m) {
+  std::unordered_map<typename C::mapped_type,
+                     std::unordered_set<typename C::key_type>>
+      m_inv;
+  for (auto const &[key, value] : m) {
+    m_inv[value].insert(key);
+  }
+  return m_inv;
+}
+
 template <typename F, typename C, typename K, typename V>
 std::unordered_map<K, V> generate_map(C const &c, F const &f) {
   static_assert(is_hashable<K>::value,
@@ -492,6 +504,13 @@ std::vector<E> as_vector(C const &c) {
   return result;
 }
 
+template <typename C>
+auto as_set(C const &c) {
+  using E = typename C::value_type;
+  std::set<E> result(c.cbegin(), c.cend());
+  return result;
+}
+
 template <typename F, typename In, typename Out>
 std::vector<Out> transform(std::vector<In> const &v, F const &f) {
   std::vector<Out> result;
@@ -553,6 +572,18 @@ std::vector<Out> flatmap(std::vector<In> const &v, F const &f) {
   return result;
 }
 
+template <typename C>
+auto pairs(C const &c) {
+  using E = typename C::value_type;
+  std::vector<std::pair<E, E>> pairs;
+  for (auto it1 = c.begin(); it1 != c.end(); ++it1) {
+    for (auto it2 = std::next(it1); it2 != c.end(); ++it2) {
+      pairs.emplace_back(*it1, *it2);
+    }
+  }
+  return pairs;
+}
+
 template <typename C, typename Enable>
 struct get_element_type {
   using type = typename C::value_type;
@@ -585,6 +616,20 @@ std::unordered_set<Out> flatmap_v2(std::unordered_set<In> const &v,
   return result;
 }
 
+template <typename C>
+void inplace_sorted(C &c) {
+  CHECK_SUPPORTS_ITERATOR_TAG(std::random_access_iterator_tag, C);
+  std::sort(c.begin(), c.end());
+}
+
+template <typename C>
+auto sorted(C const &c) {
+  using Elem = typename C::value_type;
+  std::vector<Elem> result(c.begin(), c.end());
+  inplace_sorted(result);
+  return result;
+}
+
 template <typename C, typename F, typename Elem>
 void inplace_sorted_by(C &c, F const &f) {
   CHECK_SUPPORTS_ITERATOR_TAG(std::random_access_iterator_tag, C);

lib/utils/include/utils/graph/algorithms.h

@@ -160,6 +160,10 @@ std::unordered_set<Node> get_predecessors(DiGraphView const &, Node const &);
 std::unordered_map<Node, std::unordered_set<Node>>
     get_predecessors(DiGraphView const &, std::unordered_set<Node> const &);
 
+std::unordered_set<Node> get_successors(DiGraphView const &, Node const &);
+std::unordered_map<Node, std::unordered_set<Node>>
+    get_successors(DiGraphView const &, std::unordered_set<Node> const &);
+
 Node get_src_node(MultiDiEdge const &);
 Node get_dst_node(MultiDiEdge const &);
 Node get_dst_node(InputMultiDiEdge const &);
@@ -268,6 +272,13 @@ std::vector<Node>
 std::vector<Node> get_topological_ordering(DiGraphView const &);
 std::vector<Node> get_unchecked_topological_ordering(DiGraphView const &);
 
+std::vector<Node>
+    get_topological_ordering_from_starting_node(DiGraphView const &,
+                                                Node const &);
+std::vector<Node>
+    get_unchecked_topological_ordering_from_starting_node(DiGraphView const &,
+                                                          Node const &);
+
 std::vector<DirectedEdge> get_edge_topological_ordering(DiGraphView const &);
 std::vector<MultiDiEdge>
     get_edge_topological_ordering(MultiDiGraphView const &);
@@ -313,6 +324,12 @@ std::unordered_map<Node, int> calculate_topo_rank(DiGraphView const &);
 Node get_node_with_greatest_topo_rank(std::unordered_set<Node> const &,
                                       DiGraphView const &);
 
+std::unordered_map<Node, int>
+    get_longest_path_lengths_from_source_node(DiGraphView const &g);
+std::unordered_map<Node, float>
+    get_weighted_longest_path_lengths_from_source_node(
+        DiGraphView const &g, std::unordered_map<Node, float> const &cost_map);
+
 MultiDiGraphView join(MultiDiGraphView const &lhs, MultiDiGraphView const &rhs);
 DiGraphView join(DiGraphView const &lhs, DiGraphView const &rhs);
 UndirectedGraphView join(UndirectedGraphView const &lhs,

lib/utils/include/utils/graph/serialparallel.h

@@ -8,6 +8,8 @@
 #include <vector>
 
 namespace FlexFlow {
+bool has_single_source(DiGraphView const &g);
+bool has_single_sink(DiGraphView const &g);
 
 Node find_source_node(DiGraphView const &);
 Node find_sink_node(DiGraphView const &);
@@ -22,6 +24,8 @@ struct Serial {
 
 struct Parallel {
   std::vector<std::variant<Serial, Node>> children;
+  bool operator==(Parallel const &other) const;
+  bool operator!=(Parallel const &other) const;
 };
 
 FF_VISITABLE_STRUCT_NONSTANDARD_CONSTRUCTION(Parallel, children);
@@ -40,14 +44,77 @@ std::unordered_map<Node, Node> parallel_extend(MultiDiGraph &g,
 std::unordered_map<Node, Node> serial_extend(MultiDiGraph &g,
                                              MultiDiGraph const &ext);
 
+std::unordered_map<Node, Node> parallel_extend(DiGraph &g,
+                                               DiGraphView const &ext);
+
+std::unordered_map<Node, Node> serial_extend(DiGraph &g,
+                                             DiGraphView const &ext);
+
 MultiDiGraph serial_composition(MultiDiGraph const &g1, MultiDiGraph const &g2);
 
 MultiDiGraph parallel_composition(MultiDiGraph const &g1,
                                   MultiDiGraph const &g2);
 
+SerialParallelDecomposition parallel_composition(
+    std::vector<SerialParallelDecomposition> const &sp_compositions);
+
+SerialParallelDecomposition serial_composition(
+    std::vector<SerialParallelDecomposition> const &sp_compositions);
+
 MultiDiGraph multidigraph_from_sp_decomposition(
     SerialParallelDecomposition const &sp_decomposition);
 
+bool isempty(SerialParallelDecomposition const &sp);
+
+SerialParallelDecomposition normalize(SerialParallelDecomposition const &sp);
+
+std::unordered_map<Node, size_t>
+    node_counter(SerialParallelDecomposition const &sp);
+
+size_t node_count(SerialParallelDecomposition const &sp);
+
+std::vector<SerialParallelDecomposition>
+    to_sp_decomp(std::vector<std::variant<Parallel, Node>> const &children);
+
+std::vector<SerialParallelDecomposition>
+    to_sp_decomp(std::vector<std::variant<Serial, Node>> const &children);
+
+SerialParallelDecomposition
+    to_sp_decomp(std::variant<Parallel, Node> const &child);
+
+SerialParallelDecomposition
+    to_sp_decomp(std::variant<Serial, Node> const &child);
+
+size_t node_count(DiGraphView const &g);
+
+float compute_cost(SerialParallelDecomposition const &sp,
+                   std::unordered_map<Node, float> cost_map);
+
+float compute_cost(DiGraphView const &g,
+                   std::unordered_map<Node, float> const &cost_map);
+
+float critical_path_cost(DiGraphView const &g,
+                         std::unordered_map<Node, float> const &cost_map);
+
+float critical_path_cost(SerialParallelDecomposition const &sp,
+                         std::unordered_map<Node, float> const &cost_map);
+
+float average_parallelism_degree(
+    SerialParallelDecomposition const &sp,
+    std::unordered_map<Node, float> const &cost_map);
+
+float max_parallelism_degree(SerialParallelDecomposition const &sp,
+                             std::unordered_map<Node, float> const &cost_map);
+
+float relative_cost_increase(DiGraphView const &g,
+                             SerialParallelDecomposition const &sp,
+                             std::unordered_map<Node, float> const &cost_map);
+
+float relative_critical_path_cost_increase(
+    DiGraphView const &g,
+    SerialParallelDecomposition const &sp,
+    std::unordered_map<Node, float> const &cost_map);
+
 } // namespace FlexFlow
 
 #endif

lib/utils/include/utils/graph/sp_ization.h

@@ -0,0 +1,21 @@
+#ifndef _FLEXFLOW_UTILS_GRAPH_barrier_sync_sp_ization_H
+#define _FLEXFLOW_UTILS_GRAPH_barrier_sync_sp_ization_H
+
+#include "serialparallel.h"
+
+using namespace FlexFlow;
+
+namespace FlexFlow {
+
+SerialParallelDecomposition barrier_sync_sp_ization(DiGraphView const &g);
+SerialParallelDecomposition
+    dependency_invariant_sp_ization(DiGraphView const &g);
+SerialParallelDecomposition
+    dependency_invariant_sp_ization_with_coalescing(DiGraphView const &g);
+
+SerialParallelDecomposition cost_aware_barrier_sync_sp_ization(
+    DiGraphView const &g, std::unordered_map<Node, float> const &cost_map);
+
+} // namespace FlexFlow
+
+#endif

lib/utils/src/graph/algorithms.cc

@@ -376,6 +376,24 @@
   return get_predecessors(g, std::unordered_set<Node>{n}).at(n);
 }
 
+std::unordered_map<Node, std::unordered_set<Node>>
+    get_successors(DiGraphView const &g,
+                   std::unordered_set<Node> const &nodes) {
+
+  std::unordered_map<Node, std::unordered_set<Node>> successors;
+  for (Node const &n : nodes) {
+    successors[n];
+  }
+  for (DirectedEdge const &e : get_outgoing_edges(g, nodes)) {
+    successors.at(e.src).insert(e.dst);
+  }
+  return successors;
+}
+
+std::unordered_set<Node> get_successors(DiGraphView const &g, Node const &n) {
+  return get_successors(g, std::unordered_set<Node>{n}).at(n);
+}
+
 std::vector<Node> get_unchecked_dfs_ordering(
     DiGraphView const &g, std::unordered_set<Node> const &starting_points) {
   UncheckedDFSView dfs_view = unchecked_dfs(g, starting_points);
@@ -474,6 +492,37 @@
   return get_unchecked_topological_ordering(g);
 }
 
+std::vector<Node> get_unchecked_topological_ordering_from_starting_node(
+    DiGraphView const &g, Node const &starting_node) {
+
+  auto get_descendants = [&](DiGraphView const &g, Node const &starting_node) {
+    std::unordered_set<Node> descendants;
+    std::stack<Node> to_visit;
+    to_visit.push(starting_node);
+    while (!to_visit.empty()) {
+      Node current = to_visit.top();
+      to_visit.pop();
+      descendants.insert(current);
+      for (auto const &s :
+           filter(get_successors(g, current),
+                  [&](Node const &n) { return !contains(descendants, n); })) {
+        to_visit.push(s);
+      }
+    }
+    return descendants;
+  };
+  std::unordered_set<Node> descendants = get_descendants(g, starting_node);
+  return get_topological_ordering(get_subgraph(g, descendants));
+}
+
+std::vector<Node>
+    get_topological_ordering_from_starting_node(DiGraphView const &g,
+                                                Node const &starting_nodes) {
+  assert(is_acyclic(g));
+  return get_unchecked_topological_ordering_from_starting_node(g,
+                                                               starting_nodes);
+}
+
 std::vector<DirectedEdge> get_edge_topological_ordering(DiGraphView const &g) {
   std::vector<DirectedEdge> result;
   for (Node const &n : get_topological_ordering(g)) {
@@ -639,6 +688,62 @@
                            });
 }
 
+std::unordered_map<Node, int>
+    get_unchecked_longest_path_lengths_from_source_node(DiGraphView const &g) {
+  std::vector<Node> topo_order = get_topological_ordering(g);
+  std::unordered_map<Node, int> longest_path_lengths;
+
+  for (Node const &n : topo_order) {
+    std::unordered_set<int> predecessor_path_lengths =
+        transform(get_predecessors(g, n), [&](Node const &pred) {
+          return longest_path_lengths.at(pred);
+        });
+    longest_path_lengths[n] = (predecessor_path_lengths.size() == 0)
+                                  ? 0
+                                  : maximum(predecessor_path_lengths) + 1;
+  }
+
+  return longest_path_lengths;
+}
+
+std::unordered_map<Node, float>
+    get_unchecked_weighted_longest_path_lengths_from_source_node(
+        DiGraphView const &g,
+        std::unordered_map<Node, float> const &node_costs) {
+
+  std::vector<Node> topo_order = get_topological_ordering(g);
+  std::unordered_map<Node, float> longest_path_lengths;
+
+  for (Node const &n : topo_order) {
+    std::unordered_set<float> predecessor_path_lengths =
+        transform(get_predecessors(g, n), [&](Node const &pred) {
+          return longest_path_lengths.at(pred);
+        });
+    longest_path_lengths[n] =
+        (predecessor_path_lengths.size() == 0)
+            ? node_costs.at(n)
+            : maximum(predecessor_path_lengths) + node_costs.at(n);
+  }
+
+  return longest_path_lengths;
+}
+
+std::unordered_map<Node, float>
+    get_weighted_longest_path_lengths_from_source_node(
+        DiGraphView const &g,
+        std::unordered_map<Node, float> const &node_costs) {
+
+  assert(is_acyclic(g));
+  return get_unchecked_weighted_longest_path_lengths_from_source_node(
+      g, node_costs);
+}
+
+std::unordered_map<Node, int>
+    get_longest_path_lengths_from_source_node(DiGraphView const &g) {
+  assert(is_acyclic(g));
+  return get_unchecked_longest_path_lengths_from_source_node(g);
+}
+
 std::optional<Node>
     get_imm_post_dominator(DiGraphView const &g,
                            std::unordered_set<Node> const &nodes) {