Pass initial test for find_pattern_matches

flexflow · victorli2002 · Jun 4, 2024 · Jun 4, 2024 · Jun 10, 2024 · Jun 10, 2024
commit bcf776e56fb3862c239242db8192f985faaae3e9
diff --git a/lib/local-execution/include/local-execution/op_task_invocation.h b/lib/local-execution/include/local-execution/op_task_invocation.h
@@ -12,7 +12,7 @@
 #include "local-execution/variadic_tensor_ref.h"
 #include "op-attrs/computation_graph_op_attrs.h"
 #include "pcg/computation_graph.h"
-#include "utils/bidict.h"
+#include "utils/bidict/bidict.h"
 #include "utils/stack_map.h"
 #include <typeindex>
 #include <unordered_map>

diff --git a/lib/substitutions/include/substitutions/pcg_pattern.h b/lib/substitutions/include/substitutions/pcg_pattern.h
@@ -20,6 +20,7 @@ TensorAttributePattern get_tensor_pattern(PCGPattern const &,
                                           PatternValue const &);
 OperatorAttributePattern get_operator_pattern(PCGPattern const &,
                                               PatternNode const &);
+std::unordered_set<PatternInput> get_inputs(PCGPattern const &);
 
 bool assignment_satisfies(SubParallelComputationGraph const &,
                           PCGPattern const &,

diff --git a/lib/substitutions/include/substitutions/unlabelled/edge_splits.struct.toml b/lib/substitutions/include/substitutions/unlabelled/edge_splits.struct.toml
@@ -5,7 +5,7 @@ features = [
 ]
 
 includes = [
-  "utils/bidict.h",
+  "utils/bidict/bidict.h",
   "utils/graph.h",
   "<utility>",
 ]

diff --git a/lib/substitutions/include/substitutions/unlabelled/pattern_split_result.struct.toml b/lib/substitutions/include/substitutions/unlabelled/pattern_split_result.struct.toml
@@ -6,7 +6,7 @@ includes = [
   "substitutions/unlabelled/unlabelled_graph_pattern.dtg.h",
   "substitutions/unlabelled/pattern_value.dtg.h",
   "substitutions/unlabelled/pattern_input.dtg.h",
-  "utils/bidict.h",
+  "utils/bidict/bidict.h",
 ]
 
 [[fields]]
@@ -18,5 +18,9 @@ name = "subpattern_2"
 type = "::FlexFlow::UnlabelledGraphPattern"
 
 [[fields]]
-name = "subpattern_1_outputs_to_subpattern_2_inputs"
+name = "full_pattern_values_to_subpattern_1_inputs"
+type = "::FlexFlow::bidict<::FlexFlow::PatternValue, ::FlexFlow::PatternInput>"
+
+[[fields]]
+name = "full_pattern_values_to_subpattern_2_inputs"
 type = "::FlexFlow::bidict<::FlexFlow::PatternValue, ::FlexFlow::PatternInput>"
diff --git a/lib/substitutions/include/substitutions/unlabelled/unlabelled_dataflow_graph_pattern_match.h b/lib/substitutions/include/substitutions/unlabelled/unlabelled_dataflow_graph_pattern_match.h
@@ -12,7 +12,8 @@ UnlabelledDataflowGraphPatternMatch empty_unlabelled_pattern_match();
 std::unordered_set<Node> matched_nodes(UnlabelledDataflowGraphPatternMatch const &);
 std::optional<UnlabelledDataflowGraphPatternMatch> merge_unlabelled_dataflow_graph_pattern_matches(UnlabelledDataflowGraphPatternMatch const &subpattern_1,
                                                 UnlabelledDataflowGraphPatternMatch const &subpattern_2,
-                                                bidict<PatternValue, PatternInput> const &outputs_of_1_to_inputs_of_2);
+                                                bidict<PatternValue, PatternInput> const &merged_graph_values_to_inputs_of_1,
+                                                bidict<PatternValue, PatternInput> const &merged_graph_values_to_inputs_of_2);
 
 } // namespace FlexFlow
 

