[loco] Update index on link (#6381)

Both "link" helpers now update the index of a corresponding node.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>

diff --git a/compiler/loco/src/IR/Nodes.cpp b/compiler/loco/src/IR/Nodes.cpp
index 47b61d7..000b207 100644 (file)
--- a/compiler/loco/src/IR/Nodes.cpp
+++ b/compiler/loco/src/IR/Nodes.cpp
@@ -57,7 +57,11 @@ GraphOutputIndex Push::index(void) const
   return static_cast<GraphOutputIndex>(_index);
 }
 
-void link(GraphOutput *output, Push *push) { output->node(push); }
+void link(GraphOutput *output, Push *push)
+{
+  output->node(push);
+  push->index(output->index());
+}
 
 } // namespace loco
 
@@ -137,7 +141,11 @@ DataType Pull::dtype(void) const
   }
 }
 
-void link(GraphInput *input, Pull *pull) { input->node(pull); }
+void link(GraphInput *input, Pull *pull)
+{
+  input->node(pull);
+  pull->index(input->index());
+}
 
 } // namespace loco
 

diff --git a/compiler/loco/src/loco.test.cpp b/compiler/loco/src/loco.test.cpp
index 7496a18..9c352f0 100644 (file)
--- a/compiler/loco/src/loco.test.cpp
+++ b/compiler/loco/src/loco.test.cpp
@@ -65,15 +65,11 @@ TEST(LOCO, identity_network)
   graph_output->name("output");
   loco::link(graph_output, push_node);
 
-  // Set Input Index
-  pull_node->index(0);
-
+  // loco::link SHOULD update "index"
   ASSERT_EQ(pull_node->index(), 0);
   ASSERT_EQ(graph_input->dtype(), loco::DataType::FLOAT32);
 
-  // Set Output Index
-  push_node->index(0);
-
+  // loco::link SHOULD update "index"
   ASSERT_EQ(push_node->index(), 0);
 }