path: root/compiler/moco-tf/src/Transforms/FixShapeTransform.cpp

/*
 * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "FixShapeTransform.h"

#include "LogHelper.h"

#include "Annotations/ConcatData.h"
#include "Annotations/PadData.h"
#include "Annotations/ShapeInferenceData.h"
#include "Annotations/StrideData.h"
#include "Annotations/WindowData.h"
#include "Dialect/TFNodes.h"

#include <loco.h>
#include <loco/IR/NodeShape.h>
#include <loco/Service/ShapeInference.h>
#include <moco/Log.h>
#include <stdex/Memory.h>
#include <plier/tf/Convert.h>
#include <locoex/COpCall.h>

#include <cassert>
#include <sstream>
#include <stdexcept>

namespace
{

using namespace moco::tf;

/**
 * @brief  Return true if node has shape inference data for checking shape
 *         inference is done or not
 */
bool shape_inference_done(const loco::Node *node)
{
  auto shapedata = node->annot<ShapeInferenceData>();
  return (shapedata != nullptr);
}

/**
 * @brief  Copy ShapeInferenceData values from src to dst
 *
 * @note   T can be ShapeInferenceData or loco::Node based class having shape
 *         attributes like ConstGen, Pull and so on
 */
template <class T> void copy_shape_values(const T *src, ShapeInferenceData *dst)
{
  assert(src != nullptr);
  assert(dst != nullptr);

  uint32_t rank = src->rank();
  dst->rank(rank);
  for (uint32_t index = 0; index < rank; ++index)
  {
    if (src->dim(index).known())
      dst->dim(index) = src->dim(index).value();
    else
      dst->dim(index).unset();
  }
}

/**
 * @brief  Make copy of ShapeInferenceData from src
 *
 * @note   T can be ShapeInferenceData or loco::Node based class having shape
 *         attributes like TFConst, COpCall and so on
 */
template <class T> std::unique_ptr<ShapeInferenceData> make_shape_inference_data(const T *src)
{
  assert(src != nullptr);

  auto shape_data = stdex::make_unique<ShapeInferenceData>();

  uint32_t rank = src->rank();
  shape_data->rank(rank);
  for (uint32_t index = 0; index < rank; ++index)
  {
    if (src->dim(index).known())
      shape_data->dim(index) = src->dim(index).value();
    else
      shape_data->dim(index).unset();
  }

  return std::move(shape_data);
}

std::unique_ptr<ShapeInferenceData> make_shape_inference_data(const loco::NodeShape &src)
{
  auto shape_data = stdex::make_unique<ShapeInferenceData>();

  switch (src.domain())
  {
    case loco::Domain::Tensor:
      shape_data->tensor_shape(src.as<loco::TensorShape>());
      break;

    case loco::Domain::Feature:
      shape_data->feature_shape(src.as<loco::FeatureShape>());
      break;

    case loco::Domain::Filter:
      shape_data->filter_shape(src.as<loco::FilterShape>());
      break;

    case loco::Domain::DepthwiseFilter:
      shape_data->depthwisefilter_shape(src.as<loco::DepthwiseFilterShape>());
      break;

    case loco::Domain::Bias:
      shape_data->bias_shape(src.as<loco::BiasShape>());
      break;

    default:
      throw std::runtime_error("Unsupported Domain in make_shape_inference_data");
  }

  return std::move(shape_data);
}

loco::NodeShape as_node_shape(const ShapeInferenceData *shapedata)
{
  switch (shapedata->domain())
  {
    case loco::Domain::Tensor:
      return loco::NodeShape({shapedata->tensor_shape()});

    case loco::Domain::Feature:
      return loco::NodeShape({shapedata->feature_shape()});

    case loco::Domain::Filter:
      return loco::NodeShape({shapedata->filter_shape()});

    case loco::Domain::DepthwiseFilter:
      return loco::NodeShape({shapedata->depthwisefilter_shape()});

    case loco::Domain::Bias:
      return loco::NodeShape({shapedata->bias_shape()});
  }

  throw std::runtime_error("Unsupported Domain in as_node_shape");
}

/**
 * @brief Create a higher-rank TensorShape following NumPy broadcasting semantics
 *
 * HOW TO USE:
 *
 *   auto expanded_tensor_shape = expand(tensor_shape).to(N);
 */
class TensorShapeExpander
{
public:
  TensorShapeExpander(const loco::TensorShape &shape) : _shape{shape}
  {
    // DO NOTHING
  }

public:
  loco::TensorShape to(uint32_t output_rank)
  {
    auto const &input_shape = _shape;
    uint32_t const input_rank = input_shape.rank();

    assert(input_rank <= output_rank && "Cannot shrink rank");
    uint32_t const axis_shift = output_rank - input_rank;

    loco::TensorShape output_shape;

    output_shape.rank(output_rank);
    for (uint32_t axis = 0; axis < output_rank; ++axis)
    {
      output_shape.dim(axis) = (axis < axis_shift) ? 1 : input_shape.dim(axis - axis_shift);
    }

    return output_shape;
  }

private:
  const loco::TensorShape _shape;
};

/**
 * @breif  Expand shape x and y to same rank by align right and filling with 1
 */
