Address review comments

flexflow · Aug 31, 2024 · aecbbe6 · aecbbe6
1 parent 1a63f90
commit aecbbe6
Show file tree

Hide file tree

Showing 89 changed files with 2,070 additions and 377 deletions.
diff --git a/flake.nix b/flake.nix
@@ -101,6 +101,7 @@
               tl-expected
               doxygen
               lcov # for code coverage
+              texliveFull
             ])
             (with proj-repo.packages.${system}; [
               proj

diff --git a/lib/op-attrs/include/op-attrs/dim_ordered/enumerate.h b/lib/op-attrs/include/op-attrs/dim_ordered/enumerate.h
@@ -6,12 +6,20 @@
 
 namespace FlexFlow {
 
+/**
+ * @brief Generate a map from indices to elements of \p c.
+ *
+ * @note We return a <tt>std::map</tt> to prevent mixups of \ref ff_dim_t and
+ * \ref legion_dim_t. Note that <tt>std::map</tt> provides ordered iteration in
+ * increasing order, so iterating through the result of this function should
+ * function as expected.
+ */
 template <typename T>
-bidict<ff_dim_t, T> enumerate(FFOrdered<T> const &ff_ordered) {
-  bidict<ff_dim_t, T> result;
+std::map<ff_dim_t, T> enumerate(FFOrdered<T> const &ff_ordered) {
+  std::map<ff_dim_t, T> result;
   for (int raw_ff_dim : count(ff_ordered.size())) {
     ff_dim_t ff_dim = ff_dim_t{raw_ff_dim};
-    result.equate({ff_dim, ff_ordered.at(ff_dim)});
+    result.insert({ff_dim, ff_ordered.at(ff_dim)});
   }
   return result;
 }

diff --git a/lib/op-attrs/include/op-attrs/ops/concat.h b/lib/op-attrs/include/op-attrs/ops/concat.h
@@ -1,5 +1,5 @@
-#ifndef _FLEXFLOW_CONCAT_ATTRS_H
-#define _FLEXFLOW_CONCAT_ATTRS_H
+#ifndef _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_OPS_CONCAT_H
+#define _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_OPS_CONCAT_H
 
 #include "op-attrs/ops/concat_attrs.dtg.h"
 #include "op-attrs/ops/core.h"

diff --git a/lib/op-attrs/test/src/op-attrs/dim_ordered/enumerate.cc b/lib/op-attrs/test/src/op-attrs/dim_ordered/enumerate.cc
@@ -0,0 +1,19 @@
+#include <doctest/doctest.h>
+#include "op-attrs/dim_ordered/enumerate.h"
+
+using namespace ::FlexFlow;
+
+TEST_SUITE(FF_TEST_SUITE) {
+  TEST_CASE("enumerate(FFOrdered<T>)") {
+    FFOrdered<std::string> input = {"zero", "one", "two"};
+
+    std::map<ff_dim_t, std::string> result = enumerate(input);
+    std::map<ff_dim_t, std::string> correct = {
+      {ff_dim_t{0}, "zero"},
+      {ff_dim_t{1}, "one"},
+      {ff_dim_t{2}, "two"},
+    };
+
+    CHECK(result == correct);
+  }
+}
diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_dataflow_graph.h b/lib/pcg/include/pcg/file_format/v1/graphs/v1_dataflow_graph.h
@@ -8,7 +8,7 @@ namespace FlexFlow {
 
 V1DataflowGraph to_v1(DataflowGraphView const &);
 V1DataflowGraph to_v1(DataflowGraphView const &,
-                      std::unordered_map<Node, size_t> const &);
+                      std::unordered_map<Node, int> const &);
 
 } // namespace FlexFlow
 

diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_dataflow_graph.struct.toml b/lib/pcg/include/pcg/file_format/v1/graphs/v1_dataflow_graph.struct.toml
@@ -19,7 +19,7 @@ includes = [
 
 [[fields]]
 name = "nodes"
-type = "std::vector<size_t>"
+type = "std::vector<int>"
 
 [[fields]]
 name = "edges"

diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_graph_edge.struct.toml b/lib/pcg/include/pcg/file_format/v1/graphs/v1_graph_edge.struct.toml
@@ -11,16 +11,16 @@ features = [
 
 [[fields]]
 name = "srcNode"
-type = "size_t"
+type = "int"
 
 [[fields]]
 name = "srcIdx"
-type = "size_t"
+type = "int"
 
 [[fields]]
 name = "dstNode"
-type = "size_t"
+type = "int"
 
 [[fields]]
 name = "dstIdx"
-type = "size_t"
+type = "int"
diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_dataflow_graph.h b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_dataflow_graph.h
@@ -3,27 +3,27 @@
 
 #include "pcg/file_format/v1/graphs/v1_dataflow_graph.h"
 #include "pcg/file_format/v1/graphs/v1_labelled_dataflow_graph.dtg.h"
-#include "utils/containers/enumerate.h"
 #include "utils/containers/map_values.h"
 #include "utils/containers/transform.h"
 #include "utils/graph/dataflow_graph/algorithms.h"
 #include "utils/graph/labelled_dataflow_graph/labelled_dataflow_graph_view.h"
 #include "utils/graph/node/algorithms.h"
+#include "utils/bidict/algorithms/bidict_from_enumerating.h"
 
 namespace FlexFlow {
 
 template <typename NodeLabel, typename OutputLabel>
 V1LabelledDataflowGraph<NodeLabel, OutputLabel>
     to_v1(LabelledDataflowGraphView<NodeLabel, OutputLabel> const &g) {
 
-  bidict<size_t, Node> nodes = enumerate(get_nodes(g));
+  bidict<int, Node> nodes = bidict_from_enumerating(get_nodes(g));
 
   V1DataflowGraph unlabelled = to_v1(g, nodes.reversed());
 
-  std::unordered_map<size_t, NodeLabel> node_labels = map_values(
+  std::unordered_map<int, NodeLabel> node_labels = map_values(
       nodes.as_unordered_map(), [&](Node const &n) { return g.at(n); });
 
-  std::unordered_map<size_t, std::vector<OutputLabel>> output_labels =
+  std::unordered_map<int, std::vector<OutputLabel>> output_labels =
       map_values(nodes.as_unordered_map(), [&](Node const &n) {
         return transform(get_outputs(g, n),
                          [&](DataflowOutput const &o) { return g.at(o); });

diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_dataflow_graph.struct.toml b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_dataflow_graph.struct.toml
@@ -22,11 +22,11 @@ includes = [
 
 [[fields]]
 name = "node_labels"
-type = "std::unordered_map<size_t, NodeLabel>"
+type = "std::unordered_map<int, NodeLabel>"
 
 [[fields]]
 name = "output_labels"
-type = "std::unordered_map<size_t, std::vector<OutputLabel>>"
+type = "std::unordered_map<int, std::vector<OutputLabel>>"
 
 [[fields]]
 name = "graph"

diff --git a/lib/pcg/src/pcg/computation_graph.cc b/lib/pcg/src/pcg/computation_graph.cc
@@ -47,7 +47,7 @@ std::vector<tensor_guid_t> get_outgoing_tensors(ComputationGraph const &cg,
 
 std::vector<tensor_guid_t> get_incoming_tensors(ComputationGraph const &cg,
                                                 layer_guid_t n) {
-  return transform(get_inputs(cg.raw_graph, n.raw_node),
+  return transform(get_input_values(cg.raw_graph, n.raw_node),
                    [](DataflowOutput const &o) { return tensor_guid_t{o}; });
 }
 

diff --git a/lib/pcg/src/pcg/file_format/v1/graphs/v1_dataflow_graph.cc b/lib/pcg/src/pcg/file_format/v1/graphs/v1_dataflow_graph.cc
@@ -5,21 +5,24 @@
 #include "utils/graph/dataflow_graph/algorithms.h"
 #include "utils/graph/node/algorithms.h"
 #include "utils/integer_conversions.h"
+#include "utils/bidict/algorithms/bidict_from_enumerating.h"
 
 namespace FlexFlow {
 
 V1DataflowGraph to_v1(DataflowGraphView const &g) {
-  return to_v1(g, enumerate(get_nodes(g)).reversed());
+  bidict<int, Node> node_enumeration_bidict = bidict_from_enumerating(get_nodes(g));
+  std::unordered_map<Node, int> node_enumeration = node_enumeration_bidict.reversed().as_unordered_map();
+  return to_v1(g, node_enumeration);
 }
 
 V1DataflowGraph to_v1(DataflowGraphView const &g,
-                      std::unordered_map<Node, size_t> const &nodes) {
+                      std::unordered_map<Node, int> const &nodes) {
   std::unordered_set<V1GraphEdge> edges;
   for (DataflowEdge const &e : get_edges(g)) {
     edges.insert(V1GraphEdge{nodes.at(e.src.node),
-                             size_t_from_int(e.src.idx),
+                             e.src.idx,
                              nodes.at(e.dst.node),
-                             size_t_from_int(e.dst.idx)});
+                             e.dst.idx});
   }
 
   return V1DataflowGraph{

diff --git a/lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph.cc b/lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph.cc
@@ -44,7 +44,7 @@ std::vector<parallel_tensor_guid_t>
     get_layer_inputs(ParallelComputationGraph const &pcg,
                      parallel_layer_guid_t const &l) {
   return transform(
-      get_inputs(pcg.raw_graph, l.raw_graph_node),
+      get_input_values(pcg.raw_graph, l.raw_graph_node),
       [](DataflowOutput const &o) { return parallel_tensor_guid_t{o}; });
 }
 

diff --git a/lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph_builder.cc b/lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph_builder.cc
@@ -427,7 +427,7 @@ parallel_tensor_guid_t ParallelComputationGraphBuilder::elu(
     std::optional<std::string> const &maybe_name) {
 
   ElementUnaryAttrs attrs = ElementUnaryAttrs{
-      OperatorType::TANH,
+      OperatorType::ELU,
       std::nullopt,
   };
 

diff --git a/lib/substitutions/include/substitutions/pcg_pattern.h b/lib/substitutions/include/substitutions/pcg_pattern.h
@@ -29,7 +29,7 @@ std::vector<PatternNodeOutput> get_pattern_node_outputs(PCGPattern const &,
 
 bool assignment_satisfies(SubParallelComputationGraph const &,
                           PCGPattern const &,
-                          UnlabelledDataflowGraphPatternMatch const &);
+                          PCGPatternMatch const &);
 
 } // namespace FlexFlow
 

diff --git a/lib/substitutions/include/substitutions/pcg_pattern_match.h b/lib/substitutions/include/substitutions/pcg_pattern_match.h
@@ -6,6 +6,7 @@
 #include "substitutions/pcg_pattern_match.dtg.h"
 #include "substitutions/sub_parallel_computation_graph.dtg.h"
 #include "substitutions/unlabelled/pattern_node_output.dtg.h"
+#include "substitutions/unlabelled/unlabelled_dataflow_graph_pattern_match.dtg.h"
 
 namespace FlexFlow {
 
@@ -15,6 +16,8 @@ bidict<PatternNodeOutput, parallel_tensor_guid_t>
         PCGPattern const &pattern,
         SubParallelComputationGraph const &spcg);
 
+UnlabelledDataflowGraphPatternMatch get_unlabelled_pattern_match(PCGPatternMatch const &);
+
 } // namespace FlexFlow
 
 #endif
diff --git a/lib/substitutions/include/substitutions/sub_parallel_computation_graph.h b/lib/substitutions/include/substitutions/sub_parallel_computation_graph.h
@@ -57,7 +57,7 @@ SubParallelComputationGraphData
     get_sub_pcg_data(SubParallelComputationGraph const &);
 SubParallelComputationGraph
     sub_pcg_from_graph_data(SubParallelComputationGraphData const &);
-bool are_isomorphic(SubParallelComputationGraph const &,
+bool sub_pcgs_are_isomorphic(SubParallelComputationGraph const &,
                     SubParallelComputationGraph const &);
 
 SubParallelComputationGraph

diff --git a/lib/substitutions/include/substitutions/unlabelled/unlabelled_graph_pattern.h b/lib/substitutions/include/substitutions/unlabelled/unlabelled_graph_pattern.h
@@ -17,7 +17,7 @@ std::unordered_set<PatternValue> get_values(UnlabelledGraphPattern const &);
 std::vector<PatternNode>
     get_topological_ordering(UnlabelledGraphPattern const &);
 
-std::unordered_set<PatternInput> get_inputs(UnlabelledGraphPattern const &);
+std::unordered_set<PatternInput> get_graph_inputs(UnlabelledGraphPattern const &);
 
 std::unordered_set<PatternEdge> get_edges(UnlabelledGraphPattern const &);
 

diff --git a/lib/substitutions/src/substitutions/pcg_pattern.cc b/lib/substitutions/src/substitutions/pcg_pattern.cc
@@ -1,5 +1,6 @@
 #include "substitutions/pcg_pattern.h"
 #include "substitutions/operator_pattern/satisfies_pattern.h"
+#include "substitutions/pcg_pattern_match.h"
 #include "substitutions/sub_parallel_computation_graph.h"
 #include "substitutions/tensor_pattern/satisfies_pattern.h"
 #include "substitutions/unlabelled/pattern_value.h"
@@ -76,10 +77,10 @@ std::vector<PatternNodeOutput>
 bool assignment_satisfies(
     SubParallelComputationGraph const &pcg,
     PCGPattern const &pattern,
-    UnlabelledDataflowGraphPatternMatch const &patternMatch) {
+    PCGPatternMatch const &pattern_match) {
   return unlabelled_pattern_does_match(get_unlabelled_pattern(pattern),
                                        pcg.raw_graph,
-                                       patternMatch,
+                                       get_unlabelled_pattern_match(pattern_match),
                                        pcg_pattern_criteria(pattern, pcg));
 }
 

diff --git a/lib/substitutions/src/substitutions/pcg_pattern_match.cc b/lib/substitutions/src/substitutions/pcg_pattern_match.cc
@@ -4,6 +4,7 @@
 #include "utils/bidict/algorithms/bidict_from_keys_and_values.h"
 #include "utils/bidict/algorithms/merge_bidicts.h"
 #include "utils/containers/zip.h"
+#include "utils/containers/map_values.h"
 
 namespace FlexFlow {
 
@@ -32,4 +33,11 @@ bidict<PatternNodeOutput, parallel_tensor_guid_t>
   return result;
 }
 
+UnlabelledDataflowGraphPatternMatch get_unlabelled_pattern_match(PCGPatternMatch const &match) {
+  return UnlabelledDataflowGraphPatternMatch{
+    map_values(match.node_assignment, [](parallel_layer_guid_t const &l) { return l.raw_graph_node; }),
+    map_values(match.input_assignment, [](open_parallel_tensor_guid_t const &i) { return i.raw_open_dataflow_value; }),
+  };
+}
+
 } // namespace FlexFlow
diff --git a/lib/substitutions/src/substitutions/sub_parallel_computation_graph.cc b/lib/substitutions/src/substitutions/sub_parallel_computation_graph.cc
@@ -186,7 +186,7 @@ SubParallelComputationGraph
   };
 }
 
-bool are_isomorphic(SubParallelComputationGraph const &lhs,
+bool sub_pcgs_are_isomorphic(SubParallelComputationGraph const &lhs,
                     SubParallelComputationGraph const &rhs) {
   return find_isomorphism(without_layer_names(lhs).raw_graph,
                           without_layer_names(rhs).raw_graph)

diff --git a/lib/substitutions/src/substitutions/sub_parallel_computation_graph_edge.cc b/lib/substitutions/src/substitutions/sub_parallel_computation_graph_edge.cc
@@ -31,7 +31,7 @@ SubParallelComputationGraphEdge
 
 open_parallel_tensor_guid_t
     get_parallel_tensor(SubParallelComputationGraphEdge const &e) {
-  OpenDataflowValue raw_value = get_open_dataflow_edge_source(e.raw_edge);
+  OpenDataflowValue raw_value = get_open_dataflow_edge_src(e.raw_edge);
   return open_parallel_tensor_guid_t{raw_value};
 }
 

diff --git a/lib/substitutions/src/substitutions/unlabelled/find_pattern_matches.cc b/lib/substitutions/src/substitutions/unlabelled/find_pattern_matches.cc
@@ -37,7 +37,7 @@ static std::optional<UnlabelledDataflowGraphPatternMatch>
 
   std::vector<PatternValue> pattern_node_inputs =
       get_inputs_to_pattern_node(pattern, pattern_node);
-  std::unordered_set<PatternInput> pattern_graph_inputs = get_inputs(pattern);
+  std::unordered_set<PatternInput> pattern_graph_inputs = get_graph_inputs(pattern);
 
   assert(unordered_set_of(pattern_node_inputs) ==
          transform(pattern_graph_inputs,

diff --git a/lib/substitutions/src/substitutions/unlabelled/unlabelled_graph_pattern.cc b/lib/substitutions/src/substitutions/unlabelled/unlabelled_graph_pattern.cc
@@ -6,9 +6,10 @@
 #include "utils/graph/digraph/algorithms/get_topological_ordering.h"
 #include "utils/graph/node/algorithms.h"
 #include "utils/graph/open_dataflow_graph/algorithms/get_edges.h"
-#include "utils/graph/open_dataflow_graph/algorithms/get_inputs.h"
+#include "utils/graph/open_dataflow_graph/algorithms/get_open_dataflow_graph_inputs.h"
 #include "utils/graph/open_dataflow_graph/algorithms/get_open_dataflow_values.h"
 #include "utils/graph/open_dataflow_graph/algorithms/get_subgraph.h"
+#include "utils/graph/open_dataflow_graph/algorithms/get_inputs.h"
 #include "utils/graph/open_dataflow_graph/algorithms/get_subgraph_inputs.h"
 
 namespace FlexFlow {
@@ -31,8 +32,8 @@ std::unordered_set<PatternValue> get_values(UnlabelledGraphPattern const &p) {
                    pattern_value_from_raw_open_dataflow_value);
 }
 
-std::unordered_set<PatternInput> get_inputs(UnlabelledGraphPattern const &p) {
-  return transform(get_inputs(p.raw_graph),
+std::unordered_set<PatternInput> get_graph_inputs(UnlabelledGraphPattern const &p) {
+  return transform(get_open_dataflow_graph_inputs(p.raw_graph),
                    [](DataflowGraphInput const &i) { return PatternInput{i}; });
 }
 

diff --git a/lib/substitutions/test/src/substitutions/substitution.cc b/lib/substitutions/test/src/substitutions/substitution.cc
@@ -10,7 +10,6 @@
 #include "utils/containers/get_only.h"
 #include "utils/graph/instances/unordered_set_labelled_open_dataflow_graph.h"
 #include "utils/graph/labelled_open_dataflow_graph/algorithms/get_graph_data.h"
-#include "utils/graph/open_dataflow_graph/algorithms/are_isomorphic.h"
 #include "utils/integer_conversions.h"
 #include <doctest/doctest.h>
 
@@ -225,6 +224,6 @@ TEST_SUITE(FF_TEST_SUITE) {
     // since the new nodes produced by the substitution have new ids, it's
     // easier/more correct to check that the graphs are isomorphic rather than
     // checking their exact graph data
-    CHECK(are_isomorphic(result, correct));
+    CHECK(sub_pcgs_are_isomorphic(result, correct));
   }
 }
diff --git a/lib/substitutions/test/src/substitutions/unlabelled/pattern_split.cc b/lib/substitutions/test/src/substitutions/unlabelled/pattern_split.cc
@@ -61,7 +61,7 @@ TEST_SUITE(FF_TEST_SUITE) {
       SUBCASE("full_pattern_values_to_subpattern_2_inputs") {
         bidict<PatternValue, PatternInput> result =
             split_result.full_pattern_values_to_subpattern_2_inputs;
-        PatternInput i0 = get_only(get_inputs(split_result.subpattern_2));
+        PatternInput i0 = get_only(get_graph_inputs(split_result.subpattern_2));
         bidict<PatternValue, PatternInput> correct = {
             {pv0, i0},
         };
@@ -117,7 +117,7 @@ TEST_SUITE(FF_TEST_SUITE) {
             split_result.full_pattern_values_to_subpattern_1_inputs;
         bidict<PatternValue, PatternInput> correct = {
             {PatternValue{pi0},
-             get_only(get_inputs(split_result.subpattern_1))},
+             get_only(get_graph_inputs(split_result.subpattern_1))},
         };
         CHECK(result == correct);
       }
@@ -126,7 +126,7 @@ TEST_SUITE(FF_TEST_SUITE) {
             split_result.full_pattern_values_to_subpattern_2_inputs;
         bidict<PatternValue, PatternInput> correct = {
             {PatternValue{pi1},
-             get_only(get_inputs(split_result.subpattern_2))},
+             get_only(get_graph_inputs(split_result.subpattern_2))},
         };
         CHECK(result == correct);
       }