Get tests building again

lib/compiler/include/compiler/cost_estimator/cost_estimator.h

@@ -0,0 +1,45 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_COST_ESTIMATOR_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_COST_ESTIMATOR_H
+
+#include <vector>
+#include "compiler/cost_estimator/op_cost_estimate_key.dtg.h"
+#include "compiler/cost_estimator/tensor_set_movement.dtg.h"
+#include "op-attrs/parallel_tensor_shape.dtg.h"
+#include "op-attrs/pcg_operator_attrs.dtg.h"
+#include "pcg/machine_view.dtg.h"
+
+namespace FlexFlow {
+
+struct ICostEstimator {
+  virtual float estimate_cost(OpCostEstimateKey const &) const = 0;
+  virtual float estimate_cost(TensorSetMovement const &) const = 0;
+
+  ICostEstimator() = default;
+  ICostEstimator(ICostEstimator const &) = delete;
+  ICostEstimator &operator=(ICostEstimator const &) = delete;
+
+  virtual ~ICostEstimator() = default;
+};
+CHECK_RC_COPY_VIRTUAL_COMPLIANT(ICostEstimator);
+
+struct CostEstimator {
+  float estimate_cost(OpCostEstimateKey const &k) const;
+  float estimate_cost(TensorSetMovement const &m) const;
+
+  template <typename T, typename... Args>
+  static typename std::enable_if<std::is_base_of<ICostEstimator, T>::value,
+                                 CostEstimator>::type
+      create(Args &&...args) {
+    return CostEstimator(std::make_shared<T>(std::forward<Args>(args)...));
+  }
+
+private:
+  CostEstimator(std::shared_ptr<ICostEstimator> implementation_ptr);
+
+private:
+  std::shared_ptr<ICostEstimator> implementation_ptr;
+};
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/comm_cost_estimate_key.struct.toml → lib/compiler/include/compiler/cost_estimator/single_tensor_movement.struct.toml

@@ -1,25 +1,30 @@
 namespace = "FlexFlow"
-name = "CommCostEstimateKey"
+name = "SingleTensorMovement"
 features = [
   "eq",
-  "ord",
   "hash",
   "fmt",
 ]
 
 includes = [
   "op-attrs/parallel_tensor_shape.dtg.h",
   "pcg/machine_view.dtg.h",
+  "<unordered_set>",
+]
+
+src_includes = [
+  "utils/hash/unordered_set.h",
+  "utils/fmt/unordered_set.h",
 ]
 
 [[fields]]
 name = "parallel_tensor_shape"
 type = "::FlexFlow::ParallelTensorShape"
 
 [[fields]]
-name = "src_machine_view"
-type = "::FlexFlow::MachineView"
+name = "src_machine_views"
+type = "std::unordered_set<::FlexFlow::MachineView>"
 
 [[fields]]
 name = "dst_machine_view"
-type = "::FlexFlow::MachineView"
+type = "std::unordered_set<::FlexFlow::MachineView>"

lib/compiler/include/compiler/cost_estimator/tensor_set_movement.struct.toml

@@ -0,0 +1,21 @@
+namespace = "FlexFlow"
+name = "TensorSetMovement"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+]
+
+includes = [
+  "compiler/cost_estimator/single_tensor_movement.dtg.h",
+  "<unordered_set>",
+]
+
+src_includes = [
+  "utils/fmt/unordered_multiset.h",
+  "utils/hash/unordered_multiset.h",
+]
+
+[[fields]]
+name = "single_tensor_movements"
+type = "std::unordered_multiset<::FlexFlow::SingleTensorMovement>"

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

@@ -1,7 +1,7 @@
 #ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_ESTIMATE_LAYER_COST_H
 #define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_ESTIMATE_LAYER_COST_H
 