void expand_rank(loco::TensorShape &x, loco::TensorShape &y)
{
  auto x_rank = x.rank();
  auto y_rank = y.rank();

  if (x_rank == y_rank)
    return;

  TensorShapeExpander x_exp(x);
  TensorShapeExpander y_exp(y);

  auto xy_rank = std::max(x_rank, y_rank);

  x = x_rank > y_rank ? x : x_exp.to(xy_rank);
  y = y_rank > x_rank ? y : y_exp.to(xy_rank);
}

/**
 * @breif  Returns shape of expanded dimension of input x and y having same rank
 */
loco::TensorShape expand_dimension(const loco::TensorShape &x, const loco::TensorShape &y)
{
  assert(x.rank() == y.rank());

  auto rank = x.rank();

  loco::TensorShape output_shape;

  output_shape.rank(rank);
  for (auto axis = 0; axis < rank; ++axis)
  {
    assert(x.dim(axis).known() && y.dim(axis).known());

    auto x_dim = x.dim(axis).value();
    auto y_dim = y.dim(axis).value();

    // each dimension of x and y should be same or one must be 1 if different
    if (!((x_dim == y_dim) || (x_dim == 1 || y_dim == 1)))
      throw std::runtime_error("Cannot produce expand_dimension of two shapes");

    output_shape.dim(axis) = std::max(x_dim, y_dim);
  }

  return output_shape;
}

loco::TensorShape broadcast_shape(const loco::TensorShape &x, const loco::TensorShape &y)
{
  auto x_match = x;
  auto y_match = y;

  expand_rank(x_match, y_match);

  auto output_shape = expand_dimension(x_match, y_match);

  return output_shape;
}

/**
 * @brief  Copy ShapeInferenceData from loco::Node pointer src to dst
 */
bool copy_shapedata(const loco::Node *src, loco::Node *dst)
{
  // if dst already has ShapeInferenceData, skip
  if (shape_inference_done(dst))
    return false;

  // if src has loco::NodeShape, use it
  if (loco::shape_known(src))
  {
    auto shape_data = make_shape_inference_data(loco::shape_get(src));
    dst->annot(std::move(shape_data));

    return true;
  }

  // if src doesn't have ShapeInferenceData, skip
  if (!shape_inference_done(src))
    return false;

  auto src_shapedata = src->annot<ShapeInferenceData>();
  auto shape_data = make_shape_inference_data(src_shapedata);
  dst->annot(std::move(shape_data));

  return true;
}

/**
 * @note  This will find broadcast shape from two inputs lhs and rhs using
 *        broadcast_shape() and return that shape to dst
 */
bool copy_shapedata(const loco::Node *lhs, const loco::Node *rhs, loco::Node *dst)
{
  // if dst already has ShapeInferenceData, skip
  if (shape_inference_done(dst))
    return false;

  loco::NodeShape lhs_shape;
  loco::NodeShape rhs_shape;

  if (loco::shape_known(lhs))
  {
    lhs_shape = loco::shape_get(lhs);
  }
  else
  {
    if (!shape_inference_done(lhs))
      return false;

    lhs_shape = as_node_shape(lhs->annot<ShapeInferenceData>());
  }

  if (loco::shape_known(rhs))
  {
    rhs_shape = loco::shape_get(rhs);
  }
  else
  {
    if (!shape_inference_done(rhs))
      return false;

    rhs_shape = as_node_shape(rhs->annot<ShapeInferenceData>());
  }

  if (lhs_shape.domain() != loco::Domain::Tensor || rhs_shape.domain() != loco::Domain::Tensor)
  {
    throw std::runtime_error("copy_shapedata supports only for Tensor");
  }

  loco::TensorShape lhs_tensorshape = lhs_shape.as<loco::TensorShape>();
  loco::TensorShape rhs_tensorshape = rhs_shape.as<loco::TensorShape>();
  loco::TensorShape sum_tensorshape = broadcast_shape(lhs_tensorshape, rhs_tensorshape);

  loco::NodeShape sum_shape({sum_tensorshape});
  auto shape_data = make_shape_inference_data(sum_shape);
  dst->annot(std::move(shape_data));

  LOGGER(l);

  INFO(l) << "copy_shapedata " << lhs_tensorshape << " or " << rhs_tensorshape << " -> "
          << sum_tensorshape << std::endl;

  return true;
}

/**
 * @note  While in shape inference, Node maybe Canonical, TF dialect or other dialects
 *        This will provide common loco::NodeShape as shape information
 */
bool node_shape(const loco::Node *node, loco::NodeShape &nodeshape)
{
  if (loco::shape_known(node))
  {
    nodeshape = loco::shape_get(node);
    return true;
  }

  if (!shape_inference_done(node))
    return false;

  auto shapedata = node->annot<ShapeInferenceData>();

  switch (shapedata->domain())
  {
    case loco::Domain::Tensor:
      nodeshape.set(shapedata->tensor_shape());
      break;

    case loco::Domain::Feature:
      nodeshape.set(shapedata->feature_shape());
      break;

    case loco::Domain::Filter:
      nodeshape.set(shapedata->filter_shape());
      break;

    case loco::Domain::DepthwiseFilter:
      nodeshape.set(shapedata->depthwisefilter_shape());
      break;

    case loco::Domain::Bias:
      nodeshape.set(shapedata->bias_shape());
      break;

    default:
      throw std::runtime_error("Unsupported Domain in node_shape()");
  }
  return true;
}

