[moco-tf] Use loco::NodeShape for inputs (#7008)

* [moco-tf] Use loco::NodeShape for inputs

This will change fix_shape() to use loco::NodeShape for input's shape information

Signed-off-by: SaeHie Park <saehie.park@samsung.com>
* fix vars

diff --git a/compiler/moco-tf/src/Transforms/FixShapeTransform.cpp b/compiler/moco-tf/src/Transforms/FixShapeTransform.cpp
index 8304c61..817d9c8 100644 (file)
--- a/compiler/moco-tf/src/Transforms/FixShapeTransform.cpp
+++ b/compiler/moco-tf/src/Transforms/FixShapeTransform.cpp
@@ -827,18 +827,13 @@ bool fix_shape(moco::tf::TFAdd *node)
 {
   auto x = node->x();
   auto y = node->y();
-  if (x == nullptr || y == nullptr)
-  {
-    // this node maybe disconnected by some transformation
-    // TODO remove this block after fixing to iterate only active nodes
+  loco::NodeShape x_shape;
+  loco::NodeShape y_shape;
+
+  if (!node_shape(x, x_shape))
     return false;
-  }
-  auto x_shapedata = x->annot<ShapeInferenceData>();
-  auto y_shapedata = y->annot<ShapeInferenceData>();
-  if (x_shapedata == nullptr || y_shapedata == nullptr)
-  {
+  if (!node_shape(y, y_shape))
     return false;
-  }
   // TODO check shape difference
 
   // Output shape is same as the input
@@ -856,8 +851,8 @@ bool fix_shape(moco::tf::TFAvgPool *node)
     return false;
   }
   auto value = node->value();
-  auto value_shapedata = value->annot<ShapeInferenceData>();
-  if (value_shapedata == nullptr)
+  loco::NodeShape value_shape;
+  if (!node_shape(value, value_shape))
   {
     // input node shape inference is not ready
     return false;
@@ -869,7 +864,7 @@ bool fix_shape(moco::tf::TFAvgPool *node)
   update_stride_data(node);
   update_window_data(node);
 
-  auto value_feature_shape = as_feature_shape(*value_shapedata, node->data_layout());
+  auto value_feature_shape = as_feature_shape(value_shape, node->data_layout());
 
   auto stride_data = node->annot<StrideData>();
   assert(stride_data != nullptr);
@@ -926,9 +921,9 @@ bool fix_shape(moco::tf::TFBiasAdd *node)
 {
   auto value = node->value();
   auto bias = node->bias();
-  auto value_shapedata = value->annot<ShapeInferenceData>();
-  auto bias_shapedata = bias->annot<ShapeInferenceData>();
-  if (value_shapedata == nullptr || bias_shapedata == nullptr)
+  loco::NodeShape value_shape;
+  loco::NodeShape bias_shape;
+  if (!node_shape(value, value_shape) || !node_shape(bias, bias_shape))
   {
     return false;
   }
@@ -941,8 +936,11 @@ template <class CONST_CLASS> bool valid_scala_value(CONST_CLASS *node)
 {
   LOGGER(l);
 
-  auto shapedata = node->template annot<ShapeInferenceData>();
-  assert(shapedata != nullptr);
+  loco::NodeShape nodeshape;
+  if (!node_shape(node, nodeshape))
+  {
+    return false;
+  }
 
   if (node->dtype() != loco::DataType::S32)
   {
@@ -950,7 +948,7 @@ template <class CONST_CLASS> bool valid_scala_value(CONST_CLASS *node)
     return false;
   }
 
-  auto tensor_shape = shapedata->tensor_shape();
+  auto tensor_shape = nodeshape.as<loco::TensorShape>();
   if (!(tensor_shape.rank() == 0 || tensor_shape.rank() == 1))
   {
     INFO(l) << "valid_scala_value rank not 0/1 : " << tensor_shape.rank();
@@ -962,12 +960,15 @@ template <class CONST_CLASS> bool valid_scala_value(CONST_CLASS *node)
 
 template <class CONST_CLASS> int32_t scala_value(CONST_CLASS *node)
 {
-  auto shapedata = node->template annot<ShapeInferenceData>();
-  assert(shapedata != nullptr);
+  loco::NodeShape nodeshape;
+  if (!node_shape(node, nodeshape))
+  {
+    return false;
+  }
 
   assert(node->dtype() == loco::DataType::S32);
 
-  auto tensor_shape = shapedata->tensor_shape();
+  auto tensor_shape = nodeshape.as<loco::TensorShape>();
   assert(tensor_shape.rank() == 0 || tensor_shape.rank() == 1);
 
   return node->template at<loco::DataType::S32>(0);
@@ -989,34 +990,38 @@ bool fix_shape(moco::tf::TFConcatV2 *node)
   // Check shape inference data are all ready
   // Check shape rank are all same
   auto value_a = node->values(0);
-  auto value_a_shapedata = value_a->annot<ShapeInferenceData>();
-  if (value_a_shapedata == nullptr)
+  loco::NodeShape value_a_shape;
+  if (!node_shape(value_a, value_a_shape))
   {
     // shape inference is not ready for this value
     INFO(l) << "Fix shape TFConcatV2 value 0 shape_data not ready";
     return false;
   }
-  uint32_t a_rank = value_a_shapedata->rank();
+  assert(value_a_shape.domain() == loco::Domain::Tensor);
+  auto value_a_tensor_shape = value_a_shape.as<loco::TensorShape>();
+  uint32_t a_rank = value_a_tensor_shape.rank();
 
   uint32_t num_values = node->num_values();
   for (uint32_t ni = 1; ni < num_values; ++ni)
   {
     auto value_b = node->values(ni);
-    auto value_b_shapedata = value_b->annot<ShapeInferenceData>();
-    if (value_b_shapedata == nullptr)
+    loco::NodeShape value_b_shape;
+    if (!node_shape(value_b, value_b_shape))
     {
       // shape inference is not ready for this value
       INFO(l) << "Fix shape TFConcatV2 value " << ni << " shape_data not ready";
       return false;
     }
-    uint32_t b_rank = value_b_shapedata->rank();
+    assert(value_b_shape.domain() == loco::Domain::Tensor);
+    auto value_b_tensor_shape = value_b_shape.as<loco::TensorShape>();
+    uint32_t b_rank = value_b_tensor_shape.rank();
     assert(a_rank == b_rank);
   }
 
   // check for axis
   auto axis_node = node->axis();
-  auto axis_shapedata = axis_node->annot<ShapeInferenceData>();
-  if (axis_shapedata == nullptr)
+  loco::NodeShape axis_shape;
+  if (!node_shape(axis_node, axis_shape))
   {
     // shape inference is not ready for axis_node
     INFO(l) << "Fix shape TFConcatV2 axis shape_data not ready";
@@ -1066,18 +1071,21 @@ bool fix_shape(moco::tf::TFConcatV2 *node)
 
   for (uint32_t index = 0; index < a_rank; ++index)
   {
-    if (value_a_shapedata->dim(index).known())
+    if (value_a_tensor_shape.dim(index).known())
     {
-      uint32_t dim = value_a_shapedata->dim(index).value();
+      uint32_t dim = value_a_tensor_shape.dim(index).value();
       if (index == axis_absolute)
       {
         uint32_t dim_acc = dim;
         for (uint32_t ni = 1; ni < num_values; ++ni)
         {
           auto value_b = node->values(ni);
-          auto value_b_shapedata = value_b->annot<ShapeInferenceData>();
-          assert(value_b_shapedata->dim(index).known());
-          dim_acc += value_b_shapedata->dim(index).value();
+          loco::NodeShape value_b_shape;
+          node_shape(value_b, value_b_shape);
+          assert(value_b_shape.domain() == loco::Domain::Tensor);
+          auto value_b_tensor_shape = value_b_shape.as<loco::TensorShape>();
+          assert(value_b_tensor_shape.dim(index).known());
+          dim_acc += value_b_tensor_shape.dim(index).value();
         }
         dim = dim_acc;
       }
@@ -1127,16 +1135,16 @@ bool fix_shape(moco::tf::TFConv2D *node)
     return false;
   }
   auto ifm = node->ifm();
-  auto ifm_shapedata = ifm->annot<ShapeInferenceData>();
-  if (ifm_shapedata == nullptr)
+  loco::NodeShape ifm_shape;
+  if (!node_shape(ifm, ifm_shape))
   {
     // input node shape inference is not ready
     return false;
   }
 
   auto ker = node->ker();
-  auto ker_shapedata = ker->annot<ShapeInferenceData>();
-  if (ker_shapedata == nullptr)
+  loco::NodeShape ker_shape;
+  if (!node_shape(ker, ker_shape))
   {
     return false;
   }
@@ -1152,8 +1160,8 @@ bool fix_shape(moco::tf::TFConv2D *node)
   INFO(l) << "Fix TFConv2D strides = " << stride_data->stride()->vertical() << ", "
           << stride_data->stride()->horizontal();
 
-  auto ifm_tensor_shape = ifm_shapedata->tensor_shape(); // in NHWC
-  auto ker_tensor_shape = ker_shapedata->tensor_shape(); // in HWIO
+  auto ifm_tensor_shape = ifm_shape.as<loco::TensorShape>(); // in NHWC
+  auto ker_tensor_shape = ker_shape.as<loco::TensorShape>(); // in HWIO
   assert(ifm_tensor_shape.rank() == 4);
   assert(ker_tensor_shape.rank() == 4);
 
@@ -1219,16 +1227,16 @@ bool fix_shape(moco::tf::TFDepthwiseConv2dNative *node)
     return false;
   }
   auto ifm = node->ifm();
-  auto ifm_shapedata = ifm->annot<ShapeInferenceData>();
-  if (ifm_shapedata == nullptr)
+  loco::NodeShape ifm_shape;
+  if (!node_shape(ifm, ifm_shape))
   {
     // input node shape inference is not ready
     return false;
   }
 
   auto ker = node->ker();
-  auto ker_shapedata = ker->annot<ShapeInferenceData>();
-  if (ker_shapedata == nullptr)
+  loco::NodeShape ker_shape;
+  if (!node_shape(ker, ker_shape))
   {
     return false;
   }
@@ -1241,8 +1249,8 @@ bool fix_shape(moco::tf::TFDepthwiseConv2dNative *node)
   INFO(l) << "FixShape TFDepthwiseConv2dNative strides = " << stride_data->stride()->vertical()
           << ", " << stride_data->stride()->horizontal();
 
-  auto ifm_tensor_shape = ifm_shapedata->tensor_shape(); // in NHWC
-  auto ker_tensor_shape = ker_shapedata->tensor_shape(); // in HWCM
+  auto ifm_tensor_shape = ifm_shape.as<loco::TensorShape>(); // in NHWC
+  auto ker_tensor_shape = ker_shape.as<loco::TensorShape>(); // in HWCM
   assert(ifm_tensor_shape.rank() == 4);
   assert(ker_tensor_shape.rank() == 4);
 
@@ -1322,8 +1330,8 @@ bool fix_shape(moco::tf::TFMaxPool *node)
     return false;
   }
   auto value = node->value();
-  auto value_shapedata = value->annot<ShapeInferenceData>();
-  if (value_shapedata == nullptr)
+  loco::NodeShape value_shape;
+  if (!node_shape(value, value_shape))
   {
     // input node shape inference is not ready
     return false;
@@ -1335,13 +1343,13 @@ bool fix_shape(moco::tf::TFMaxPool *node)
   update_stride_data(node);
   update_window_data(node);
 
-  auto value_feature_shape = as_feature_shape(*value_shapedata, node->data_layout());
-
   auto stride_data = node->annot<StrideData>();
   assert(stride_data != nullptr);
   auto window_data = node->annot<WindowData>();
   assert(window_data != nullptr);
 
+  auto value_feature_shape = as_feature_shape(value_shape, node->data_layout());
+
   uint32_t input_height = value_feature_shape.height().value();
   uint32_t input_width = value_feature_shape.width().value();
   uint32_t stride_height = stride_data->stride()->vertical();
@@ -1392,12 +1400,13 @@ bool fix_shape(moco::tf::TFMul *node)
 {
   auto x = node->x();
   auto y = node->y();
-  auto x_shapedata = x->annot<ShapeInferenceData>();
-  auto y_shapedata = y->annot<ShapeInferenceData>();
-  if (x_shapedata == nullptr || y_shapedata == nullptr)
-  {
+  loco::NodeShape x_shape;
+  loco::NodeShape y_shape;
+
+  if (!node_shape(x, x_shape))
+    return false;
+  if (!node_shape(y, y_shape))
     return false;
-  }
   // TODO check shape difference
 
   // Output shape is same as the input
@@ -1408,18 +1417,14 @@ bool fix_shape(moco::tf::TFRealDiv *node)
 {
   auto x = node->x();
   auto y = node->y();
-  if (x == nullptr || y == nullptr)
-  {
-    // this node maybe disconnected by some transformation
-    // TODO remove this block after fixing to iterate only active nodes
+  loco::NodeShape x_shape;
+  loco::NodeShape y_shape;
+
+  if (!node_shape(x, x_shape))
     return false;
-  }
-  auto x_shapedata = x->annot<ShapeInferenceData>();
-  auto y_shapedata = y->annot<ShapeInferenceData>();
-  if (x_shapedata == nullptr || y_shapedata == nullptr)
-  {
+  if (!node_shape(y, y_shape))
     return false;
-  }
+
   // TODO check shape difference
 
   // Output shape is same as the input
@@ -1518,18 +1523,19 @@ bool fix_shape(moco::tf::TFShape *node)
   }
 
   auto input = node->input();
-  auto input_shape = input->annot<ShapeInferenceData>();
-  if (input_shape == nullptr)
+  loco::NodeShape input_shape;
+  if (!node_shape(input, input_shape))
   {
     // Input shape is required for TFShape shape inference
     return false;
   }
+  loco::TensorShape input_tensor_shape = input_shape.as<loco::TensorShape>();
 
   loco::TensorShape node_shape;
 
   // Note that input shape becomes node(TFShape)'s value
   node_shape.rank(1);
-  node_shape.dim(0) = input_shape->rank();
+  node_shape.dim(0) = input_tensor_shape.rank();
 
   auto shape_annot = stdex::make_unique<ShapeInferenceData>();
   shape_annot->tensor_shape(node_shape);
@@ -1572,8 +1578,8 @@ bool fix_shape(moco::tf::TFSqueeze *node)
   }
 
   auto input = node->input();
-  auto input_shape = input->annot<ShapeInferenceData>();
-  if (input_shape == nullptr)
+  loco::NodeShape input_shape;
+  if (!node_shape(input, input_shape))
   {
     // Input shape is required for TFSqueeze shape inference
     return false;
@@ -1581,7 +1587,7 @@ bool fix_shape(moco::tf::TFSqueeze *node)
 
   // TODO Not sure Squeeze only get input as Tensor
   // Note that tensor_shape() has assertion in it
-  auto input_tensor_shape = input_shape->tensor_shape();
+  auto input_tensor_shape = input_shape.as<loco::TensorShape>();
 
   auto squeeze_dims_vec = node->squeeze_dims();
   std::set<int64_t> squeeze_dims(squeeze_dims_vec.cbegin(), squeeze_dims_vec.cend());
@@ -1680,18 +1686,14 @@ bool fix_shape(moco::tf::TFSub *node)
 {
   auto x = node->x();
   auto y = node->y();
-  if (x == nullptr || y == nullptr)
-  {
-    // this node maybe disconnected by some transformation
-    // TODO remove this block after fixing to iterate only active nodes
+  loco::NodeShape x_shape;
+  loco::NodeShape y_shape;
+
+  if (!node_shape(x, x_shape))
     return false;
-  }
-  auto x_shapedata = x->annot<ShapeInferenceData>();
-  auto y_shapedata = y->annot<ShapeInferenceData>();
-  if (x_shapedata == nullptr || y_shapedata == nullptr)
-  {
+  if (!node_shape(y, y_shape))
     return false;
-  }
+
   // TODO check shape difference
 
   // Output shape is same as the input