Initial implementation of NonZero op. (#437)

Initial implementation of NonZero op.

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<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 6, Clip)>());
@@ -499,6 +503,10 @@ void RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
   kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, string, TfIdfVectorizer)>());
   kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, int32_t, TfIdfVectorizer)>());
   kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, int64_t, TfIdfVectorizer)>());
+  kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, bool, NonZero)>());
+  kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, float, NonZero)>());
+  kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, int32_t, NonZero)>());
+  kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, int64_t, NonZero)>());
 }
 
 // Forward declarations of ml op kernels

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 <cassert>
+#include <vector>
+
+#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<type>()), \
+      NonZero<type>)
+
+#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<int64_t>* 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 <typename T>
+Status NonZero<T>::Compute(OpKernelContext* context) const {
+  const auto X = context->Input<Tensor>(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<int64_t> 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<T>());
+    if (value != T{}) {
+      non_zero_indices_buffer.push_back(0);
+    }
+  } else {
+    std::vector<int64_t> coordinate(coordinate_size, 0);
+    for (const T& value : X->DataAsSpan<T>()) {
+      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<int64_t> 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<int64_t> y_matrix{
+      Y->MutableData<int64_t>(),
+      coordinate_size, num_non_zero_values};
+  y_matrix = non_zero_indices_matrix.transpose();
+
+  return Status::OK();
+}  // namespace onnxruntime
+}  // namespace onnxruntime

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 <typename T>
+class NonZero : public OpKernel {
+ public:
+  explicit NonZero(const OpKernelInfo& info) : OpKernel{info} {}
+
+  Status Compute(OpKernelContext* context) const override;
+};
+}  // namespace onnxruntime

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"}};
 

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 <typename TTarget, typename TNarrow = int8_t>
+void NonZeroBasicNumericTest() {
+  OpTester test{kOpName, kOpVersion};
+
+  std::vector<int64_t> X_dims{1, 2, 3};
+  std::vector<TNarrow> X{0, 1, 2,
+                         0, 3, 4};
+  test.AddInput<TTarget>("X", X_dims, std::vector<TTarget>{X.begin(), X.end()});
+  test.AddOutput<int64_t>(
+      "Y", {3, 4},
+      {0, 0, 0, 0,
+       0, 0, 1, 1,
+       1, 2, 1, 2});
+  test.Run();
+}
+}  // namespace
+
+TEST(NonZeroOpTest, BasicNumeric) {
+  NonZeroBasicNumericTest<int32_t>();
+  NonZeroBasicNumericTest<int64_t>();
+  NonZeroBasicNumericTest<float>();
+}
+
+TEST(NonZeroOpTest, BasicBool) {
+  OpTester test{kOpName, kOpVersion};
+  test.AddInput<bool>(
+      "X", {2, 3},
+      {true, false, false,
+       false, false, true});
+  test.AddOutput<int64_t>(
+      "Y", {2, 2},
+      {0, 1,
+       0, 2});
+  test.Run();
+}
+
+TEST(NonZeroOpTest, Scalar) {
+  {
+    OpTester test{kOpName, kOpVersion};
+    test.AddInput<int32_t>("X", {}, {0});
+    test.AddOutput<int64_t>("Y", {1, 0}, {});
+    test.Run();
+  }
+  {
+    OpTester test{kOpName, kOpVersion};
+    test.AddInput<int32_t>("X", {}, {1});
+    test.AddOutput<int64_t>("Y", {1, 1}, {0});
+    test.Run();
+  }
+}
+
+TEST(NonZeroOpTest, EmptyInput) {
+  OpTester test{kOpName, kOpVersion};
+  test.AddInput<int32_t>(
+      "X", {1, 0, 2},
+      {});
+  test.AddOutput<int64_t>(
+      "Y", {3, 0},
+      {});
+  test.Run();
+}
+
+}  // namespace test
+}  // namespace onnxruntime

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.*'