loco::FeatureShape as_feature_shape(const loco::NodeShape &nodeshape,
                                    const TFDataLayout &data_layout)
{
  if (nodeshape.domain() == loco::Domain::Feature)
    return nodeshape.as<loco::FeatureShape>();

  loco::FeatureShape feature_shape;

  // only convert from tensor to feature
  if (nodeshape.domain() != loco::Domain::Tensor)
  {
    throw std::runtime_error("as_feature_shape: domain is not tensor");
  }

  loco::TensorShape tensor_shape = nodeshape.as<loco::TensorShape>();

  if (tensor_shape.rank() != 4)
  {
    throw std::runtime_error("as_feature_shape: rank is not 4");
  }

  if (data_layout == "NHWC")
  {
    feature_shape.count() = tensor_shape.dim(0);
    feature_shape.height() = tensor_shape.dim(1);
    feature_shape.width() = tensor_shape.dim(2);
    feature_shape.depth() = tensor_shape.dim(3);
  }
  else if (data_layout == "NCHW")
  {
    feature_shape.count() = tensor_shape.dim(0);
    feature_shape.depth() = tensor_shape.dim(1);
    feature_shape.height() = tensor_shape.dim(2);
    feature_shape.width() = tensor_shape.dim(3);
  }
  else
  {
    // TODO support for other data_layout if needed
    throw std::runtime_error("as_feature_shape: only supports NHWC or NCHW");
  }

  return feature_shape;
}

struct FixPadContext
{
  uint32_t input_height;
  uint32_t input_width;
  uint32_t output_height;
  uint32_t output_width;
  uint32_t stride_height;
  uint32_t stride_width;
  uint32_t effective_window_height;
  uint32_t effective_window_width;
};

PadData calc_paddata(const FixPadContext &ctx)
{
  assert(ctx.output_height > 0);
  assert(ctx.output_width > 0);

  // calculate padding height, width
  int64_t i_height = (int64_t)(ctx.output_height - 1) * (int64_t)ctx.stride_height +
                     (int64_t)ctx.effective_window_height - (int64_t)ctx.input_height;
  int64_t i_width = (int64_t)(ctx.output_width - 1) * (int64_t)ctx.stride_width +
                    (int64_t)ctx.effective_window_width - (int64_t)ctx.input_width;
  uint32_t pad_height = i_height >= 0 ? (uint32_t)i_height : 0U;
  uint32_t pad_width = i_width >= 0 ? (uint32_t)i_width : 0U;

  PadData pad_data;

  pad_data.pad()->top(pad_height / 2);
  pad_data.pad()->bottom(pad_height - pad_data.pad()->top());
  pad_data.pad()->left(pad_width / 2);
  pad_data.pad()->right(pad_width - pad_data.pad()->left());

  return pad_data;
}

template <class T> void calc_annot_paddata(T *node, const FixPadContext &ctx)
{
  assert(node != nullptr);

  PadData pd = calc_paddata(ctx);

  // annotation of pad data
  auto pad_data = stdex::make_unique<PadData>(pd);

  node->annot(std::move(pad_data));

  assert(node->template annot<PadData>() != nullptr);
}

template <class T> void update_stride_data(T *node)
{
  auto stride_data = stdex::make_unique<StrideData>();
  auto strides = node->strides();
  auto data_layout = plier::tf::as_data_layout(node->data_layout());
  if (data_layout == plier::tf::DataLayout::NHWC)
  {
    stride_data->stride()->vertical(strides[1]);
    stride_data->stride()->horizontal(strides[2]);
  }
  else if (data_layout == plier::tf::DataLayout::NCHW)
  {
    stride_data->stride()->vertical(strides[2]);
    stride_data->stride()->horizontal(strides[3]);
  }
  node->annot(std::move(stride_data));
}

template <class T> void update_window_data(T *node)
{
  auto window_data = stdex::make_unique<WindowData>();
  auto ksize = node->ksize();
  auto data_layout = plier::tf::as_data_layout(node->data_layout());
  if (data_layout == plier::tf::DataLayout::NHWC)
  {
    window_data->window()->vertical(ksize[1]);
    window_data->window()->horizontal(ksize[2]);
  }
  else if (data_layout == plier::tf::DataLayout::NCHW)
  {
    window_data->window()->vertical(ksize[2]);
    window_data->window()->horizontal(ksize[3]);
  }
  node->annot(std::move(window_data));
}

bool fix_shape(loco::Pull *node)
{
  if (shape_inference_done(node))
    return false;

  // Pull itself has shape information, copy them
  auto shape_data = make_shape_inference_data(node);
  node->annot(std::move(shape_data));

  return true;
}

bool fix_shape(loco::Push *node)
{
  // Output shape is same as the from
  auto from = node->from();
  return copy_shapedata(from, node);
}

bool fix_shape(moco::tf::TFAdd *node)
{
  auto x = node->x();
  auto y = node->y();
  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
  return copy_shapedata(x, node);
}

