Unity device mapping algorithm

lib/compiler/include/compiler/graph_optimize_result.struct.toml

@@ -0,0 +1,16 @@
+namespace = "FlexFlow"
+name = "GraphOptimizeResult"
+features = [ ]
+
+includes = [ 
+  "compiler/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"
+#include "compiler/machine_mapping.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

lib/compiler/include/compiler/machine_mapping.h

@@ -3,7 +3,7 @@
 
 #include "compiler/machine_mapping.dtg.h"
 #include "compiler/optimal_cost_state.dtg.h"
-#include "cost_estimate.h"
+#include "cost_estimator.h"
 #include "pcg/machine_specification.dtg.h"
 #include "pcg/machine_specification.h"
 #include "pcg/machine_view.h"
@@ -29,10 +29,6 @@ struct OptimalCostResult {
 };
 FF_VISITABLE_STRUCT(OptimalCostResult, runtime, machine_mapping);
 
-struct OptimalCostRuntimeCmp {
-  bool operator()(OptimalCostResult const &, OptimalCostResult const &);
-};
-
 class OptimalCostCache {
 public:
   OptimalCostCache() = default;
@@ -55,15 +51,4 @@ OptimalCostResult optimal_cost(
 
 } // namespace FlexFlow
 
-// namespace std {
-//
-// template <>
-// struct hash<std::unordered_map<FlexFlow::Node, FlexFlow::MachineMapping>> {
-//   size_t operator()(
-//       std::unordered_map<FlexFlow::Node, FlexFlow::MachineMapping> const &g)
-//       const;
-// };
-
-// }; // namespace std
-
 #endif

lib/compiler/include/compiler/optimal_cost_state.struct.toml

@@ -13,10 +13,10 @@ includes = [
   "utils/graph/serial_parallel/serial_parallel_decomposition.dtg.h",
   "pcg/machine_specification.dtg.h",
   "pcg/machine_view.dtg.h",
-  "utils/graph/node/node.dtg.h",
   "utils/graph/open_dataflow_graph/open_dataflow_edge.dtg.h",
+  "utils/graph/open_dataflow_graph/open_dataflow_value.dtg.h",
+  "utils/hash/unordered_map.h",
   "utils/fmt/unordered_map.h", 
-  "utils/hash/unordered_map.h", 
 ]
 
 [[fields]]
@@ -28,9 +28,5 @@ name = "resource"
 type = "::FlexFlow::MachineSpecification"
 
 [[fields]]
-name = "given_machine_views"
-type = "std::unordered_map<::FlexFlow::Node, ::FlexFlow::MachineView>"
-
-[[fields]]
-name = "frontier_machine_views"
-type = "std::unordered_map<::FlexFlow::OpenDataflowEdge, ::FlexFlow::MachineView>"
+name = "fixed_machine_views"
+type = "std::unordered_map<::FlexFlow::OpenDataflowValue, ::FlexFlow::MachineView>"

lib/compiler/include/compiler/optimizer_config.struct.toml

@@ -0,0 +1,26 @@
+namespace = "FlexFlow"
+name = "OptimizerConfig"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+]
+
+includes = [
+]
+
+[[fields]]
+name = "alpha"
+type = "float"
+
+[[fields]]
+name = "budget"
+type = "int"
+
+[[fields]]
+name = "threshold"
+type = "float"
+
+[[fields]]
+name = "max_num_ops"
+type = "int"

lib/compiler/include/compiler/unity_algorithm.h

@@ -1,39 +1,18 @@
 #ifndef _FLEXFLOW_COMPILER_UNITY_ALGORITHM_H
 #define _FLEXFLOW_COMPILER_UNITY_ALGORITHM_H
 
+#include "compiler/graph_optimize_result.dtg.h"
 #include "compiler/machine_mapping.h"
-#include "cost_estimate.h"
-#include "machine_mapping.h"
+#include "cost_estimator.h"
+#include "optimizer_config.dtg.h"
 #include "pcg/computation_graph.h"
 #include "pcg/machine_specification.dtg.h"
 #include "substitutions/sub_parallel_computation_graph.h"
-namespace FlexFlow {
-
-struct Strategy {
-  ParallelComputationGraph pcg;
-  MachineMapping machine_mapping;
-  req<float> runtime;
-  friend bool operator!=(Strategy const &lhs, Strategy const &rhs) {
-    return (lhs.machine_mapping != rhs.machine_mapping) ||
-           (lhs.runtime != rhs.runtime);
-  }
-};
-
-FF_VISITABLE_STRUCT(Strategy, pcg, machine_mapping, runtime);
 
-struct StrategyRuntimeCmp {
-  bool operator()(Strategy const &, Strategy const &);
-};
-
-struct OptimizerConfig {
-  float alpha;
-  int budget;
-  float threshold;
-  int max_num_ops;
-};
+namespace FlexFlow {
 
-Strategy graph_optimize(
-    ComputationGraph &cg,
+GraphOptimizeResult graph_optimize(
+    ParallelComputationGraph &pcg,
     CostEstimator const &cost_estimator,
     MachineSpecification const &resources,
     std::function<std::unordered_set<MachineView>(

lib/compiler/src/graph_optimize_state.cc

@@ -0,0 +1,80 @@
+#include "compiler/graph_optimize_state.h"
+
+namespace FlexFlow {
+
+GraphOptimizeState::GraphOptimizeState(
+    GraphOptimizeResult const &graph_optimize_result, float runtime)
+    : graph_optimize_result(graph_optimize_result), runtime(runtime) {}
+
+bool GraphOptimizeState::operator==(GraphOptimizeState const &other) const {
+  auto layers1 = topological_ordering(graph_optimize_result.pcg);
+  auto layers2 = topological_ordering(other.graph_optimize_result.pcg);
+  if (layers1.size() != layers2.size()) {
+    return false;
+  }
+  std::unordered_map<parallel_tensor_guid_t, parallel_tensor_guid_t> mapping;
+  for (size_t i = 0; i < layers1.size(); ++i) {
+    if (get_parallel_layer_attrs(graph_optimize_result.pcg, layers1[i]) !=
+        get_parallel_layer_attrs(other.graph_optimize_result.pcg, layers2[i])) {
+      return false;
+    }
+    auto inputs1 = get_layer_inputs(graph_optimize_result.pcg, layers1[i]);
+    auto inputs2 =
+        get_layer_inputs(other.graph_optimize_result.pcg, layers2[i]);
+    if (inputs1.size() != inputs2.size()) {
+      return false;
+    }
+    for (size_t j = 0; j < inputs1.size(); ++j) {
+      if (inputs1[j] != mapping.at(inputs2[j])) {
+        return false;
+      }
+    }
+    auto outputs1 = get_layer_outputs(graph_optimize_result.pcg, layers1[i]);
+    auto outputs2 =
+        get_layer_outputs(other.graph_optimize_result.pcg, layers2[i]);
+    if (outputs1.size() != outputs2.size()) {
+      return false;
+    }
+    for (size_t j = 0; j < outputs1.size(); ++j) {
+      mapping.emplace(outputs2[j], outputs1[j]);
+    }
+  }
+  return true;
+}
+
+bool GraphOptimizeState::operator!=(GraphOptimizeState const &other) const {
+  return !(*this == other);
+}
+
+bool GraphOptimizeState::operator<(GraphOptimizeState const &other) const {
+  return runtime < other.runtime;
+}
+
+} // namespace FlexFlow
+
+namespace std {
+
+size_t hash<::FlexFlow::GraphOptimizeState>::operator()(
+    ::FlexFlow::GraphOptimizeState const &state) const {
+  size_t seed = 0;
+  auto layers = topological_ordering(state.graph_optimize_result.pcg);
+  ::FlexFlow::hash_combine(seed, layers.size());
+  for (auto layer : layers) {
+    ::FlexFlow::hash_combine(
+        seed, get_parallel_layer_attrs(state.graph_optimize_result.pcg, layer));
+    auto inputs = get_layer_inputs(state.graph_optimize_result.pcg, layer);
+    ::FlexFlow::hash_combine(seed, inputs.size());
+    for (auto input : inputs) {
+      for (size_t i = 0; i < layers.size(); ++i) {
+        if (get_source_layer(state.graph_optimize_result.pcg, input) ==
+            layers[i]) {
+          ::FlexFlow::hash_combine(seed, i);
+          break;
+        }
+      }
+    }
+  }
+  return seed;
+}
+
+} // namespace std