Task Simulator

lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.h

@@ -0,0 +1,24 @@
+#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/device_id_t.dtg.h"
+#include "pcg/machine_specification.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);
+
+std::unordered_set<device_id_t>
+    get_devices_from_op_key(OpCostEstimateKey const &op_key,
+                            ParallelComputationGraph const &pcg,
+                            MachineSpecification const &machine_spec);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/cost_estimator/tensor_set_movement.h

@@ -0,0 +1,19 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_TENSOR_SET_MOVEMENT_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_TENSOR_SET_MOVEMENT_H
+
+#include "compiler/cost_estimator/tensor_set_movement.dtg.h"
+#include "pcg/machine_view.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph_edge.dtg.h"
+
+namespace FlexFlow {
+
+TensorSetMovement get_tensor_set_movement_from_pcg_edge(
+    ParallelComputationGraphEdge const &edge,
+    ParallelComputationGraph const &pcg,
+    MachineView const &src_mv,
+    MachineView const &dst_mv);
+
+} // namespace FlexFlow
+
+#endif

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 {
 

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

@@ -0,0 +1,18 @@
+#ifndef _FLEXFLOW_COMPILER_MACHINE_MAPPING_UNSTRUCTURED_DEVICE_MAPPING_H
+#define _FLEXFLOW_COMPILER_MACHINE_MAPPING_UNSTRUCTURED_DEVICE_MAPPING_H
+
+#include "compiler/machine_mapping/machine_mapping.dtg.h"
+#include "compiler/machine_mapping/unstructured_device_mapping.dtg.h"
+#include "pcg/machine_specification.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
+
+namespace FlexFlow {
+
+UnstructuredDeviceMapping
+    get_unstructured_device_mapping(MachineMapping const &machine_mapping,
+                                    MachineSpecification const &machine_spec,
+                                    ParallelComputationGraph const &pcg);
+
+} // namespace FlexFlow
+
+#endif

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

@@ -0,0 +1,26 @@
+namespace = "FlexFlow"
+name = "UnstructuredDeviceMapping"
+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>>"

lib/compiler/include/compiler/task_graph_simulator/in_progress_task.struct.toml

@@ -0,0 +1,26 @@
+namespace = "FlexFlow"
+name = "InProgressTask"
+
+features = [
+  "eq",
+  "hash",
+  "fmt",
+  "ord"
+]
+
+includes = [
+  "utils/graph/node/node.dtg.h"
+]
+
+
+[[fields]]
+name = "start_time"
+type = "float"
+
+[[fields]]
+name = "end_time"
+type = "float"
+
+[[fields]]
+name = "node"
+type = "::FlexFlow::Node"

lib/compiler/include/compiler/task_graph_simulator/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/task_graph_simulator/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

lib/compiler/include/compiler/task_graph_simulator/pcg_task.variant.toml

@@ -0,0 +1,20 @@
+namespace = "FlexFlow"
+name = "PCGTask"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+]
+
+includes = [
+  "compiler/cost_estimator/op_cost_estimate_key.dtg.h",
+  "compiler/cost_estimator/tensor_set_movement.dtg.h",
+]
+
+[[values]]
+type = "::FlexFlow::OpCostEstimateKey"
+key = "operator"
+
+[[values]]
+type = "::FlexFlow::TensorSetMovement"
+key = "tensor_movement"

lib/compiler/include/compiler/task_graph_simulator/pcg_task_graph.h