bool fix_shape(moco::tf::TFAvgPool *node)
{
  LOGGER(l);

  if (shape_inference_done(node))
    return false;

  auto value = node->value();
  loco::NodeShape value_shape;
  if (!node_shape(value, value_shape))
  {
    // input node shape inference is not ready
    return false;
  }

  auto padding = node->padding();
  assert(padding == "VALID" || padding == "SAME");

  update_stride_data(node);
  update_window_data(node);

  auto value_feature_shape = as_feature_shape(value_shape, node->data_layout());

  auto stride_data = node->annot<StrideData>();
  assert(stride_data != nullptr);
  auto window_data = node->annot<WindowData>();
  assert(window_data != nullptr);

  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();
  uint32_t stride_width = stride_data->stride()->horizontal();
  uint32_t window_height = window_data->window()->vertical();
  uint32_t window_width = window_data->window()->horizontal();
  uint32_t dilation_height = 1; // dilation is 1
  uint32_t dilation_width = 1;
  uint32_t effective_window_height = dilation_height * (window_height - 1) + 1;
  uint32_t effective_window_width = dilation_width * (window_width - 1) + 1;
  uint32_t output_height;
  uint32_t output_width;

  if (padding == "VALID")
  {
    output_height = (input_height + stride_height - effective_window_height) / stride_height;
    output_width = (input_width + stride_width - effective_window_width) / stride_width;
  }
  else if (padding == "SAME")
  {
    output_height = (input_height + stride_height - 1) / stride_height;
    output_width = (input_width + stride_width - 1) / stride_width;
  }

  loco::FeatureShape ofm_feature_shape;
  ofm_feature_shape.count() = value_feature_shape.count();
  ofm_feature_shape.height() = output_height;
  ofm_feature_shape.width() = output_width;
  ofm_feature_shape.depth() = value_feature_shape.depth();

  auto shape_data = stdex::make_unique<ShapeInferenceData>();
  as_tensor_shape(*shape_data.get(), ofm_feature_shape, node->data_layout());
  node->annot(std::move(shape_data));

  FixPadContext ctx = {
      input_height,  input_width,  output_height,           output_width,
      stride_height, stride_width, effective_window_height, effective_window_width};

  calc_annot_paddata(node, ctx);

  INFO(l) << "Fix TFAvgPool shape = ifm" << value_feature_shape << " --> ofm" << ofm_feature_shape;
  INFO(l) << "              pad = " << *node->annot<PadData>();

  return true;
}

bool fix_shape(moco::tf::TFBiasAdd *node)
{
  auto value = node->value();
  auto bias = node->bias();
  loco::NodeShape value_shape;
  loco::NodeShape bias_shape;
  if (!node_shape(value, value_shape) || !node_shape(bias, bias_shape))
  {
    return false;
  }

  // Output shape is same as the value shape
  return copy_shapedata(value, node);
}

template <class CONST_CLASS> bool valid_scala_value(CONST_CLASS *node)
{
  LOGGER(l);

  loco::NodeShape nodeshape;
  if (!node_shape(node, nodeshape))
  {
    return false;
  }

  if (node->dtype() != loco::DataType::S32)
  {
    INFO(l) << "valid_scala_value not S32";
    return false;
  }

  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();
    return false;
  }

  return true;
}

template <class CONST_CLASS> int32_t scala_value(CONST_CLASS *node)
{
  loco::NodeShape nodeshape;
  if (!node_shape(node, nodeshape))
  {
    return false;
  }

  assert(node->dtype() == loco::DataType::S32);

  auto tensor_shape = nodeshape.as<loco::TensorShape>();
  assert(tensor_shape.rank() == 0 || tensor_shape.rank() == 1);

  return node->template at<loco::DataType::S32>(0);
}

bool fix_shape(moco::tf::TFConcatV2 *node)
{
  LOGGER(l);

  if (shape_inference_done(node))
  {
    INFO(l) << "Fix shape TFConcatV2 already done";
    return false;
  }
  // ConcatData should be null
  assert(node->annot<ConcatData>() == nullptr);

  // Check shape inference data are all ready
  // Check shape rank are all same
  auto value_a = node->values(0);
  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;
  }
  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);
    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;
    }
    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();
  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";
    return false;
  }

  int32_t axis_value = 0;
  bool axis_available = false;
  {
    // check for axis is TFConst
    auto tfconst = dynamic_cast<moco::tf::TFConst *>(axis_node);
    if (tfconst != nullptr)
    {
      if (valid_scala_value(tfconst))
      {
        axis_value = scala_value(tfconst);
        axis_available = true;
      }
    }
  }
  {
    // check for axis is ConstGen
    auto constgen = dynamic_cast<loco::ConstGen *>(axis_node);
    if (constgen != nullptr)
    {
      if (valid_scala_value(constgen))
      {
        axis_value = scala_value(constgen);
        axis_available = true;
      }
    }
  }
  if (!axis_available)
  {
    // we cannot find a valid axis value
    INFO(l) << "Fix shape TFConcatV2 axis_available false";
    return false;
  }

  auto concat_data = stdex::make_unique<ConcatData>(axis_value);
  node->annot(std::move(concat_data));

  uint32_t axis_absolute = (axis_value >= 0) ? axis_value : (int32_t)a_rank + axis_value;

  auto shape_data = stdex::make_unique<ShapeInferenceData>();
  shape_data->rank(a_rank);

  for (uint32_t index = 0; index < a_rank; ++index)
  {
    if (value_a_tensor_shape.dim(index).known())
    {
      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);
          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;
      }
      shape_data->dim(index) = dim;
    }
    else
      shape_data->dim(index).unset();
  }
  node->annot(std::move(shape_data));

  INFO(l) << "Fix TFConcat shape = " << node->annot<ShapeInferenceData>();

  return true;
}

