Top-level loop for compiler

lib/compiler/include/compiler/algorithm_config.variant.toml

@@ -0,0 +1,18 @@
+namespace = "FlexFlow"
+name = "AlgorithmConfig"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+]
+
+includes = [
+  "compiler/data_parallelism/data_parallelism_config.dtg.h",
+  "compiler/unity_algorithm/unity_search_config.dtg.h",
+]
+
+[[values]]
+type = "::FlexFlow::DataParallelismConfig"
+
+[[values]]
+type = "::FlexFlow::UnitySearchConfig"

lib/compiler/include/compiler/compiler.h

@@ -1,42 +1,22 @@
 #ifndef _FLEXFLOW_COMPILER_COMPILER_H
 #define _FLEXFLOW_COMPILER_COMPILER_H
 
-#include "pcg/cost_values.h"
-#include "pcg/machine_view.h"
-#include "pcg/parallel_computation_graph/parallel_computation_graph.h"
-#include "pcg/tensor_mapping.h"
+#include "compiler/algorithm_config.dtg.h"
+#include "compiler/cost_estimator/cost_estimator.h"
+#include "compiler/search_result.dtg.h"
+#include "pcg/machine_specification.dtg.h"
 
 namespace FlexFlow {
 
 enum class SearchAlgorithm {
   DATA_PARALLEL,
-};
-
-using SearchAlgorithmConfig = std::variant<>;
-using SearchSolution = std::variant<>;
-
-struct SearchResult {
-  ParallelComputationGraph pcg;
-  TensorMapping tensor_mapping;
-  SearchSolution solution;
-  CostValues cost_values;
+  UNITY,
 };
 
 SearchResult optimize(ComputationGraph const &,
                       MachineSpecification const &,
                       CostEstimator const &,
-                      SearchAlgorithm,
-                      optional<AlgorithmConfig> const &);
-
-// struct SearchSolution {
-//   LabelledMultiDiGraph<PCGOperatorAttrs, ParallelTensorShape> optimized_pcg;
-//   std::unordered_map<Node, MachineView> device_assignments;
-//   /* std::unordered_map<tensor_guid_t,
-//   std::unordered_set<parallel_tensor_guid_t>> tensor_mappings; */
-// };
-//
-// SearchSolution run_data_parallelize(ComputationGraph const &,
-// MachineSpecification const &);
+                      AlgorithmConfig const &);
 
 } // namespace FlexFlow
 

lib/compiler/include/compiler/data_parallelism/data_parallelism_config.struct.toml

@@ -0,0 +1,14 @@
+namespace = "FlexFlow"
+name = "DataParallelismConfig"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+]
+
+includes = [
+]
+
+[[fields]]
+name = "degree"
+type = "int"

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

@@ -2,6 +2,8 @@
 #define _FLEXFLOW_COMPILER_MACHINE_MAPPING_H
 
 #include "compiler/machine_mapping/machine_mapping.dtg.h"
