From 9e975f64c33a7f86c40fea179fbc38196a1f34cc Mon Sep 17 00:00:00 2001 From: Bowen Bao Date: Wed, 2 Oct 2019 14:34:57 -0700 Subject: [PATCH] Add scatter_nd cpu (#1603) * Add ScatterND impl * remove skipped tests * add rank check for input tensor. --- .../providers/cpu/cpu_execution_provider.cc | 2 + .../core/providers/cpu/tensor/scatter_nd.cc | 157 +++++++++++++++++ .../core/providers/cpu/tensor/scatter_nd.h | 47 ++++++ .../cpu/tensor/scatter_nd_op_test.cc | 159 ++++++++++++++++++ .../test/python/onnx_backend_test_series.py | 1 - 5 files changed, 365 insertions(+), 1 deletion(-) create mode 100644 onnxruntime/core/providers/cpu/tensor/scatter_nd.cc create mode 100644 onnxruntime/core/providers/cpu/tensor/scatter_nd.h create mode 100644 onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc index f3d9f008d0..17f4d0679b 100644 --- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc +++ b/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, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, }; diff --git a/onnxruntime/core/providers/cpu/tensor/scatter_nd.cc b/onnxruntime/core/providers/cpu/tensor/scatter_nd.cc new file mode 100644 index 0000000000..98a31ea17b --- /dev/null +++ b/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()), + ScatterND); + +template +Status ScatterNDBase::PrepareForCompute(OpKernelContext* context, Prepare& p) const { + + auto input_tensor = context->Input(0); + auto indice_tensor = context->Input(1); + auto update_tensor = context->Input(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(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(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(); + + // 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(); + const std::string* str_end = str_begin + input_shape.Size(); + auto* dst = output_tensor->template MutableData(); + std::copy(str_begin, str_end, dst); + } else { + memcpy(dst_base, src_base, input_tensor->SizeInBytes()); + } + } + + std::vector 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(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()) { + p.input_str_base = static_cast(update_tensor->DataRaw()); + p.output_str_base = static_cast(output_tensor->MutableDataRaw()); + } else { + p.input_base = static_cast(update_tensor->DataRaw()); + p.output_base = static_cast(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(OpKernelContext*, Prepare&) const; + +Status ScatterND::Compute(OpKernelContext* context) const { + Prepare p; + ORT_RETURN_IF_ERROR(PrepareForCompute(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(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(p.element_offsets.size()); ++i) { + for (int64_t j = 0; j < static_cast(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(); +} + +} \ No newline at end of file diff --git a/onnxruntime/core/providers/cpu/tensor/scatter_nd.h b/onnxruntime/core/providers/cpu/tensor/scatter_nd.h new file mode 100644 index 0000000000..3c91441bbe --- /dev/null +++ b/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 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 + 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 \ No newline at end of file diff --git a/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc b/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc new file mode 100644 index 0000000000..145fbacdb3 --- /dev/null +++ b/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("data", {2,2}, {"h","h","o","z"}); + test1.AddInput("indices", {2}, {0,1}); + test1.AddInput("updates", {}, {"k"}); + test1.AddOutput("output", {2,2}, {"h","k","o","z"}); + test1.Run(); + + OpTester test2("ScatterND", 11); + test2.AddInput("data", {6}, {"h","k","o","o","l","t"}); + test2.AddInput("indices", {1}, {3}); + test2.AddInput("updates", {}, {"z"}); + test2.AddOutput("output", {6}, {"h","k","o","z","l","t"}); + test2.Run(); + + OpTester test3("ScatterND", 11); + test3.AddInput("data", {3,2}, {"h","k","o","z","l","z"}); + test3.AddInput("indices", {2}, {2,1}); + test3.AddInput("updates", {}, {"t"}); + test3.AddOutput("output", {3,2}, {"h","k","o","z","l","t"}); + test3.Run(); +} + +TEST(ScatterNDOpTest, ScatterND_matrice_int64_int64) { + OpTester test("ScatterND", 11); + test.AddInput ("data", {2,2}, {1LL,1LL,2LL,2LL}); + test.AddInput ("indices", {2,2}, {0LL,0LL,1LL,1LL}); + test.AddInput("updates", {2}, {0LL,3LL}); + test.AddOutput("output", {2,2}, {0LL,1LL,2LL,3LL}); + test.Run(); +} + +TEST(ScatterNDOpTest, ScatterND_matrice_string_int64) { + OpTester test1("ScatterND", 11); + test1.AddInput("data", {2,2,2}, {"egg","dance","bob","air","smart","terry","laugh","kite"}); + test1.AddInput("indices", {2,1,2}, {0,1,1,0}); + test1.AddInput("updates", {2,1,2}, {"air","bob","terry","smart"}); + test1.AddOutput("output", {2,2,2}, {"egg","dance","air","bob","terry","smart","laugh","kite"}); + test1.Run(); + + OpTester test2("ScatterND", 11); + test2.AddInput("data", {3,3}, {"egg","","air","","terry","smart","laugh","","hop"}); + test2.AddInput("indices", {3,2}, {2,1,1,0,0,1}); + test2.AddInput("updates", {3}, {"kite","bob","dance"}); + test2.AddOutput("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("data", {2,2}, {0.0f,0.1f,0.1f,0.1f}); + test.AddInput("indices", {2,1}, {1LL,0LL}); + test.AddInput("updates", {2,2}, {0.2f,0.3f,0.0f,0.1f}); + test.AddOutput("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("data", {2,2}, {0.0f,0.1f,0.1f,0.1f}); + test.AddInput("indices", {2,1}, {1LL,0LL}); + test.AddInput("updates", {2,2}, {0.2f,0.3f,0.0f,0.1f}); + test.AddOutput("output", {2,2}, {0.0f,0.1f,0.2f,0.3f}); + test.Run(); +} + +TEST(ScatterNDOpTest, ScatterND_3tensor_int64) { + OpTester test1("ScatterND", 11); + test1.AddInput("data", {2,2,2}, {0LL,1LL,1LL,1LL,1LL,1LL,6LL,7LL}); + test1.AddInput("indices", {2,2}, {0LL,1LL,1LL,0LL}); + test1.AddInput("updates", {2,2}, {2LL,3LL,4LL,5LL}); + test1.AddOutput("output", {2,2,2}, {0LL,1LL,2LL,3LL,4LL,5LL,6LL,7LL}); + test1.Run(); + + OpTester test2("ScatterND", 11); + test2.AddInput("data", {2,2,2}, {0,0,2,3,4,0,6,7}); + test2.AddInput("indices", {2,3}, {0,0,1,1,0,1}); + test2.AddInput("updates", {2}, {1,5}); + test2.AddOutput("output", {2,2,2}, {0,1,2,3,4,5,6,7}); + test2.Run(); + + OpTester test3("ScatterND", 11); + test3.AddInput("data", {2,2,2}, {0,1,2,3,0,1,2,3}); + test3.AddInput("indices", {1,1}, {1LL}); + test3.AddInput("updates", {1,2,2}, {4,5,6,7}); + test3.AddOutput("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("data", {2,2}, {2LL,3LL,2LL,3LL}); + test.AddInput("indices", {2,1,2}, {0LL,0LL,0LL,1LL}); + test.AddInput("updates", {2,1}, {0LL,1LL}); + test.AddOutput("output", {2,2}, {0LL,1LL,2LL,3LL}); + test.Run(); +} + +TEST(ScatterNDOpTest, ScatterND_batched_index_bool_int64) { + OpTester test("ScatterND", 11); + test.AddInput("data", {2,2}, {false,true,false,true}); + test.AddInput("indices", {2,1,2}, {0LL,0LL,0LL,1LL}); + test.AddInput("updates", {2,1}, {true,false}); + test.AddOutput("output", {2,2}, {true,false,false,true}); + test.Run(); +} + +TEST(ScatterNDOpTest, ScatterND_sliced_index_int64) { + OpTester test("ScatterND", 11); + test.AddInput("data", {2,2}, {0LL,0LL,0LL,0LL}); + test.AddInput("indices", {2,1,1}, {1LL,0LL}); + test.AddInput("updates", {2,1,2}, {2LL,3LL,0LL,1LL}); + test.AddOutput("output", {2,2}, {0LL,1LL,2LL,3LL}); + test.Run(); +} + +TEST(ScatterNDOpTest, ScatterND_sliced_index_string_int64) { + OpTester test("ScatterND", 11); + test.AddInput("data", {2,2}, {"","","",""}); + test.AddInput("indices", {2,1,1}, {1LL,0LL}); + test.AddInput("updates", {2,1,2}, {"f","ghi","ab","cde"}); + test.AddOutput("output", {2,2}, {"ab","cde","f","ghi"}); + test.Run(); +} + +TEST(ScatterNDOpTest, ScatterND_batched_3tensor_int64) { + OpTester test1("ScatterND", 11); + test1.AddInput("data", {2,2,2}, {0,0,0,0,0,0,0,0}); + test1.AddInput("indices", {2,2,2}, {0LL,1LL,1LL,0LL,0LL,0LL,1LL,1LL}); + test1.AddInput("updates", {2,2,2}, {2,3,4,5,0,1,6,7}); + test1.AddOutput("output", {2,2,2}, {0,1,2,3,4,5,6,7}); + test1.Run(); + + OpTester test2("ScatterND", 11); + test2.AddInput("data", {2,2,2}, {0,0,2,0,4,0,0,7}); + test2.AddInput("indices", {2,2,3}, {0,0,1,1,0,1,0,1,1,1,1,0}); + test2.AddInput("updates", {2,2}, {1,5,3,6}); + test2.AddOutput("output", {2,2,2}, {0,1,2,3,4,5,6,7}); + test2.Run(); + + OpTester test3("ScatterND", 11); + test3.AddInput("data", {2,2,2}, {1LL,0LL,0LL,0LL,0LL,0LL,0LL,0LL}); + test3.AddInput("indices", {2,1,1}, {1,0}); + test3.AddInput("updates", {2,1,2,2}, {4LL,5LL,6LL,7LL,0LL,1LL,2LL,3LL}); + test3.AddOutput("output", {2,2,2}, {0LL,1LL,2LL,3LL,4LL,5LL,6LL,7LL}); + test3.Run(); +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/python/onnx_backend_test_series.py b/onnxruntime/test/python/onnx_backend_test_series.py index 053a9c4d14..4c1db26639 100644 --- a/onnxruntime/test/python/onnx_backend_test_series.py +++ b/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_*',