bool fix_shape(moco::tf::TFConst *node)
{
  if (shape_inference_done(node))
    return false;

  // TFConst itself has shape information, copy them
  auto shape_data = make_shape_inference_data(node);
  node->annot(std::move(shape_data));

  {
    LOGGER(l);
    auto shapedata = node->annot<ShapeInferenceData>();
    assert(shapedata != nullptr);
    INFO(l) << "Fix TFConst shape = " << shapedata->tensor_shape();
  }

  return true;
}

bool fix_shape(moco::tf::TFConv2D *node)
{
  LOGGER(l);

  if (shape_inference_done(node))
    return false;

  auto ifm = node->input();
  loco::NodeShape ifm_shape;
  if (!node_shape(ifm, ifm_shape))
  {
    // input node shape inference is not ready
    return false;
  }

  auto ker = node->filter();
  loco::NodeShape ker_shape;
  if (!node_shape(ker, ker_shape))
  {
    return false;
  }

  auto padding = node->padding();
  assert(padding == "VALID" || padding == "SAME");

  update_stride_data(node);

  auto stride_data = node->annot<StrideData>();
  assert(stride_data != nullptr);
  // TODO add and use 'stride_data->stride()' stream out
  INFO(l) << "Fix TFConv2D strides = " << stride_data->stride()->vertical() << ", "
          << stride_data->stride()->horizontal();

  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);

  uint32_t input_height = ifm_tensor_shape.dim(1).value();
  uint32_t input_width = ifm_tensor_shape.dim(2).value();
  uint32_t stride_height = stride_data->stride()->vertical();
  uint32_t stride_width = stride_data->stride()->horizontal();
  uint32_t ker_height = ker_tensor_shape.dim(0).value();
  uint32_t ker_width = ker_tensor_shape.dim(1).value();
  uint32_t dilation_height = 1; // TODO Consider dilation
  uint32_t dilation_width = 1;
  uint32_t effective_ker_height = dilation_height * (ker_height - 1) + 1;
  uint32_t effective_ker_width = dilation_width * (ker_width - 1) + 1;
  uint32_t output_height;
  uint32_t output_width;

  if (padding == "VALID")
  {
    output_height = (input_height + stride_height - effective_ker_height) / stride_height;
    output_width = (input_width + stride_width - effective_ker_width) / stride_width;
  }
  else if (padding == "SAME")
  {
    output_height = (input_height + stride_height - 1) / stride_height;
    output_width = (input_width + stride_width - 1) / stride_width;
  }
  else
  {
    assert(false && "Unknown padding in fix_shape for TFConv2D");
  }

  loco::TensorShape ofm_tensor_shape;
  ofm_tensor_shape.rank(4);
  ofm_tensor_shape.dim(0) = ifm_tensor_shape.dim(0);
  ofm_tensor_shape.dim(1) = output_height;
  ofm_tensor_shape.dim(2) = output_width;
  ofm_tensor_shape.dim(3) = ker_tensor_shape.dim(3);

  auto shape_data = stdex::make_unique<ShapeInferenceData>();
  shape_data->tensor_shape(ofm_tensor_shape);
  node->annot(std::move(shape_data));

  FixPadContext ctx = {input_height,  input_width,  output_height,        output_width,
                       stride_height, stride_width, effective_ker_height, effective_ker_width};

  calc_annot_paddata(node, ctx);

  INFO(l) << "Fix TFConv2D shape = ifm" << ifm_tensor_shape << " ker" << ker_tensor_shape
          << " --> ofm" << ofm_tensor_shape;
  INFO(l) << "             pad = " << *node->annot<PadData>();

  return true;
}

