Adding task-based simulator for PCGs

lib/compiler/include/compiler/machine_mapping.h

@@ -69,6 +69,13 @@ OptimalCostResult
                  MachineSpecification const &resources,
                  OptimalCostCache &cached_subgraph_costs);
 
+float parallel_estimate_cost(
+    SubParallelComputationGraphView const &g,
+    CostEstimator const &estimator,
+    MachineMapping const &device_mapping,
+    std::unordered_map<OpenMultiDiEdge, MachineView> const
+        &frontier_machine_views);
+
 } // namespace FlexFlow
 
 namespace std {

lib/compiler/src/machine_mapping.cc

@@ -5,6 +5,9 @@
 #include "utils/exception.h"
 #include "utils/graph/serialparallel.h"
 
+#include "utils/deduplicated_priority_queue.h"
+#include <algorithm>
+
 namespace FlexFlow {
 
 MachineMapping MachineMapping::combine(MachineMapping const &s1,
@@ -110,6 +113,100 @@ float estimate_cost(SubParallelComputationGraphView const &g,
   return cost;
 }
 
+// Computes estimated execution cost for a single node
+float node_estimate_cost(Node node,
+                         SubParallelComputationGraphView const &g,
+                         CostEstimator const &estimator,
+                         MachineMapping const &device_mapping) {
+  std::unordered_set<UpwardOpenMultiDiEdge> incoming_edges =
+      get_incoming_edges(g, node);
+  std::vector<ParallelTensorShape> inputs = transform(
+      as_vector(incoming_edges), [&](UpwardOpenMultiDiEdge const &input_edge) {
+        return g.at(input_edge).get_shape();
+      });
+  float cost = estimator.estimate_cost(
+      g.at(node).attrs, inputs, device_mapping.machine_views.at(node));
+  return cost;
+}
+
+struct TimedNode { // Node and associated finishing time
+  Node node;
+  float endtime;
+};
+
+struct TimeComparison {
+  bool operator()(TimedNode const &lhs, TimedNode const &rhs) const {
+    return (lhs.endtime < rhs.endtime);
+  }
+};
+
+float parallel_estimate_cost(
+    SubParallelComputationGraphView const &g,
+    CostEstimator const &estimator,
+    MachineMapping const &device_mapping,
+    std::unordered_map<OpenMultiDiEdge, MachineView> const
+        &frontier_machine_views) {
+
+  std::unordered_set<Node>
+      frontier; // nodes whose dependencies (previous nodes) have been met, and
+                // are waiting to be processed.
+  DeduplicatedPriorityQueue<TimedNode, std::vector<TimedNode>, TimeComparison>
+      processing; // nodes currently being processed.
+  std::unordered_set<Node>
+      processed; // set of nodes that have already been processed
+  std::unordered_map<int, bool>
+      occupied; // keeps track of the devices that are currently occupied
+
+  // Filling the frontier
+  for (auto &[edge, _] : frontier_machine_views) {
+    Node node = get_dst_node(edge);
+    frontier.insert(node);
+  }
+
+  float current_time = 0;
+  while (!frontier.empty() || !processing.empty()) {
+
+    // Processing new nodes
+    std::unordered_set<Node> copy(frontier);
+    for (Node const &node : copy) {
+      std::vector<int> devices =
+          device_mapping.machine_views.at(node).machine_ids();
+      if (std::all_of(devices.begin(), devices.end(), [&occupied](int d) {
+            return occupied[d] == false
+          })) {
+        float cost = node_estimate_cost(node, g, estimator, device_mapping);
+        processing.push({node, current_time + cost});
+        for (int d : devices) {
+          occupied[d] = true;
+        }
+        frontier.erase(node);
+      }
+    }
+
+    // Finish processing one node
+    TimedNode finished = processing.pop();
+    std::vector<int> devices =
+        device_mapping.machine_views.at(finished.node).machine_ids();
+    for (int d : devices) { // free devices
+      occupied[d] = false;
+    }
+    processed.insert(finished.node);
+    current_time = finished.endtime;
+
+    // Adding candidates to the frontier
+    for (Node const &successor :
+         get_successors(g, finished.node)) { // All nodes depending on finished
+      std::unordered_set<Node> predecessors = get_predecessors(g, successor);
+      if (std::all_of(predecessors.begin(),
+                      predecessors.end(),
+                      [&processed](int p) { return processed.contains(p) })) {
+        frontier.insert(successor);
+      }
+    }
+  }
+  return current_time;
+}
+
 void minimize_runtime(OptimalCostResult &m1, OptimalCostResult const &m2) {
   minimize(m1, m2, OptimalCostRuntimeCmp{});
 }

lib/compiler/test/src/test_parallel_cost_estimator.cpp

@@ -0,0 +1,50 @@
+#include "compiler/machine_mapping.h"
+#include "doctest/doctest.h"
+#include "doctest/test_cost_estimator.h#include "compiler/cost_estimate.h"
+
+namespace FlexFlow {
+
+struct TestCostEstimator : public ICostEstimator {
+  float estimate_cost(PCGOperatorAttrs const &op,
+                      std::vector<ParallelTensorShape> const &inputs,
+                      MachineView const &mv) const override {
+    return 0.1;
+  }"
+#include "rapidcheck.h"
+
+
+using namespace FlexFlow;
+
+
+
+TEST_SUITE(FF_TEST_SUITE) {
+  TEST_CASE("parallel_estimate_cost: linear graph") {
+    //Straight line example
+    SubParallelComputationGraphView g =
+        OutputLabelledOpenMultiDiGraphView<Operator, ParallelTensor>::create<
+            UnorderedOutputLabelledOpenMultiDiGraph<Operator, ParallelTensor>>();
+    Node n0 = g.add_node(Operator{InputAttrs{}, "n0"});
+    Node n1 = g.add_node(Operator{InputAttrs{}, "n1"});
+    Node n2 = g.add_node(Operator{InputAttrs{}, "n2"});
+
+    NodePort p0 = g.add_node_port();
+    NodePort p1 = g.add_node_port();
+    NodePort p2 = g.add_node_port();
+
+    // MultiDiEdge: dst, dstport, src, srcport
+    MultiDiEdge e0{n1, p1, n0, p0};
+    MultiDiEdge e1{n2, p2, n1, p1};
+
+    g.add_edge(e0);
+    g.add_edge(e1);
+    g.add_edge(e2);
+    g.add_label(e0, 10);
+    g.add_label(e1, 11);
+
+    CostEstimator estimator = CostEstimator::create<TestCostEstimator>();
+    MachineMapping device_mapping = std::unordered_set<MachineView>{
+          make_1d_machine_view(gpu_id_t(1), gpu_id_t(2))};
+    views
+    parallel_estimate_cost(g, estimator, device_mapping, frontier_machine_views);
+  }
+}

lib/utils/include/utils/graph/algorithms.h

@@ -160,6 +160,10 @@ std::unordered_set<Node> get_predecessors(DiGraphView const &, Node const &);
 std::unordered_map<Node, std::unordered_set<Node>>
     get_predecessors(DiGraphView const &, std::unordered_set<Node> const &);
 
+std::unordered_set<Node> get_successors(DiGraphView const &, Node const &);
+std::unordered_map<Node, std::unordered_set<Node>>
+    get_successors(DiGraphView const &, std::unordered_set<Node> const &);
+
 Node get_src_node(MultiDiEdge const &);
 Node get_dst_node(MultiDiEdge const &);
 Node get_dst_node(InputMultiDiEdge const &);

lib/utils/src/graph/algorithms.cc

@@ -376,6 +376,24 @@ std::unordered_set<Node> get_predecessors(DiGraphView const &g, Node const &n) {
   return get_predecessors(g, std::unordered_set<Node>{n}).at(n);
 }
 
+std::unordered_map<Node, std::unordered_set<Node>>
+    get_successors(DiGraphView const &g,
+                   std::unordered_set<Node> const &nodes) {
+
+  std::unordered_map<Node, std::unordered_set<Node>> successors;
+  for (Node const &n : nodes) {
+    successors[n];
+  }
+  for (DirectedEdge const &e : get_outgoing_edges(g, nodes)) {
+    successors.at(e.src).insert(e.dst);
+  }
+  return successors;
+}
+
+std::unordered_set<Node> get_successors(DiGraphView const &g, Node const &n) {
+  return get_successors(g, std::unordered_set<Node>{n}).at(n);
+}
+
 std::vector<Node> get_unchecked_dfs_ordering(
     DiGraphView const &g, std::unordered_set<Node> const &starting_points) {
   UncheckedDFSView dfs_view = unchecked_dfs(g, starting_points);