Task Simulator

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

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

lib/compiler/include/compiler/cost_estimator/tasks_state_tracker.struct.toml

@@ -0,0 +1,32 @@
+namespace = "FlexFlow"
+name = "TasksStateTracker"
+
+features = [
+]
+
+includes = [
+  "utils/deduplicated_priority_queue.h",
+  "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h",
+  "compiler/cost_estimator/timed_component.dtg.h",
+  "compiler/cost_estimator/timed_component_comparator.h",
+  "<unordered_set>" 
+]
+
+src_includes = [
+  "utils/hash/unordered_set.h", 
+  "utils/fmt/unordered_set.h",
+  "utils/fmt/vector.h", 
+  "utils/hash/vector.h"
+]
+
+[[fields]]
+name = "ready_layers"
+type = "std::unordered_set<::FlexFlow::parallel_layer_guid_t>"
+
+[[fields]]
+name = "component_processing"
+type = "::FlexFlow::DeduplicatedPriorityQueue<::FlexFlow::TimedComponent, std::vector<::FlexFlow::TimedComponent>, ::FlexFlow::TimedComponentComparator>"
+
+[[fields]]
+name = "processed_components"
+type = "std::unordered_set<::FlexFlow::TimedComponent>"

lib/compiler/include/compiler/cost_estimator/timed_component.variant.toml

@@ -0,0 +1,19 @@
+namespace = "FlexFlow"
+name = "TimedComponent"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+  "ord"
+]
+
+includes = [ 
+  "compiler/cost_estimator/timed_dependency.dtg.h",
+  "compiler/cost_estimator/timed_layer.dtg.h",
+]
+
+[[values]]
+type = "::FlexFlow::TimedLayer"
+
+[[values]]
+type = "::FlexFlow::TimedDependency"

lib/compiler/include/compiler/cost_estimator/timed_component_comparator.h

@@ -0,0 +1,14 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_TIMED_COMPONENT_COMPARATOR_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_TIMED_COMPONENT_COMPARATOR_H
+
+#include "compiler/cost_estimator/timed_component.dtg.h"
+
+namespace FlexFlow {
+
+struct TimedComponentComparator {
+  bool operator()(TimedComponent const &lhs, TimedComponent const &rhs) const;
+};
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/cost_estimator/timed_dependency.struct.toml

@@ -0,0 +1,20 @@
+namespace = "FlexFlow"
+name = "TimedDependency"
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "fmt",
+]
+
+includes = [ 
+  "pcg/parallel_computation_graph/parallel_computation_graph_edge.dtg.h",
+]
+
+[[fields]]
+name = "endtime"
+type = "float"
+
+[[fields]]
+name = "raw_edge"
+type = "::FlexFlow::ParallelComputationGraphEdge"

lib/compiler/include/compiler/cost_estimator/timed_layer.struct.toml

@@ -0,0 +1,20 @@
+namespace = "FlexFlow"
+name = "TimedLayer"
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "fmt",
+]
+
+includes = [ 
+  "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h",
+]
+
+[[fields]]
+name = "endtime"
+type = "float"
+
+[[fields]]
+name = "layer"
+type = "::FlexFlow::parallel_layer_guid_t"

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

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

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>>"

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

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