A lot of simplifying and modularizing of unity dp code

wmdi · Sep 26, 2024 · 3a35951 · 3a35951
1 parent a112225
commit 3a35951
Show file tree

Hide file tree

Showing 105 changed files with 2,370 additions and 939 deletions.
diff --git a/.proj.toml b/.proj.toml
@@ -22,11 +22,11 @@ test_targets = [
   "utils-tests",
   "op-attrs-tests",
   "pcg-tests",
-  "substitutions-tests",
+  # "substitutions-tests",
   "compiler-tests",
-  "substitution-generator-tests",
-  "local-execution-tests",
-  "models-tests",
+  # "substitution-generator-tests",
+  # "local-execution-tests",
+  # "models-tests",
 ]
 
 [cmake_flags_extra]

diff --git a/lib/compiler/include/compiler/comm_cost_estimate_key.struct.toml b/lib/compiler/include/compiler/comm_cost_estimate_key.struct.toml
@@ -0,0 +1,25 @@
+namespace = "FlexFlow"
+name = "CommCostEstimateKey"
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "fmt",
+]
+
+includes = [
+  "op-attrs/parallel_tensor_shape.dtg.h",
+  "pcg/machine_view.dtg.h",
+]
+
+[[fields]]
+name = "parallel_tensor_shape"
+type = "::FlexFlow::ParallelTensorShape"
+
+[[fields]]
+name = "src_machine_view"
+type = "::FlexFlow::MachineView"
+
+[[fields]]
+name = "dst_machine_view"
+type = "::FlexFlow::MachineView"
diff --git a/lib/compiler/include/compiler/cost_estimator.h b/lib/compiler/include/compiler/cost_estimator.h
@@ -1,19 +1,18 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_H
 
-#ifndef _FLEXFLOW_COMPILER_COST_ESTIMATE_H
-#define _FLEXFLOW_COMPILER_COST_ESTIMATE_H
-
-#include "op-attrs/operator_attrs.h"
-#include "op-attrs/parallel_tensor_shape.h"
-#include "pcg/machine_view.h"
-#include "pcg/parallel_computation_graph/parallel_tensor_attrs.dtg.h"
+#include <vector>
+#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(PCGOperatorAttrs const &op,
                               std::vector<ParallelTensorShape> const &inputs,