diff --git a/...ions/include/substitutions/unlabelled/unlabelled_dataflow_graph_pattern_match.struct.toml b/...ions/include/substitutions/unlabelled/unlabelled_dataflow_graph_pattern_match.struct.toml
@@ -8,7 +8,7 @@ features = [
 ]
 
 includes = [
-  "utils/bidict.h",
+  "utils/bidict/bidict.h",
   "utils/graph/node/node.dtg.h",
   "utils/graph/open_dataflow_graph/open_dataflow_value.dtg.h",
   "substitutions/unlabelled/pattern_input.dtg.h",

diff --git a/lib/substitutions/include/substitutions/unlabelled/unlabelled_graph_pattern.h b/lib/substitutions/include/substitutions/unlabelled/unlabelled_graph_pattern.h
@@ -6,6 +6,7 @@
 #include "substitutions/unlabelled/pattern_node.dtg.h"
 #include "substitutions/unlabelled/unlabelled_graph_pattern.dtg.h"
 #include "substitutions/unlabelled/pattern_input.dtg.h"
+#include "substitutions/unlabelled/unlabelled_graph_pattern_subgraph_result.dtg.h"
 
 namespace FlexFlow {
 
@@ -26,8 +27,8 @@ std::vector<PatternValue>
 std::vector<PatternValue>
     get_outputs_from_pattern_node(UnlabelledGraphPattern const &, PatternNode const &);
 
-UnlabelledGraphPattern get_subgraph(UnlabelledGraphPattern const &,
-                                    std::unordered_set<PatternNode> const &);
+UnlabelledGraphPatternSubgraphResult get_subgraph(UnlabelledGraphPattern const &,
+                                                  std::unordered_set<PatternNode> const &);
 
 } // namespace FlexFlow
 

diff --git a/...ons/include/substitutions/unlabelled/unlabelled_graph_pattern_subgraph_result.struct.toml b/...ons/include/substitutions/unlabelled/unlabelled_graph_pattern_subgraph_result.struct.toml
@@ -0,0 +1,18 @@
+namespace = "FlexFlow"
+name = "UnlabelledGraphPatternSubgraphResult"
+features = [ ]
+
+includes = [
+  "substitutions/unlabelled/unlabelled_graph_pattern.dtg.h",
+  "substitutions/unlabelled/pattern_value.dtg.h",
+  "substitutions/unlabelled/pattern_input.dtg.h",
+  "utils/bidict/bidict.h",
+]
+
+[[fields]]
+name = "subpattern"
+type = "::FlexFlow::UnlabelledGraphPattern"
+
+[[fields]]
+name = "full_pattern_values_to_subpattern_inputs"
+type = "::FlexFlow::bidict<::FlexFlow::PatternValue, ::FlexFlow::PatternInput>"
diff --git a/lib/substitutions/src/substitutions/unlabelled/find_pattern_matches.cc b/lib/substitutions/src/substitutions/unlabelled/find_pattern_matches.cc
@@ -76,7 +76,7 @@ std::vector<UnlabelledDataflowGraphPatternMatch>
     }
   } else {
     PatternSplit split = find_even_split(pattern);
-    auto subpatterns = apply_split(pattern, split);
+    PatternSplitResult subpatterns = apply_split(pattern, split);
     std::vector<UnlabelledDataflowGraphPatternMatch> prefix_matches =
         find_pattern_matches(subpatterns.subpattern_1, graph, additional_criterion);
     std::vector<UnlabelledDataflowGraphPatternMatch> postfix_matches =
@@ -87,7 +87,8 @@ std::vector<UnlabelledDataflowGraphPatternMatch>
         std::optional<UnlabelledDataflowGraphPatternMatch> unsplit =
           merge_unlabelled_dataflow_graph_pattern_matches(prefix_match,
                                                           postfix_match,
-                                                          subpatterns.subpattern_1_outputs_to_subpattern_2_inputs);
+                                                          subpatterns.full_pattern_values_to_subpattern_1_inputs,
+                                                          subpatterns.full_pattern_values_to_subpattern_2_inputs);
         if (unsplit.has_value() && unlabelled_pattern_does_match(pattern, graph, unsplit.value(), additional_criterion)) {
           matches.push_back(unsplit.value());
         }

