Unity device mapping algorithm (#1459)

* pass existing tests * unity algorithm builds * fmt * fix * refactor machine mapping * add unit tests * fmt * add more tests * fmt * fix * refactor get_optimal_machine_mapping a bit and improve the tests * remove debug codes * A lot of simplifying and modularizing of unity dp code * Get tests building again * Get all the new testcases working * Move over to ProblemTree/ResultTree framework for machine mapping * Settle on ProblemTree/BinaryTreePath-indexed-MachineMappingResult for machine mapping * More code cleanup and PR prep * Get tests building again * Pass some basic tests of get_optimal_machine_mapping * Migrate over to use type-erased binary tree * Move back to templated FullBinaryTree * Get all existing tests passing again * Fix tests and format * Move graph_optimize_state.cc to correct location --------- Co-authored-by: Mengdi Wu <[email protected]> Co-authored-by: Colin Unger <[email protected]> Co-authored-by: Colin Unger <[email protected]>
flexflow · Oct 8, 2024 · b77598c · b77598c
1 parent cf96db6
commit b77598c
Show file tree

Hide file tree

Showing 318 changed files with 8,579 additions and 3,627 deletions.
diff --git a/.github/workflows/per-lib-check.yml b/.github/workflows/per-lib-check.yml
@@ -123,9 +123,9 @@ jobs:
         run: |
           test_target.sh substitutions
 
-      # - name: Test compiler
-      #   run: |
-      #     test_target.sh compiler
+      - name: Test compiler
+        run: |
+          test_target.sh compiler
 
       - name: Test substitution-generator
         run: |

diff --git a/bin/export-model-arch/include/export_model_arch/json_sp_model_export.struct.toml b/bin/export-model-arch/include/export_model_arch/json_sp_model_export.struct.toml
@@ -3,24 +3,22 @@ name = "JsonSPModelExport"
 features = [
   "eq",
   "hash",
-  "json",
   "fmt",
+  "json", 
 ]
 
 includes = [
   "pcg/file_format/v1/v1_computation_graph.dtg.h",
-  "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree.h",
+  "pcg/file_format/v1/v1_binary_sp_decomposition/v1_binary_sp_decomposition.dtg.h",
 ]
 
 src_includes = [
-  "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/hash.h",
-  "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/fmt.h",
-  "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/json.h",
+  "pcg/file_format/v1/v1_binary_sp_decomposition/json.h",
 ]
 
 [[fields]]
 name = "sp_decomposition"
-type = "::FlexFlow::GenericBinarySPDecompositionTree<int>"
+type = "::FlexFlow::V1BinarySPDecomposition"
 
 [[fields]]
 name = "computation_graph"

diff --git a/bin/export-model-arch/src/export_model_arch.cc b/bin/export-model-arch/src/export_model_arch.cc
@@ -1,5 +1,5 @@
-#include "compiler/series_parallel/computation_graph_binary_sp_decomposition.h"
-#include "compiler/series_parallel/get_computation_graph_series_parallel_decomposition.h"
+#include "compiler/series_parallel/computation_graph/computation_graph_binary_sp_decomposition.h"
+#include "compiler/series_parallel/computation_graph/get_computation_graph_series_parallel_decomposition.h"
 #include "export_model_arch/json_sp_model_export.dtg.h"
 #include "models/bert/bert.h"
 #include "models/candle_uno/candle_uno.h"
@@ -13,7 +13,6 @@
 #include "utils/cli/cli_parse.h"
 #include "utils/cli/cli_parse_result.h"
 #include "utils/cli/cli_spec.h"
-#include "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/transform.h"
 #include "utils/graph/series_parallel/binary_sp_decomposition_tree/right_associative_binary_sp_tree_from_nary.h"
 #include "utils/graph/series_parallel/get_series_parallel_decomposition.h"
 
@@ -105,9 +104,8 @@ tl::expected<JsonSPModelExport, std::string>
       to_v1_including_node_numbering(computation_graph);
   V1ComputationGraph v1_cg = v1_result.first;
   bidict<int, layer_guid_t> layer_numbering = v1_result.second;
-  GenericBinarySPDecompositionTree<int> v1_sp_decomposition =
-      transform(sp_decomposition.raw_tree,
-                [&](layer_guid_t const &l) { return layer_numbering.at_r(l); });
+  V1BinarySPDecomposition v1_sp_decomposition =
+      to_v1(sp_decomposition, layer_numbering);
 
   return JsonSPModelExport{
       v1_sp_decomposition,

diff --git a/flake.lock b/flake.lock
diff --git a/lib/compiler/include/compiler/cost_estimate.h b/lib/compiler/include/compiler/cost_estimate.h
diff --git a/lib/compiler/include/compiler/cost_estimator/cost_estimator.h b/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 "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"
+#include <vector>
+
+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
diff --git a/lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.struct.toml b/lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.struct.toml
@@ -0,0 +1,40 @@
+namespace = "FlexFlow"
+name = "OpCostEstimateKey"
+features = [
+  "eq",
+  "ord",
+  "fmt",
+  "hash",
+]
+
+includes = [
+  "op-attrs/pcg_operator_attrs.dtg.h",
+  "op-attrs/parallel_tensor_shape.dtg.h",
+  "<vector>",
+  "pcg/machine_view.dtg.h",
+]
+
+src_includes = [
+  "utils/hash/vector.h",
+  "utils/fmt/vector.h",
+]
+
+[[fields]]
+name = "op_attrs"
+type = "::FlexFlow::PCGOperatorAttrs"
+
+[[fields]]
+name = "input_shapes"
+type = "std::vector<::FlexFlow::ParallelTensorShape>"
+
+[[fields]]
+name = "weight_shapes"
+type = "std::vector<::FlexFlow::ParallelTensorShape>"
+
+[[fields]]
+name = "output_shapes"
+type = "std::vector<::FlexFlow::ParallelTensorShape>"
+
+[[fields]]
+name = "machine_view"
+type = "::FlexFlow::MachineView"
diff --git a/lib/compiler/include/compiler/cost_estimator/single_tensor_movement.struct.toml b/lib/compiler/include/compiler/cost_estimator/single_tensor_movement.struct.toml
@@ -0,0 +1,30 @@
+namespace = "FlexFlow"
+name = "SingleTensorMovement"
+features = [
+  "eq",
+  "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_views"
+type = "std::unordered_set<::FlexFlow::MachineView>"
+
+[[fields]]
+name = "dst_machine_views"
+type = "std::unordered_set<::FlexFlow::MachineView>"
diff --git a/lib/compiler/include/compiler/cost_estimator/tensor_set_movement.struct.toml b/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>"
diff --git a/lib/compiler/include/compiler/graph_optimize_result.struct.toml b/lib/compiler/include/compiler/graph_optimize_result.struct.toml
@@ -0,0 +1,16 @@
+namespace = "FlexFlow"
+name = "GraphOptimizeResult"
+features = [ ]
+
+includes = [ 
+  "compiler/machine_mapping/machine_mapping.dtg.h",
+  "pcg/parallel_computation_graph/parallel_computation_graph.h"
+]
+
+[[fields]]
+name = "pcg"
+type = "::FlexFlow::ParallelComputationGraph"
+
+[[fields]]
+name = "machine_mapping"
+type = "::FlexFlow::MachineMapping"
diff --git a/lib/compiler/include/compiler/graph_optimize_state.h b/lib/compiler/include/compiler/graph_optimize_state.h
@@ -0,0 +1,31 @@
+#ifndef _FLEXFLOW_COMPILER_MCMC_STATE_H
+#define _FLEXFLOW_COMPILER_MCMC_STATE_H
+
+#include "compiler/graph_optimize_result.dtg.h"
+
+namespace FlexFlow {
+
+struct GraphOptimizeState {
+  GraphOptimizeState(GraphOptimizeResult const &graph_optimize_result,
+                     float runtime);
+
+  GraphOptimizeResult graph_optimize_result;
+  float runtime;
+
+  bool operator==(GraphOptimizeState const &other) const;
+  bool operator!=(GraphOptimizeState const &other) const;
+  bool operator<(GraphOptimizeState const &other) const;
+};
+
+} // namespace FlexFlow
+
+namespace std {
+
+template <>
+struct hash<::FlexFlow::GraphOptimizeState> {
+  size_t operator()(::FlexFlow::GraphOptimizeState const &) const;
+};
+
+} // namespace std
+
+#endif
diff --git a/lib/compiler/include/compiler/graph_utils.h b/lib/compiler/include/compiler/graph_utils.h