-                              std::vector<ParallelTensorAttrs> const &weights,
-                              std::vector<ParallelTensorAttrs> const &outputs,
+                              std::vector<ParallelTensorShape> const &weights,
+                              std::vector<ParallelTensorShape> const &outputs,
                               MachineView const &mv) const = 0;
   virtual float estimate_cost(ParallelTensorShape const &tensor_shape,
                               MachineView const &src,
@@ -30,8 +29,8 @@ CHECK_RC_COPY_VIRTUAL_COMPLIANT(ICostEstimator);
 struct CostEstimator {
   float estimate_cost(PCGOperatorAttrs const &op,
                       std::vector<ParallelTensorShape> const &inputs,
-                      std::vector<ParallelTensorAttrs> const &weights,
-                      std::vector<ParallelTensorAttrs> const &outputs,
+                      std::vector<ParallelTensorShape> const &weights,
+                      std::vector<ParallelTensorShape> const &outputs,
                       MachineView const &mv) const {
     return this->implementation_ptr->estimate_cost(
         op, inputs, weights, outputs, mv);

diff --git a/lib/compiler/include/compiler/machine_mapping/estimate_cost_across_split.h b/lib/compiler/include/compiler/machine_mapping/estimate_cost_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.h"
+#include "compiler/machine_mapping/transitive_reduced_pcg.dtg.h"
+#include "pcg/machine_view.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h"
+
+namespace FlexFlow {
+
+float estimate_cost_across_split(TransitiveReducedPCG const &,
+                                 CostEstimator const &,
+                                 std::unordered_map<parallel_layer_guid_t, MachineView> const &,
+                                 std::unordered_map<parallel_layer_guid_t, MachineView> const &);
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/compiler/include/compiler/machine_mapping/estimate_layer_cost.h b/lib/compiler/include/compiler/machine_mapping/estimate_layer_cost.h
@@ -0,0 +1,16 @@
+#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 "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h"
+namespace FlexFlow {
+
+float estimate_layer_cost(ParallelComputationGraph const &pcg,
+                          CostEstimator const &cost_estimator,
+                          parallel_layer_guid_t const &layer,
+                          MachineView const &machine_view);
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/compiler/include/compiler/machine_mapping/get_allowed_machine_views_list.h b/lib/compiler/include/compiler/machine_mapping/get_allowed_machine_views_list.h
@@ -3,6 +3,8 @@
 
 #include "compiler/machine_mapping/machine_mapping_context.dtg.h"
 #include "pcg/machine_specification.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_tensor_guid_t.dtg.h"
 #include <vector>
 
 namespace FlexFlow {

diff --git a/lib/compiler/include/compiler/machine_mapping/get_machine_resource_splits.h b/lib/compiler/include/compiler/machine_mapping/get_machine_resource_splits.h
@@ -0,0 +1,15 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_GET_MACHINE_RESOURCE_SPLITS_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_GET_MACHINE_RESOURCE_SPLITS_H
+
+#include "pcg/machine_specification.dtg.h"
+#include <utility>
+#include <unordered_set>
+
+namespace FlexFlow {
+
+std::unordered_set<std::pair<MachineSpecification, MachineSpecification>>
+    get_machine_resource_splits(MachineSpecification const &resource);
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/compiler/include/compiler/machine_mapping/get_optimal_machine_mapping.h b/lib/compiler/include/compiler/machine_mapping/get_optimal_machine_mapping.h
@@ -4,6 +4,9 @@
 #include "compiler/machine_mapping/machine_mapping.h"
 #include "compiler/machine_mapping/machine_mapping_cache.h"
 #include "compiler/machine_mapping/machine_mapping_context.dtg.h"
+#include "compiler/machine_mapping/partial_machine_mapping.dtg.h"
+#include "compiler/series_parallel/pcg_binary_parallel_split.dtg.h"
+#include "compiler/series_parallel/pcg_binary_series_split.dtg.h"
 #include "pcg/machine_specification.h"
 #include "pcg/machine_view.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph.h"
@@ -21,36 +24,37 @@ MachineMappingResult get_optimal_machine_mapping(
     MachineMappingCache &cached_subgraph_results);
 
 MachineMappingResult
-    get_optimal_machine_mapping_internal(MachineMappingContext &context,
+    get_optimal_machine_mapping_internal(MachineMappingCache &result_cache,
+                                         MachineMappingContext const &context,
                                          MachineSpecification const &resources);
 
 MachineMappingResult get_optimal_machine_mapping_internal(
-    MachineMappingContext &context,
-    SerialParallelDecomposition const &decompn,
-    MachineSpecification const &resource,
-    std::unordered_map<parallel_tensor_guid_t, MachineView> const
-        &fixed_machine_views);
+    MachineMappingCache &result_cache,
+    MachineMappingContext const &context,
+    PCGBinarySPDecomposition const &sp_decomposition,
+    MachineSpecification const &resources,
+    PartialMachineMapping const &);
 
 MachineMappingResult get_optimal_machine_mapping_internal(
-    MachineMappingContext &context,
-    SerialSplit const &serial,
-    MachineSpecification const &resource,
-    std::unordered_map<parallel_tensor_guid_t, MachineView> const
-        &fixed_machine_views);
+    MachineMappingCache &result_cache,
+    MachineMappingContext const &context,
+    PCGBinarySeriesSplit const &series,
+    MachineSpecification const &resources,
+    PartialMachineMapping const &);
 
 MachineMappingResult get_optimal_machine_mapping_internal(
-    MachineMappingContext &context,
-    ParallelSplit const &parallel,
-    MachineSpecification const &resource,
-    std::unordered_map<parallel_tensor_guid_t, MachineView> const
-        &fixed_machine_views);
+    MachineMappingCache &result_cache,
+    MachineMappingContext const &context,
+    PCGBinaryParallelSplit const &parallel,
+    MachineSpecification const &resources,
+    PartialMachineMapping const &);
 
 MachineMappingResult get_optimal_machine_mapping_internal(
-    MachineMappingContext &context,
-    Node const &node,
-    MachineSpecification const &resource,
-    std::unordered_map<parallel_tensor_guid_t, MachineView> const
-        &fixed_machine_views);
+    MachineMappingCache &result_cache,
+    MachineMappingContext const &,
+    parallel_layer_guid_t const &,
+    MachineSpecification const &,
+    PartialMachineMapping const &);
 
 } // namespace FlexFlow
 

diff --git a/lib/compiler/include/compiler/machine_mapping/machine_mapping.struct.toml b/lib/compiler/include/compiler/machine_mapping/machine_mapping.struct.toml
@@ -21,4 +21,4 @@ src_includes = [
 
 [[fields]]
 name = "machine_views"
-type = "std::unordered_map<::FlexFlow::parallel_layer_guid_t, ::FlexFlow::MachineView>"
+type = "std::unordered_map<::FlexFlow::parallel_layer_guid_t, ::FlexFlow::MachineView>"
diff --git a/lib/compiler/include/compiler/machine_mapping/machine_mapping_context.h b/lib/compiler/include/compiler/machine_mapping/machine_mapping_context.h
@@ -0,0 +1,23 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_MACHINE_MAPPING_CONTEXT_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_MACHINE_MAPPING_CONTEXT_H
+
+#include "compiler/machine_mapping/machine_mapping_context.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph_edge.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_tensor_guid_t.dtg.h"
+
+namespace FlexFlow {
+
+std::unordered_set<MachineView> get_allowed_machine_views_for_tensor(MachineMappingContext const &,
+                                                                     parallel_tensor_guid_t const &);
+std::unordered_set<MachineView> get_allowed_machine_views_for_layer(MachineMappingContext const &,
+                                                                    parallel_layer_guid_t const &);
+
+MachineMappingContext make_machine_mapping_context(ParallelComputationGraph const &pcg,
+                                                   CostEstimator const &cost_estimator,
+                                                   std::function<std::unordered_set<MachineView>(
+                                                     ParallelLayerAttrs const &, MachineSpecification const &)> const &allowed_machine_views);
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/compiler/include/compiler/machine_mapping/machine_mapping_context.struct.toml b/lib/compiler/include/compiler/machine_mapping/machine_mapping_context.struct.toml
@@ -1,20 +1,17 @@
 namespace = "FlexFlow"
 name = "MachineMappingContext"
-features = [
-]
+features = []
 
 includes = [
-  "compiler/machine_mapping/machine_mapping.dtg.h",
-  "pcg/parallel_computation_graph/parallel_computation_graph.h",
   "compiler/cost_estimator.h",
-  "pcg/machine_view.h",
+  "pcg/machine_view.dtg.h",
   "pcg/machine_specification.dtg.h",
-  "compiler/machine_mapping/machine_mapping_cache.h"
+  "compiler/machine_mapping/transitive_reduced_pcg.dtg.h",
 ]
 
 [[fields]]
-name = "pcg"
-type = "::FlexFlow::ParallelComputationGraph"
+name = "transitive_reduced_pcg"
+type = "::FlexFlow::TransitiveReducedPCG"
 
 [[fields]]
 name = "cost_estimator"
@@ -23,7 +20,3 @@ type = "::FlexFlow::CostEstimator"
 [[fields]]
 name = "allowed_machine_views"
 type = "std::function<std::unordered_set<::FlexFlow::MachineView>(::FlexFlow::ParallelLayerAttrs const &, ::FlexFlow::MachineSpecification const &)>"
-
-[[fields]]
-name = "cached_subgraph_results"
-type = "::FlexFlow::MachineMappingCache"
diff --git a/lib/compiler/include/compiler/machine_mapping/machine_mapping_result.h b/lib/compiler/include/compiler/machine_mapping/machine_mapping_result.h
@@ -6,6 +6,7 @@
 namespace FlexFlow {
 
 MachineMappingResult sequential_combine(MachineMappingResult const &s1,
+                                        float comm_cost,
                                         MachineMappingResult const &s2);
 MachineMappingResult parallel_combine(MachineMappingResult const &s1,
                                       MachineMappingResult const &s2);

diff --git a/lib/compiler/include/compiler/machine_mapping/machine_mapping_result.struct.toml b/lib/compiler/include/compiler/machine_mapping/machine_mapping_result.struct.toml
@@ -2,6 +2,7 @@ namespace = "FlexFlow"
 name = "MachineMappingResult"
 features = [
   "eq",
+  "hash",
   "fmt",
 ]
 

diff --git a/lib/compiler/include/compiler/machine_mapping/machine_mapping_state.struct.toml b/lib/compiler/include/compiler/machine_mapping/machine_mapping_state.struct.toml
@@ -7,25 +7,19 @@ features = [
 ]
 
 includes = [
-  "utils/graph/serial_parallel/serial_parallel_decomposition.dtg.h",
   "pcg/machine_specification.dtg.h",
-  "pcg/machine_view.dtg.h",
-  "pcg/parallel_computation_graph/parallel_tensor_guid_t.dtg.h",
-]
-
-src_includes = [
-  "utils/hash/unordered_map.h",
-  "utils/fmt/unordered_map.h",
+  "compiler/series_parallel/pcg_binary_sp_decomposition.dtg.h",
+  "compiler/machine_mapping/partial_machine_mapping.dtg.h",
 ]
 
 [[fields]]
 name = "subgraph"
-type = "::FlexFlow::SerialParallelDecomposition"
+type = "::FlexFlow::PCGBinarySPDecomposition"
 
 [[fields]]
 name = "resource"
 type = "::FlexFlow::MachineSpecification"
 
 [[fields]]
-name = "fixed_machine_views"
-type = "std::unordered_map<::FlexFlow::parallel_tensor_guid_t, ::FlexFlow::MachineView>"
+name = "partial_solution"
+type = "::FlexFlow::PartialMachineMapping"
diff --git a/lib/compiler/include/compiler/machine_mapping/partial_machine_mapping.h b/lib/compiler/include/compiler/machine_mapping/partial_machine_mapping.h
@@ -0,0 +1,25 @@
+#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_context.dtg.h"
+#include "compiler/machine_mapping/partial_machine_mapping.dtg.h"
+
+namespace FlexFlow {
+
+PartialMachineMapping get_unconstrained_solution();
+
+PartialMachineMapping get_sub_solution(MachineMappingContext const &ctx, 
+                                PartialMachineMapping const &partial_solution, 
+                                PCGBinarySPDecomposition const &sub_problem);
+
+PartialMachineMapping with_additional_tensor_machine_views(MachineMappingContext const &ctx,
+                                                           PartialMachineMapping const &partial_solution,
+                                                           std::unordered_map<parallel_tensor_guid_t, MachineView> const &additional);
+
+PartialMachineMapping with_additional_layer_machine_views(MachineMappingContext const &ctx,
+                                                          PartialMachineMapping const &partial_solution,
+                                                          std::unordered_map<parallel_layer_guid_t, MachineView> const &additional);
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/compiler/include/compiler/machine_mapping/partial_machine_mapping.struct.toml b/lib/compiler/include/compiler/machine_mapping/partial_machine_mapping.struct.toml
@@ -0,0 +1,21 @@
+namespace = "FlexFlow"
+name = "PartialMachineMapping"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+]
+
+includes = [
+  "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h",
+  "pcg/machine_view.dtg.h",
+]  
+
+src_includes = [   
+  "utils/hash/unordered_map.h",
+  "utils/fmt/unordered_map.h",
+]
+
+[[fields]]
+name = "machine_views"
+type = "std::unordered_map<::FlexFlow::parallel_layer_guid_t, ::FlexFlow::MachineView>"
diff --git a/lib/compiler/include/compiler/machine_mapping/split_sp_decomposition.h b/lib/compiler/include/compiler/machine_mapping/split_sp_decomposition.h