diff --git a/lib/substitutions/src/substitutions/unlabelled/pattern_matching.cc b/lib/substitutions/src/substitutions/unlabelled/pattern_matching.cc
@@ -61,8 +61,8 @@ OpenDataflowSubgraphResult subgraph_matched(OpenDataflowGraphView const &g,
 //   return true;
 // }
 
-struct ConcreteFromPattern {
-  ConcreteFromPattern(UnlabelledDataflowGraphPatternMatch const &match, 
+struct SubgraphConcreteFromPattern {
+  SubgraphConcreteFromPattern(UnlabelledDataflowGraphPatternMatch const &match, 
                       bidict<OpenDataflowValue, DataflowGraphInput> const &full_graph_values_to_subgraph_inputs)
     : match(match), full_graph_values_to_subgraph_inputs(full_graph_values_to_subgraph_inputs)
   { }
@@ -124,7 +124,7 @@ bool pattern_matches_subgraph_under(UnlabelledGraphPattern const &pattern,
                                     bidict<OpenDataflowValue, DataflowGraphInput> const &full_graph_values_to_subgraph_inputs,
                                     UnlabelledDataflowGraphPatternMatch const &match,
                                     MatchAdditionalCriterion const &additional_criterion) {
-  ConcreteFromPattern concrete_from_pattern{match, full_graph_values_to_subgraph_inputs};
+  SubgraphConcreteFromPattern concrete_from_pattern{match, full_graph_values_to_subgraph_inputs};
 
   std::unordered_set<Node> concrete_nodes = get_nodes(subgraph);
   std::unordered_set<Node> concrete_nodes_from_match = transform(get_nodes(pattern), concrete_from_pattern);
@@ -174,7 +174,20 @@ bool unlabelled_pattern_does_match(
   assert (keys(match.node_assignment) == get_nodes(pattern));
   assert (keys(match.node_assignment.reversed()) == get_nodes(matched_subgraph));
 
-  return pattern_matches_subgraph_under(pattern, matched_subgraph, subgraph_result.full_graph_values_to_subgraph_inputs, match, additional_criterion);
+  MatchAdditionalCriterion through_subgraph_operation = MatchAdditionalCriterion{
+    additional_criterion.node_criterion,
+    [&](PatternValue const &pv, OpenDataflowValue const &v) {
+      return v.visit<bool>(overload {
+        [&](DataflowOutput const &) { return additional_criterion.value_criterion(pv, v); },
+        [&](DataflowGraphInput const &subgraph_input) { 
+          OpenDataflowValue full_graph_value = subgraph_result.full_graph_values_to_subgraph_inputs.at_r(subgraph_input);
+          return additional_criterion.value_criterion(pv, full_graph_value);
+        }
+      });
+    },
+  };
+
+  return pattern_matches_subgraph_under(pattern, matched_subgraph, subgraph_result.full_graph_values_to_subgraph_inputs, match, through_subgraph_operation);
 }
 
 } // namespace FlexFlow
diff --git a/lib/substitutions/src/substitutions/unlabelled/pattern_split.cc b/lib/substitutions/src/substitutions/unlabelled/pattern_split.cc
@@ -37,20 +37,14 @@ PatternSplit find_even_split(UnlabelledGraphPattern const &pattern) {
 
 PatternSplitResult
     apply_split(UnlabelledGraphPattern const &p, PatternSplit const &s) {
-  OpenDataflowSubgraphResult raw_second_subgraph_result = get_subgraph(p.raw_graph, transform(s.second, [](PatternNode const &pn) { return pn.raw_node; }));
-
-  bidict<PatternValue, PatternInput> subpattern_1_outputs_to_subpattern_2_inputs;
-  for (auto const &kv : raw_second_subgraph_result.full_graph_values_to_subgraph_inputs) {
-    OpenDataflowValue open_dataflow_value = kv.first;
-    DataflowGraphInput dataflow_graph_input = kv.second;
-    subpattern_1_outputs_to_subpattern_2_inputs.equate(
-                                                       pattern_value_from_raw_open_dataflow_value(open_dataflow_value), PatternInput{dataflow_graph_input});
-  }
-
+  UnlabelledGraphPatternSubgraphResult first_subgraph_result = get_subgraph(p, s.first);
+  UnlabelledGraphPatternSubgraphResult second_subgraph_result = get_subgraph(p, s.second);
+
   return PatternSplitResult{
-    get_subgraph(p, s.first),
-    UnlabelledGraphPattern{raw_second_subgraph_result.graph},
-    subpattern_1_outputs_to_subpattern_2_inputs,
+    first_subgraph_result.subpattern,
+    second_subgraph_result.subpattern,
+    first_subgraph_result.full_pattern_values_to_subpattern_inputs,
+    second_subgraph_result.full_pattern_values_to_subpattern_inputs
   };
 }
 

diff --git a/lib/substitutions/src/substitutions/unlabelled/unlabelled_dataflow_graph_pattern_match.cc b/lib/substitutions/src/substitutions/unlabelled/unlabelled_dataflow_graph_pattern_match.cc
@@ -1,5 +1,8 @@
 #include "substitutions/unlabelled/unlabelled_dataflow_graph_pattern_match.h"
 #include "utils/containers.h"
+#include "utils/containers/filtermap_keys.h"
+#include "utils/bidict/try_merge_nondisjoint_bidicts.h"
+#include "utils/containers/try_merge_nondisjoint_unordered_maps.h"
 
 namespace FlexFlow {
 
@@ -10,34 +13,11 @@ UnlabelledDataflowGraphPatternMatch empty_unlabelled_pattern_match() {
   };
 }
 
-template <typename L, typename R>
-std::optional<bidict<L, R>> try_merge_nondisjoint_bidicts(bidict<L, R> const &d1,
-                                                          bidict<L, R> const &d2) {
-  bidict<L, R> result;
-  for (L const &l : set_union(keys(d1), keys(d2))) {
-    if (d1.contains_l(l) && d2.contains_l(l)) {
-      if (d1.at_l(l) == d2.at_l(l)) {
-        result.equate(l, d1.at_l(l));
-      } else {
-        return std::nullopt;
-      }
-    } else if (d1.contains_l(l)) {
-      result.equate(l, d1.at_l(l));
-    } else {
-      assert (d2.contains_l(l));
-
-      result.equate(l, d2.at_l(l));
-    }
-  }
-
-  return result;
-}
-
-
 std::optional<UnlabelledDataflowGraphPatternMatch>
   merge_unlabelled_dataflow_graph_pattern_matches(UnlabelledDataflowGraphPatternMatch const &subpattern_1,
                                                 UnlabelledDataflowGraphPatternMatch const &subpattern_2,
-                                                bidict<PatternValue, PatternInput> const &outputs_of_1_to_inputs_of_2) {
+                                                bidict<PatternValue, PatternInput> const &merged_graph_values_to_inputs_of_1,
+                                                bidict<PatternValue, PatternInput> const &merged_graph_values_to_inputs_of_2) {
   bidict<PatternNode, Node> merged_node_assignment = ({
     std::optional<bidict<PatternNode, Node>> result = try_merge_nondisjoint_bidicts(
       subpattern_1.node_assignment, subpattern_2.node_assignment);
@@ -47,11 +27,37 @@ std::optional<UnlabelledDataflowGraphPatternMatch>
     result.value();
   });
 
-  assert (all_of(keys(subpattern_2.input_assignment), [&](PatternInput const &i) { return outputs_of_1_to_inputs_of_2.contains_r(i); }));
+  std::unordered_map<PatternInput, OpenDataflowValue> merged_input_assignment = ({
+    std::unordered_map<PatternValue, OpenDataflowValue> lifted_input_assignment_1 = map_keys(
+      subpattern_1.input_assignment, 
+      [&](PatternInput const &pi1) {
+        return merged_graph_values_to_inputs_of_1.at_r(pi1);
+      }
+    );
+    std::unordered_map<PatternValue, OpenDataflowValue> lifted_input_assignment_2 = map_keys(
+      subpattern_2.input_assignment, 
+      [&](PatternInput const &pi2) {
+        return merged_graph_values_to_inputs_of_2.at_r(pi2);
+      }
+    );
+    std::optional<std::unordered_map<PatternValue, OpenDataflowValue>> merged = try_merge_nondisjoint_unordered_maps(
+      lifted_input_assignment_1, lifted_input_assignment_2);
+    if (!merged.has_value()) {
+      return std::nullopt;
+    }
+    filtermap_keys(merged.value(), 
+                   [](PatternValue const &v) -> std::optional<PatternInput> {
+                     if (v.has<PatternInput>()) {
+                       return v.get<PatternInput>();  
+                     } else {
+                       return std::nullopt;
+                     }
+                   });
+  });
 
   return UnlabelledDataflowGraphPatternMatch{
     merged_node_assignment,
-    subpattern_1.input_assignment,
+    merged_input_assignment,
   };
 }
 

diff --git a/lib/substitutions/src/substitutions/unlabelled/unlabelled_graph_pattern.cc b/lib/substitutions/src/substitutions/unlabelled/unlabelled_graph_pattern.cc
@@ -52,12 +52,22 @@ std::vector<PatternValue>
                    [](DataflowOutput const &o) { return pattern_value_from_raw_open_dataflow_value(OpenDataflowValue{o}); });
 }
 
-UnlabelledGraphPattern get_subgraph(UnlabelledGraphPattern const &p,
-                                    std::unordered_set<PatternNode> const &n) {
-  OpenDataflowGraphView raw_subgraph = 
-    get_subgraph(p.raw_graph, transform(n, [](PatternNode const &pn) { return pn.raw_node; })).graph;
-  return UnlabelledGraphPattern{
-    raw_subgraph,
+UnlabelledGraphPatternSubgraphResult get_subgraph(UnlabelledGraphPattern const &p,
+                                                  std::unordered_set<PatternNode> const &n) {
+  OpenDataflowSubgraphResult raw_result =  
+    get_subgraph(p.raw_graph, transform(n, [](PatternNode const &pn) { return pn.raw_node; }));
+  bidict<PatternValue, PatternInput> full_pattern_values_to_subpattern_inputs = transform(
+    raw_result.full_graph_values_to_subgraph_inputs,
+    [](OpenDataflowValue const &v, DataflowGraphInput const &i) {
+      return std::make_pair(
+        pattern_value_from_raw_open_dataflow_value(v),
+        PatternInput{i}
+      );
+    }
+  );
+  return UnlabelledGraphPatternSubgraphResult{
+    UnlabelledGraphPattern{raw_result.graph},
+    full_pattern_values_to_subpattern_inputs ,
   };
 }
 

diff --git a/lib/substitutions/test/src/substitutions/pcg_pattern.cc b/lib/substitutions/test/src/substitutions/pcg_pattern.cc
@@ -101,33 +101,33 @@ TEST_SUITE(FF_TEST_SUITE) {
     PCGPattern pattern = PCGPattern{g};
 
     std::unordered_set<UnlabelledDataflowGraphPatternMatch> result = without_order(find_pattern_matches(pattern, sub_pcg_from_full_pcg(pcg)));
-    //
-    // UnlabelledDataflowGraphPatternMatch match1 = UnlabelledDataflowGraphPatternMatch{
-    //   bidict<PatternNode, Node>{
-    //     {op_pattern_1_node, x_matmul.raw_graph_node},
-    //     {op_pattern_2_node, y_matmul.raw_graph_node},
-    //   },
-    //   bidict<PatternInput, OpenDataflowValue>{
-    //     {PatternInput{pt_a}, OpenDataflowValue{a_tensor.raw_graph_output}},
-    //     {PatternInput{pt_b}, OpenDataflowValue{x_weights.raw_graph_output}},
-    //     {PatternInput{pt_c}, OpenDataflowValue{y_weights.raw_graph_output}},
-    //   }
-    // };
-    //
-    // UnlabelledDataflowGraphPatternMatch match2 = UnlabelledDataflowGraphPatternMatch{
-    //   bidict<PatternNode, Node>{
-    //     {op_pattern_1_node, y_matmul.raw_graph_node},
-    //     {op_pattern_2_node, x_matmul.raw_graph_node},
-    //   },
-    //   bidict<PatternInput, OpenDataflowValue>{
-    //     {PatternInput{pt_a}, OpenDataflowValue{a_tensor.raw_graph_output}},
-    //     {PatternInput{pt_b}, OpenDataflowValue{y_weights.raw_graph_output}},
-    //     {PatternInput{pt_c}, OpenDataflowValue{x_weights.raw_graph_output}},
-    //   }
-    // };
-    //
-    // std::unordered_set<UnlabelledDataflowGraphPatternMatch> correct = {match1, match2};
-    //
-    // CHECK(result == correct);
+
+    UnlabelledDataflowGraphPatternMatch match1 = UnlabelledDataflowGraphPatternMatch{
+      bidict<PatternNode, Node>{
+        {op_pattern_1_node, x_matmul.raw_graph_node},
+        {op_pattern_2_node, y_matmul.raw_graph_node},
+      },
+      bidict<PatternInput, OpenDataflowValue>{
+        {PatternInput{pt_a}, OpenDataflowValue{a_tensor.raw_graph_output}},
+        {PatternInput{pt_b}, OpenDataflowValue{x_weights.raw_graph_output}},
+        {PatternInput{pt_c}, OpenDataflowValue{y_weights.raw_graph_output}},
+      }
+    };
+
+    UnlabelledDataflowGraphPatternMatch match2 = UnlabelledDataflowGraphPatternMatch{
+      bidict<PatternNode, Node>{
+        {op_pattern_1_node, y_matmul.raw_graph_node},
+        {op_pattern_2_node, x_matmul.raw_graph_node},
+      },
+      bidict<PatternInput, OpenDataflowValue>{
+        {PatternInput{pt_a}, OpenDataflowValue{a_tensor.raw_graph_output}},
+        {PatternInput{pt_b}, OpenDataflowValue{y_weights.raw_graph_output}},
+        {PatternInput{pt_c}, OpenDataflowValue{x_weights.raw_graph_output}},
+      }
+    };
+
+    std::unordered_set<UnlabelledDataflowGraphPatternMatch> correct = {match1, match2};
+
+    CHECK(result == correct);
   }
 }