Top-level loop for compiler

bin/export-model-arch/src/export-model-arch/main.cc

@@ -1,6 +1,6 @@
 #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 "export-model-arch/json_sp_model_export.dtg.h"
 #include "models/bert/bert.h"
 #include "models/candle_uno/candle_uno.h"
 #include "models/dlrm/dlrm.h"

lib/compiler/include/compiler/algorithm_config.dtg.toml

@@ -0,0 +1,19 @@
+namespace = "FlexFlow"
+name = "AlgorithmConfig"
+type = "variant"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+]
+
+includes = [
+  "compiler/data_parallelism/data_parallelism_config.dtg.h",
+  "compiler/unity_algorithm/unity_search_config.dtg.h",
+]
+
+[[values]]
+type = "::FlexFlow::DataParallelismConfig"
+
+[[values]]
+type = "::FlexFlow::UnitySearchConfig"

lib/compiler/include/compiler/compiler.h

@@ -0,0 +1,19 @@
+#ifndef _FLEXFLOW_COMPILER_COMPILER_H
+#define _FLEXFLOW_COMPILER_COMPILER_H
+
+#include "compiler/algorithm_config.dtg.h"
+#include "compiler/cost_estimator/cost_estimator.h"
+#include "compiler/search_result.dtg.h"
+#include "pcg/computation_graph.h"
+#include "pcg/machine_specification.dtg.h"
+
+namespace FlexFlow {
+
+SearchResult optimize(ComputationGraph const &,
+                      MachineSpecification const &,
+                      CostEstimator const &,
+                      AlgorithmConfig const &);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/data_parallelism/data_parallelism_config.dtg.toml

@@ -0,0 +1,15 @@
+namespace = "FlexFlow"
+name = "DataParallelismConfig"
+type = "struct"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+]
+
+includes = [
+]
+
+[[fields]]
+name = "degree"
+type = "int"

lib/compiler/include/compiler/machine_mapping/abstracted_tensor_set_movement/abstracted_device.dtg.toml

