renaming

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

@@ -1,7 +1,7 @@
 #ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_COST_ESTIMATOR_H
 #define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_COST_ESTIMATOR_H
 
-#include "compiler/cost_estimator/cost_metric.dtg.h"
+#include "compiler/cost_estimator/op_cost_metrics.dtg.h"
 #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"
@@ -12,10 +12,8 @@
 namespace FlexFlow {
 
 struct ICostEstimator {
-  virtual float estimate_cost(OpCostEstimateKey const &) const = 0;
+  virtual OpCostMetrics estimate_cost(OpCostEstimateKey const &) const = 0;
   virtual float estimate_cost(TensorSetMovement const &) const = 0;
-  virtual CostMetric
-      estimate_cost_with_memory(OpCostEstimateKey const &) const = 0;
 
   ICostEstimator() = default;
   ICostEstimator(ICostEstimator const &) = delete;
@@ -26,9 +24,8 @@ struct ICostEstimator {
 CHECK_RC_COPY_VIRTUAL_COMPLIANT(ICostEstimator);
 
 struct CostEstimator {
-  float estimate_cost(OpCostEstimateKey const &k) const;
+  OpCostMetrics estimate_cost(OpCostEstimateKey const &) const;
   float estimate_cost(TensorSetMovement const &m) const;
-  CostMetric estimate_cost_with_memory(OpCostEstimateKey const &k) const;
 
   template <typename T, typename... Args>
   static typename std::enable_if<std::is_base_of<ICostEstimator, T>::value,

lib/compiler/include/compiler/cost_estimator/cost_metric.struct.toml → lib/compiler/include/compiler/cost_estimator/op_cost_metrics.struct.toml

@@ -1,5 +1,5 @@
 namespace = "FlexFlow"
-name = "CostMetric"
+name = "OpCostMetrics"
 features = [
   "eq",
   "fmt",

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

@@ -2,7 +2,6 @@
 #define _FLEXFLOW_COMPILER_MACHINE_MAPPING_GET_OPTIMAL_MACHINE_MAPPING_H
 
 #include "compiler/machine_mapping/machine_mapping_cache.dtg.h"
-#include "compiler/machine_mapping/machine_mapping_config.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_constraints.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_context.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_problem_tree/machine_mapping_problem_tree.dtg.h"

lib/compiler/include/compiler/machine_mapping/memory_optimization/get_optimal_machine_mapping_with_memory.h

@@ -2,43 +2,42 @@
 #define _FLEXFLOW_COMPILER_MACHINE_MAPPING_MEMORY_OPTIMIZATION_GET_OPTIMAL_MACHINE_MAPPING_WITH_MEMORY_H
 
 #include "compiler/machine_mapping/machine_mapping_cache.dtg.h"
-#include "compiler/machine_mapping/machine_mapping_config.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_constraints.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_context.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_problem_tree/machine_mapping_problem_tree.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_problem_tree/mm_problem_tree_parallel_split.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_problem_tree/mm_problem_tree_series_split.dtg.h"
-#include "compiler/machine_mapping/memory_optimization/machine_mapping_cache_with_memory.dtg.h"
+#include "compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_cache.dtg.h"
 #include "compiler/machine_mapping/parallel_split_transformation.dtg.h"
 #include "pcg/machine_specification.dtg.h"
 
 namespace FlexFlow {
 
-MachineMappingResultWithMemory get_optimal_machine_mapping_with_memory(
-    MachineMappingCacheWithMemory &result_cache,
+MachineMappingWithMemoryResult get_optimal_machine_mapping_with_memory(
+    MachineMappingWithMemoryCache &result_cache,
     MachineMappingContext const &context,
     MachineMappingProblemTree const &problem_tree,
     MachineSpecification const &resources,
     MachineMappingConstraints const &constraints);
 
-MachineMappingResultWithMemory get_optimal_machine_mapping_with_memory(
-    MachineMappingCacheWithMemory &result_cache,
+MachineMappingWithMemoryResult get_optimal_machine_mapping_with_memory(
+    MachineMappingWithMemoryCache &result_cache,
     MachineMappingContext const &context,
     MMProblemTreeSeriesSplit const &series_split,
     MachineSpecification const &resources,
     MachineMappingConstraints const &constraints,
     std::optional<ParallelSplitTransformation> const
         &parallel_split_transformation);
 
-MachineMappingResultWithMemory get_optimal_machine_mapping_with_memory(
-    MachineMappingCacheWithMemory &result_cache,
+MachineMappingWithMemoryResult get_optimal_machine_mapping_with_memory(
+    MachineMappingWithMemoryCache &result_cache,
     MachineMappingContext const &context,
     MMProblemTreeParallelSplit const &parallel_split,
     MachineSpecification const &resources,
     MachineMappingConstraints const &constraints);
 
-MachineMappingResultWithMemory get_optimal_machine_mapping_with_memory(
-    MachineMappingCacheWithMemory &result_cache,
+MachineMappingWithMemoryResult get_optimal_machine_mapping_with_memory(
+    MachineMappingWithMemoryCache &result_cache,
     MachineMappingContext const &,
     UnmappedOpCostEstimateKey const &leaf,
     MachineSpecification const &resources,

lib/compiler/include/compiler/machine_mapping/memory_optimization/machine_mapping_cache_with_memory.h → lib/compiler/include/compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_cache.h

@@ -1,18 +1,18 @@
 #ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_MEMORY_OPTIMIZATION_MACHINE_MAPPING_CACHE_WITH_MEMORY_H
 #define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_MEMORY_OPTIMIZATION_MACHINE_MAPPING_CACHE_WITH_MEMORY_H
 
-#include "compiler/machine_mapping/memory_optimization/machine_mapping_cache_with_memory.dtg.h"
+#include "compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_cache.dtg.h"
 
 namespace FlexFlow {
 
-MachineMappingCacheWithMemory empty_machine_mapping_cache_with_memory();
-std::optional<MachineMappingResultWithMemory>
-    machine_mapping_cache_with_memory_load(
-        MachineMappingCacheWithMemory const &, MachineMappingState const &);
-void machine_mapping_cache_with_memory_save(
-    MachineMappingCacheWithMemory &,
+MachineMappingWithMemoryCache empty_machine_mapping_with_memory_cache();
+std::optional<MachineMappingWithMemoryResult>
+    machine_mapping_with_memory_cache_load(
+        MachineMappingWithMemoryCache const &, MachineMappingState const &);
+void machine_mapping_with_memory_cache_save(
+    MachineMappingWithMemoryCache &,
     MachineMappingState const &,
-    MachineMappingResultWithMemory const &);
+    MachineMappingWithMemoryResult const &);
 
 } // namespace FlexFlow
 

lib/compiler/include/compiler/machine_mapping/memory_optimization/machine_mapping_cache_with_memory.struct.toml → lib/compiler/include/compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_cache.struct.toml

@@ -1,5 +1,5 @@
 namespace = "FlexFlow"
-name = "MachineMappingCacheWithMemory"
+name = "MachineMappingWithMemoryCache"
 features = [
   "eq",
   "hash",
@@ -9,7 +9,7 @@ features = [
 includes = [
   "<unordered_map>",
   "compiler/machine_mapping/machine_mapping_state.dtg.h",
-  "compiler/machine_mapping/memory_optimization/machine_mapping_result_with_memory.dtg.h",
+  "compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_result.dtg.h",
 ]
 
 src_includes = [
@@ -19,4 +19,4 @@ src_includes = [
 
 [[fields]]
 name = "raw_map"
-type = "std::unordered_map<::FlexFlow::MachineMappingState, ::FlexFlow::MachineMappingResultWithMemory>"
+type = "std::unordered_map<::FlexFlow::MachineMappingState, ::FlexFlow::MachineMappingWithMemoryResult>"

lib/compiler/include/compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_result.h

@@ -0,0 +1,41 @@
+#ifndef _FLEXFLOW_COMPILER_MACHINE_MAPPING_MEMORY_OPTIMIZATION_MACHINE_MAPPING_RESULT_WITH_MEMORY_H
+#define _FLEXFLOW_COMPILER_MACHINE_MAPPING_MEMORY_OPTIMIZATION_MACHINE_MAPPING_RESULT_WITH_MEMORY_H
+
+#include "compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_result.dtg.h"
+#include "compiler/machine_mapping/parallel_split_transformation.dtg.h"
+#include <optional>
+
+namespace FlexFlow {
+
+[[nodiscard]] MachineMappingWithMemoryResult
+    empty_machine_mapping_with_memory_result();
+[[nodiscard]] bool is_empty(MachineMappingWithMemoryResult const &);
+
+[[nodiscard]] MachineMappingWithMemoryResult get_mapping_with_minimal_runtime(
+    std::unordered_set<MachineMappingWithMemoryResult> const &);
+
+[[nodiscard]] MachineMappingWithMemoryResult
+    remove_non_pareto_optimal_machine_mapping_result(
+        MachineMappingWithMemoryResult const &);
+
+[[nodiscard]] MachineMappingWithMemoryResult
+    series_combine(float comm_cost,
+                   MachineMappingWithMemoryResult const &pre_result,
+                   MachineMappingWithMemoryResult const &post_result,
+                   std::optional<ParallelSplitTransformation> const
+                       &parallel_split_transformation);
+[[nodiscard]] MachineMappingWithMemoryResult
+    parallel_combine(MachineMappingWithMemoryResult const &lhs_result,
+                     MachineMappingWithMemoryResult const &rhs_result);
+
+[[nodiscard]] MachineMappingWithMemoryResult
+    minimize_runtime(MachineMappingWithMemoryResult const &m1,
+                     MachineMappingWithMemoryResult const &m2);
+
+[[nodiscard]] MachineMappingWithMemoryResult
+    make_singleton_machine_mapping_with_memory_result(
+        OpCostMetrics cost, MachineView const &machine_view);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/machine_mapping/memory_optimization/machine_mapping_result_with_memory.struct.toml → lib/compiler/include/compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_result.struct.toml

@@ -1,5 +1,5 @@
 namespace = "FlexFlow"
-name = "MachineMappingResultWithMemory"
+name = "MachineMappingWithMemoryResult"
 features = [
   "eq",
   "hash",

lib/compiler/include/compiler/machine_mapping/memory_optimization/single_machine_mapping.struct.toml

@@ -8,12 +8,12 @@ features = [
 
 includes = [
   "compiler/machine_mapping/parallel_layer_guid_oblivious_machine_mapping.dtg.h",
-  "compiler/cost_estimator/cost_metric.dtg.h",
+  "compiler/cost_estimator/op_cost_metrics.dtg.h",
 ]
 
 [[fields]]
 name = "cost"
-type = "::FlexFlow::CostMetric"
+type = "::FlexFlow::OpCostMetrics"
 
 [[fields]]
 name = "machine_mapping"

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

@@ -5,17 +5,12 @@ namespace FlexFlow {
 CostEstimator::CostEstimator(std::shared_ptr<ICostEstimator> implementation_ptr)
     : implementation_ptr(implementation_ptr) {}
 
-float CostEstimator::estimate_cost(OpCostEstimateKey const &k) const {
+OpCostMetrics 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);
 }
 
-CostMetric
-    CostEstimator::estimate_cost_with_memory(OpCostEstimateKey const &k) const {
-  return this->implementation_ptr->estimate_cost_with_memory(k);
-}
-
 } // namespace FlexFlow

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

@@ -240,7 +240,7 @@ MachineMappingResult
   auto get_mapping_result = [&](MachineView const &machine_view) {
     OpCostEstimateKey mapped =
         map_unmapped_op_cost_estimate_key(leaf, machine_view);
-    float cost = context.cost_estimator.estimate_cost(mapped);
+    float cost = context.cost_estimator.estimate_cost(mapped).runtime;
 
     return make_singleton_machine_mapping_result(cost, machine_view);
   };