bool fix_shape(moco::tf::TFDepthwiseConv2dNative *node)
{
  LOGGER(l);

  if (shape_inference_done(node))
    return false;

  auto ifm = node->input();
  loco::NodeShape ifm_shape;
  if (!node_shape(ifm, ifm_shape))
  {
    // input node shape inference is not ready
    return false;
  }

  auto ker = node->filter();
  loco::NodeShape ker_shape;
  if (!node_shape(ker, ker_shape))
  {
    return false;
  }

  update_stride_data(node);

  auto stride_data = node->annot<StrideData>();
  assert(stride_data != nullptr);

  INFO(l) << "FixShape TFDepthwiseConv2dNative strides = " << stride_data->stride()->vertical()
          << ", " << stride_data->stride()->horizontal();

  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);

  uint32_t input_height = ifm_tensor_shape.dim(1).value();
  uint32_t input_width = ifm_tensor_shape.dim(2).value();
  uint32_t stride_height = stride_data->stride()->vertical();
  uint32_t stride_width = stride_data->stride()->horizontal();
  uint32_t ker_height = ker_tensor_shape.dim(0).value();
  uint32_t ker_width = ker_tensor_shape.dim(1).value();
  uint32_t dilation_height = 1; // TODO Consider dilation
  uint32_t dilation_width = 1;
  uint32_t effective_ker_height = dilation_height * (ker_height - 1) + 1;
  uint32_t effective_ker_width = dilation_width * (ker_width - 1) + 1;
  uint32_t output_height;
  uint32_t output_width;

  auto padding = node->padding();
  assert(padding == "VALID" || padding == "SAME");

  if (padding == "VALID")
  {
    output_height = (input_height + stride_height - effective_ker_height) / stride_height;
    output_width = (input_width + stride_width - effective_ker_width) / stride_width;
  }
  else // padding == "SAME"
  {
    output_height = (input_height + stride_height - 1) / stride_height;
    output_width = (input_width + stride_width - 1) / stride_width;
  }

  loco::TensorShape ofm_tensor_shape;
  ofm_tensor_shape.rank(4);
  ofm_tensor_shape.dim(0) = ifm_tensor_shape.dim(0);
  ofm_tensor_shape.dim(1) = output_height;
  ofm_tensor_shape.dim(2) = output_width;
  ofm_tensor_shape.dim(3) =
      loco::Dimension(ker_tensor_shape.dim(2).value() * ker_tensor_shape.dim(3).value());

  auto shape_data = stdex::make_unique<ShapeInferenceData>();
  shape_data->tensor_shape(ofm_tensor_shape);
  node->annot(std::move(shape_data));

  FixPadContext ctx = {input_height,  input_width,  output_height,        output_width,
                       stride_height, stride_width, effective_ker_height, effective_ker_width};

  calc_annot_paddata(node, ctx);

  INFO(l) << "Fix TFDepthwiseConv2dNative shape = ifm" << ifm_tensor_shape << " ker"
          << ker_tensor_shape << " --> ofm" << ofm_tensor_shape;
  INFO(l) << "                            pad = " << *node->annot<PadData>();

  return true;
}

bool fix_shape(moco::tf::TFFusedBatchNorm *node)
{
  // Output shape is same as the input
  auto input = node->input();
  return copy_shapedata(input, node);
}

bool fix_shape(moco::tf::TFIdentity *node)
{
  // Output shape is same as the input
  auto input = node->input();
  return copy_shapedata(input, node);
}

bool fix_shape(moco::tf::TFMaxPool *node)
{
  LOGGER(l);

  if (shape_inference_done(node))
    return false;

  auto value = node->value();
  loco::NodeShape value_shape;
  if (!node_shape(value, value_shape))
  {
    // input node shape inference is not ready
    return false;
  }

  auto padding = node->padding();
  assert(padding == "VALID" || padding == "SAME");

  update_stride_data(node);
  update_window_data(node);

  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();
  uint32_t stride_width = stride_data->stride()->horizontal();
  uint32_t window_height = window_data->window()->vertical();
  uint32_t window_width = window_data->window()->horizontal();
  uint32_t dilation_height = 1; // dilation for MaxPool is 1
  uint32_t dilation_width = 1;
  uint32_t effective_window_height = dilation_height * (window_height - 1) + 1;
  uint32_t effective_window_width = dilation_width * (window_width - 1) + 1;
  uint32_t output_height;
  uint32_t output_width;

  if (padding == "VALID")
  {
    output_height = (input_height + stride_height - effective_window_height) / stride_height;
    output_width = (input_width + stride_width - effective_window_width) / stride_width;
  }
  else if (padding == "SAME")
  {
    output_height = (input_height + stride_height - 1) / stride_height;
    output_width = (input_width + stride_width - 1) / stride_width;
  }

  loco::FeatureShape ofm_feature_shape;
  ofm_feature_shape.count() = value_feature_shape.count();
  ofm_feature_shape.height() = output_height;
  ofm_feature_shape.width() = output_width;
  ofm_feature_shape.depth() = value_feature_shape.depth();

  auto shape_data = stdex::make_unique<ShapeInferenceData>();
  as_tensor_shape(*shape_data.get(), ofm_feature_shape, node->data_layout());
  node->annot(std::move(shape_data));

  FixPadContext ctx = {
      input_height,  input_width,  output_height,           output_width,
      stride_height, stride_width, effective_window_height, effective_window_width};

  calc_annot_paddata(node, ctx);

  INFO(l) << "Fix TFMaxPool shape = ifm" << value_feature_shape << " --> ofm" << ofm_feature_shape;
  INFO(l) << "              pad = " << *node->annot<PadData>();

  return true;
}

bool fix_shape(moco::tf::TFMul *node)
{
  auto x = node->x();
  auto y = node->y();
  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
  return copy_shapedata(x, node);
}