+#include "compiler/machine_mapping/machine_mapping_result.h"
+#include "compiler/series_parallel/pcg/pcg_binary_sp_decomposition.dtg.h"
 
 namespace FlexFlow {
 
@@ -10,6 +12,9 @@ MachineMapping combine_disjoint_mappings(MachineMapping const &,
 
 bool nodes_are_disjoint(MachineMapping const &m1, MachineMapping const &m2);
 
+MachineMapping get_machine_mapping_from_machine_mapping_result(
+    PCGBinarySPDecomposition const &, MachineMappingResult const &);
+
 } // namespace FlexFlow
 
 #endif

lib/compiler/include/compiler/search_result.struct.toml

@@ -0,0 +1,17 @@
+namespace = "FlexFlow"
+name = "SearchResult"
+features = [
+]
+
+includes = [
+  "pcg/parallel_computation_graph/parallel_computation_graph.h",
+  "compiler/machine_mapping/machine_mapping.h",
+]
+
+[[fields]]
+name = "pcg"
+type = "::FlexFlow::ParallelComputationGraph"
+
+[[fields]]
+name = "machine_mapping"
+type = "::FlexFlow::MachineMapping"

lib/compiler/include/compiler/series_parallel/pcg/get_pcg_balanced_binary_sp_decomposition.h

@@ -1,6 +1,8 @@
 #ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_SERIES_PARALLEL_GET_PCG_BALANCED_BINARY_SP_DECOMPOSITION_H
 #define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_SERIES_PARALLEL_GET_PCG_BALANCED_BINARY_SP_DECOMPOSITION_H
 
+#include "compiler/series_parallel/pcg/pcg_binary_sp_decomposition.dtg.h"
+
 namespace FlexFlow {
 
 std::optional<PCGBinarySPDecomposition>

lib/compiler/include/compiler/graph_optimize_state.h → lib/compiler/include/compiler/unity_algorithm/graph_optimize_state.h

@@ -1,16 +1,16 @@
 #ifndef _FLEXFLOW_COMPILER_MCMC_STATE_H
 #define _FLEXFLOW_COMPILER_MCMC_STATE_H
 
-#include "compiler/graph_optimize_result.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.h"
 
 namespace FlexFlow {
 
 struct GraphOptimizeState {
-  GraphOptimizeState(GraphOptimizeResult const &graph_optimize_result,
-                     float runtime);
+  GraphOptimizeState(ParallelComputationGraph const &pcg,
+                     float runtime_with_optimal_mm);
 
-  GraphOptimizeResult graph_optimize_result;
-  float runtime;
+  ParallelComputationGraph pcg;
+  float runtime_with_optimal_mm;
 
   bool operator==(GraphOptimizeState const &other) const;
   bool operator!=(GraphOptimizeState const &other) const;

lib/compiler/include/compiler/unity_algorithm/unity_algorithm.h

@@ -0,0 +1,21 @@
+#ifndef _FLEXFLOW_COMPILER_UNITY_ALGORITHM_H
+#define _FLEXFLOW_COMPILER_UNITY_ALGORITHM_H
+
+#include "compiler/cost_estimator/cost_estimator.h"
+#include "compiler/search_result.dtg.h"
+#include "compiler/unity_algorithm/unity_search_config.dtg.h"
+#include "pcg/machine_specification.dtg.h"
+#include "substitutions/substitution.h"
+
+namespace FlexFlow {
+
+SearchResult graph_optimize(ParallelComputationGraph &pcg,
+                            CostEstimator const &cost_estimator,
+                            MachineSpecification const &resources,
+                            std::vector<Substitution> const &substitutions,
+                            UnitySearchConfig const &search_config,
+                            DeviceType device_type);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/optimizer_config.struct.toml → lib/compiler/include/compiler/unity_algorithm/unity_search_config.struct.toml

@@ -1,5 +1,5 @@
 namespace = "FlexFlow"
-name = "OptimizerConfig"
+name = "UnitySearchConfig"
 features = [
   "eq",
   "hash",

lib/compiler/src/compiler/compiler.cc

@@ -0,0 +1,31 @@
+#include "compiler/compiler.h"
+#include "compiler/unity_algorithm/unity_algorithm.h"
+#include "utils/overload.h"
+
+namespace FlexFlow {
+
+SearchResult optimize(ComputationGraph const &computation_graph,
+                      MachineSpecification const &machine_specification,
+                      CostEstimator const &cost_estimator,
+                      AlgorithmConfig const &search_config) {
+  return search_config.visit<SearchResult>(overload{
+    [&](DataParallelismConfig const &config) -> SearchResult {
+      throw std::runtime_error(
+          "Data parallel search algorithm is not implemented yet");
+    },
+    [&](UnitySearchConfig const &config) {
+      ParallelComputationGraph pcg =
+          parallel_computation_graph_from_computation_graph(computation_graph);
+      std::vector<Substitution> substitutions; // TODO: Implement this
+      return graph_optimize(pcg,
+                            cost_estimator,
+                            machine_specification,
+                            substitutions,
+                            config,
+                            DeviceType::GPU);
+
+    },
+  });
+}
+
+} // namespace FlexFlow

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

@@ -1,7 +1,10 @@
 #include "compiler/machine_mapping/machine_mapping.h"
+#include "compiler/series_parallel/pcg/pcg_binary_sp_decomposition.h"
 #include "utils/containers/are_disjoint.h"
 #include "utils/containers/keys.h"
 #include "utils/containers/merge_maps.h"
+#include "utils/containers/map_keys.h"
+#include "utils/graph/series_parallel/binary_sp_decomposition_tree/binary_sp_decomposition_tree.h"
 
 namespace FlexFlow {
 
@@ -14,4 +17,38 @@ bool nodes_are_disjoint(MachineMapping const &m1, MachineMapping const &m2) {
   return are_disjoint(keys(m1.machine_views), keys(m2.machine_views));
 }
 
+MachineMapping get_machine_mapping_from_machine_mapping_result(
+    PCGBinarySPDecomposition const &sp_decomposition,
+    MachineMappingResult const &mm_result) {
+
+  BinarySPDecompositionTree sp_tree =
+      binary_sp_tree_from_pcg_sp_tree(sp_decomposition);
+
+  auto get_layer_from_path =
+      [&](BinaryTreePath const &path) -> parallel_layer_guid_t {
+    std::optional<BinarySPDecompositionTree> subtree_optional =
+        binary_sp_decomposition_tree_get_subtree_at_path(sp_tree, path);
+    if (!subtree_optional.has_value()) {
+      throw std::runtime_error(fmt::format("Invalid tree path {}", path));
+    }
+    BinarySPDecompositionTree subtree = subtree_optional.value();
+    if (!subtree.is_node()) {
+      throw std::runtime_error(fmt::format("Invalid tree path to a leaf: found {} instead", subtree));
+    }
+    return parallel_layer_guid_t{
+        subtree.require_node(),
+    };
+  };
+
+  std::unordered_map<parallel_layer_guid_t, MachineView> mm;
+
+  if (mm_result.raw_result) {
+    FeasibleMachineMappingResult const &feasible_mm_result =
+        mm_result.raw_result.value();
+    mm = map_keys(feasible_mm_result.machine_mapping.raw_mapping, get_layer_from_path);
+  }
+
+  return MachineMapping{mm};
+}
+
 } // namespace FlexFlow

lib/compiler/src/compiler/allowed_machine_views.cc → lib/compiler/src/compiler/unity_algorithm/allowed_machine_views.cc

@@ -1,4 +1,4 @@
-#include "compiler/allowed_machine_views.h"
+#include "compiler/unity_algorithm/allowed_machine_views.h"
 #include "pcg/machine_specification.h"
 #include "pcg/machine_view.h"
 #include "pcg/multi_dimensional_stride.dtg.h"