Unity device mapping algorithm

.github/workflows/per-lib-check.yml

@@ -108,9 +108,9 @@ jobs:
         run: |
           test_libs.sh substitutions
 
-      # - name: Test compiler
-      #   run: |
-      #     test_libs.sh compiler
+      - name: Test compiler
+        run: |
+          test_libs.sh compiler
 
       - name: Test substitution-generator
         run: |

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

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

@@ -0,0 +1,56 @@
+#ifndef _FLEXFLOW_COMPILER_MACHINE_MAPPING_GET_OPTIMAL_MACHINE_MAPPING_H
+#define _FLEXFLOW_COMPILER_MACHINE_MAPPING_GET_OPTIMAL_MACHINE_MAPPING_H
+
+#include "machine_mapping.h"
+#include "machine_mapping_cache.h"
+#include "machine_mapping_context.dtg.h"
+#include "pcg/machine_specification.h"
+#include "pcg/machine_view.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.h"
+
+namespace FlexFlow {
+
+MachineMappingResult get_optimal_machine_mapping(
+    ParallelComputationGraph const &pcg,
+    std::function<std::unordered_set<MachineView>(
+        ParallelLayerAttrs const &, MachineSpecification const &)> const
+        &allowed_machine_views,
+    CostEstimator const &cost_estimator,
+    MachineSpecification const &resources,
+    MachineMappingCache &cached_subgraph_results);
+
+MachineMappingResult
+    get_optimal_machine_mapping_internal(MachineMappingContext &context,
+                                MachineSpecification const &resources);
+
+MachineMappingResult get_optimal_machine_mapping_internal(
+    MachineMappingContext &context,
+    SerialParallelDecomposition const &decompn,
+    MachineSpecification const &resource,
+    std::unordered_map<OpenDataflowValue, MachineView> const
+        &fixed_machine_views);
+
+MachineMappingResult get_optimal_machine_mapping_internal(
+    MachineMappingContext &context,
+    SerialSplit const &serial,
+    MachineSpecification const &resource,
+    std::unordered_map<OpenDataflowValue, MachineView> const
+        &fixed_machine_views);
+
+MachineMappingResult get_optimal_machine_mapping_internal(
+    MachineMappingContext &context,
+    ParallelSplit const &parallel,
+    MachineSpecification const &resource,
+    std::unordered_map<OpenDataflowValue, MachineView> const
+        &fixed_machine_views);
+
+MachineMappingResult get_optimal_machine_mapping_internal(
+    MachineMappingContext &context,
+    Node const &node,
+    MachineSpecification const &resource,
+    std::unordered_map<OpenDataflowValue, MachineView> const
+        &fixed_machine_views);
+
+} // namespace FlexFlow
+
+#endif

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

@@ -0,0 +1,14 @@
+#ifndef _FLEXFLOW_COMPILER_MACHINE_MAPPING_H
+#define _FLEXFLOW_COMPILER_MACHINE_MAPPING_H
+
+#include "machine_mapping.dtg.h"
+
+namespace FlexFlow {
+
+MachineMapping combine(MachineMapping const &, MachineMapping const &);
+
+bool nodes_are_disjoint(MachineMapping const &m1, MachineMapping const &m2);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/machine_mapping/machine_mapping_cache.h

@@ -0,0 +1,23 @@
+#ifndef _FLEXFLOW_COMPILER_MACHINE_MAPPING_DP_CACHE_H
+#define _FLEXFLOW_COMPILER_MACHINE_MAPPING_DP_CACHE_H
+
+#include "machine_mapping_result.dtg.h"
+#include "machine_mapping_state.dtg.h"
+#include "utils/optional.h"
+
+namespace FlexFlow {
+
+class MachineMappingCache {
+public:
+  MachineMappingCache() = default;
+
+  std::optional<MachineMappingResult> load(MachineMappingState const &) const;
+  void save(MachineMappingState const &, MachineMappingResult const &);
+
+private:
+  std::unordered_map<MachineMappingState, MachineMappingResult> cache;
+};
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/machine_mapping/machine_mapping_context.struct.toml

@@ -0,0 +1,29 @@
+namespace = "FlexFlow"
+name = "MachineMappingContext"
+features = [
+]
+
+includes = [
+  "machine_mapping.dtg.h",
+  "pcg/parallel_computation_graph/parallel_computation_graph.h",
+  "compiler/cost_estimator.h",
+  "pcg/machine_view.h",
+  "pcg/machine_specification.dtg.h",
+  "compiler/machine_mapping/machine_mapping_cache.h"
+]
+
+[[fields]]
+name = "pcg"
+type = "::FlexFlow::ParallelComputationGraph"
+
+[[fields]]
+name = "cost_estimator"
+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"

lib/compiler/include/compiler/machine_mapping/machine_mapping_result.h

@@ -0,0 +1,18 @@
+#ifndef _FLEXFLOW_COMPILER_MACHINE_MAPPING_MACHINE_MAPPING_RESULT_H
+#define _FLEXFLOW_COMPILER_MACHINE_MAPPING_MACHINE_MAPPING_RESULT_H
+
+#include "machine_mapping_result.dtg.h"
+
+namespace FlexFlow {
+
+MachineMappingResult sequential_combine(MachineMappingResult const &s1,
+                                        MachineMappingResult const &s2);
+MachineMappingResult parallel_combine(MachineMappingResult const &s1,
+                                      MachineMappingResult const &s2);
+MachineMappingResult get_infinity_machine_mapping_result();
+
+void minimize_runtime(MachineMappingResult &m1, MachineMappingResult const &m2);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/machine_mapping/machine_mapping_result.struct.toml

@@ -0,0 +1,18 @@
+namespace = "FlexFlow"
+name = "MachineMappingResult"
+features = [
+  "eq",
+  "fmt",
+]
+
+includes = [
+  "machine_mapping.dtg.h",
+]
+
+[[fields]]
+name = "runtime"
+type = "float"
+
+[[fields]]
+name = "machine_mapping"
+type = "::FlexFlow::MachineMapping"

lib/compiler/include/compiler/optimal_cost_state.struct.toml → lib/compiler/include/compiler/machine_mapping/machine_mapping_state.struct.toml

@@ -1,22 +1,22 @@
 namespace = "FlexFlow"
-name = "OptimalCostState"
+name = "MachineMappingState"
 features = [
   "eq",
-  # "ord",
   "hash",
-  # "json",
-  # "rapidcheck",
   "fmt",
 ]
 
 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/fmt/unordered_map.h", 
-  "utils/hash/unordered_map.h", 
+  "utils/graph/open_dataflow_graph/open_dataflow_value.dtg.h",
+]
+
+src_includes = [
+  "utils/hash/unordered_map.h",
+  "utils/fmt/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/machine_mapping/split_sp_decomposition.h

@@ -0,0 +1,16 @@
+#ifndef _FLEXFLOW_COMPILER_MACHINE_MAPPING_SPLIT_SP_DECOMPOSITION_H
+#define _FLEXFLOW_COMPILER_MACHINE_MAPPING_SPLIT_SP_DECOMPOSITION_H
+
+#include "utils/graph/serial_parallel/serial_parallel_decomposition.dtg.h"
+
+namespace FlexFlow {
+
+std::pair<SerialParallelDecomposition, SerialParallelDecomposition>
+    split_sp_decomposition(SerialSplit const &serial);
+
+std::pair<SerialParallelDecomposition, SerialParallelDecomposition>
+    split_sp_decomposition(ParallelSplit const &parallel);
+
+}
+
+#endif