bool fix_shape(moco::tf::TFMean *node)
{
  if (shape_inference_done(node))
    return false;

  LOGGER(l);

  auto input = node->input();
  auto reduction_indices = node->reduction_indices();
  loco::NodeShape input_shape;
  loco::NodeShape reduction_indices_shape;

  if (!node_shape(input, input_shape) || !node_shape(reduction_indices, reduction_indices_shape))
  {
    // Input and reduction_indices shape are required for TFMean shape inference
    return false;
  }

  // Get constant values if reduction_indeces is const
  std::vector<int32_t> reduction_values;
  if (auto tfconst = dynamic_cast<moco::tf::TFConst *>(reduction_indices))
  {
    assert(tfconst->dtype() == loco::DataType::S32);
    auto const_size = tfconst->size<loco::DataType::S32>();
    for (uint32_t i = 0; i < const_size; ++i)
    {
      int32_t axis = tfconst->at<loco::DataType::S32>(i);
      if (axis < 0)
        axis += input_shape.as<loco::TensorShape>().rank();
      reduction_values.push_back(axis);
    }
  }
  else
  {
    // we cannot find a valid reduction indices value
    INFO(l) << "Fix shape TFMean fail : reduction indeces are not constant or not valid";
    return false;
  }

  loco::TensorShape shape_data;
  loco::TensorShape input_tensor_shape = input_shape.as<loco::TensorShape>();

  if (node->keep_dims())
  {
    shape_data.rank(input_tensor_shape.rank());
    for (uint32_t i = 0; i < input_tensor_shape.rank(); ++i)
      shape_data.dim(i) = input_tensor_shape.dim(i);
    for (uint32_t i = 0; i < reduction_values.size(); ++i)
      shape_data.dim(reduction_values.at(i)) = 1;
  }
  else
  {
    std::vector<bool> check_reduce(input_tensor_shape.rank(), false);
    for (uint32_t i = 0; i < reduction_values.size(); ++i)
      check_reduce.at(reduction_values.at(i)) = true;

    uint32_t reduce_cnt = 0;
    for (uint32_t i = 0; i < check_reduce.size(); ++i)
      if (check_reduce.at(i))
        ++reduce_cnt;

    shape_data.rank(input_tensor_shape.rank() - reduce_cnt);
    for (uint32_t i = 0, j = 0; i < check_reduce.size(); ++i)
      if (check_reduce.at(i) == false)
        shape_data.dim(j++) = i;
  }

  auto shape_annot = stdex::make_unique<ShapeInferenceData>();
  shape_annot->tensor_shape(shape_data);
  node->annot(std::move(shape_annot));

  return true;
}

bool fix_shape(moco::tf::TFRealDiv *node)
{
  auto x = node->x();
  auto y = node->y();
  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
  return copy_shapedata(x, node);
}

bool fix_shape(moco::tf::TFRelu *node)
{
  // Output shape is same as the features
  auto features = node->features();
  return copy_shapedata(features, node);
}

bool fix_shape(moco::tf::TFRelu6 *node)
{
  // Output shape is same as the features
  auto features = node->features();
  return copy_shapedata(features, node);
}

bool fix_shape(moco::tf::TFReshape *node)
{
  if (shape_inference_done(node))
    return false;

  // For now, we only consider Fixed Reshape, i.e. Reshape with determined
  //      'shape' input. So here we only support case when 'shape' input of
  //      TFReshape is TFConst. If 'shape' input is not TFConst, another
  //      transform (e.g. constant folding) should be done beforehand to make
  //      it TFConst.
  // TODO Support dynamic Reshape
  // Note that 'shape()' here is 'shape' input, not node's shape information
  auto const_shape_input = dynamic_cast<moco::tf::TFConst *>(node->shape());
  if (!const_shape_input)
  {
    // 'shape' input of TFReshape is not TFConst, try next time when it becomes TFConst
    return false;
  }

  // 'Shape' input should be integer tensor of rank 1, e.g. [2, 3, 4] or [3, -1]
  assert(const_shape_input->dtype() == loco::DataType::S32);
  assert(const_shape_input->rank() == 1);

  auto shape_rank = const_shape_input->dim(0).value();
  assert(shape_rank > 0);

  loco::TensorShape shape_data;
  shape_data.rank(shape_rank);
  for (uint32_t axis = 0; axis < shape_rank; ++axis)
  {
    auto shape_dim = const_shape_input->at<loco::DataType::S32>(axis);
    if (shape_dim == -1)
    {
      // Reshape's new shape has wildcard dimension, i.e. dynamic reshape
      return false;
    }
    assert(shape_dim >= 1);
    shape_data.dim(axis) = shape_dim;
  }

  // TODO Compare 'tensor' input and validate coherency?
  // Not sure this is appropriate stage for this task.

  auto shape_annot = stdex::make_unique<ShapeInferenceData>();
  shape_annot->tensor_shape(shape_data);
  node->annot(std::move(shape_annot));

  {
    LOGGER(l);
    auto shapedata = node->annot<ShapeInferenceData>();
    assert(shapedata != nullptr);
    INFO(l) << "Fix TFReshape shape = " << shapedata->tensor_shape();
  }

  return true;
}

bool fix_shape(moco::tf::TFRsqrt *node)
{
  // Output shape is same as the input x
  auto x = node->x();
  return copy_shapedata(x, node);
}

bool fix_shape(moco::tf::TFShape *node)
{
  if (shape_inference_done(node))
    return false;

  auto input = node->input();
  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_tensor_shape.rank();

  auto shape_annot = stdex::make_unique<ShapeInferenceData>();
  shape_annot->tensor_shape(node_shape);
  node->annot(std::move(shape_annot));

  LOGGER(l);
  INFO(l) << "Fix TFShape shape = " << node_shape;

  return true;
}

bool fix_shape(moco::tf::TFSqrt *node)
{
  // Output shape is same as the input x
  auto x = node->x();
  return copy_shapedata(x, node);
}

