Add scatter_nd cpu (#1603)

* Add ScatterND impl

* remove skipped tests

* add rank check for input tensor.

onnxruntime/core/providers/cpu/cpu_execution_provider.cc

@@ -384,6 +384,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Lp
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Conv);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, ConvTranspose);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, If);
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, ScatterND);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Gemm);
 
 void RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
@@ -753,6 +754,7 @@ void RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Conv)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, ConvTranspose)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, If)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, ScatterND)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Gemm)>,
   };
 

onnxruntime/core/providers/cpu/tensor/scatter_nd.cc

@@ -0,0 +1,157 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "scatter_nd.h"
+
+namespace onnxruntime {
+
+ONNX_CPU_OPERATOR_KERNEL(
+    ScatterND,
+    11,
+    KernelDefBuilder()
+        .TypeConstraint("T",    DataTypeImpl::AllTensorTypes())
+        .TypeConstraint("Tind", DataTypeImpl::GetTensorType<int64_t>()),
+    ScatterND);
+
+template<typename Tind>
+Status ScatterNDBase::PrepareForCompute(OpKernelContext* context, Prepare& p) const {
+
+  auto input_tensor  = context->Input<Tensor>(0);
+  auto indice_tensor = context->Input<Tensor>(1);
+  auto update_tensor = context->Input<Tensor>(2);
+  ORT_ENFORCE(input_tensor  != nullptr);
+  ORT_ENFORCE(indice_tensor != nullptr);
+  ORT_ENFORCE(update_tensor != nullptr);
+
+  auto input_shape   = input_tensor->Shape();
+  auto indice_shape  = indice_tensor->Shape();
+  auto update_shape  = update_tensor->Shape();
+  if (indice_shape.NumDimensions() == 0 || input_shape.NumDimensions() == 0) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+      "input tensor and indices tensor must has rank larger than 0. ",
+      "input shape: ", input_shape, ", indices shape: ", indice_shape);
+  }
+
+  auto indice_rank = indice_shape.NumDimensions();
+  auto last_indice_dimension = indice_shape[indice_rank - 1];
+  if (last_indice_dimension > static_cast<int64_t>(input_shape.NumDimensions())) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+      "last dimension of indices must not be larger than rank of input tensor");
+  }
+
+  bool is_update_shape_invalid = [&](){
+    auto update_rank = update_shape.NumDimensions();
+    auto input_rank = input_shape.NumDimensions();
+    if (update_rank < indice_rank - 1) {
+      return true;
+    }
+    if ((update_rank >= indice_rank - 1) &&
+        (indice_rank - 1 >= 0) &&
+        (indice_shape.Slice(0, indice_rank - 1) != update_shape.Slice(0, indice_rank - 1))) {
+      return true;
+    }
+    if ((static_cast<int64_t>(input_rank) > last_indice_dimension) &&
+        (update_rank >= indice_rank - 1) &&
+        (input_shape.Slice(last_indice_dimension) != update_shape.Slice(indice_rank - 1))) {
+      return true;
+    }
+    return false;
+  }();
+  if (is_update_shape_invalid) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+      "updates tensor should have shape equal to indices.shape[:-1] + data.shape[indices.shape[-1]:]. ",
+      "updates shape: ", update_shape, ", indices shape: ", indice_shape, ", data shape: ", input_shape);
+  }
+
+  auto output_tensor = context->Output(0, TensorShape(input_shape));
+
+  const auto* src_base = input_tensor->DataRaw();
+  auto* dst_base = output_tensor->MutableDataRaw();
+  bool is_string_type = input_tensor->DataType() == DataTypeImpl::GetType<std::string>();
+
+  // Re-use input for output. If input/output Tensor* are the same, do not copy.
+  if (src_base != dst_base) {
+    if (is_string_type) {
+      const auto* str_begin = input_tensor->template Data<std::string>();
+      const std::string* str_end = str_begin + input_shape.Size();
+      auto* dst = output_tensor->template MutableData<std::string>();
+      std::copy(str_begin, str_end, dst);
+    } else {
+      memcpy(dst_base, src_base, input_tensor->SizeInBytes());
+    }
+  }
+
+  std::vector<int64_t> element_counts(last_indice_dimension, 0LL); // Number of elements for each input dimension
+
+#ifdef USE_OPENMP
+#pragma omp parallel for
+#endif
+  for (int64_t i = 0; i < last_indice_dimension; ++i) {
+    element_counts[i] = input_shape.SizeFromDimension(i + 1);
+  }
+
+  int64_t err_indice = 0;
+  p.element_bytes    = input_tensor->DataType()->Size();
+  p.element_to_copy  = input_shape.SizeFromDimension(last_indice_dimension);
+  p.bytes_to_copy    = p.element_bytes * p.element_to_copy;
+  auto indice_offset = static_cast<const Tind*>(indice_tensor->DataRaw());
+  auto offset_count  = indice_shape.Size() / last_indice_dimension; // Times to copy
+  p.element_offsets.assign(offset_count, 0LL);
+
+  if (input_tensor->DataType() == DataTypeImpl::GetType<std::string>()) {
+    p.input_str_base  = static_cast<const std::string*>(update_tensor->DataRaw());
+    p.output_str_base = static_cast<std::string*>(output_tensor->MutableDataRaw());
+  } else {
+    p.input_base      = static_cast<const uint8_t*>(update_tensor->DataRaw());
+    p.output_base     = static_cast<uint8_t*>(output_tensor->MutableDataRaw());
+  }
+
+#ifdef USE_OPENMP
+#pragma omp parallel for
+#endif
+  for (int64_t i = 0; i < offset_count; ++i) {
+    for (int64_t j = 0; j < last_indice_dimension; ++j) {
+      auto indice = *(indice_offset + i * last_indice_dimension + j);
+      if (indice < 0 || indice >= input_shape[j]) {
+        err_indice = indice;
+      }
+      p.element_offsets[i] += indice * element_counts[j];
+    }
+  }
+  return err_indice == 0 ? Status::OK() :
+    ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "invalid indice found, indice = ", err_indice);
+}
+
+template Status ScatterNDBase::PrepareForCompute<int64_t>(OpKernelContext*, Prepare&) const;
+
+Status ScatterND::Compute(OpKernelContext* context) const {
+  Prepare p;
+  ORT_RETURN_IF_ERROR(PrepareForCompute<int64_t>(context, p));
+  return nullptr == p.input_str_base ? ScatterNumber(p) : ScatterString(p);
+}
+
+Status ScatterND::ScatterNumber(const Prepare& p) const {
+#ifdef USE_OPENMP
+#pragma omp parallel for
+#endif
+  for (int64_t i = 0; i < static_cast<int64_t>(p.element_offsets.size()); ++i) {
+    memcpy(p.output_base + p.element_offsets[i] * p.element_bytes,
+           p.input_base + i * p.bytes_to_copy,
+           p.bytes_to_copy);
+  }
+  return Status::OK();
+}
+
+Status ScatterND::ScatterString(const Prepare& p) const {
+#ifdef USE_OPENMP
+#pragma omp parallel for
+#endif
+  for (int64_t i = 0; i < static_cast<int64_t>(p.element_offsets.size()); ++i) {
+    for (int64_t j = 0; j < static_cast<int64_t>(p.element_to_copy); ++j) {
+      p.output_str_base[p.element_offsets[i] + j] = p.input_str_base[i * p.element_to_copy + j];
+    }
+  }
+  return Status::OK();
+}
+
+}

