[locop] Introduce push_shape/pull_shape helper (#6398)

Let's replace tensor_shape with push_shape/pull_shape as the first step
to remove shape attribtues from Pull/Push node.

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

diff --git a/compiler/locop/src/CanonicalNodeSummaryBuilder.cpp b/compiler/locop/src/CanonicalNodeSummaryBuilder.cpp
index 7335a66..9b762b2 100644 (file)
--- a/compiler/locop/src/CanonicalNodeSummaryBuilder.cpp
+++ b/compiler/locop/src/CanonicalNodeSummaryBuilder.cpp
@@ -62,6 +62,34 @@ loco::TensorShape tensor_shape(const loco::NodeMixin<loco::NodeTrait::TensorShap
   return res;
 }
 
+loco::TensorShape push_shape(const loco::Push *m)
+{
+  loco::TensorShape res;
+
+  res.rank(m->rank());
+
+  for (uint32_t axis = 0; axis < m->rank(); ++axis)
+  {
+    res.dim(axis) = m->dim(axis);
+  }
+
+  return res;
+}
+
+loco::TensorShape pull_shape(const loco::Pull *m)
+{
+  loco::TensorShape res;
+
+  res.rank(m->rank());
+
+  for (uint32_t axis = 0; axis < m->rank(); ++axis)
+  {
+    res.dim(axis) = m->dim(axis);
+  }
+
+  return res;
+}
+
 std::ostream &operator<<(std::ostream &os, const loco::Dimension &d)
 {
   os << (d.known() ? std::to_string(d.value()) : std::string{"?"});
@@ -193,7 +221,7 @@ public:
 
     res.arg("index", node->indexed() ? pp::fmt(node->index()) : pp::fmt('?'));
     res.arg("from", symbol_lookup(*_symtbl, node->from()));
-    res.arg("shape", pp::fmt(pretty(tensor_shape(node))));
+    res.arg("shape", pp::fmt(pretty(push_shape(node))));
     res.state(NodeDesc::State::Complete);
 
     return res;
@@ -205,7 +233,7 @@ public:
 
     res.arg("index", node->indexed() ? pp::fmt(node->index()) : pp::fmt('?'));
     // TODO Print dtype
-    res.arg("shape", pp::fmt(pretty(tensor_shape(node))));
+    res.arg("shape", pp::fmt(pretty(pull_shape(node))));
     res.state(NodeDesc::State::PartiallyKnown);
 
     return res;