Fix is_valid_machine_mapping_problem_tree, add hacky printing for pro…

…blem trees

lib/compiler/include/compiler/machine_mapping/machine_mapping_problem_tree/machine_mapping_problem_tree.h

@@ -28,6 +28,9 @@ std::optional<MachineMappingProblemTree>
     mm_problem_tree_get_subtree_at_path(MachineMappingProblemTree const &,
                                         BinaryTreePath const &);
 
+std::string as_dot(MachineMappingProblemTree const &);
+void debug_print_dot(MachineMappingProblemTree const &);
+
 } // namespace FlexFlow
 
 #endif

lib/compiler/src/compiler/machine_mapping/machine_mapping_problem_tree/get_machine_mapping_problem_tree.cc

@@ -7,6 +7,7 @@
 #include "compiler/series_parallel/pcg/pcg_binary_sp_decomposition.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph.h"
 #include "utils/overload.h"
+#include "utils/containers/all_of.h"
 
 namespace FlexFlow {
 
@@ -16,12 +17,15 @@ bool is_valid_machine_mapping_problem_tree(
       [&](MMProblemTreeSeriesSplit const &series_split) {
         AbstractedTensorSetMovement tensor_movement =
             series_split.tensor_set_movement;
-        for (BinaryTreePath const &path : get_src_layers(tensor_movement)) {
-          if (!mm_problem_tree_get_subtree_at_path(problem_tree, path)) {
-            return false;
-          }
-        }
-        return is_valid_machine_mapping_problem_tree(
+
+        auto contains_paths = [](MachineMappingProblemTree const &t,
+                                 std::unordered_set<BinaryTreePath> const &paths) {
+          return all_of(paths, [&](BinaryTreePath const &p) { return mm_problem_tree_get_subtree_at_path(t, p).has_value(); });
+        };
+
+        return contains_paths(series_split.get_left_child(), get_src_layers(tensor_movement)) &&
+               contains_paths(series_split.get_right_child(), get_dst_layers(tensor_movement)) &&
+               is_valid_machine_mapping_problem_tree(
                    series_split.get_left_child()) &&
                is_valid_machine_mapping_problem_tree(
                    series_split.get_right_child());
@@ -50,26 +54,32 @@ MachineMappingProblemTree get_machine_mapping_problem_tree(
         [&](PCGBinarySeriesSplit const &series) {
           AbstractedTensorSetMovement tensor_movement =
               get_abstracted_tensor_set_movement_across_split(tr_pcg, series);
-          return MachineMappingProblemTree{
-              MMProblemTreeSeriesSplit{
-                  /*tensor_set_movement=*/tensor_movement,
-                  /*lhs=*/to_problem_tree(series.get_left_child()),
-                  /*rhs=*/to_problem_tree(series.get_right_child()),
-              },
+          MachineMappingProblemTree result = MachineMappingProblemTree{
+            MMProblemTreeSeriesSplit{
+              /*tensor_set_movement=*/tensor_movement,
+              /*lhs=*/to_problem_tree(series.get_left_child()),
+              /*rhs=*/to_problem_tree(series.get_right_child()),
+            },
           };
+          assert (is_valid_machine_mapping_problem_tree(result));
+          return result;
         },
         [&](PCGBinaryParallelSplit const &parallel) {
-          return MachineMappingProblemTree{
-              MMProblemTreeParallelSplit{
-                  to_problem_tree(parallel.get_left_child()),
-                  to_problem_tree(parallel.get_right_child()),
-              },
+          MachineMappingProblemTree result = MachineMappingProblemTree{
+            MMProblemTreeParallelSplit{
+                to_problem_tree(parallel.get_left_child()),
+                to_problem_tree(parallel.get_right_child()),
+            },
           };
+          assert (is_valid_machine_mapping_problem_tree(result));
+          return result;
         },
         [&](parallel_layer_guid_t const &leaf) {
-          return MachineMappingProblemTree{
+          MachineMappingProblemTree result = MachineMappingProblemTree{
               get_unmapped_op_cost_estimate_key_for_layer(pcg, leaf),
           };
+          assert (is_valid_machine_mapping_problem_tree(result));
+          return result;
         },
     });
   };

lib/compiler/src/compiler/machine_mapping/machine_mapping_problem_tree/machine_mapping_problem_tree.cc

@@ -1,7 +1,9 @@
 #include "compiler/machine_mapping/machine_mapping_problem_tree/machine_mapping_problem_tree.h"
+#include "compiler/machine_mapping/abstracted_tensor_set_movement/abstracted_tensor_set_movement.h"
 #include "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/get_all_leaf_paths.h"
 #include "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/get_leaves.h"
 #include "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/get_subtree_at_path.h"
+#include "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/as_dot.h"
 
 namespace FlexFlow {
 
@@ -88,4 +90,44 @@ std::optional<MachineMappingProblemTree>
       tree, generic_binary_sp_impl_for_mm_problem_tree(), path);
 }
 
+std::string as_dot(MachineMappingProblemTree const &tree) {
+  std::function<std::string(MMProblemTreeSeriesSplit const &)> get_series_label 
+    = [](MMProblemTreeSeriesSplit const &series) -> std::string {
+      auto path_as_dot = [](BinaryTreePath const &path) -> std::string {
+        return "(" + join_strings(path.entries,
+                            ", ",
+                            [](BinaryTreePathEntry const &entry) -> std::string {
+                              if (entry == BinaryTreePathEntry::LEFT_CHILD) {
+                                return "l";
+                              } else {
+                                assert (entry == BinaryTreePathEntry::RIGHT_CHILD);
+                                return "r";
+                              }
+                            }) + ")";
+      };
+
+      auto path_set_as_dot = [&](std::unordered_set<BinaryTreePath> const &path_set) -> std::string {
+        return "(" + join_strings(path_set, ", ", path_as_dot) + ")";
+      };
+
+      return fmt::format("srcs={} dsts={}", path_set_as_dot(get_src_layers(series.tensor_set_movement)), path_set_as_dot(get_dst_layers(series.tensor_set_movement)));
+    };
+
+  std::function<std::string(MMProblemTreeParallelSplit const &)> get_parallel_label 
+    = [](MMProblemTreeParallelSplit const &parallel) -> std::string {
+      return "P";
+    };
+
+  std::function<std::string(UnmappedOpCostEstimateKey const &)> get_leaf_label 
+    = [](UnmappedOpCostEstimateKey const &leaf) -> std::string {
+      return "";
+    };
+
+  return as_dot(tree, generic_binary_sp_impl_for_mm_problem_tree(), get_series_label, get_parallel_label, get_leaf_label);
+}
+
+void debug_print_dot(MachineMappingProblemTree const &tree) {
+  std::cout << as_dot(tree) << std::endl;
+}
+
 } // namespace FlexFlow

lib/utils/include/utils/full_binary_tree/as_dot.h

@@ -0,0 +1,70 @@
+#ifndef _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_FULL_BINARY_TREE_AS_DOT_H
+#define _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_FULL_BINARY_TREE_AS_DOT_H
+
+#include "utils/containers/get_only.h"
+#include "utils/dot_file.h"
+#include "utils/full_binary_tree/full_binary_tree_implementation.dtg.h"
+#include "utils/full_binary_tree/full_binary_tree_visitor.dtg.h"
+#include "utils/graph/dataflow_graph/dataflow_graph.h"
+#include "utils/graph/dataflow_graph/dataflow_graph_view.h"
+#include "utils/graph/digraph/digraph_view.h"
+#include "utils/graph/instances/adjacency_digraph.h"
+#include "utils/graph/instances/unordered_set_dataflow_graph.h"
+#include "utils/graph/instances/unordered_set_labelled_open_dataflow_graph.h"
+#include "utils/graph/labelled_dataflow_graph/algorithms/view_as_labelled_open_dataflow_graph.h"
+#include "utils/graph/labelled_dataflow_graph/labelled_dataflow_graph.h"
+#include <string>
+#include <functional>
+#include <sstream>
+#include "utils/graph/labelled_open_dataflow_graph/algorithms/as_dot.h"
+#include "utils/full_binary_tree/visit.h"
+
+namespace FlexFlow {
+
+template <typename Tree, typename Parent, typename Leaf, typename NodeLabel>
+LabelledDataflowGraph<NodeLabel, std::monostate> as_labelled_dataflow_graph(Tree const &tree,
+                             FullBinaryTreeImplementation<Tree, Parent, Leaf> const &impl,
+                             std::function<NodeLabel(Parent const &)> const &get_parent_label,
+                             std::function<NodeLabel(Leaf const &)> const &get_leaf_label) {
+  auto g = LabelledDataflowGraph<NodeLabel, std::monostate>
+    ::template create<UnorderedSetLabelledOpenDataflowGraph<NodeLabel, std::monostate>>();
+
+  FullBinaryTreeVisitor<DataflowOutput, Tree, Parent, Leaf> visitor = FullBinaryTreeVisitor<DataflowOutput, Tree, Parent, Leaf>{
+    [&](Parent const &parent) -> DataflowOutput {
+      DataflowOutput left_child_output = visit(impl.get_left_child(parent), impl, visitor);
+      DataflowOutput right_child_output = visit(impl.get_right_child(parent), impl, visitor);
+      NodeLabel parent_label = get_parent_label(parent);
+      NodeAddedResult parent_added = g.add_node(parent_label, {left_child_output, right_child_output}, {std::monostate{}});
+      return get_only(parent_added.outputs);
+    },
+    [&](Leaf const &leaf) -> DataflowOutput {
+      NodeLabel leaf_label = get_leaf_label(leaf);
+      NodeAddedResult leaf_added = g.add_node(leaf_label, {}, {std::monostate{}});
+      return get_only(leaf_added.outputs);
+    },
+  };
+
+  visit(tree, impl, visitor);
+
+  return g;
+}
+
+template <typename Tree, typename Parent, typename Leaf>
+std::string as_dot(Tree const &tree,
+                   FullBinaryTreeImplementation<Tree, Parent, Leaf> const &impl,
+                   std::function<std::string(Parent const &)> const &get_parent_label,
+                   std::function<std::string(Leaf const &)> const &get_leaf_label) {
+
+  LabelledDataflowGraphView<std::string, std::monostate> g = as_labelled_dataflow_graph(tree, impl, get_parent_label, get_leaf_label);
+
+  std::function<std::string(std::string const &)> get_node_label = [](std::string const &s) { return s; };
+  std::function<std::string(std::monostate const &)> get_input_label = [](std::monostate const &) { return ""; };
+
+  return as_dot(view_as_labelled_open_dataflow_graph(g), 
+                get_node_label,
+                get_input_label);
+}
+
+} // namespace FlexFlow
+
+#endif

lib/utils/include/utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/as_dot.h

@@ -0,0 +1,37 @@
+#ifndef _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_GRAPH_SERIES_PARALLEL_BINARY_SP_DECOMPOSITION_TREE_GENERIC_BINARY_SP_DECOMPOSITION_TREE_AS_DOT_H
+#define _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_GRAPH_SERIES_PARALLEL_BINARY_SP_DECOMPOSITION_TREE_GENERIC_BINARY_SP_DECOMPOSITION_TREE_AS_DOT_H
+
+#include "utils/full_binary_tree/as_dot.h"
+#include "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree_implementation.dtg.h"
+#include "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree_implementation.h"
+#include "utils/overload.h"
+
+namespace FlexFlow {
+
+template <typename Tree, typename Series, typename Parallel, typename Leaf>
+std::string as_dot(Tree const &tree, 
+                   GenericBinarySPDecompositionTreeImplementation<Tree, Series, Parallel, Leaf> const &impl,
+                   std::function<std::string(Series const &)> const &get_series_label,
+                   std::function<std::string(Parallel const &)> const &get_parallel_label,
+                   std::function<std::string(Leaf const &)> const &get_leaf_label) {
+  FullBinaryTreeImplementation<Tree, std::variant<Series, Parallel>, Leaf> full_binary_tree_impl 
+    = get_full_binary_impl_from_generic_sp_impl(impl);
+
+  std::function<std::string(std::variant<Series, Parallel> const &)> get_parent_label 
+    = [&](std::variant<Series, Parallel> const &parent) -> std::string {
+      return std::visit(overload {
+                          [&](Series const &series) -> std::string {
+                            return get_series_label(series);
+                          },
+                          [&](Parallel const &parallel) -> std::string {
+                            return get_parallel_label(parallel);
+                          },
+                        }, parent);
+    };
+
+  return as_dot(tree, full_binary_tree_impl, get_parent_label, get_leaf_label);
+}
+
+} // namespace FlexFlow
+
+#endif

lib/utils/src/utils/full_binary_tree/as_dot.cc

@@ -0,0 +1,16 @@
+#include "utils/full_binary_tree/as_dot.h"
+#include "utils/archetypes/value_type.h"
+
+namespace FlexFlow {
+
+using Tree = value_type<0>;
+using Parent = value_type<1>;
+using Leaf = value_type<2>;
+
+template
+  std::string as_dot(Tree const &,
+                     FullBinaryTreeImplementation<Tree, Parent, Leaf> const &,
+                     std::function<std::string(Parent const &)> const &,
+                     std::function<std::string(Leaf const &)> const &);
+
+} // namespace FlexFlow

lib/utils/src/utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/as_dot.cc

@@ -0,0 +1,18 @@
+#include "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/as_dot.h"
+#include "utils/archetypes/value_type.h"
+
+namespace FlexFlow {
+
+using Tree = value_type<0>;
+using Series = value_type<1>;
+using Parallel = value_type<2>;
+using Leaf = value_type<3>;
+
+template
+  std::string as_dot(Tree const &, 
+                     GenericBinarySPDecompositionTreeImplementation<Tree, Series, Parallel, Leaf> const &,
+                     std::function<std::string(Series const &)> const &,
+                     std::function<std::string(Parallel const &)> const &,
+                     std::function<std::string(Leaf const &)> const &);
+
+} // namespace FlexFlow