flexflow · wmdi · Aug 2, 2024 · Jun 4, 2024 · Jun 4, 2024 · Jun 10, 2024
diff --git a/.proj.toml b/.proj.toml
@@ -9,7 +9,7 @@ build_targets = [
   "kernels",
   "pcg",
   "substitutions",
-  # "compiler",
+  "compiler",
   "substitution-generator",
   "local-execution", 
 ]
@@ -19,7 +19,7 @@ test_targets = [
   "op-attrs-tests",
   "pcg-tests",
   "substitutions-tests",
-  # "compiler-tests",
+  "compiler-tests",
   "substitution-generator-tests",
 ]
 

diff --git a/lib/compiler/include/compiler/compiler.h b/lib/compiler/include/compiler/compiler.h
@@ -3,7 +3,7 @@
 
 #include "pcg/cost_values.h"
 #include "pcg/machine_view.h"
-#include "pcg/parallel_computation_graph.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.h"
 #include "pcg/tensor_mapping.h"
 
 namespace FlexFlow {

diff --git a/lib/compiler/include/compiler/cost_estimate.h b/lib/compiler/include/compiler/cost_estimate.h
@@ -5,7 +5,7 @@
 #include "op-attrs/operator_attrs.h"
 #include "op-attrs/parallel_tensor_shape.h"
 #include "pcg/machine_view.h"
-#include "pcg/parallel_tensor_attrs.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_tensor_attrs.dtg.h"
 
 namespace FlexFlow {
 

diff --git a/lib/compiler/src/graph_utils.h → lib/compiler/include/compiler/graph_utils.h b/lib/compiler/src/graph_utils.h → lib/compiler/include/compiler/graph_utils.h
@@ -2,15 +2,18 @@
 #define _FLEXFLOW_COMPILER_GRAPH_UTILS_H
 
 #include "compiler/unity_algorithm.h"
+#include "pcg/computation_graph.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
+#include "substitutions/sub_parallel_computation_graph.dtg.h"
+#include "utils/graph/serial_parallel/serial_parallel_decomposition.dtg.h"
 
 namespace FlexFlow {
 
 SerialParallelDecomposition
     get_serial_parallel_decomposition(ParallelComputationGraph const &pcg);
 
 ParallelComputationGraph cg_to_pcg(ComputationGraph const &g);
-SubParallelComputationGraphView
-    pcg_to_subpcg(ParallelComputationGraph const &g);
+SubParallelComputationGraph pcg_to_subpcg(ParallelComputationGraph const &g);
 
 // NOTE(@wmdi): I think we should have the following interfaces in the graph
 // library eventually.

diff --git a/lib/compiler/include/compiler/machine_mapping.h b/lib/compiler/include/compiler/machine_mapping.h
@@ -1,37 +1,21 @@
 #ifndef _FLEXFLOW_COMPILER_MACHINE_MAPPING_H
 #define _FLEXFLOW_COMPILER_MACHINE_MAPPING_H
 
+#include "compiler/machine_mapping.dtg.h"
+#include "compiler/optimal_cost_state.dtg.h"
 #include "cost_estimate.h"
+#include "pcg/machine_specification.dtg.h"
 #include "pcg/machine_specification.h"
 #include "pcg/machine_view.h"
-#include "pcg/parallel_computation_graph.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.h"
 #include "substitutions/sub_parallel_computation_graph.h"
+#include "utils/graph/serial_parallel/serial_parallel_decomposition.dtg.h"
 
 namespace FlexFlow {
 
-using SubParallelComputationGraphView =
-    OutputLabelledOpenMultiDiGraphView<Operator, ParallelTensor>;
+MachineMapping combine(MachineMapping const &, MachineMapping const &);
 
-struct MachineMapping {
-  static MachineMapping combine(MachineMapping const &, MachineMapping const &);
-  static bool nodes_are_disjoint(MachineMapping const &m1,
-                                 MachineMapping const &m2);
-
-  req<std::unordered_map<Node, MachineView>> machine_views;
-};
-FF_VISITABLE_STRUCT(MachineMapping, machine_views);
-
-struct OptimalCostState {
-  SerialParallelDecomposition subgraph;
-  MachineSpecification resource;
-  std::unordered_map<Node, MachineView> given_machine_views;
-  req<std::unordered_map<OpenMultiDiEdge, MachineView>> frontier_machine_views;
-};
-FF_VISITABLE_STRUCT(OptimalCostState,
-                    subgraph,
-                    resource,
-                    given_machine_views,
-                    frontier_machine_views);
+bool nodes_are_disjoint(MachineMapping const &m1, MachineMapping const &m2);
 
 struct OptimalCostResult {
   static OptimalCostResult sequential_combine(OptimalCostResult const &s1,
@@ -60,26 +44,26 @@ class OptimalCostCache {
   std::unordered_map<OptimalCostState, OptimalCostResult> cache;
 };
 
-OptimalCostResult
-    optimal_cost(ParallelComputationGraph const &g,
-                 std::function<std::unordered_set<MachineView>(
-                     Operator const &, MachineSpecification const &)> const
-                     &allowed_machine_views,
-                 CostEstimator const &cost_estimator,
-                 MachineSpecification const &resources,
-                 OptimalCostCache &cached_subgraph_costs);
+OptimalCostResult optimal_cost(
+    ParallelComputationGraph const &g,
+    std::function<std::unordered_set<MachineView>(
+        ParallelLayerAttrs const &, MachineSpecification const &)> const
+        &allowed_machine_views,
+    CostEstimator const &cost_estimator,
+    MachineSpecification const &resources,
+    OptimalCostCache &cached_subgraph_costs);
 
 } // 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 {
+//
+// 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
+// }; // namespace std
 
 #endif
diff --git a/lib/compiler/include/compiler/machine_mapping.struct.toml b/lib/compiler/include/compiler/machine_mapping.struct.toml
@@ -0,0 +1,21 @@
+namespace = "FlexFlow"
+name = "MachineMapping"
+features = [
+  "eq",
+  # "ord",
+  "hash",
+  # "json",
+  # "rapidcheck",
+  "fmt",
+]
+
+includes = [ 
+  "utils/graph/node/node.dtg.h",
+  "pcg/machine_view.dtg.h",
+  "utils/hash/unordered_map.h",
+  "utils/fmt/unordered_map.h", 
+]
+
+[[fields]]
+name = "machine_views"
+type = "std::unordered_map<::FlexFlow::Node, ::FlexFlow::MachineView>"
diff --git a/lib/compiler/include/compiler/optimal_cost_state.struct.toml b/lib/compiler/include/compiler/optimal_cost_state.struct.toml
@@ -0,0 +1,36 @@
+namespace = "FlexFlow"
+name = "OptimalCostState"
+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", 
+]
+
+[[fields]]
+name = "subgraph"
+type = "::FlexFlow::SerialParallelDecomposition"
+
+[[fields]]
+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>"
diff --git a/lib/compiler/include/compiler/unity_algorithm.h b/lib/compiler/include/compiler/unity_algorithm.h
@@ -1,17 +1,22 @@
 #ifndef _FLEXFLOW_COMPILER_UNITY_ALGORITHM_H
 #define _FLEXFLOW_COMPILER_UNITY_ALGORITHM_H
 
+#include "compiler/machine_mapping.h"
 #include "cost_estimate.h"
 #include "machine_mapping.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);
@@ -27,14 +32,14 @@ struct OptimizerConfig {
   int max_num_ops;
 };
 
-Strategy
-    graph_optimize(ComputationGraph &cg,
-                   CostEstimator const &cost_estimator,
-                   MachineSpecification const &resources,
-                   std::function<std::unordered_set<MachineView>(
-                       Operator const &, MachineSpecification const &)> const
-                       &allowed_machine_views,
-                   OptimizerConfig const &opt_config);
+Strategy graph_optimize(
+    ComputationGraph &cg,
+    CostEstimator const &cost_estimator,
+    MachineSpecification const &resources,
+    std::function<std::unordered_set<MachineView>(
+        ParallelLayerAttrs const &, MachineSpecification const &)> const
+        &allowed_machine_views,
+    OptimizerConfig const &opt_config);
 
 } // namespace FlexFlow