-#include "compiler/cost_estimator.h"
+#include "compiler/cost_estimator/cost_estimator.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
 #include "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h"
 namespace FlexFlow {

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

@@ -0,0 +1,18 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_ESTIMATE_COST_ACROSS_SPLIT_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_ESTIMATE_COST_ACROSS_SPLIT_H
+
+#include "compiler/cost_estimator/tensor_set_movement.dtg.h"
+#include "compiler/machine_mapping/partial_machine_mapping.dtg.h"
+#include "compiler/machine_mapping/transitive_reduced_pcg.dtg.h"
+#include "compiler/series_parallel/pcg_binary_series_split.dtg.h"
+
+namespace FlexFlow {
+
+TensorSetMovement get_tensor_set_movement_across_split(TransitiveReducedPCG const &transitive_reduced_pcg,
+                                                       PCGBinarySeriesSplit const &split,
+                                                       PartialMachineMapping const &pre_mapping,
+                                                       PartialMachineMapping const &post_mapping);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/machine_mapping/include_unconstrained.struct.toml

@@ -0,0 +1,16 @@
+namespace = "FlexFlow"
+name = "IncludeUnconstrained"
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "fmt",
+  "rapidcheck",
+  "json",
+]
+
+includes = []
+
+[[fields]]
+name = "raw_bool"
+type = "bool"

lib/compiler/include/compiler/machine_mapping/machine_mapping_context.struct.toml

@@ -3,7 +3,7 @@ name = "MachineMappingContext"
 features = []
 
 includes = [
-  "compiler/cost_estimator.h",
+  "compiler/cost_estimator/cost_estimator.h",
   "pcg/machine_view.dtg.h",
   "pcg/machine_specification.dtg.h",
   "compiler/machine_mapping/transitive_reduced_pcg.dtg.h",

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

@@ -1,25 +1,31 @@
 #ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_PARTIAL_MACHINE_MAPPING_H
 #define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_PARTIAL_MACHINE_MAPPING_H
 
+#include "compiler/machine_mapping/machine_mapping.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_context.dtg.h"
 #include "compiler/machine_mapping/partial_machine_mapping.dtg.h"
+#include "compiler/machine_mapping/include_unconstrained.dtg.h"
+#include "compiler/series_parallel/pcg_binary_sp_decomposition.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_tensor_guid_t.dtg.h"
 
 namespace FlexFlow {
 
-PartialMachineMapping get_unconstrained_solution();
+PartialMachineMapping get_unconstrained_solution_for_layers(std::unordered_set<parallel_layer_guid_t> const &);
 
-PartialMachineMapping get_sub_solution(MachineMappingContext const &ctx, 
-                                PartialMachineMapping const &partial_solution, 
-                                PCGBinarySPDecomposition const &sub_problem);
+std::unordered_set<parallel_layer_guid_t> get_all_layers(PartialMachineMapping const &, 
+                                                         IncludeUnconstrained const &);
 
-PartialMachineMapping with_additional_tensor_machine_views(MachineMappingContext const &ctx,
-                                                           PartialMachineMapping const &partial_solution,
-                                                           std::unordered_map<parallel_tensor_guid_t, MachineView> const &additional);
+std::optional<MachineView> get_machine_view_for_layer(PartialMachineMapping const &,
+                                                      parallel_layer_guid_t const &);
 
-PartialMachineMapping with_additional_layer_machine_views(MachineMappingContext const &ctx,
-                                                          PartialMachineMapping const &partial_solution,
+PartialMachineMapping get_sub_solution(PartialMachineMapping const &partial_solution, 
+                                       PCGBinarySPDecomposition const &sub_problem);
+
+PartialMachineMapping with_additional_layer_machine_views(PartialMachineMapping const &partial_solution,
                                                           std::unordered_map<parallel_layer_guid_t, MachineView> const &additional);
 
+MachineMapping require_complete_mapping(PartialMachineMapping const &);
+
 } // namespace FlexFlow
 
 #endif

lib/compiler/include/compiler/machine_mapping/partial_machine_mapping.struct.toml

@@ -9,13 +9,15 @@ features = [
 includes = [
   "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h",
   "pcg/machine_view.dtg.h",
+  "<optional>",
 ]  
 
 src_includes = [   
   "utils/hash/unordered_map.h",
   "utils/fmt/unordered_map.h",
+  "utils/fmt/optional.h",
 ]
 
 [[fields]]
 name = "machine_views"
-type = "std::unordered_map<::FlexFlow::parallel_layer_guid_t, ::FlexFlow::MachineView>"
+type = "std::unordered_map<::FlexFlow::parallel_layer_guid_t, std::optional<::FlexFlow::MachineView>>"

lib/compiler/include/compiler/series_parallel/pcg_binary_sp_decomposition.h

@@ -17,9 +17,9 @@ std::unordered_multiset<parallel_layer_guid_t>
 
 SPDecompositionTreeNodeType get_node_type(PCGBinarySPDecomposition const &);
 
-PCGBinarySeriesSplit make_pcg_series_split(PCGBinarySPDecomposition const &, PCGBinarySPDecomposition const &);
-PCGBinarySeriesSplit make_pcg_parallel_split(PCGBinarySPDecomposition const &, PCGBinarySPDecomposition const &);
-PCGBinarySeriesSplit make_pcg_leaf_node(parallel_layer_guid_t const &);
+PCGBinarySPDecomposition make_pcg_series_split(PCGBinarySPDecomposition const &, PCGBinarySPDecomposition const &);
+PCGBinarySPDecomposition make_pcg_parallel_split(PCGBinarySPDecomposition const &, PCGBinarySPDecomposition const &);
+PCGBinarySPDecomposition make_pcg_leaf_node(parallel_layer_guid_t const &);
 
 PCGBinarySeriesSplit require_series(PCGBinarySPDecomposition const &);
 PCGBinaryParallelSplit require_parallel(PCGBinarySPDecomposition const &);

lib/compiler/include/compiler/unity_algorithm.h

@@ -2,7 +2,7 @@
 #define _FLEXFLOW_COMPILER_UNITY_ALGORITHM_H
 
 #include "compiler/graph_optimize_result.dtg.h"
-#include "cost_estimator.h"
+#include "compiler/cost_estimator/cost_estimator.h"
 #include "optimizer_config.dtg.h"
 #include "pcg/computation_graph.h"
 #include "pcg/machine_specification.dtg.h"

lib/compiler/src/compiler/cost_estimator/cost_estimator.cc

@@ -0,0 +1,16 @@
+#include "compiler/cost_estimator/cost_estimator.h"
+
+namespace FlexFlow {
+
+CostEstimator::CostEstimator(std::shared_ptr<ICostEstimator> implementation_ptr)
+    : implementation_ptr(implementation_ptr) {}
+
+float CostEstimator::estimate_cost(OpCostEstimateKey const &k) const {
+  return this->implementation_ptr->estimate_cost(k);
+}
+
+float CostEstimator::estimate_cost(TensorSetMovement const &m) const {
+  return this->implementation_ptr->estimate_cost(m);
+}
+
+} // namespace FlexFlow

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

@@ -17,11 +17,15 @@ float estimate_layer_cost(ParallelComputationGraph const &pcg,
   std::vector<parallel_tensor_guid_t> weight_tensors = get_incoming_weights(pcg, layer);
   std::vector<parallel_tensor_guid_t> output_tensors = get_layer_outputs(pcg, layer);
 
-  return cost_estimator.estimate_cost(op_attrs,
-                                      transform(input_tensors, get_tensor_shape),
-                                      transform(weight_tensors, get_tensor_shape),
-                                      transform(output_tensors, get_tensor_shape),
-                                      machine_view);
+  OpCostEstimateKey key = OpCostEstimateKey{
+    /*op_attrs=*/op_attrs,
+    /*input_shapes=*/transform(input_tensors, get_tensor_shape),
+    /*weight_shapes=*/transform(weight_tensors, get_tensor_shape),
+    /*output_shapes=*/transform(output_tensors, get_tensor_shape),
+    /*machine_view=*/machine_view,
+  };
+
+  return cost_estimator.estimate_cost(key);
 }
 
 } // namespace FlexFlow