text: // Copyright (C) 2018-2022 Intel Corporation
  // SPDX-License-Identifier: Apache-2.0
  //
  
  #include "vpu/ngraph/operations/dynamic_shape_resolver.hpp"
  
  #include "ngraph/opsets/opset3.hpp"
  
  namespace ngraph { namespace vpu { namespace op {
  
  DynamicShapeResolver::DynamicShapeResolver(
          const Output<Node>& tensorWithData,
          const Output<Node>& tensorWithDims,
          const DynamicShapeResolverMode& mode)
      : Op(OutputVector{tensorWithData, tensorWithDims}), m_mode(mode) {
      constructor_validate_and_infer_types();
  }
  
  std::shared_ptr<Node> DynamicShapeResolver::clone_with_new_inputs(const OutputVector& new_args) const {
      check_new_args_count(this, new_args);
      return std::make_shared<DynamicShapeResolver>(new_args.at(0), new_args.at(1), m_mode);
  }
  
  void DynamicShapeResolver::validate_and_infer_types() {
      NODE_VALIDATION_CHECK(this, get_input_size() == 2, "(", get_friendly_name(), ") supports only ", 2, " inputs, but ", get_input_size(), " provided");
      NODE_VALIDATION_CHECK(this, get_input_partial_shape(1).is_static(), "(", get_friendly_name(), ") does not support dynamic shape for dims tensor");
  
      const auto& dataElementType = get_input_element_type(0);
      NODE_VALIDATION_CHECK(this, dataElementType.is_static(), "(", get_friendly_name(), ") does not support dynamic element type for data tensor");
      const auto& dimsElementType = get_input_element_type(1);
      NODE_VALIDATION_CHECK(this, dimsElementType.is_static() && (dimsElementType.compatible(ngraph::element::i64) ||
                                                                  dimsElementType.compatible(ngraph::element::i32)),
          "(", get_friendly_name(), ") supports only i64 and i32 number type for dims tensor, but ", dimsElementType, " provided");
  
      const auto& dimsShape = get_input_shape(1);
  
      if (m_mode == DynamicShapeResolverMode::INFER_UPPER_BOUND_SHAPE) {
          NODE_VALIDATION_CHECK(this, get_input_partial_shape(0).is_static(), "(", get_friendly_name(), ") does not support dynamic shape for data tensor");
  
          const auto& dataShape = get_input_shape(0);
          NODE_VALIDATION_CHECK(this, dimsShape.size() == 1 && dimsShape.front() == dataShape.size(), "(", get_friendly_name(), ") inputs shapes mismatch: first "
              "input shape = ", dataShape, " second input shape = ", dimsShape, " but ", dataShape, " and ", Shape{dataShape.size()}, " are expected");
  
          set_output_type(0, dataElementType, dataShape);
      } else if (m_mode == DynamicShapeResolverMode::INFER_DYNAMIC_SHAPE) {
          NODE_VALIDATION_CHECK(this, get_input_partial_shape(0).rank() == dimsShape.front(),
                  "(", get_friendly_name(), ") data and shape ranks must be equal, provided: ",
                  get_input_partial_shape(0).rank(), " vs ", dimsShape.front());
  
          set_output_type(0, dataElementType,
                          ngraph::PartialShape::dynamic(get_input_partial_shape(0).rank()));
      } else {
          NGRAPH_UNREACHABLE(this, "Unknown DynamicShapeResolverMode value, expected one of: INFER_UPPER_BOUND_SHAPE, INFER_DYNAMIC_SHAPE");
      }
  }
  
  bool DynamicShapeResolver::visit_attributes(ngraph::AttributeVisitor&) {
      return true;
  }
  
  namespace dyn_shape {
  namespace {
  
  template<element::Type_t ET>
  bool getShapeFromHostTensorData(const HostTensorPtr& data, Shape& result) {
      using T = typename element_type_traits<ET>::value_type;
      T *dataPtr = data->get_data_ptr<ET>();
      if (!dataPtr) {
          return false;
      }
      if (data->get_shape().size() != 1) {
          return false;
      }
      size_t outputRank = data->get_shape()[0];
  
      for (size_t i = 0; i < outputRank; i++) {
          result.push_back(dataPtr[i]);
      }
  
      return true;
  }
  
  bool getShapeFromHostTensorData(const HostTensorPtr& data, Shape& shape) {
      switch (data->get_element_type()) {
          case element::Type_t::i8:
              return getShapeFromHostTensorData<element::Type_t::i8>(data, shape);
          case element::Type_t::i16:
              return getShapeFromHostTensorData<element::Type_t::i16>(data, shape);
          case element::Type_t::i32:
              return getShapeFromHostTensorData<element::Type_t::i32>(data, shape);
          case element::Type_t::i64:
              return getShapeFromHostTensorData<element::Type_t::i64>(data, shape);
          case element::Type_t::u8:
              return getShapeFromHostTensorData<element::Type_t::u8>(data, shape);
          case element::Type_t::u16:
              return getShapeFromHostTensorData<element::Type_t::u16>(data, shape);
          case element::Type_t::u32:
              return getShapeFromHostTensorData<element::Type_t::u32>(data, shape);
          case element::Type_t::u64:
              return getShapeFromHostTensorData<element::Type_t::u64>(data, shape);
          default:
              return false;
      }
      return true;
  }
  
  template<element::Type_t DataType>
  bool evaluate(const HostTensorPtr& inputTensor,
                const HostTensorPtr& inputShapeTensor,
                const HostTensorPtr& outputTensor) {
      Shape inputShape = inputTensor->get_shape();
      Shape outputShape;
      if (!getShapeFromHostTensorData(inputShapeTensor, outputShape)) {
          return false;
      }
  
      if (!ngraph::PartialShape(outputShape).refines(outputTensor->get_partial_shape())) {
          return false;
      }
  
      outputTensor->set_shape(outputShape);
  
      using T = typename element_type_traits<DataType>::value_type;
      T *inputPtr = inputTensor->get_data_ptr<DataType>();
      T *outputPtr = outputTensor->get_data_ptr<DataType>();
  
      const auto inTotalDimSize = shape_size(inputShape);
      const auto stridesByElements = row_major_strides(inputShape);
  
      const auto inLineSize = inputShape[inputShape.size() - 1];
      const auto outLineSize = outputShape[outputShape.size() - 1];
  
      for (size_t inElementOffset = 0, outElementOffset = 0; inElementOffset < inTotalDimSize; inElementOffset += inLineSize) {
          auto offset = inElementOffset;
          bool isGarbageLine = false;
          for (size_t dim = 0; dim < stridesByElements.size() - 1; ++dim) {
              const auto coordAlongDim = offset / stridesByElements[dim];
              if (coordAlongDim > outputShape[dim] - 1) {
                  isGarbageLine = true;
                  break;
              }
  
              offset %= stridesByElements[dim];
          }
          if (!isGarbageLine) {
              std::copy_n(inputPtr + inElementOffset, outLineSize, outputPtr + outElementOffset);
              outElementOffset += outLineSize;
          }
      }
      return true;
  }
  
  bool evaluateDynamicShapeResolver(const HostTensorPtr& inputTensor,
                                    const HostTensorPtr& inputShapeTensor,
                                    const HostTensorPtr& outputTensor) {
      bool rc = true;
  
      switch (inputTensor->get_element_type()) {
          TYPE_CASE(i8)(inputTensor, inputShapeTensor, outputTensor);
              break;
          TYPE_CASE(i16)(inputTensor, inputShapeTensor, outputTensor);
              break;
          TYPE_CASE(i32)(inputTensor, inputShapeTensor, outputTensor);
              break;
          TYPE_CASE(i64)(inputTensor, inputShapeTensor, outputTensor);
              break;
          TYPE_CASE(u8)(inputTensor, inputShapeTensor, outputTensor);
              break;
          TYPE_CASE(u16)(inputTensor, inputShapeTensor, outputTensor);
              break;
          TYPE_CASE(u32)(inputTensor, inputShapeTensor, outputTensor);
              break;
          TYPE_CASE(u64)(inputTensor, inputShapeTensor, outputTensor);
              break;
          TYPE_CASE(bf16)(inputTensor, inputShapeTensor, outputTensor);
              break;
          TYPE_CASE(f32)(inputTensor, inputShapeTensor, outputTensor);
              break;
          TYPE_CASE(f64)(inputTensor, inputShapeTensor, outputTensor);
              break;
          TYPE_CASE(boolean)(inputTensor, inputShapeTensor, outputTensor);
              break;
          default:
              rc = false;
              break;
      }
  
      return rc;
  }
  
  }  // namespace
  }  // namespace dyn_shape
  
  bool DynamicShapeResolver::evaluate(const HostTensorVector& outputs,
                                      const HostTensorVector& inputs) const {
      return dyn_shape::evaluateDynamicShapeResolver(inputs[0], inputs[1], outputs[0]);
  }
  
  }  // namespace op
  }  // namespace vpu
  }  // namespace ngraph
  