Per-operator stream assignment.

It works by traversing the graph in depth-first fashion, from outputs
to inputs. Each node has a set of "ready" streams which are merged
when the graph converges. When a stream is reused in a subsequent
node, the stream id is removed from the producer's ready set
and the removal happens recursively, but only once for each stream id.

Typical complexity is O(V+E), pessimistic complexity is O(V+E^2).

Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>

dali/pipeline/executor/executor2/stream_assignment.h

@@ -38,8 +38,6 @@ enum class StreamPolicy : int {
   Single,       //< There's just one stream that's used by all operators
   PerBackend,   //< Operators are scheduled on a stream specific to their backend (mixed or GPU)
   PerOperator   //< Independent operators are executed on separate streams.
-
-  // TODO(michalz): implement minimal assignment for PerOperator policy
 };
 
 
@@ -164,6 +162,159 @@ class StreamAssignment<StreamPolicy::PerBackend> {
   bool has_mixed_ = false;
 };
 
+/** Implements per-operator stream assignment.
+ *
+ * This policy implements stream assingment such that independent GPU/Mixed operators get
+ * separate streams. When there's a dependency then one dependent operator shares the stream of
+ * its predecessor.
+ *
+ * Example - numbers are stream indices, "X" means no stream, "s" means synchronization
+ * ```
+ * CPU(X) ---- GPU(0) --- GPU(0) -- GPU(0) -- output 0
+ *              \                    s
+ *               \                  /
+ *                ----- GPU(1) ----
+ *                        \
+ *                         \
+ * CPU(X) --- GPU(2) ----s GPU(1) ----------s output 1
+ * ```
+ */
+template <>
+class StreamAssignment<StreamPolicy::PerOperator> {
+ public:
+  explicit StreamAssignment(ExecGraph &graph) {
+    Assign(graph);
+  }
+
+  std::optional<int> operator[](const ExecNode *node) const {
+    auto it = node_meta_.find(node);
+    assert(it != node_meta_.end());
+    return it->second.stream_id;
+  }
+
+  /** Gets the total number of streams required to run independent operators on separate streams. */
+  int NumStreams() const {
+    return total_streams_;
+  }
+
+ private:
+  struct NodeMeta {
+    int index;  // index in the sorted_nodes_
+    bool needs_stream = false;     // whether this exact node needs a stream
+    bool gpu_contributor = false;  // whether the node contributes to a GPU operator
+    std::optional<int> stream_id;
+    std::set<int> ready_streams;
+  };
+
+  void Assign(ExecGraph &graph) {
+    int num_nodes = graph.Nodes().size();
+
+    // the nodes in the graph must be sorted topologically
+    sorted_nodes_.reserve(num_nodes);
+    for (auto &node : graph.Nodes()) {
+      int idx = sorted_nodes_.size();
+      NodeMeta &meta = node_meta_.insert({ &node, {} }).first->second;
+      meta.index = idx;
+      sorted_nodes_.push_back({ &node, &meta });
+    }
+
+    assert(static_cast<size_t>(num_nodes) == sorted_nodes_.size());
+
+    FindGPUContributors(graph);
+    RunAssignment(graph);
+    ClearCPUStreams();
+  }
+
+
+  void ClearCPUStreams() {
+    for (auto &[node, meta] : node_meta_)
+      if (!meta.needs_stream)
+        meta.stream_id = std::nullopt;
+  }
+
+
+  void FindGPUContributors(ExecGraph &graph) {
+    // Run DFS, output to input, and find nodes which contribute to any node that requires a stream
+    graph::ClearVisitMarkers(graph.Nodes());
+    for (auto it = graph.Nodes().rbegin(); it != graph.Nodes().rend(); ++it) {
+      auto &node = *it;
+      FindGPUContributors(&node, false);
+    }
+  }
+
+  void FindGPUContributors(const ExecNode *node, bool is_gpu_contributor) {
+    graph::Visit v(node);
+    if (!v)
+      return;
+    auto &meta = node_meta_[node];
+    meta.needs_stream = NeedsStream(node);
+    if (!is_gpu_contributor)
+      is_gpu_contributor = meta.needs_stream;
+    if (is_gpu_contributor)
+      meta.gpu_contributor = true;
+    for (auto *inp : node->inputs)
+      FindGPUContributors(inp->producer, is_gpu_contributor);
+  }
+
+  void RunAssignment(ExecGraph &graph) {
+    graph::ClearVisitMarkers(graph.Nodes());
+    for (int i = sorted_nodes_.size() - 1; i >= 0; i--) {
+      ProcessNode(sorted_nodes_[i].first, sorted_nodes_[i].second);
+    }
+  }
+
+  void ProcessNode(const ExecNode *node, NodeMeta *meta) {
+    graph::Visit v(node);
+    if (!v)
+      return;
+
+    auto &stream_id = meta->stream_id;
+    assert(!stream_id.has_value());
+
+    for (auto &e : node->inputs) {
+      auto &prod_meta = node_meta_[e->producer];
+      ProcessNode(e->producer, &prod_meta);
+      if (meta->gpu_contributor && !prod_meta.ready_streams.empty()) {
+        if (!stream_id.has_value() || *stream_id > *prod_meta.ready_streams.begin()) {
+          stream_id = *prod_meta.ready_streams.begin();
+        }
+      }
+      meta->ready_streams.insert(prod_meta.ready_streams.begin(), prod_meta.ready_streams.end());
+    }
+
+    if (stream_id.has_value()) {
+      for (auto &e : node->inputs) {
+        EraseReady(e->producer, *stream_id, node);
+      }
+    } else {
+      if (meta->needs_stream) {
+        stream_id = total_streams_++;
+        meta->ready_streams.insert(*stream_id);
+      }
+    }
+
+    assert(!stream_id.has_value() || meta->ready_streams.count(*stream_id));
+  }
+
+  void EraseReady(const ExecNode *node, int id, const ExecNode *sentinel) {
+    if (node == sentinel)
+      return;
+    auto &meta = node_meta_[node];
+    if (meta.ready_streams.erase(id)) {
+      for (auto &e : node->inputs)
+        EraseReady(e->producer, id, sentinel);
+      for (auto &out : node->outputs)
+        for (auto &e : out.consumers)
+          EraseReady(e->consumer, id, sentinel);
+    }
+  }
+
+  int total_streams_ = 0;
+
+  std::unordered_map<const ExecNode *, NodeMeta> node_meta_;
+  std::vector<std::pair<const ExecNode *, NodeMeta *>> sorted_nodes_;
+};
+
 }  // namespace exec2
 }  // namespace dali