@@ -0,0 +1,17 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_PCG_TASK_GRAPH_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_PCG_TASK_GRAPH_H
+
+#include "compiler/machine_mapping/machine_mapping.dtg.h"
+#include "compiler/task_graph_simulator/pcg_task_graph.dtg.h"
+#include "pcg/machine_specification.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
+
+namespace FlexFlow {
+
+PCGTaskGraph get_pcg_task_graph(ParallelComputationGraph const &pcg,
+                                MachineMapping const &machine_mapping,
+                                MachineSpecification const &machine_spec);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/task_graph_simulator/pcg_task_graph.struct.toml

@@ -0,0 +1,19 @@
+namespace = "FlexFlow"
+name = "PCGTaskGraph"
+
+features = [
+]
+
+includes = [
+  "utils/graph/digraph/digraph_view.h",
+  "utils/bidict/bidict.h",
+  "compiler/task_graph_simulator/pcg_task.dtg.h"
+]
+
+[[fields]]
+name = "graph"
+type = "::FlexFlow::DiGraphView"
+
+[[fields]]
+name = "node_map"
+type = "::FlexFlow::bidict<::FlexFlow::Node, ::FlexFlow::PCGTask>"

lib/compiler/include/compiler/task_graph_simulator/simulate_task_graph_execution.h

@@ -0,0 +1,17 @@
+#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/task_graph_simulator/task_execution_constraint.dtg.h"
+#include "compiler/task_graph_simulator/task_graph_execution_trace.dtg.h"
+#include "utils/graph/digraph/digraph_view.h"
+#include <functional>
+namespace FlexFlow {
+
+TaskGraphExecutionTrace simulate_task_graph_execution(
+    DiGraphView const &task_graph,
+    std::function<float(Node const &)> cost_function,
+    TaskExecutionConstraint const &constraint);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/task_graph_simulator/task_execution_constraint.struct.toml

@@ -0,0 +1,15 @@
+namespace = "FlexFlow"
+name = "TaskExecutionConstraint"
+features = [
+]
+
+includes = [
+  "utils/graph/node/node.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 &)>"

lib/compiler/include/compiler/task_graph_simulator/task_graph_execution_state.struct.toml

@@ -0,0 +1,40 @@
+namespace = "FlexFlow"
+name = "TaskGraphExecutionState"
+
+features = [
+]
+
+includes = [
+  "utils/deduplicated_priority_queue.h",
+  "utils/graph/node/node.dtg.h",
+  "compiler/task_graph_simulator/in_progress_task.dtg.h", 
+  "compiler/task_graph_simulator/in_progress_task_comparator.h",
+  "<unordered_set>",
+  "<set>",
+  "<functional>"
+]
+
+src_includes = [
+  "utils/hash/unordered_set.h", 
+  "utils/fmt/unordered_set.h",
+  "utils/hash/set.h", 
+  "utils/fmt/set.h",
+  "utils/fmt/vector.h", 
+  "utils/hash/vector.h"
+]
+
+[[fields]]
+name = "ready_tasks"
+type = "std::set<::FlexFlow::Node>"
+
+[[fields]]
+name = "in_progress_tasks"
+type = "::FlexFlow::DeduplicatedPriorityQueue<::FlexFlow::InProgressTask, std::vector<::FlexFlow::InProgressTask>, ::FlexFlow::InProgressTaskComparator>"
+
+[[fields]]
+name = "finished_tasks"
+type = "std::unordered_set<::FlexFlow::Node>"
+
+[[fields]]
+name = "current_time"
+type = "float"

lib/compiler/include/compiler/task_graph_simulator/task_graph_execution_trace.h

@@ -0,0 +1,12 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_TASK_GRAPH_SIMULATOR_TASK_GRAPH_EXECUTION_TRACE_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_TASK_GRAPH_SIMULATOR_TASK_GRAPH_EXECUTION_TRACE_H
+
+#include "compiler/task_graph_simulator/task_graph_execution_trace.dtg.h"
+
+namespace FlexFlow {
+
+float get_total_execution_time(TaskGraphExecutionTrace const &trace);
+
+} // namespace FlexFlow
+
+#endif // _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_TASK_GRAPH_SIMULATOR_TASK_GRAPH_EXECUTION_TRACE_H

lib/compiler/include/compiler/task_graph_simulator/task_graph_execution_trace.struct.toml

@@ -0,0 +1,23 @@
+namespace = "FlexFlow"
+name = "TaskGraphExecutionTrace"
+
+features = [
+  "hash",
+  "fmt",
+  "eq"
+]
+
+includes = [
+  "compiler/task_graph_simulator/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>"

lib/compiler/include/compiler/task_graph_simulator/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"

lib/compiler/include/compiler/task_graph_simulator/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

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);