From 6722a82c4b6ed5c95193cc61d6a7a896c4b74d00 Mon Sep 17 00:00:00 2001
From: =?utf8?q?=EC=9D=B4=ED=95=9C=EC=A2=85/On-Device=20Lab=28SR=29/Enginee?=
 =?utf8?q?r/=EC=82=BC=EC=84=B1=EC=A0=84=EC=9E=90?=
 <hanjoung.lee@samsung.com>
Date: Wed, 17 Jul 2019 19:34:51 +0900
Subject: [PATCH] Chnage DotDumper implementation (#5661)

This commit changes `DotDumper`'s dumping procedure so we can change
attributes of `Node` after creation.

Signed-off-by: Hanjoung Lee <hanjoung.lee@samsung.com>
---
 runtimes/neurun/core/src/dumper/dot/DotDumper.cc | 37 +++++++++++++++++-------
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/runtimes/neurun/core/src/dumper/dot/DotDumper.cc b/runtimes/neurun/core/src/dumper/dot/DotDumper.cc
index 2cfd988..7c8bc67 100644
--- a/runtimes/neurun/core/src/dumper/dot/DotDumper.cc
+++ b/runtimes/neurun/core/src/dumper/dot/DotDumper.cc
@@ -15,6 +15,7 @@
  */
 
 #include <fstream>
+#include <unordered_map>
 
 #include "DotDumper.h"
 #include "DotBuilder.h"
@@ -41,18 +42,21 @@ void DotDumper::dumpIfNeeded(const std::string &tag)
   auto &operations = _graph.operations();
   auto &operands = _graph.operands();
 
-  operations.iterate([&](const model::OperationIndex &index, const model::Operation &node) {
-    neurun::dumper::dot::OperationNode node_info(index, node);
+  std::unordered_map<model::OperationIndex, std::unique_ptr<OperationNode>> operation_nodes;
+  std::unordered_map<model::OperandIndex, std::unique_ptr<OperandNode>> operand_nodes;
 
-    for (auto output : node.getOutputs())
+  operations.iterate([&](const model::OperationIndex &index, const model::Operation &op) {
+    auto node = nnfw::cpp14::make_unique<OperationNode>(index, op);
+
+    for (auto output : op.getOutputs())
     {
       using neurun::dumper::dot::OperandNode;
       auto child = std::make_shared<OperandNode>(output, OperandNode::Type::MODEL_OUTPUT,
                                                  _graph.getLowerInfo(output));
-      node_info.addEdge(child);
+      node->addEdge(child);
     }
 
-    dot_builder.update(node_info);
+    operation_nodes.emplace(index, std::move(node));
   });
 
   util::Set<model::OperandIndex> shown_operand_set;
@@ -85,16 +89,16 @@ void DotDumper::dumpIfNeeded(const std::string &tag)
       }();
 
       auto lower_info = _graph.getLowerInfo(index);
-      neurun::dumper::dot::OperandNode operand_info(index, type, lower_info);
+      auto node = nnfw::cpp14::make_unique<OperandNode>(index, type, lower_info);
 
       for (auto operation_index : object.getUses().list())
       {
-        auto &node = operations.at(operation_index);
-        auto child = std::make_shared<neurun::dumper::dot::OperationNode>(operation_index, node);
-        operand_info.addEdge(child);
+        auto &operation = operations.at(operation_index);
+        auto child = std::make_shared<OperationNode>(operation_index, operation);
+        node->addEdge(child);
       }
 
-      dot_builder.update(operand_info);
+      operand_nodes.emplace(index, std::move(node));
     }
   });
 
@@ -104,9 +108,22 @@ void DotDumper::dumpIfNeeded(const std::string &tag)
     subg_ctx->iterate([&](const model::SubgraphIndex &index, const model::Subgraph &subgraph) {
       DotSubgraphInfo subgraph_info{index, subgraph, shown_operand_set};
       dot_builder.addSubgraph(subgraph_info);
+      for (const auto &op : subgraph.operations())
+      {
+        auto found = operation_nodes.find(op.index);
+        if (found != operation_nodes.end())
+        {
+          found->second->setAttribute("fillcolor", "lightgrey");
+        }
+      }
     });
   }
 
+  for (const auto &e : operation_nodes)
+    dot_builder.update(*e.second);
+  for (const auto &e : operand_nodes)
+    dot_builder.update(*e.second);
+
   // Dump to file
   {
     std::string file_name;
-- 
2.7.4