Add unit tests for lib/utils/graph

lib/CMakeLists.txt

@@ -1,8 +1,8 @@
-add_subdirectory(pcg)
-add_subdirectory(compiler)
+#add_subdirectory(pcg)
+#add_subdirectory(compiler)
 # add_subdirectory(runtime)
-add_subdirectory(op-attrs)
-add_subdirectory(kernels)
+#add_subdirectory(op-attrs)
+#add_subdirectory(kernels)
 add_subdirectory(utils)
 # add_subdirectory(ffi)
-add_subdirectory(substitutions)
+#add_subdirectory(substitutions)

lib/utils/CMakeLists.txt

@@ -36,4 +36,4 @@ define_ff_vars(${project_target})
 ff_set_cxx_properties(${project_target})
 
 add_subdirectory(ffi)
-# add_subdirectory(test)
+add_subdirectory(test)

lib/utils/include/utils/containers.h

@@ -4,6 +4,7 @@
 #include "bidict.h"
 #include "invoke.h"
 #include "optional.h"
+#include "utils/exception.h"
 #include <algorithm>
 #include <cassert>
 #include <functional>
@@ -137,7 +138,7 @@ template <typename K,
 std::unordered_map<K2, V> map_keys(std::unordered_map<K, V> const &m,
                                    F const &f) {
   std::unordered_map<K2, V> result;
-  for (auto const &kv : f) {
+  for (auto const &kv : m) {
     result.insert({f(kv.first), kv.second});
   }
   return result;
@@ -149,7 +150,7 @@ template <typename K,
           typename K2 = decltype(std::declval<F>()(std::declval<K>()))>
 bidict<K2, V> map_keys(bidict<K, V> const &m, F const &f) {
   bidict<K2, V> result;
-  for (auto const &kv : f) {
+  for (auto const &kv : m) {
     result.equate(f(kv.first), kv.second);
   }
   return result;
@@ -159,14 +160,25 @@ template <typename K, typename V, typename F>
 std::unordered_map<K, V> filter_keys(std::unordered_map<K, V> const &m,
                                      F const &f) {
   std::unordered_map<K, V> result;
-  for (auto const &kv : f) {
+  for (auto const &kv : m) {
     if (f(kv.first)) {
       result.insert(kv);
     }
   }
   return result;
 }
 
+template <typename K, typename V, typename F>
+std::unordered_map<K, V> filter_values(bidict<K, V> const &m, F const &f) {
+  std::unordered_map<K, V> result;
+  for (auto const &kv : m) {
+    if (f(kv.second)) {
+      result.insert(kv);
+    }
+  }
+  return result;
+}
+
 template <typename K,
           typename V,
           typename F,
@@ -255,6 +267,20 @@ std::unordered_set<T> intersection(std::unordered_set<T> const &l,
   return result;
 }
 
+template <typename T>
+std::unordered_set<T>
+    intersection(std::vector<std::unordered_set<T>> const &v) {
+  if (v.empty()) {
+    throw mk_runtime_error("cannot take intersection of no sets");
+  }
+
+  std::unordered_set<T> result = v[0];
+  for (int i = 1; i < v.size(); i++) {
+    result = intersection(result, v[i]);
+  }
+  return result;
+}
+
 template <typename T>
 bool are_disjoint(std::unordered_set<T> const &l,
                   std::unordered_set<T> const &r) {
@@ -311,12 +337,12 @@ std::function<V(K const &)> lookup_in(std::unordered_map<K, V> const &m) {
 
 template <typename L, typename R>
 std::function<R(L const &)> lookup_in_l(bidict<L, R> const &m) {
-  return [&m](L const &l) -> L { return m.at_l(l); };
+  return [&m](L const &l) -> R { return m.at_l(l); };
 }
 
 template <typename L, typename R>
 std::function<L(R const &)> lookup_in_r(bidict<L, R> const &m) {
-  return [&m](R const &r) -> R { return m.at_r(r); };
+  return [&m](R const &r) -> L { return m.at_r(r); };
 }
 
 template <typename T>
@@ -509,13 +535,6 @@ std::unordered_set<Out> flatmap_v2(std::unordered_set<In> const &v,
   return result;
 }
 
-template <typename T, typename F>
-std::vector<T> sorted_by(std::unordered_set<T> const &s, F const &f) {
-  std::vector<T> result(s.begin(), s.end());
-  inplace_sorted_by(s, f);
-  return result;
-}
-
 template <typename T, typename F>
 void inplace_sorted_by(std::vector<T> &v, F const &f) {
   auto custom_comparator = [&](T const &lhs, T const &rhs) -> bool {
@@ -524,6 +543,13 @@ void inplace_sorted_by(std::vector<T> &v, F const &f) {
   std::sort(v.begin(), v.end(), custom_comparator);
 }
 
+template <typename T, typename F>
+std::vector<T> sorted_by(std::unordered_set<T> const &s, F const &f) {
+  std::vector<T> result(s.begin(), s.end());
+  inplace_sorted_by(result, f);
+  return result;
+}
+
 template <typename C, typename F>
 C filter(C const &v, F const &f) {
   C result(v);
@@ -559,27 +585,27 @@ std::pair<std::vector<T>, std::vector<T>> vector_split(std::vector<T> const &v,
 
 template <typename C>
 typename C::value_type maximum(C const &v) {
-  return std::max_element(v.begin(), v.end());
+  return *std::max_element(v.begin(), v.end());
 }
 
 template <typename T>
 T reversed(T const &t) {
   T r;
-  for (auto i = t.cend() - 1; i >= t.begin(); i++) {
+  for (auto i = t.cend() - 1; i >= t.begin(); i--) {
     r.push_back(*i);
   }
   return r;
 }
 
 template <typename T>
 std::vector<T> value_all(std::vector<optional<T>> const &v) {
-  std::vector<T> result;
-
-  for (auto const &element : v) {
-    result.push_back(element.value());
-  }
-
-  return result;
+  return transform(v, [](optional<T> const &t) {
+    if (t == nullopt) {
+      throw mk_runtime_error("value is nullopt");
+    } else {
+      return t.value();
+    }
+  });
 }
 
 template <typename T>
@@ -602,7 +628,7 @@ std::vector<T> subvec(std::vector<T> const &v,
     begin_iter += resolve_loc(maybe_start.value());
   }
   if (maybe_end.has_value()) {
-    end_iter = v.cbegin() + resolve_loc(maybe_start.value());
+    end_iter = v.cbegin() + resolve_loc(maybe_end.value());
   }
 
   std::vector<T> output(begin_iter, end_iter);
@@ -688,4 +714,4 @@ reversed_container_t<C> reversed_container(C const &c) {
 
 } // namespace FlexFlow
 
-#endif
+#endif

lib/utils/include/utils/fmt.h

@@ -14,6 +14,12 @@ template <typename T>
 struct is_fmtable<T, void_t<decltype(fmt::to_string(std::declval<T>()))>>
     : std::true_type {};
 
+// template <typename T>
+// typename std::enable_if<is_fmtable<T>::value, std::ostream &>::type
+// operator<<(std::ostream &s, T const &t) {
+//     return s << fmt::to_string(t);
+// }
+
 } // namespace FlexFlow
 
 #endif

lib/utils/include/utils/graph/adjacency_digraph.h

@@ -9,6 +9,7 @@ namespace FlexFlow {
 
 class AdjacencyDiGraph : public IDiGraph {
 public:
+  AdjacencyDiGraph() = default;
   Node add_node() override;
   void add_node_unsafe(Node const &) override;
   void remove_node_unsafe(Node const &) override;
@@ -28,8 +29,8 @@ class AdjacencyDiGraph : public IDiGraph {
 private:
   using ContentsType = std::unordered_map<Node, std::unordered_set<Node>>;
 
-  AdjacencyDiGraph(std::size_t, ContentsType);
-
+  AdjacencyDiGraph(std::size_t next_node_idx, ContentsType adjacency)
+      : next_node_idx(next_node_idx), adjacency(adjacency) {}
   std::size_t next_node_idx = 0;
   ContentsType adjacency;
 };

lib/utils/include/utils/graph/adjacency_multidigraph.h

@@ -9,16 +9,20 @@ namespace FlexFlow {
 
 class AdjacencyMultiDiGraph : public IMultiDiGraph {
 public:
+  AdjacencyMultiDiGraph() = default;
   Node add_node() override;
   void add_node_unsafe(Node const &) override;
+  NodePort add_node_port() override;
+  void add_node_port_unsafe(NodePort const &) override;
   void remove_node_unsafe(Node const &) override;
   void add_edge(Edge const &) override;
   void remove_edge(Edge const &) override;
   std::unordered_set<Edge> query_edges(EdgeQuery const &) const override;
   std::unordered_set<Node> query_nodes(NodeQuery const &) const override;
 
   AdjacencyMultiDiGraph *clone() const override {
-    return new AdjacencyMultiDiGraph(this->next_node_idx, this->adjacency);
+    return new AdjacencyMultiDiGraph(
+        this->next_node_idx, this->next_node_port, this->adjacency);
   }
 
 private:
@@ -28,9 +32,12 @@ class AdjacencyMultiDiGraph : public IMultiDiGraph {
           Node,
           std::unordered_map<NodePort, std::unordered_set<NodePort>>>>;
 
-  AdjacencyMultiDiGraph(std::size_t, ContentsType const &);
+  AdjacencyMultiDiGraph(std::size_t next_node_idx,
+                        std::size_t next_node_port,
+                        ContentsType const &adjacency)
+      : next_node_idx(next_node_idx), next_node_port(next_node_port),
+        adjacency(adjacency) {}
 
-private:
   std::size_t next_node_idx = 0;
   std::size_t next_node_port = 0;
   ContentsType adjacency;

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

@@ -15,8 +15,7 @@
 #include <vector>
 
 namespace FlexFlow {
-
-std::vector<Node> add_nodes(Graph &, int);
+std::vector<Node> add_nodes(DiGraph &, int);
 std::unordered_set<Node> get_nodes(GraphView const &);
 std::unordered_set<NodePort> get_node_ports(MultiDiGraphView const &);
 
@@ -144,6 +143,10 @@ std::unordered_map<Node, std::unordered_set<Node>>
 std::unordered_map<Node, std::unordered_set<Node>>
     get_predecessors(DiGraphView const &, std::unordered_set<Node> const &);
 
+std::unordered_set<Node> get_neighbors(UndirectedGraphView const &,
+                                       Node const &);
+std::unordered_set<Node> get_neighbors(DiGraphView const &, Node const &);
+std::unordered_set<Node> get_neighbors(MultiDiGraphView const &, Node const &);
 std::unordered_set<Node> get_sources(DiGraphView const &);
 std::unordered_set<Node> get_sources(MultiDiGraphView const &);
 
@@ -230,6 +233,10 @@ MultiDiGraphView get_subgraph(MultiDiGraphView const &,
 OpenMultiDiGraphView get_subgraph(OpenMultiDiGraphView const &,
                                   std::unordered_set<Node> const &);
 
+std::unordered_map<Node, int> calculate_topo_rank(DiGraphView const &);
+Node get_node_with_greatest_topo_rank(std::unordered_set<Node> const &,
+                                      DiGraphView const &);
+
 MultiDiGraphView join(MultiDiGraphView const &lhs, MultiDiGraphView const &rhs);
 DiGraphView join(DiGraphView const &lhs, DiGraphView const &rhs);
 UndirectedGraphView join(UndirectedGraphView const &lhs,
@@ -241,6 +248,7 @@ DiGraphView with_added_edges(DiGraphView const &,
                              std::unordered_set<DirectedEdge> const &);
 
 UndirectedGraphView as_undirected(DiGraphView const &);
+UndirectedGraphView as_undirected(MultiDiGraphView const &);
 MultiDiGraphView as_multidigraph(DiGraphView const &);
 DiGraphView as_digraph(MultiDiGraphView const &);
 MultiDiGraphView as_multidigraph(OpenMultiDiGraphView const &);

lib/utils/include/utils/graph/digraph.h

@@ -14,15 +14,16 @@ struct DirectedEdge {
   Node src;
   Node dst;
 };
+
+std::ostream &operator<<(std::ostream &s, DirectedEdge const &e);
+
 FF_VISITABLE_STRUCT(DirectedEdge, src, dst);
 
 struct DirectedEdgeQuery {
   query_set<Node> srcs;
   query_set<Node> dsts;
 
-  static DirectedEdgeQuery all() {
-    NOT_IMPLEMENTED();
-  }
+  static DirectedEdgeQuery all();
 };
 FF_VISITABLE_STRUCT(DirectedEdgeQuery, srcs, dsts);
 
@@ -38,7 +39,7 @@ struct IDiGraphView : public IGraphView {
   IDiGraphView &operator=(IDiGraphView const &) = delete;
 
   virtual std::unordered_set<Edge> query_edges(EdgeQuery const &) const = 0;
-  virtual ~IDiGraphView();
+  virtual ~IDiGraphView() = default;
 
 protected:
   IDiGraphView() = default;
@@ -56,15 +57,9 @@ struct DiGraphView {
 
   friend void swap(DiGraphView &, DiGraphView &);
 
-  bool operator==(DiGraphView const &) const;
-  bool operator!=(DiGraphView const &) const;
-
   std::unordered_set<Node> query_nodes(NodeQuery const &) const;
   std::unordered_set<Edge> query_edges(EdgeQuery const &) const;
-
-  IDiGraphView const *unsafe() const {
-    return this->ptr.get();
-  }
+  friend bool is_ptr_equal(DiGraphView const &, DiGraphView const &);
 
   template <typename T, typename... Args>
   static typename std::enable_if<std::is_base_of<IDiGraphView, T>::value,
@@ -73,18 +68,19 @@ struct DiGraphView {
     return DiGraphView(std::make_shared<T>(std::forward<Args>(args)...));
   }
 
-private:
-  DiGraphView(std::shared_ptr<IDiGraphView const>);
+  static DiGraphView
+      unsafe_create_without_ownership(IDiGraphView const &graphView);
 
-  friend DiGraphView unsafe(IDiGraphView const &);
+  DiGraphView(std::shared_ptr<IDiGraphView const> const &ptr,
+              should_only_be_used_internally_tag_t const &tag)
+      : DiGraphView(ptr) {}
 
 private:
+  DiGraphView(std::shared_ptr<IDiGraphView const> ptr) : ptr(ptr) {}
   std::shared_ptr<IDiGraphView const> ptr;
 };
 CHECK_WELL_BEHAVED_VALUE_TYPE_NO_EQ(DiGraphView);
 
-DiGraphView unsafe(IDiGraphView const &);
-
 struct IDiGraph : public IDiGraphView, public IGraph {
   virtual void add_edge(Edge const &) = 0;
   virtual void remove_edge(Edge const &) = 0;
@@ -101,7 +97,15 @@ struct DiGraph {
   DiGraph(DiGraph const &) = default;
   DiGraph &operator=(DiGraph const &) = default;
 
-  operator DiGraphView() const;
+  operator DiGraphView() const {
+    return DiGraphView(this->ptr.get(), should_only_be_used_internally_tag_t{});
+  }
+
+  operator Graph() const {
+    return Graph(std::static_pointer_cast<IGraph>(this->ptr.get_mutable()),
+                 should_only_be_used_internally_tag_t{});
+  } // Note(lambda):because ptr is cow_ptr_t<IDiGraph>, but ptr in Graph is
+    // cow_ptr_t<IGraph>, so we need to static_pointer_cast
 
   friend void swap(DiGraph &, DiGraph &);
 
@@ -123,9 +127,7 @@ struct DiGraph {
   }
 
 private:
-  DiGraph(std::unique_ptr<IDiGraph>);
-
-private:
+  DiGraph(std::unique_ptr<IDiGraph> ptr);
   cow_ptr_t<IDiGraph> ptr;
 };
 CHECK_WELL_BEHAVED_VALUE_TYPE_NO_EQ(DiGraph);