@@ -6,6 +6,7 @@ features = [
   "ord",
   "hash",
   "fmt",
+  "json",
 ]
 
 includes = [

lib/compiler/include/compiler/machine_mapping/abstracted_tensor_set_movement/abstracted_single_tensor_communication_edge.dtg.toml

@@ -6,6 +6,7 @@ features = [
   "ord",
   "hash",
   "fmt",
+  "json",
 ]
 
 includes = [

lib/compiler/include/compiler/machine_mapping/abstracted_tensor_set_movement/abstracted_single_tensor_movement.dtg.toml

@@ -5,6 +5,7 @@ features = [
   "eq",
   "hash",
   "fmt",
+  "json",
 ]
 
 includes = [

lib/compiler/include/compiler/machine_mapping/abstracted_tensor_set_movement/abstracted_tensor_set_movement.dtg.toml

@@ -5,6 +5,7 @@ features = [
   "eq",
   "hash",
   "fmt",
+  "json",
 ]
 
 includes = [

lib/compiler/include/compiler/allowed_machine_views.h → lib/compiler/include/compiler/machine_mapping/allowed_machine_views.h

@@ -1,18 +1,19 @@
-#ifndef _FLEXFLOW_COMPILER_ALLOWED_MACHINE_VIEWS_H
-#define _FLEXFLOW_COMPILER_ALLOWED_MACHINE_VIEWS_H
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_ALLOWED_MACHINE_VIEWS_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_ALLOWED_MACHINE_VIEWS_H
 
 #include "compiler/machine_mapping/machine_view.dtg.h"
 #include "op-attrs/operator_task_space.dtg.h"
+#include "pcg/machine_compute_resource_slice.dtg.h"
 #include "pcg/machine_compute_specification.dtg.h"
 
 namespace FlexFlow {
 
 bool is_valid_machine_view(MachineView const &mv,
                            OperatorTaskSpace const &task,
-                           MachineComputeSpecification const &ms);
+                           MachineComputeResourceSlice const &ms);
 
 std::unordered_set<MachineView>
-    get_allowed_machine_views(MachineComputeSpecification const &machine_spec,
+    get_allowed_machine_views(MachineComputeResourceSlice const &machine_spec,
                               OperatorTaskSpace const &task,
                               DeviceType device_type);
 

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

@@ -1,14 +1,14 @@
 #ifndef _FLEXFLOW_COMPILER_MACHINE_MAPPING_GET_OPTIMAL_MACHINE_MAPPING_H
 #define _FLEXFLOW_COMPILER_MACHINE_MAPPING_GET_OPTIMAL_MACHINE_MAPPING_H
 
-#include "compiler/machine_mapping/machine_compute_resource_slice.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_cache.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_constraints.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_context.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_problem_tree/machine_mapping_problem_tree.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_problem_tree/mm_problem_tree_parallel_split.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_problem_tree/mm_problem_tree_series_split.dtg.h"
 #include "compiler/machine_mapping/parallel_split_transformation.dtg.h"
+#include "pcg/machine_compute_resource_slice.dtg.h"
 
 namespace FlexFlow {
 

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

@@ -2,6 +2,8 @@
 #define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_MACHINE_MAPPING_H
 
 #include "compiler/machine_mapping/machine_mapping.dtg.h"
+#include "compiler/machine_mapping/machine_mapping_result.h"
+#include "compiler/series_parallel/pcg/pcg_binary_sp_decomposition.dtg.h"
 #include "pcg/device_id_t.dtg.h"
 #include "pcg/machine_specification.dtg.h"
 #include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.dtg.h"
@@ -18,6 +20,9 @@ MappedParallelComputationGraph
     mapped_pcg_from_pcg_and_mapping(ParallelComputationGraph const &,
                                     MachineMapping const &);
 
+std::optional<MachineMapping> get_machine_mapping_from_machine_mapping_result(
+    PCGBinarySPDecomposition const &, MachineMappingResult const &);
+
 } // namespace FlexFlow
 
 #endif

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

@@ -14,8 +14,13 @@ MachineMappingConstraints get_unconstrained_solution_for_layers(
     std::unordered_set<BinaryTreePath> const &);
 
 std::unordered_set<BinaryTreePath>
-    get_all_layers(MachineMappingConstraints const &,
-                   IncludeUnconstrained const &);
+    get_unconstrained_layers(MachineMappingConstraints const &);
+
+std::unordered_set<BinaryTreePath>
+    get_constrained_layers(MachineMappingConstraints const &);
+
+std::unordered_set<BinaryTreePath>
+    get_all_layers(MachineMappingConstraints const &);
 
 std::optional<MachineView>
     get_machine_view_for_layer(MachineMappingConstraints const &,

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

@@ -6,7 +6,7 @@ features = []
 includes = [
   "compiler/cost_estimator/runtime_only_cost_estimator.h",
   "compiler/machine_mapping/machine_view.dtg.h",
-  "compiler/machine_mapping/machine_compute_resource_slice.dtg.h",
+  "pcg/machine_compute_resource_slice.dtg.h",
   "compiler/machine_mapping/machine_mapping_problem_tree/unmapped_runtime_only_op_cost_estimate_key.dtg.h",
 ]
 

lib/compiler/include/compiler/machine_mapping/machine_mapping_problem_tree/get_machine_mapping_problem_tree.h

@@ -9,6 +9,9 @@
 
 namespace FlexFlow {
 
+bool is_valid_machine_mapping_problem_tree(
+    MachineMappingProblemTree const &problem_tree);
+
 MachineMappingProblemTree
     get_machine_mapping_problem_tree(ParallelComputationGraph const &pcg,
                                      PCGBinarySPDecomposition const &sp);

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

@@ -1,7 +1,6 @@
 #ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_MACHINE_MAPPING_PROBLEM_TREE_H
 #define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_MACHINE_MAPPING_PROBLEM_TREE_H
 
-#include "compiler/machine_mapping/abstracted_tensor_set_movement/machine_space_stencil.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_problem_tree/machine_mapping_problem_tree.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_problem_tree/mm_problem_tree_parallel_split.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_problem_tree/mm_problem_tree_series_split.dtg.h"
@@ -32,6 +31,9 @@ std::optional<MachineMappingProblemTree>
 std::unordered_map<BinaryTreePath, UnmappedRuntimeOnlyOpCostEstimateKey>
     mm_problem_tree_get_path_to_leaf_map(MachineMappingProblemTree const &);
 
+std::string as_dot(MachineMappingProblemTree const &);
+void debug_print_dot(MachineMappingProblemTree const &);
+
 } // namespace FlexFlow
 
 #endif

lib/compiler/include/compiler/machine_mapping/machine_mapping_problem_tree/unmapped_runtime_only_op_cost_estimate_key.h

@@ -3,6 +3,7 @@
 
 #include "compiler/cost_estimator/runtime_only_op_cost_estimate_key.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_problem_tree/unmapped_runtime_only_op_cost_estimate_key.dtg.h"
+#include "op-attrs/operator_task_space.dtg.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
 #include "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h"
 
@@ -17,6 +18,9 @@ RuntimeOnlyOpCostEstimateKey map_unmapped_runtime_only_op_cost_estimate_key(
     UnmappedRuntimeOnlyOpCostEstimateKey const &unmapped,
     MachineView const &machine_view);
 
+OperatorTaskSpace get_operator_task_space_for_runtime_only_op_cost_estimate_key(
+    UnmappedRuntimeOnlyOpCostEstimateKey const &unmapped);
+
 } // namespace FlexFlow
 
 #endif

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

@@ -35,6 +35,9 @@ FeasibleMachineMappingResult require_feasible(MachineMappingResult const &);
     make_singleton_machine_mapping_result(milliseconds_t runtime,
                                           MachineView const &machine_view);
 
+[[nodiscard]] milliseconds_t
+    get_runtime_cost(MachineMappingResult const &mm_result);
+
 } // namespace FlexFlow
 
 #endif

lib/compiler/include/compiler/machine_mapping/machine_mapping_state.dtg.toml

@@ -9,7 +9,7 @@ features = [
 
 includes = [
   "compiler/machine_mapping/machine_mapping_constraints.dtg.h",
-  "compiler/machine_mapping/machine_compute_resource_slice.dtg.h",
+  "pcg/machine_compute_resource_slice.dtg.h",
   "compiler/machine_mapping/machine_mapping_problem_tree/machine_mapping_problem_tree.dtg.h",
 ]
 

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

@@ -1,10 +1,10 @@
 #ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_MACHINE_RESOURCE_SPLIT_H
 #define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MACHINE_MAPPING_MACHINE_RESOURCE_SPLIT_H
 
-#include "compiler/machine_mapping/machine_compute_resource_slice.dtg.h"
 #include "compiler/machine_mapping/machine_resource_split.dtg.h"
 #include "compiler/machine_mapping/machine_view.dtg.h"
 #include "compiler/machine_mapping/parallel_layer_guid_oblivious_machine_mapping.dtg.h"
+#include "pcg/machine_compute_resource_slice.dtg.h"
 
 namespace FlexFlow {
 

lib/compiler/include/compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_context.dtg.toml

@@ -6,7 +6,7 @@ features = []
 includes = [
   "compiler/cost_estimator/cost_estimator.h",
   "compiler/machine_mapping/machine_view.dtg.h",
-  "compiler/machine_mapping/machine_compute_resource_slice.dtg.h",
+  "pcg/machine_compute_resource_slice.dtg.h",
   "compiler/machine_mapping/machine_mapping_problem_tree/unmapped_runtime_only_op_cost_estimate_key.dtg.h",
   "pcg/optimizer_attrs.dtg.h",
 ]

lib/compiler/include/compiler/search_result.dtg.toml

@@ -0,0 +1,18 @@
+namespace = "FlexFlow"
+name = "SearchResult"
+type = "struct"
+features = [
+]
+
+includes = [
+  "pcg/parallel_computation_graph/parallel_computation_graph.h",
+  "compiler/machine_mapping/machine_mapping.h",
+]
+
+[[fields]]
+name = "pcg"
+type = "::FlexFlow::ParallelComputationGraph"
+
+[[fields]]
+name = "machine_mapping"
+type = "::FlexFlow::MachineMapping"

lib/compiler/include/compiler/search_result.h

@@ -0,0 +1,16 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_SEARCH_RESULT_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_SEARCH_RESULT_H
+
+#include "compiler/search_result.dtg.h"
+
+namespace FlexFlow {
+
+MappedParallelComputationGraph
+    get_mapped_pcg_from_search_result(SearchResult const &);
+
+std::string format_as(SearchResult const &);
+std::ostream &operator<<(std::ostream &, SearchResult const &);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/series_parallel/pcg/pcg_binary_sp_decomposition.h

@@ -22,11 +22,13 @@ GenericBinarySPDecompositionTreeImplementation<PCGBinarySPDecomposition,
 BinarySPDecompositionTree
     binary_sp_tree_from_pcg_sp_tree(PCGBinarySPDecomposition const &);
 
-std::optional<PCGBinarySPDecomposition>
-    get_pcg_balanced_binary_sp_decomposition(ParallelComputationGraph const &);
 std::unordered_multiset<parallel_layer_guid_t>
     get_parallel_layers(PCGBinarySPDecomposition const &);
 
+PCGBinarySPDecomposition
+    pcg_binary_sp_decomposition_from_binary_sp_decomposition_tree(
+        BinarySPDecompositionTree const &);
+
 SPDecompositionTreeNodeType get_node_type(PCGBinarySPDecomposition const &);
 
 std::unordered_set<BinaryTreePath>
@@ -36,9 +38,6 @@ std::unordered_set<BinaryTreePath>
     find_paths_to_leaf(PCGBinarySPDecomposition const &,
                        parallel_layer_guid_t const &);
 
-PCGBinarySPDecomposition pcg_binary_sp_decomposition_from_binary_sp_tree(
-    BinarySPDecompositionTree const &spd_tree);
-
 std::unordered_map<BinaryTreePath, parallel_layer_guid_t>
     pcg_sp_tree_get_path_to_leaf_map(PCGBinarySPDecomposition const &);