From ae2e783c3166bf0102a5220e11f62a2c64d27669 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=EC=96=91=EC=A2=85=EC=9B=90?= <yjw963@gmail.com>
Date: Thu, 19 Sep 2024 10:28:40 +0900
Subject: [PATCH] [luci/service] Support CircleConst as reshape's shape

This commit enhances robustness in handling CircleConst as reshape's
shape.

ONE-DCO-1.0-Signed-off-by: Jongwon Yang <yjw963@gmail.com>
---
 .../luci/service/src/Nodes/CircleReshape.cpp     | 16 +++++++++++++---
 .../service/src/Nodes/CircleReshape.test.cpp     |  4 ++--
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/compiler/luci/service/src/Nodes/CircleReshape.cpp b/compiler/luci/service/src/Nodes/CircleReshape.cpp
index 0de10960b51..bac8b874e91 100644
--- a/compiler/luci/service/src/Nodes/CircleReshape.cpp
+++ b/compiler/luci/service/src/Nodes/CircleReshape.cpp
@@ -65,6 +65,13 @@ luci::CircleNode *CloneNodeLet<CN::OPQR>::visit(const luci::CircleReshape *node)
 namespace sinf
 {
 
+/**
+ * @note  CircleReshape always has two inputs: `tensor` and `shape`.
+ *        The `shape` can be CircleConst, CircleOutputDummy, or CircleNode.
+ *        - If the `shape` is CircleConst, the shape is inferred from the constant.
+ *        - Else, the shape is inferred from the node iteself.
+ *        - TODO support CircleOutputDummy and CircleNode
+ */
 loco::TensorShape Algorithm::visit(const luci::CircleReshape *node)
 {
   LOGGER(l);
@@ -77,8 +84,7 @@ loco::TensorShape Algorithm::visit(const luci::CircleReshape *node)
 
     // Only support node's shape() is CircleConst with S32
     // TODO support other node with other types
-    auto const_shape_node = dynamic_cast<luci::CircleConst *>(node->shape());
-    if (const_shape_node != nullptr)
+    if (auto const_shape_node = dynamic_cast<luci::CircleConst *>(node->shape()))
     {
       LUCI_ASSERT(const_shape_node->dtype() == S32, "Only support int32 CircleConst");
 
@@ -87,6 +93,10 @@ loco::TensorShape Algorithm::visit(const luci::CircleReshape *node)
       for (uint32_t axis = 0; axis < shape_by_input.rank(); ++axis)
       {
         shape_by_input.dim(axis) = const_shape_node->at<S32>(axis);
+        if (const_shape_node->at<S32>(axis) < 0)
+        {
+          shape_by_input.dim(axis).unset();
+        }
       }
     }
     else
@@ -139,7 +149,7 @@ loco::TensorShape Algorithm::visit(const luci::CircleReshape *node)
   for (uint32_t dim_index = 0; dim_index < output_shape.rank(); ++dim_index)
   {
     const uint32_t dim_value = output_shape.dim(dim_index).value();
-    if (static_cast<int>(dim_value) == -1)
+    if (output_shape.dim(dim_index).known() == false)
     {
       LUCI_ASSERT(unknown_dim_index == UINT32_MAX, "More than one unknown dimension");
       unknown_dim_index = dim_index;
diff --git a/compiler/luci/service/src/Nodes/CircleReshape.test.cpp b/compiler/luci/service/src/Nodes/CircleReshape.test.cpp
index a6ae6735500..50ce3aaa3f6 100644
--- a/compiler/luci/service/src/Nodes/CircleReshape.test.cpp
+++ b/compiler/luci/service/src/Nodes/CircleReshape.test.cpp
@@ -39,7 +39,7 @@ TEST(CloneNodeTest, clone_Reshape)
   ASSERT_EQ(node_reshape->newShape()->dim(1), cloned_reshape->newShape()->dim(1));
 }
 
-TEST(ShapeRuleTest, reshape_by_input_const_static)
+TEST(ShapeRuleTest, reshape_by_const_static)
 {
   auto g = loco::make_graph();
   auto node_reshape = g->nodes()->create<luci::CircleReshape>();
@@ -71,7 +71,7 @@ TEST(ShapeRuleTest, reshape_by_input_const_static)
   ASSERT_EQ(4, output_shape.dim(1).value());
 }
 
-TEST(ShapeRuleTest, reshape_by_input_const_dynamic)
+TEST(ShapeRuleTest, reshape_by_const_dynamic)
 {
   auto g = loco::make_graph();
   auto node_reshape = g->nodes()->create<luci::CircleReshape>();