[onert] Add iterateOpSeqs into PostDfsIterator (Samsung#1961)

This commit adds iterateOpSeqs into PostDfsIterator. - iterateOpSeqs : A method for iterating OpSeqeunces in post dfs Signed-off-by: ragmani <[email protected]>
chxckj · Jun 8, 2020 · d1310e2 · d1310e2
1 parent 1fa3e50
commit d1310e2
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 1 deletion.
diff --git a/runtime/onert/core/src/ir/GraphIterator.cc b/runtime/onert/core/src/ir/GraphIterator.cc
@@ -17,7 +17,7 @@
 #include "GraphIterator.h"
 
 #include "ir/OperationIndexMap.h"
-#include "ir/Graph.h"
+#include "ir/LoweredGraph.h"
 
 namespace onert
 {
@@ -80,5 +80,39 @@ void PostDfsIterator<is_const>::iterate(GraphRef graph, const IterFn &fn) const
                      [](const std::pair<const OperationIndex, bool> &v) { return v.second; }));
 }
 
+template <bool is_const>
+void PostDfsIterator<is_const>::iterateOpSeqs(LoweredGraphRef lowered_graph,
+                                              const OpSeqIterFn &fn) const
+{
+  std::unordered_map<OpSequenceIndex, bool> visited;
+  lowered_graph.op_seqs().iterate(
+      [&](const OpSequenceIndex &index, OpSequenceRef) { visited[index] = false; });
+
+  std::function<void(const OpSequenceIndex &, OpSequenceRef)> dfs_recursive =
+      [&](const OpSequenceIndex &index, OpSequenceRef op_seq) -> void {
+    if (visited[index])
+      return;
+    visited[index] = true;
+
+    for (const auto output : op_seq.getOutputs() | Remove::DUPLICATED)
+    {
+      const auto &operand = lowered_graph.graph().operands().at(output);
+      for (const auto &use : operand.getUses())
+      {
+        const auto use_op_seq_index = lowered_graph.op_seqs().getOperation(use);
+        dfs_recursive(use_op_seq_index, lowered_graph.op_seqs().at(use_op_seq_index));
+      }
+    }
+
+    fn(index, op_seq);
+  };
+
+  lowered_graph.op_seqs().iterate(dfs_recursive);
+
+  // All of the operations(nodes) must have been visited.
+  assert(std::all_of(visited.begin(), visited.end(),
+                     [](const std::pair<const OpSequenceIndex, bool> &v) { return v.second; }));
+}
+
 } // namespace ir
 } // namespace onert
diff --git a/runtime/onert/core/src/ir/GraphIterator.h b/runtime/onert/core/src/ir/GraphIterator.h
@@ -28,6 +28,8 @@ namespace ir
 
 class Graph;
 class Operation;
+class LoweredGraph;
+class OpSequence;
 
 template <bool is_const> class Iterator
 {
@@ -62,9 +64,15 @@ template <bool is_const = false> class PostDfsIterator final : public Iterator<i
   using IndexRef = typename Iterator<is_const>::IndexRef;
   using NodeRef = typename Iterator<is_const>::NodeRef;
   using IterFn = typename Iterator<is_const>::IterFn;
+  using LoweredGraphRef =
+      typename std::conditional<is_const, const LoweredGraph &, LoweredGraph &>::type;
+  using OpSequenceRef = typename std::conditional<is_const, const OpSequence &, OpSequence &>::type;
+  using OpSeqIndexRef = const OpSequenceIndex &;
+  using OpSeqIterFn = std::function<void(OpSeqIndexRef, OpSequenceRef &)>;
 
 public:
   void iterate(GraphRef graph, const IterFn &fn) const;
+  void iterateOpSeqs(LoweredGraphRef lowered_graph, const OpSeqIterFn &f) const;
 };
 using PostDfsConstIterator = PostDfsIterator<true>;