Unity device mapping algorithm

.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: |

.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]

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/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"

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>"

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/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"

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