flexflow · lockshaw · Jan 24, 2025 · Dec 27, 2024 · Jan 10, 2025 · Jan 12, 2025
diff --git a/lib/compiler/include/compiler/cost_estimator/in_progress_task.struct.toml b/lib/compiler/include/compiler/cost_estimator/in_progress_task.struct.toml
@@ -0,0 +1,22 @@
+namespace = "FlexFlow"
+name = "InProgressTask"
+
+features = [
+  "eq",
+  "hash",
+  "fmt",
+  "ord"
+]
+
+includes = [
+  "utils/graph/node/node.dtg.h"
+]
+
+
+[[fields]]
+name = "endtime"
+type = "float"
+
+[[fields]]
+name = "node"
+type = "::FlexFlow::Node"
diff --git a/lib/compiler/include/compiler/cost_estimator/in_progress_task_comparator.h b/lib/compiler/include/compiler/cost_estimator/in_progress_task_comparator.h
@@ -0,0 +1,13 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_IN_PROGRESS_TASK_COMPARATOR_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_IN_PROGRESS_TASK_COMPARATOR_H
+
+#include "compiler/cost_estimator/in_progress_task.dtg.h"
+#include <tuple>
+
+namespace FlexFlow {
+struct InProgressTaskComparator {
+  bool operator()(InProgressTask const &lhs, InProgressTask const &rhs) const;
+};
+} // namespace FlexFlow
+
+#endif // _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_IN_PROGRESS_TASK_COMPARATOR_H
diff --git a/lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.h b/lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.h
@@ -0,0 +1,17 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_OP_COST_ESTIMATE_KEY_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_OP_COST_ESTIMATE_KEY_H
+
+#include "compiler/cost_estimator/op_cost_estimate_key.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h"
+
+namespace FlexFlow {
+
+OpCostEstimateKey get_mapped_op_cost_estimate_key_for_layer(
+    ParallelComputationGraph const &pcg,
+    parallel_layer_guid_t const &layer,
+    MachineView const &machine_view);
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/compiler/include/compiler/cost_estimator/pcg_task.variant.toml b/lib/compiler/include/compiler/cost_estimator/pcg_task.variant.toml
@@ -0,0 +1,21 @@
+namespace = "FlexFlow"
+name = "PCGTask"
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "fmt",
+]
+
+includes = [
+  "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h",
+  "pcg/parallel_computation_graph/parallel_computation_graph_edge.dtg.h",
+]
+
+[[values]]
+type = "::FlexFlow::parallel_layer_guid_t"
+key = "layer"
+
+[[values]]
+type = "::FlexFlow::ParallelComputationGraphEdge"
+key = "edge"
diff --git a/lib/compiler/include/compiler/cost_estimator/simulate_task_graph_execution.h b/lib/compiler/include/compiler/cost_estimator/simulate_task_graph_execution.h
@@ -0,0 +1,16 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_SIMULATE_TASK_GRAPH_EXECUTION_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_SIMULATE_TASK_GRAPH_EXECUTION_H
+
+#include "compiler/cost_estimator/task_constraint.dtg.h"
+#include "compiler/cost_estimator/task_graph.dtg.h"
+#include "compiler/cost_estimator/task_graph_execution_trace.dtg.h"
+
+namespace FlexFlow {
+
+TaskGraphExecutionTrace
+    simulate_task_graph_execution(TaskGraph const &task_graph,
+                                  TaskConstraint const &constraint);
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/compiler/include/compiler/cost_estimator/task_constraint.struct.toml b/lib/compiler/include/compiler/cost_estimator/task_constraint.struct.toml
@@ -0,0 +1,16 @@
+namespace = "FlexFlow"
+name = "TaskConstraint"
+features = [
+]
+
+includes = [
+  "utils/graph/node/node.dtg.h",
+  "compiler/cost_estimator/task_graph.dtg.h",
+  "<functional>",
+  "<unordered_set>"
+]
+
+
+[[fields]]
+name = "is_satisfied"
+type = "std::function<bool(Node const &, std::unordered_set<Node> const &, std::unordered_set<Node> const &)>"
diff --git a/lib/compiler/include/compiler/cost_estimator/task_graph.struct.toml b/lib/compiler/include/compiler/cost_estimator/task_graph.struct.toml
@@ -0,0 +1,24 @@
+namespace = "FlexFlow"
+name = "TaskGraph"
+features = [
+]
+
+includes = [
+  "utils/graph/digraph/digraph_view.h",
+  "<unordered_map>",
+  "<functional>",
+  "compiler/cost_estimator/task_graph_execution_state.dtg.h"
+]
+
+src_includes = [
+  "utils/hash/unordered_map.h",
+  "utils/fmt/unordered_map.h",
+]
+
+[[fields]]
+name = "graph"
+type = "::FlexFlow::DiGraphView"
+
+[[fields]]
+name = "cost_map"
+type = "std::unordered_map<::FlexFlow::Node, float>"
diff --git a/lib/compiler/include/compiler/cost_estimator/task_graph_execution_state.struct.toml b/lib/compiler/include/compiler/cost_estimator/task_graph_execution_state.struct.toml
@@ -0,0 +1,37 @@
+namespace = "FlexFlow"
+name = "TaskGraphExecutionState"
+
+features = [
+]
+
+includes = [
+  "utils/deduplicated_priority_queue.h",
+  "utils/graph/node/node.dtg.h",
+  "compiler/cost_estimator/in_progress_task.dtg.h", 
+  "compiler/cost_estimator/in_progress_task_comparator.h",
+  "<unordered_set>",
+  "<functional>"
+]
+
+src_includes = [
+  "utils/hash/unordered_set.h", 
+  "utils/fmt/unordered_set.h",
+  "utils/fmt/vector.h", 
+  "utils/hash/vector.h"
+]
+
+[[fields]]
+name = "ready_tasks"
+type = "std::unordered_set<::FlexFlow::Node>"
+
+[[fields]]
+name = "tasks_processing"
+type = "::FlexFlow::DeduplicatedPriorityQueue<::FlexFlow::InProgressTask, std::vector<::FlexFlow::InProgressTask>, ::FlexFlow::InProgressTaskComparator>"
+
+[[fields]]
+name = "processed_tasks"
+type = "std::unordered_set<::FlexFlow::Node>"
+
+[[fields]]
+name = "current_time"
+type = "float"
diff --git a/lib/compiler/include/compiler/cost_estimator/task_graph_execution_trace.struct.toml b/lib/compiler/include/compiler/cost_estimator/task_graph_execution_trace.struct.toml
@@ -0,0 +1,27 @@
+namespace = "FlexFlow"
+name = "TaskGraphExecutionTrace"
+
+features = [
+  "hash",
+  "fmt",
+  "eq"
+]
+
+includes = [
+  "compiler/cost_estimator/task_profile.dtg.h",
+  "<unordered_set>"
+]
+
+src_includes = [
+  "utils/fmt/unordered_set.h",
+  "utils/hash/unordered_set.h"
+]
+
+
+[[fields]]
+name = "task_profiles"
+type = "std::unordered_set<::FlexFlow::TaskProfile>"
+
+[[fields]]
+name = "end_time" 
+type = "float"
diff --git a/lib/compiler/include/compiler/cost_estimator/task_profile.struct.toml b/lib/compiler/include/compiler/cost_estimator/task_profile.struct.toml
@@ -0,0 +1,26 @@
+namespace = "FlexFlow"
+name = "TaskProfile"
+
+features = [
+  "eq",
+  "hash",
+  "fmt",
+  "ord"
+]
+
+includes = [
+  "utils/graph/node/node.dtg.h"
+]
+
+
+[[fields]]
+name = "node"
+type = "::FlexFlow::Node"
+
+[[fields]]
+name = "start_time"
+type = "float"
+
+[[fields]]
+name = "end_time"
+type = "float"
diff --git a/lib/compiler/include/compiler/cost_estimator/task_simulator.h b/lib/compiler/include/compiler/cost_estimator/task_simulator.h
@@ -0,0 +1,18 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_TASK_SIMULATOR_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_TASK_SIMULATOR_H
+
+#include "compiler/cost_estimator/cost_estimator.h"
+#include "compiler/machine_mapping/machine_mapping.dtg.h"
+#include "pcg/machine_specification.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
+
+namespace FlexFlow {
+float task_simulator_estimate_forward_pass_time(
+    ParallelComputationGraph const &pcg,
+    CostEstimator const &estimator,
+    MachineMapping const &machine_mapping,
+    MachineSpecification const &machine_spec);
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/compiler/include/compiler/machine_mapping/device_mapping.cc b/lib/compiler/include/compiler/machine_mapping/device_mapping.cc
@@ -0,0 +1,28 @@
+
+#include "compiler/machine_mapping/device_mapping.h"
+#include "compiler/machine_mapping/device_mapping.dtg.h"
+#include "pcg/machine_specification.h"
+#include "pcg/machine_view.h"
+#include "pcg/operator_task_space.dtg.h"
+#include "pcg/operator_task_space.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.h"
+#include "utils/containers/keys.h"
+#include "utils/containers/map_values.h"
+
+namespace FlexFlow {
+
+DeviceMapping get_device_mapping(MachineMapping const &machine_mapping,
+                                 MachineSpecification const &machine_spec,
+                                 ParallelComputationGraph const &pcg) {
+  std::unordered_map<parallel_layer_guid_t, std::unordered_set<device_id_t>>
+      device_mapping;
+  for (auto const &[layer, machine_view] : machine_mapping.machine_views) {
+    OperatorTaskSpace op =
+        get_operator_task_space(pcg, layer));
+    device_mapping.insert(
+        {layer, get_device_ids(op, machine_view, machine_spec)});
+  }
+  return DeviceMapping{device_mapping};
+}
+
+} // namespace FlexFlow
diff --git a/lib/compiler/include/compiler/machine_mapping/device_mapping.h b/lib/compiler/include/compiler/machine_mapping/device_mapping.h
@@ -0,0 +1,17 @@
+#ifndef _FLEXFLOW_COMPILER_DEVICE_MAPPING_H
+#define _FLEXFLOW_COMPILER_DEVICE_MAPPING_H
+
+#include "compiler/machine_mapping/device_mapping.dtg.h"
+#include "compiler/machine_mapping/machine_mapping.dtg.h"
+#include "pcg/machine_specification.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
+
+namespace FlexFlow {
+
+DeviceMapping get_device_mapping(MachineMapping const &machine_mapping,
+                                 MachineSpecification const &machine_spec,
+                                 ParallelComputationGraph const &pcg);
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/compiler/include/compiler/machine_mapping/device_mapping.struct.toml b/lib/compiler/include/compiler/machine_mapping/device_mapping.struct.toml
@@ -0,0 +1,26 @@
+namespace = "FlexFlow"
+name = "DeviceMapping"
+features = [
+  "eq",
+  # "ord",
+  "hash",
+  # "json",
+  # "rapidcheck",
+  "fmt",
+]
+
+includes = [
+  "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h",
+  "pcg/device_id_t.dtg.h"
+]  
+
+src_includes = [   
+  "utils/hash/unordered_map.h",
+  "utils/fmt/unordered_map.h",  
+  "utils/hash/unordered_set.h",  
+  "utils/fmt/unordered_set.h" 
+]
+
+[[fields]]
+name = "raw_device_map"
+type = "std::unordered_map<::FlexFlow::parallel_layer_guid_t, std::unordered_set<::FlexFlow::device_id_t>>"
diff --git a/lib/compiler/include/compiler/machine_mapping/machine_mapping.h b/lib/compiler/include/compiler/machine_mapping/machine_mapping.h
@@ -2,6 +2,10 @@
 #define _FLEXFLOW_COMPILER_MACHINE_MAPPING_H
 
 #include "compiler/machine_mapping/machine_mapping.dtg.h"
+#include "pcg/device_id_t.dtg.h"
+#include "pcg/machine_specification.dtg.h"
+#include "pcg/operator_task_space.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
 
 namespace FlexFlow {
 

diff --git a/lib/compiler/src/compiler/allowed_machine_views.cc b/lib/compiler/src/compiler/allowed_machine_views.cc
@@ -24,6 +24,10 @@ namespace FlexFlow {
 bool is_valid_machine_view(MachineView const &mv,
                            OperatorTaskSpace const &task,
                            MachineSpecification const &ms) {
+  if (num_dims(mv) != num_dims(task)) {
+    return false;
+  }
+
   std::optional<MachineSpaceCoordinate> maximum_device_coord =
       get_machine_space_coordinate(
           task, mv, get_task_space_maximum_coordinate(task), ms);

diff --git a/lib/compiler/src/compiler/cost_estimator/in_progress_task_comparator.cc b/lib/compiler/src/compiler/cost_estimator/in_progress_task_comparator.cc
@@ -0,0 +1,11 @@
+#include "compiler/cost_estimator/in_progress_task_comparator.h"
+#include <tuple>
+
+namespace FlexFlow {
+
+bool InProgressTaskComparator::operator()(InProgressTask const &lhs,
+                                          InProgressTask const &rhs) const {
+  return std::tie(lhs.endtime, lhs.node) > std::tie(rhs.endtime, rhs.node);
+}
+
+} // namespace FlexFlow
diff --git a/lib/compiler/src/compiler/cost_estimator/op_cost_estimate_key.cc b/lib/compiler/src/compiler/cost_estimator/op_cost_estimate_key.cc
@@ -0,0 +1,14 @@
+#include "compiler/cost_estimator/op_cost_estimate_key.h"
+#include "compiler/machine_mapping/machine_mapping_problem_tree/unmapped_op_cost_estimate_key.h"
+
+namespace FlexFlow {
+
+OpCostEstimateKey get_mapped_op_cost_estimate_key_for_layer(
+    ParallelComputationGraph const &pcg,
+    parallel_layer_guid_t const &layer,
+    MachineView const &machine_view) {
+  return map_unmapped_op_cost_estimate_key(
+      get_unmapped_op_cost_estimate_key_for_layer(pcg, layer), machine_view);
+}
+
+} // namespace FlexFlow