adjust scatter_nd input order and cast mindir float64 input

mindspore/ccsrc/plugin/device/cpu/kernel/nnacl/fp32/scatter_nd_fp32.c → mindspore/ccsrc/plugin/device/cpu/kernel/nnacl/base/scatter_nd_base.c

@@ -1,5 +1,5 @@
 /**
- * Copyright 2020 Huawei Technologies Co., Ltd
+ * Copyright 2020-2022 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "nnacl/fp32/scatter_nd_fp32.h"
+#include "nnacl/base/scatter_nd_base.h"
 #include <string.h>
 #include "nnacl/errorcode.h"
 

mindspore/ccsrc/plugin/device/cpu/kernel/nnacl/fp32/scatter_nd_fp32.h → mindspore/ccsrc/plugin/device/cpu/kernel/nnacl/base/scatter_nd_base.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2020 Huawei Technologies Co., Ltd
+ * Copyright 2020-2022 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 
-#ifndef MINDSPORE_NNACL_FP32_SCATTER_ND_FP32_H_
-#define MINDSPORE_NNACL_FP32_SCATTER_ND_FP32_H_
+#ifndef MINDSPORE_NNACL_BASE_SCATTER_ND_BASE_H_
+#define MINDSPORE_NNACL_BASE_SCATTER_ND_BASE_H_
 
 #include "nnacl/op_base.h"
 
@@ -35,4 +35,4 @@ int DoScatterND(void *output, const void *update, int *output_unit_offsets, cons
 }
 #endif
 
-#endif  // MINDSPORE_NNACL_FP32_SCATTER_ND_FP32_H_
+#endif  // MINDSPORE_NNACL_BASE_SCATTER_ND_BASE_H_

mindspore/ccsrc/plugin/device/cpu/kernel/nnacl/infer/scatter_nd_infer.c

@@ -1,5 +1,5 @@
 /**
- * Copyright 2021 Huawei Technologies Co., Ltd
+ * Copyright 2021-2022 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -24,11 +24,11 @@ int ScatterNdInferShape(const TensorC *const *inputs, size_t inputs_size, Tensor
     return check_ret;
   }
 
-  const TensorC *shape = inputs[0];
+  const TensorC *shape = inputs[THIRD_INPUT];
   if (shape->data_ == NULL) {
     return NNACL_INFER_INVALID;
   }
-  const TensorC *update = inputs[2];
+  const TensorC *update = inputs[FIRST_INPUT];
   TensorC *output = outputs[0];
 
   SetDataTypeFormat(output, update);

mindspore/lite/src/ops/populate/scatter_nd_populate.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2019-2021 Huawei Technologies Co., Ltd
+ * Copyright 2019-2022 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 #include "src/ops/populate/populate_register.h"
-#include "nnacl/fp32/scatter_nd_fp32.h"
+#include "nnacl/base/scatter_nd_base.h"
 using mindspore::schema::PrimitiveType_ScatterNd;
 
 namespace mindspore {

mindspore/lite/src/ops/populate/scatter_nd_update_populate.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2021 Huawei Technologies Co., Ltd
+ * Copyright 2021-2022 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 #include "src/ops/populate/populate_register.h"
-#include "nnacl/fp32/scatter_nd_fp32.h"
+#include "nnacl/base/scatter_nd_base.h"
 using mindspore::schema::PrimitiveType_ScatterNdUpdate;
 
 namespace mindspore {

mindspore/lite/src/ops/populate/v0/scatter_nd_populate_v0.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2021 Huawei Technologies Co., Ltd
+ * Copyright 2021-2022 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,7 +16,7 @@
 
 #include "schema/model_v0_generated.h"
 #include "src/ops/populate/populate_register.h"
-#include "nnacl/fp32/scatter_nd_fp32.h"
+#include "nnacl/base/scatter_nd_base.h"
 
 namespace mindspore {
 namespace lite {

mindspore/lite/src/runtime/kernel/arm/fp32/scatter_nd_fp32.cc → mindspore/lite/src/runtime/kernel/arm/base/scatter_nd_base.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2020 Huawei Technologies Co., Ltd
+ * Copyright 2020-2022 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "src/runtime/kernel/arm/fp32/scatter_nd_fp32.h"
+#include "src/runtime/kernel/arm/base/scatter_nd_base.h"
 #include <cstring>
 #include <vector>
 #include "schema/model_generated.h"
@@ -29,9 +29,9 @@ using mindspore::schema::PrimitiveType_ScatterNd;
 
 namespace mindspore::kernel {
 namespace {
-constexpr int kScatterShapeIndex = 0;
-constexpr int kScatterIndicesIndex = 1;
-constexpr int kScatterUpdateIndex = 2;
+constexpr int kScatterIndicesIndex = 0;
+constexpr int kScatterUpdateIndex = 1;
+constexpr int kScatterShapeIndex = 2;
 }  // namespace
 int ScatterNDCPUKernel::Prepare() {
   CHECK_LESS_RETURN(in_tensors_.size(), DIMENSION_3D);
@@ -46,45 +46,28 @@ int ScatterNDCPUKernel::ReSize() {
   auto indices = in_tensors_[kScatterIndicesIndex];
   auto update = in_tensors_[kScatterUpdateIndex];
   auto shape = in_tensors_[kScatterShapeIndex];
-  CHECK_NULL_RETURN(indices);
-  CHECK_NULL_RETURN(update);
-  CHECK_NULL_RETURN(shape);
-
-  // check indices shape
-  auto shape_rank = shape->ElementsNum();
+  auto indices_shape = indices->shape();
+  auto update_shape = update->shape();
   auto shape_data = reinterpret_cast<int *>(shape->data());
   CHECK_NULL_RETURN(shape_data);
-  auto indice_unit_rank = indices->shape().back();
-  if (indice_unit_rank > shape_rank) {
-    MS_LOG(ERROR) << "Value of last dimension of indices is greater than shape rank.";
-    return RET_ERROR;
-  }
-
-  if (indices->shape().size() < 2) {
-    MS_LOG(ERROR) << "Indices dimension smaller than 2.";
-    return RET_ERROR;
-  }
+  int indices_rank = static_cast<int>(indices->shape().size());
+  int update_rank = static_cast<int>(update->shape().size());
+  auto shape_rank = shape->ElementsNum();
+  int indice_unit_rank = indices->shape().back();
 
+  // check indices shape
+  MS_CHECK_TRUE_MSG(indices_rank >= DIMENSION_2D, RET_ERROR, "The rank of indices must be greater equal than 2.");
+  MS_CHECK_TRUE_MSG(indice_unit_rank <= shape_rank, RET_ERROR,
+                    "The value of indices' last dimension must be less equal than the input rank.");
+  MS_CHECK_TRUE_MSG(update_rank == indices_rank - 1 + shape_rank - indice_unit_rank, RET_ERROR,
+                    "The rank of update is illegal.");
   // check consistency of the shape indices and shape
-  auto update_rank = static_cast<int>(update->shape().size());
-  auto indices_shape = indices->shape();
-  if (update_rank != static_cast<int>(indices->shape().size() - 1 + shape_rank - indice_unit_rank)) {
-    MS_LOG(ERROR) << "Update, shape rank and indices rank inconsistent.";
-    return RET_ERROR;
-  }
-  // check update shape
-  auto update_shape = update->shape();
-  for (size_t i = 0; i < indices_shape.size() - 1; i++) {
-    if (update_shape.at(i) != indices_shape.at(i)) {
-      MS_LOG(ERROR) << "Value of " << i << " th dimension of indices is not equal to that of update.";
-      return RET_ERROR;
+  for (int i = 0; i < update_rank; i++) {
+    if (i < indices_rank - 1) {
+      MS_CHECK_TRUE_MSG(update_shape[i] == indices_shape[i], RET_ERROR, "the shape of update tensor is illegal.");
     }
-  }
-  for (size_t i = 0; i < shape_rank - (indices_shape.size() - 1); i++) {
-    if (update_shape.at(i + indices_shape.size() - 1) != shape_data[i + indices_shape.size() - 1]) {
-      MS_LOG(ERROR) << "Value of " << i + indices_shape.size() - 1
-                    << " th dimension of indices is not equal to the corresbonding dimension of shape.";
-      return RET_ERROR;
+    if (i >= indice_unit_rank) {
+      MS_CHECK_TRUE_MSG(update_shape[i] == shape_data[i], RET_ERROR, "the shape of update tensor is illegal.");
     }
   }
 
@@ -96,29 +79,17 @@ int ScatterNDCPUKernel::ReSize() {
 
   // calculate offsets
   int out_stride = 1;
-  std::vector<int> out_strides;
-  out_strides.push_back(1);
-  for (int i = indice_unit_rank - 2; i >= 0; i--) {
+  out_strides_.push_back(1);
+  for (int i = indice_unit_rank - C2NUM; i >= 0; i--) {
     out_stride *= shape_data[i + 1];
-    out_strides.push_back(out_stride);
+    out_strides_.push_back(out_stride);
   }
 
   param_->num_unit = 1;
   param_->num_unit *= update_shape.at(indices_shape.size() - C2NUM);
-  for (int i = indices_shape.size() - 3; i >= 0; i--) {
+  for (int i = indices_shape.size() - C3NUM; i >= 0; i--) {
     param_->num_unit *= update_shape.at(i);
   }
-
-  int *indices_ptr = reinterpret_cast<int *>(indices->data());
-  CHECK_NULL_RETURN(indices_ptr);
-  output_unit_offsets_.clear();
-  for (int i = 0; i < param_->num_unit; i++) {
-    int tmp_stride = 0;
-    for (int j = 0; j < indice_unit_rank; j++) {
-      tmp_stride += indices_ptr[i * indice_unit_rank + j] * out_strides.at(j) * param_->unit_size;
-    }
-    output_unit_offsets_.push_back(tmp_stride);
-  }
   return RET_OK;
 }
 
@@ -139,6 +110,17 @@ int ScatterNDRun(void *cdata, int task_id, float lhs_scale, float rhs_scale) {
 }
 
 int ScatterNDCPUKernel::Run() {
+  auto indices = in_tensors_[kScatterIndicesIndex];
+  auto indice_unit_rank = indices->shape().back();
+  int *indices_ptr = reinterpret_cast<int *>(indices->data());
+  output_unit_offsets_.clear();
+  for (int i = 0; i < param_->num_unit; i++) {
+    int tmp_stride = 0;
+    for (int j = 0; j < indice_unit_rank; j++) {
+      tmp_stride += indices_ptr[i * indice_unit_rank + j] * out_strides_.at(j) * param_->unit_size;
+    }
+    output_unit_offsets_.push_back(tmp_stride);
+  }
   auto ret = ParallelLaunch(ms_context_, ScatterNDRun, this, op_parameter_->thread_num_);
   if (ret != RET_OK) {
     MS_LOG(ERROR) << "ScatterNDRun failed, ret: " << ret;
@@ -146,8 +128,5 @@ int ScatterNDCPUKernel::Run() {
   return ret;
 }
 
-REG_KERNEL(kCPU, kNumberTypeFloat32, PrimitiveType_ScatterNd, LiteKernelCreator<ScatterNDCPUKernel>)
-#ifdef ENABLE_FP16
-REG_KERNEL(kCPU, kNumberTypeFloat16, PrimitiveType_ScatterNd, LiteKernelCreator<ScatterNDCPUKernel>)
-#endif
+REG_KERNEL(kCPU, kNumberTypeInt32, PrimitiveType_ScatterNd, LiteKernelCreator<ScatterNDCPUKernel>)
 }  // namespace mindspore::kernel

mindspore/lite/src/runtime/kernel/arm/fp32/scatter_nd_fp32.h → mindspore/lite/src/runtime/kernel/arm/base/scatter_nd_base.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2020 Huawei Technologies Co., Ltd
+ * Copyright 2020-2022 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,7 +19,7 @@
 
 #include <vector>
 #include "src/inner_kernel.h"
-#include "nnacl/fp32/scatter_nd_fp32.h"
+#include "nnacl/base/scatter_nd_base.h"
 
 namespace mindspore::kernel {
 class ScatterNDCPUKernel : public InnerKernel {
@@ -39,6 +39,7 @@ class ScatterNDCPUKernel : public InnerKernel {
  private:
   ScatterNDParameter *param_ = nullptr;
   std::vector<int> output_unit_offsets_;
+  std::vector<int> out_strides_;
 };
 }  // namespace mindspore::kernel
 

mindspore/lite/src/runtime/kernel/arm/fp32/scatter_nd_update_fp32.cc

@@ -31,7 +31,6 @@ namespace {
 constexpr int kScatterUpdateInputIndex = 0;
 constexpr int kScatterIndicesIndex = 1;
 constexpr int kScatterUpdateIndex = 2;
-constexpr size_t kScatterIndicesDims = 2;
 }  // namespace
 int ScatterNdUpdateCPUKernel::Prepare() {
   CHECK_LESS_RETURN(in_tensors_.size(), DIMENSION_3D);
@@ -46,24 +45,31 @@ int ScatterNdUpdateCPUKernel::ReSize() {
   auto input = in_tensors_.at(kScatterUpdateInputIndex);
   auto indices = in_tensors_.at(kScatterIndicesIndex);
   auto update = in_tensors_.at(kScatterUpdateIndex);
-
-  // check indices shape
+  auto input_shape = input->shape();
+  auto indices_shape = indices->shape();
+  auto update_shape = update->shape();
   int input_rank = static_cast<int>(input->shape().size());
+  int indices_rank = static_cast<int>(indices->shape().size());
+  int update_rank = static_cast<int>(update->shape().size());
   int indice_unit_rank = indices->shape().back();
-  if (indice_unit_rank > input_rank) {
-    MS_LOG(ERROR) << "Value of last dimension of indices is greater than input rank.";
-    return RET_ERROR;
-  }
 
-  if (indices->shape().size() < kScatterIndicesDims) {
-    MS_LOG(ERROR) << "Indices dimension smaller than 2.";
-    return RET_ERROR;
+  // check indices shape
+  MS_CHECK_TRUE_MSG(indices_rank >= DIMENSION_2D, RET_ERROR, "The rank of indices must be greater equal than 2.");
+  MS_CHECK_TRUE_MSG(indice_unit_rank <= input_rank, RET_ERROR,
+                    "The value of indices' last dimension must be less equal than the input rank.");
+  MS_CHECK_TRUE_MSG(update_rank == indices_rank - 1 + input_rank - indice_unit_rank, RET_ERROR,
+                    "The rank of update is illegal.");
+  // check consistency of the shape indices and shape
+  for (int i = 0; i < update_rank; i++) {
+    if (i < indices_rank - 1) {
+      MS_CHECK_TRUE_MSG(update_shape[i] == indices_shape[i], RET_ERROR, "the shape of update tensor is illegal.");
+    }
+    if (i >= indice_unit_rank) {
+      MS_CHECK_TRUE_MSG(update_shape[i] == input_shape[i], RET_ERROR, "the shape of update tensor is illegal.");
+    }
   }
 
-  // check consistency of the shape indices and shape
-  int update_rank = static_cast<int>(update->shape().size());
-  auto indices_shape = indices->shape();
-  auto update_shape = update->shape();
+  // calculate unit_size
   param_->unit_size = 1;
   for (int i = indices_shape.size() - 1; i < update_rank; i++) {
     param_->unit_size *= update_shape.at(i);
@@ -73,15 +79,15 @@ int ScatterNdUpdateCPUKernel::ReSize() {
   int out_stride = 1;
   std::vector<int> out_strides;
   out_strides.push_back(1);
-  for (int i = indice_unit_rank - 2; i >= 0; i--) {
-    out_stride *= input->shape()[i + 1];
+  for (int i = indice_unit_rank - C2NUM; i >= 0; i--) {
+    out_stride *= input_shape[i + 1];
     out_strides.push_back(out_stride);
   }
   std::reverse(out_strides.begin(), out_strides.end());
 
   param_->num_unit = 1;
   param_->num_unit *= update_shape.at(indices_shape.size() - C2NUM);
-  for (int i = indices_shape.size() - 3; i >= 0; i--) {
+  for (int i = indices_shape.size() - C3NUM; i >= 0; i--) {
     param_->num_unit *= update_shape.at(i);
   }
 

mindspore/lite/src/runtime/kernel/arm/fp32/scatter_nd_update_fp32.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2021 Huawei Technologies Co., Ltd
+ * Copyright 2021-2022 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,7 +19,7 @@
 
 #include <vector>
 #include "src/inner_kernel.h"
-#include "nnacl/fp32/scatter_nd_fp32.h"
+#include "nnacl/base/scatter_nd_base.h"
 
 namespace mindspore::kernel {
 

mindspore/lite/tools/converter/import/mindir_adjust.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2020-2021 Huawei Technologies Co., Ltd
+ * Copyright 2020-2022 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -181,6 +181,24 @@ int ConvertQuantParam(const PrimitivePtr &prim, const std::vector<AnfNodePtr> &i
 }
 }  // namespace
 
+int MindirAdjust::AdjustInputDataType(AnfNodePtr anf_node) {
+  MS_CHECK_TRUE_MSG(anf_node != nullptr, RET_ERROR, "anf_node is nullptr");
+  auto param_node = anf_node->cast<ParameterPtr>();
+  MS_CHECK_TRUE_MSG(param_node != nullptr, RET_ERROR, "param_node is nullptr");
+  auto abstract_tensor = param_node->abstract()->cast<abstract::AbstractTensorPtr>();
+  MS_CHECK_TRUE_MSG(abstract_tensor != nullptr, RET_ERROR, "param node has no abstract tensor.");
+  auto tensor_element = abstract_tensor->element();
+  MS_CHECK_TRUE_MSG(tensor_element != nullptr, RET_ERROR, "abstract tensor's element is null.");
+  auto type_ptr = abstract_tensor->element()->GetTypeTrack();
+  MS_CHECK_TRUE_MSG(type_ptr != nullptr, RET_ERROR, "Type pointer is null.");
+  auto org_type = type_ptr->type_id();
+  if (!param_node->has_default() && (org_type == kNumberTypeInt64 || org_type == kNumberTypeFloat64)) {
+    TypeId dst_type = org_type == kNumberTypeInt64 ? kNumberTypeInt32 : kNumberTypeFloat32;
+    tensor_element->set_type(TypeIdToType(dst_type));
+  }
+  return RET_OK;
+}
+
 int MindirAdjust::ValueNodeInt64Convert(AnfNodePtr anf_node) {
   MS_CHECK_TRUE_MSG(anf_node != nullptr, RET_ERROR, "anf_node is nullptr");
   if (!utils::isa<ValueNodePtr>(anf_node)) {
@@ -327,6 +345,8 @@ bool MindirAdjust::Run(const FuncGraphPtr &func_graph) {
         if (status == RET_OK || status == RET_NO_CHANGE) {
           status = UpdateConv2DTransposeInput(node->cast<CNodePtr>());
         }
+      } else if (utils::isa<ParameterPtr>(node)) {
+        status = AdjustInputDataType(node);
       } else if (utils::isa<ValueNodePtr>(node)) {
         status = ValueNodeInt64Convert(node);
       }

mindspore/lite/tools/converter/import/mindir_adjust.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2020-2021 Huawei Technologies Co., Ltd
+ * Copyright 2020-2022 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -34,6 +34,7 @@ class MindirAdjust {
   bool Run(const FuncGraphPtr &graph);
 
  private:
+  int AdjustInputDataType(AnfNodePtr anf_node);
   int ValueNodeInt64Convert(AnfNodePtr anf_node);
   int ComputeQuantParams(AnfNodePtr anf_node);
   int UpdateConv2DTransposeInput(const CNodePtr &cnode);