bool fix_shape(moco::tf::TFSoftmax *node)
{
  // Output shape is same as the input x
  auto logits = node->logits();
  return copy_shapedata(logits, node);
}

bool fix_shape(moco::tf::TFSquaredDifference *node)
{
  // Output shape is same as the input x
  auto x = node->x();
  return copy_shapedata(x, node);
}

bool fix_shape(moco::tf::TFSqueeze *node)
{
  if (shape_inference_done(node))
    return false;

  auto input = node->input();
  loco::NodeShape input_shape;
  if (!node_shape(input, input_shape))
  {
    // Input shape is required for TFSqueeze shape inference
    return false;
  }

  // TODO Not sure Squeeze only get input as Tensor
  // Note that tensor_shape() has assertion in it
  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());

  loco::TensorShape node_shape;
  uint32_t node_rank = 0;

  if (squeeze_dims.empty())
  {
    // Remove all dimensions whose value is 1
    for (uint32_t axis = 0; axis < input_tensor_shape.rank(); ++axis)
    {
      assert(input_tensor_shape.dim(axis).known());
      auto dim = input_tensor_shape.dim(axis).value();
      if (dim != 1)
      {
        assert(dim > 1);
        node_shape.rank(++node_rank);
        node_shape.dim(node_rank - 1) = dim;
      }
    }
  }
  else
  {
    uint32_t input_rank = input_tensor_shape.rank();

    // Sanity check for 'squeeze_dims'
    auto is_valid_squeeze_dims = [&squeeze_dims, &input_rank]() {
      if (!(squeeze_dims.size() < input_rank))
        return false;
      for (auto squeeze_dim : squeeze_dims)
      {
        if (!(squeeze_dim >= -(int64_t)input_rank))
          return false;
        if (!(squeeze_dim < (int64_t)input_rank))
          return false;
      }
      return true;
    };

    if (!is_valid_squeeze_dims())
    {
      throw std::runtime_error("Fix shape for TFSqueeze: invalid squeeze dimension");
    }

    // Resolve negative squeeze dimension
    std::set<int64_t> resolved_squeeze_dims;
    for (auto squeeze_dim : squeeze_dims)
    {
      if (squeeze_dim < 0)
        resolved_squeeze_dims.insert(squeeze_dim + (int64_t)input_rank);
      else
        resolved_squeeze_dims.insert(squeeze_dim);
    }

    // Remove squeeze dimensions only
    for (uint32_t axis = 0; axis < input_rank; ++axis)
    {
      assert(input_tensor_shape.dim(axis).known());
      auto dim = input_tensor_shape.dim(axis).value();
      if (resolved_squeeze_dims.find((int64_t)axis) == resolved_squeeze_dims.cend())
      {
        // Not squeeze dim
        node_shape.rank(++node_rank);
        node_shape.dim(node_rank - 1) = dim;
      }
      else
      {
        // Is squeeze dim
        assert(dim == 1);
        // DO NOTHING
      }
    }
  }

  assert(node_shape.rank() > 0);

  auto shape_annot = stdex::make_unique<ShapeInferenceData>();
  shape_annot->tensor_shape(node_shape);
  node->annot(std::move(shape_annot));

  LOGGER(l);
  INFO(l) << "Fix TFSqueeze shape = " << node_shape;

  return true;
}

bool fix_shape(moco::tf::TFStopGradient *node)
{
  // Output shape is same as the input
  auto input = node->input();
  return copy_shapedata(input, node);
}

bool fix_shape(moco::tf::TFSub *node)
{
  auto x = node->x();
  auto y = node->y();
  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
  return copy_shapedata(x, node);
}

bool fix_shape(moco::tf::TFTanh *node)
{
  // Output shape is same as the input
  auto x = node->x();
  return copy_shapedata(x, node);
}

bool fix_shape(locoex::COpCall *node)
{
  if (shape_inference_done(node))
    return false;

  auto shape_data = make_shape_inference_data(node);
  node->annot(std::move(shape_data));

  return true;
}

} // namespace

namespace moco
{
namespace tf
{

bool FixShapeTransform::run(loco::Graph *graph)
{
  bool changed = false;
  for (auto node : loco::active_nodes(loco::output_nodes(graph)))
  {
// clang-format off
// TODO remove this block after Pull, Push is not used in import
#define CANONICAL_NODE(TYPE_NAME)                     \
    if (as<loco::TYPE_NAME>(node))              \
    {                                           \
      if (fix_shape(as<loco::TYPE_NAME>(node))) \
        changed = true;                         \
    }                                           \
    else
CANONICAL_NODE(Pull)
CANONICAL_NODE(Push)
#undef CANONICAL_NODE

#define TENSORFLOW_NODE(OPCODE,CLASS)           \
    if (as<moco::tf::CLASS>(node))              \
    {                                           \
      if (fix_shape(as<moco::tf::CLASS>(node))) \
        changed = true;                         \
    }                                           \
    else
#include "Dialect/TFNodes.lst"
#undef TENSORFLOW_NODE
    // clang-format on

    if (as<locoex::COpCall>(node))
    {
      if (fix_shape(as<locoex::COpCall>(node)))
        changed = true;
    }
    else
    {
      // Skip nodes that are not interested
    }
  }

  return changed;
}

} // namespace tf
} // namespace moco