diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc index de896fa7b1..955d545300 100644 --- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc +++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc @@ -253,6 +253,10 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, Sca class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, string, TfIdfVectorizer); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, int32_t, TfIdfVectorizer); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, int64_t, TfIdfVectorizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, bool, NonZero); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, float, NonZero); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, int32_t, NonZero); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, int64_t, NonZero); void RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { kernel_registry.Register(BuildKernelCreateInfo()); @@ -499,6 +503,10 @@ void RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { kernel_registry.Register(BuildKernelCreateInfo()); kernel_registry.Register(BuildKernelCreateInfo()); kernel_registry.Register(BuildKernelCreateInfo()); + kernel_registry.Register(BuildKernelCreateInfo()); + kernel_registry.Register(BuildKernelCreateInfo()); + kernel_registry.Register(BuildKernelCreateInfo()); + kernel_registry.Register(BuildKernelCreateInfo()); } // Forward declarations of ml op kernels diff --git a/onnxruntime/core/providers/cpu/tensor/nonzero_op.cc b/onnxruntime/core/providers/cpu/tensor/nonzero_op.cc new file mode 100644 index 0000000000..90951ba3b3 --- /dev/null +++ b/onnxruntime/core/providers/cpu/tensor/nonzero_op.cc @@ -0,0 +1,107 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cpu/tensor/nonzero_op.h" + +#include +#include + +#include "core/util/math_cpuonly.h" + +namespace onnxruntime { +// kernel builder functions +#define NONZERO_TYPED_KERNEL_WITH_TYPE_NAME(type, type_name) \ + ONNX_CPU_OPERATOR_TYPED_KERNEL( \ + NonZero, \ + 9, \ + type_name, \ + KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + NonZero) + +#define NONZERO_TYPED_KERNEL(type) \ + NONZERO_TYPED_KERNEL_WITH_TYPE_NAME(type, type) + +// start with a subset of types, enable more as needed... +NONZERO_TYPED_KERNEL(bool) +//NONZERO_TYPED_KERNEL(uint8_t) +//NONZERO_TYPED_KERNEL(uint16_t) +//NONZERO_TYPED_KERNEL(uint32_t) +//NONZERO_TYPED_KERNEL(uint64_t) +//NONZERO_TYPED_KERNEL(int8_t) +//NONZERO_TYPED_KERNEL(int16_t) +NONZERO_TYPED_KERNEL(int32_t) +NONZERO_TYPED_KERNEL(int64_t) +//NONZERO_TYPED_KERNEL(MLFloat16) +//NONZERO_TYPED_KERNEL(BFloat16) +NONZERO_TYPED_KERNEL(float) +//NONZERO_TYPED_KERNEL(double) +//NONZERO_TYPED_KERNEL_WITH_TYPE_NAME(std::string, string) + +#undef NONZERO_TYPED_KERNEL_WITH_TYPE_NAME +#undef NONZERO_TYPED_KERNEL + +namespace { +void IncrementCoordinate(const TensorShape& shape, std::vector* coordinate) { + assert(coordinate->size() == shape.NumDimensions()); + + size_t i = 0; + const size_t i_end = coordinate->size(); + for (; i < i_end; ++i) { + const size_t i_from_back = i_end - i - 1; + if ((*coordinate)[i_from_back] != shape[i_from_back] - 1) break; + (*coordinate)[i_from_back] = 0; + } + + if (i < i_end) { + ++(*coordinate)[i_end - i - 1]; + } +} +} // namespace + +template +Status NonZero::Compute(OpKernelContext* context) const { + const auto X = context->Input(0); + ORT_ENFORCE(X, "X input is required!"); + + const auto X_shape = X->Shape(); + assert(X_shape.Size() >= 0); + + const int64_t coordinate_size = X_shape.IsScalar() ? 1 : X_shape.NumDimensions(); + std::vector non_zero_indices_buffer{}; + // reserve enough space for indices for every element of X + non_zero_indices_buffer.reserve(X_shape.Size() * coordinate_size); + + if (X_shape.IsScalar()) { + const T& value = *(X->Data()); + if (value != T{}) { + non_zero_indices_buffer.push_back(0); + } + } else { + std::vector coordinate(coordinate_size, 0); + for (const T& value : X->DataAsSpan()) { + if (value != T{}) { + non_zero_indices_buffer.insert(non_zero_indices_buffer.end(), + coordinate.begin(), coordinate.end()); + } + IncrementCoordinate(X_shape, &coordinate); + } + } + + const int64_t num_non_zero_values = non_zero_indices_buffer.size() / coordinate_size; + + // transpose result for output + ConstEigenMatrixMapRowMajor non_zero_indices_matrix{ + non_zero_indices_buffer.data(), + num_non_zero_values, coordinate_size}; + + Tensor* const Y = context->Output(0, TensorShape{coordinate_size, num_non_zero_values}); + ORT_ENFORCE(Y, "failed to get first output!"); + + EigenMatrixMapRowMajor y_matrix{ + Y->MutableData(), + coordinate_size, num_non_zero_values}; + y_matrix = non_zero_indices_matrix.transpose(); + + return Status::OK(); +} // namespace onnxruntime +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/tensor/nonzero_op.h b/onnxruntime/core/providers/cpu/tensor/nonzero_op.h new file mode 100644 index 0000000000..5d43724915 --- /dev/null +++ b/onnxruntime/core/providers/cpu/tensor/nonzero_op.h @@ -0,0 +1,17 @@ +// 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 { +template +class NonZero : public OpKernel { + public: + explicit NonZero(const OpKernelInfo& info) : OpKernel{info} {} + + Status Compute(OpKernelContext* context) const override; +}; +} // namespace onnxruntime diff --git a/onnxruntime/test/onnx/main.cc b/onnxruntime/test/onnx/main.cc index be3f1fb79b..faa2527e20 100644 --- a/onnxruntime/test/onnx/main.cc +++ b/onnxruntime/test/onnx/main.cc @@ -325,7 +325,6 @@ int real_main(int argc, char* argv[]) { {"cast_FLOAT_to_STRING", "Cast opset 9 not supported yet"}, {"cast_FLOAT_to_FLOAT16", "Cast opset 9 not supported yet"}, {"cast_FLOAT16_to_DOUBLE", "Cast opset 9 not supported yet"}, - {"nonzero_example", "NonZero opset 9 not supported yet"}, {"tf_inception_resnet_v2", "Cast opset 9 not supported yet"}, {"tf_inception_v4", "Cast opset 9 not supported yet"}}; diff --git a/onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc b/onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc new file mode 100644 index 0000000000..75afeca8a7 --- /dev/null +++ b/onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc @@ -0,0 +1,78 @@ +// 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 { + +namespace { +constexpr char kOpName[] = "NonZero"; +constexpr int kOpVersion = 9; + +template +void NonZeroBasicNumericTest() { + OpTester test{kOpName, kOpVersion}; + + std::vector X_dims{1, 2, 3}; + std::vector X{0, 1, 2, + 0, 3, 4}; + test.AddInput("X", X_dims, std::vector{X.begin(), X.end()}); + test.AddOutput( + "Y", {3, 4}, + {0, 0, 0, 0, + 0, 0, 1, 1, + 1, 2, 1, 2}); + test.Run(); +} +} // namespace + +TEST(NonZeroOpTest, BasicNumeric) { + NonZeroBasicNumericTest(); + NonZeroBasicNumericTest(); + NonZeroBasicNumericTest(); +} + +TEST(NonZeroOpTest, BasicBool) { + OpTester test{kOpName, kOpVersion}; + test.AddInput( + "X", {2, 3}, + {true, false, false, + false, false, true}); + test.AddOutput( + "Y", {2, 2}, + {0, 1, + 0, 2}); + test.Run(); +} + +TEST(NonZeroOpTest, Scalar) { + { + OpTester test{kOpName, kOpVersion}; + test.AddInput("X", {}, {0}); + test.AddOutput("Y", {1, 0}, {}); + test.Run(); + } + { + OpTester test{kOpName, kOpVersion}; + test.AddInput("X", {}, {1}); + test.AddOutput("Y", {1, 1}, {0}); + test.Run(); + } +} + +TEST(NonZeroOpTest, EmptyInput) { + OpTester test{kOpName, kOpVersion}; + test.AddInput( + "X", {1, 0, 2}, + {}); + test.AddOutput( + "Y", {3, 0}, + {}); + test.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 0c0e4523a4..829975e641 100644 --- a/onnxruntime/test/python/onnx_backend_test_series.py +++ b/onnxruntime/test/python/onnx_backend_test_series.py @@ -27,7 +27,6 @@ backend_test.exclude(r'(' '|^test_eyelike_populate_off_main_diagonal_cpu.*' '|^test_eyelike_with_dtype_cpu.*' '|^test_eyelike_without_dtype_cpu.*' -'|^test_nonzero_example_cpu.*' '|^test_scatter_with_axis_cpu.*' '|^test_scatter_without_axis_cpu.*' '|^test_shrink_hard_cpu.*'