onnxruntime/core/providers/cpu/tensor/scatter_nd.h

@@ -0,0 +1,47 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+
+namespace onnxruntime {
+
+class ScatterNDBase
+{
+protected:
+  struct Prepare {
+    const uint8_t*        input_base;
+    const std::string*    input_str_base;
+    uint8_t*              output_base;
+    std::string*          output_str_base;
+    uint64_t              bytes_to_copy;
+    uint64_t              element_bytes;
+    uint64_t              element_to_copy;
+    std::vector<uint64_t> element_offsets;
+
+    Prepare(): input_base      (nullptr),
+               input_str_base  (nullptr),
+               output_base     (nullptr),
+               output_str_base (nullptr),
+               bytes_to_copy   (0),
+               element_bytes   (0),
+               element_to_copy (0),
+               element_offsets (0) {}
+  }; // struct Prepare
+
+  template<typename Tind>
+  Status PrepareForCompute(OpKernelContext* context, Prepare& p) const;
+}; // class ScatterNDBase
+
+class ScatterND final : public OpKernel, protected ScatterNDBase {
+public:
+  explicit ScatterND(const OpKernelInfo& info) : OpKernel(info) {}
+  Status Compute(OpKernelContext* context) const override;
+private:
+  Status ScatterNumber(const Prepare& p) const;
+  Status ScatterString(const Prepare& p) const;
+};
+
+} // namespace onnxruntime

onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc

@@ -0,0 +1,159 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "gtest/gtest.h"
+#include "test/providers/provider_test_utils.h"
+
+namespace onnxruntime {
+namespace test {
+
+TEST(ScatterNDOpTest, ScatterND_scaler_string_int64) {
+  OpTester test1("ScatterND", 11);
+  test1.AddInput<std::string>("data", {2,2}, {"h","h","o","z"});
+  test1.AddInput<int64_t>("indices", {2}, {0,1});
+  test1.AddInput<std::string>("updates", {}, {"k"});
+  test1.AddOutput<std::string>("output", {2,2}, {"h","k","o","z"});
+  test1.Run();
+
+  OpTester test2("ScatterND", 11);
+  test2.AddInput<std::string>("data", {6}, {"h","k","o","o","l","t"});
+  test2.AddInput<int64_t>("indices", {1}, {3});
+  test2.AddInput<std::string>("updates", {}, {"z"});
+  test2.AddOutput<std::string>("output", {6}, {"h","k","o","z","l","t"});
+  test2.Run();
+
+  OpTester test3("ScatterND", 11);
+  test3.AddInput<std::string>("data", {3,2}, {"h","k","o","z","l","z"});
+  test3.AddInput<int64_t>("indices", {2}, {2,1});
+  test3.AddInput<std::string>("updates", {}, {"t"});
+  test3.AddOutput<std::string>("output", {3,2}, {"h","k","o","z","l","t"});
+  test3.Run();
+}
+
+TEST(ScatterNDOpTest, ScatterND_matrice_int64_int64) {
+  OpTester test("ScatterND", 11);
+  test.AddInput<int64_t> ("data", {2,2}, {1LL,1LL,2LL,2LL});
+  test.AddInput<int64_t> ("indices", {2,2}, {0LL,0LL,1LL,1LL});
+  test.AddInput<int64_t>("updates", {2}, {0LL,3LL});
+  test.AddOutput<int64_t>("output", {2,2}, {0LL,1LL,2LL,3LL});
+  test.Run();
+}
+
+TEST(ScatterNDOpTest, ScatterND_matrice_string_int64) {
+  OpTester test1("ScatterND", 11);
+  test1.AddInput<std::string>("data", {2,2,2}, {"egg","dance","bob","air","smart","terry","laugh","kite"});
+  test1.AddInput<int64_t>("indices", {2,1,2}, {0,1,1,0});
+  test1.AddInput<std::string>("updates", {2,1,2}, {"air","bob","terry","smart"});
+  test1.AddOutput<std::string>("output", {2,2,2}, {"egg","dance","air","bob","terry","smart","laugh","kite"});
+  test1.Run();
+
+  OpTester test2("ScatterND", 11);
+  test2.AddInput<std::string>("data", {3,3}, {"egg","","air","","terry","smart","laugh","","hop"});
+  test2.AddInput<int64_t>("indices", {3,2}, {2,1,1,0,0,1});
+  test2.AddInput<std::string>("updates", {3}, {"kite","bob","dance"});
+  test2.AddOutput<std::string>("output", {3,3}, {"egg","dance","air","bob","terry","smart","laugh","kite","hop"});
+  test2.Run();
+}
+
+TEST(ScatterNDOpTest, ScatterND_slice_float_int64_t) {
+  OpTester test("ScatterND", 11);
+  test.AddInput<float>("data", {2,2}, {0.0f,0.1f,0.1f,0.1f});
+  test.AddInput<int64_t>("indices", {2,1}, {1LL,0LL});
+  test.AddInput<float>("updates", {2,2}, {0.2f,0.3f,0.0f,0.1f});
+  test.AddOutput<float>("output", {2,2}, {0.0f,0.1f,0.2f,0.3f});
+  test.Run();
+}
+
+TEST(ScatterNDOpTest, ScatterND_slice_double_int64_t) {
+  OpTester test("ScatterND", 11);
+  test.AddInput<double>("data", {2,2}, {0.0f,0.1f,0.1f,0.1f});
+  test.AddInput<int64_t>("indices", {2,1}, {1LL,0LL});
+  test.AddInput<double>("updates", {2,2}, {0.2f,0.3f,0.0f,0.1f});
+  test.AddOutput<double>("output", {2,2}, {0.0f,0.1f,0.2f,0.3f});
+  test.Run();
+}
+
+TEST(ScatterNDOpTest, ScatterND_3tensor_int64) {
+  OpTester test1("ScatterND", 11);
+  test1.AddInput<int64_t>("data", {2,2,2}, {0LL,1LL,1LL,1LL,1LL,1LL,6LL,7LL});
+  test1.AddInput<int64_t>("indices", {2,2}, {0LL,1LL,1LL,0LL});
+  test1.AddInput<int64_t>("updates", {2,2}, {2LL,3LL,4LL,5LL});
+  test1.AddOutput<int64_t>("output", {2,2,2}, {0LL,1LL,2LL,3LL,4LL,5LL,6LL,7LL});
+  test1.Run();
+
+  OpTester test2("ScatterND", 11);
+  test2.AddInput<int8_t>("data", {2,2,2}, {0,0,2,3,4,0,6,7});
+  test2.AddInput<int64_t>("indices", {2,3}, {0,0,1,1,0,1});
+  test2.AddInput<int8_t>("updates", {2}, {1,5});
+  test2.AddOutput<int8_t>("output", {2,2,2}, {0,1,2,3,4,5,6,7});
+  test2.Run();
+
+  OpTester test3("ScatterND", 11);
+  test3.AddInput<int16_t>("data", {2,2,2}, {0,1,2,3,0,1,2,3});
+  test3.AddInput<int64_t>("indices", {1,1}, {1LL});
+  test3.AddInput<int16_t>("updates", {1,2,2}, {4,5,6,7});
+  test3.AddOutput<int16_t>("output", {2,2,2}, {0,1,2,3,4,5,6,7});
+  test3.Run();
+}
+
+TEST(ScatterNDOpTest, ScatterND_batched_index_int64) {
+  OpTester test("ScatterND", 11);
+  test.AddInput<int64_t>("data", {2,2}, {2LL,3LL,2LL,3LL});
+  test.AddInput<int64_t>("indices", {2,1,2}, {0LL,0LL,0LL,1LL});
+  test.AddInput<int64_t>("updates", {2,1}, {0LL,1LL});
+  test.AddOutput<int64_t>("output", {2,2}, {0LL,1LL,2LL,3LL});
+  test.Run();
+}
+
+TEST(ScatterNDOpTest, ScatterND_batched_index_bool_int64) {
+  OpTester test("ScatterND", 11);
+  test.AddInput<bool>("data", {2,2}, {false,true,false,true});
+  test.AddInput<int64_t>("indices", {2,1,2}, {0LL,0LL,0LL,1LL});
+  test.AddInput<bool>("updates", {2,1}, {true,false});
+  test.AddOutput<bool>("output", {2,2}, {true,false,false,true});
+  test.Run();
+}
+
+TEST(ScatterNDOpTest, ScatterND_sliced_index_int64) {
+  OpTester test("ScatterND", 11);
+  test.AddInput<int64_t>("data", {2,2}, {0LL,0LL,0LL,0LL});
+  test.AddInput<int64_t>("indices", {2,1,1}, {1LL,0LL});
+  test.AddInput<int64_t>("updates", {2,1,2}, {2LL,3LL,0LL,1LL});
+  test.AddOutput<int64_t>("output", {2,2}, {0LL,1LL,2LL,3LL});
+  test.Run();
+}
+
+TEST(ScatterNDOpTest, ScatterND_sliced_index_string_int64) {
+  OpTester test("ScatterND", 11);
+  test.AddInput<std::string>("data", {2,2}, {"","","",""});
+  test.AddInput<int64_t>("indices", {2,1,1}, {1LL,0LL});
+  test.AddInput<std::string>("updates", {2,1,2}, {"f","ghi","ab","cde"});
+  test.AddOutput<std::string>("output", {2,2}, {"ab","cde","f","ghi"});
+  test.Run();
+}
+
+TEST(ScatterNDOpTest, ScatterND_batched_3tensor_int64) {
+  OpTester test1("ScatterND", 11);
+  test1.AddInput<uint32_t>("data", {2,2,2}, {0,0,0,0,0,0,0,0});
+  test1.AddInput<int64_t>("indices", {2,2,2}, {0LL,1LL,1LL,0LL,0LL,0LL,1LL,1LL});
+  test1.AddInput<uint32_t>("updates", {2,2,2}, {2,3,4,5,0,1,6,7});
+  test1.AddOutput<uint32_t>("output", {2,2,2}, {0,1,2,3,4,5,6,7});
+  test1.Run();
+
+  OpTester test2("ScatterND", 11);
+  test2.AddInput<uint32_t>("data", {2,2,2}, {0,0,2,0,4,0,0,7});
+  test2.AddInput<int64_t>("indices", {2,2,3}, {0,0,1,1,0,1,0,1,1,1,1,0});
+  test2.AddInput<uint32_t>("updates", {2,2}, {1,5,3,6});
+  test2.AddOutput<uint32_t>("output", {2,2,2}, {0,1,2,3,4,5,6,7});
+  test2.Run();
+
+  OpTester test3("ScatterND", 11);
+  test3.AddInput<int64_t>("data", {2,2,2}, {1LL,0LL,0LL,0LL,0LL,0LL,0LL,0LL});
+  test3.AddInput<int64_t>("indices", {2,1,1}, {1,0});
+  test3.AddInput<int64_t>("updates", {2,1,2,2}, {4LL,5LL,6LL,7LL,0LL,1LL,2LL,3LL});
+  test3.AddOutput<int64_t>("output", {2,2,2}, {0LL,1LL,2LL,3LL,4LL,5LL,6LL,7LL});
+  test3.Run();
+}
+
+}  // namespace test
+}  // namespace onnxruntime

onnxruntime/test/python/onnx_backend_test_series.py

@@ -142,7 +142,6 @@ def create_backend_test(testname=None):
                                  '^test_resize_upsample_sizes_nearest_cpu.*',
                                  '^test_resize_upsample_sizes_nearest_floor_align_corners_cpu.*',
                                  '^test_resize_upsample_sizes_nearest_round_prefer_ceil_asymmetric_cpu.*',
-                                 '^test_scatternd_cpu.*',
                                  '^test_sequence_*',
                                  '^test_scatter_*',
                                  '^test_onehot_*',