!4219 Add Unique CPU kernel

Merge pull request !4219 from huanghui/add-op-unique

mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.cc

@@ -0,0 +1,78 @@
+/**
+ * Copyright 2020 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.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "backend/kernel_compiler/cpu/unique_cpu_kernel.h"
+#include "runtime/device/cpu/cpu_device_address.h"
+
+namespace mindspore {
+namespace kernel {
+void UniqueCPUKernel::InitKernel(const CNodePtr &kernel_node) {
+  CheckParam(kernel_node);
+  auto input_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
+  n_ = input_shape[0];
+  dtype_ = AnfAlgo::GetPrevNodeOutputInferDataType(kernel_node, 0);
+}
+
+bool UniqueCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
+                             const std::vector<kernel::AddressPtr> & /*workspace*/,
+                             const std::vector<kernel::AddressPtr> &outputs) {
+  if (dtype_ == kNumberTypeInt32) {
+    LaunchKernel<int>(inputs, outputs);
+  } else if (dtype_ == kNumberTypeFloat32) {
+    LaunchKernel<float>(inputs, outputs);
+  } else if (dtype_ == kNumberTypeInt64) {
+    LaunchKernel<int64_t>(inputs, outputs);
+  }
+  return true;
+}
+
+template <typename T>
+void UniqueCPUKernel::LaunchKernel(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &outputs) {
+  auto x_addr = reinterpret_cast<T *>(inputs[0]->addr);
+  auto y_addr = reinterpret_cast<T *>(outputs[0]->addr);
+  auto idx_addr = reinterpret_cast<int *>(outputs[1]->addr);
+
+  std::unordered_map<T, int> uniq;
+  int n = SizeToInt(n_);
+  uniq.reserve(n * 2);
+  for (int i = 0, j = 0; i < n; ++i) {
+    auto it = uniq.emplace(x_addr[i], j);
+    idx_addr[i] = it.first->second;
+    if (it.second) {
+      ++j;
+    }
+  }
+  for (const auto &it : uniq) {
+    y_addr[it.second] = it.first;
+  }
+}
+
+void UniqueCPUKernel::CheckParam(const CNodePtr &kernel_node) {
+  auto input_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
+  if (input_shape.size() != 1) {
+    MS_LOG(EXCEPTION) << "Input dims is " << input_shape.size() << ", but UniqueCPUKernel only support 1d.";
+  }
+  size_t input_num = AnfAlgo::GetInputTensorNum(kernel_node);
+  if (input_num != 1) {
+    MS_LOG(EXCEPTION) << "Input number is " << input_num << ", but UniqueCPUKernel needs 1 input.";
+  }
+  size_t output_num = AnfAlgo::GetOutputTensorNum(kernel_node);
+  if (output_num != 2) {
+    MS_LOG(EXCEPTION) << "Output number is " << output_num << ", but UniqueCPUKernel needs 2 output.";
+  }
+}
+}  // namespace kernel
+}  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.h

@@ -0,0 +1,61 @@
+/**
+ * Copyright 2020 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.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_UNIQUE_CPU_KERNEL_H_
+#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_UNIQUE_CPU_KERNEL_H_
+#include <vector>
+#include <memory>
+#include <unordered_map>
+#include "backend/kernel_compiler/cpu/cpu_kernel.h"
+#include "backend/kernel_compiler/cpu/cpu_kernel_factory.h"
+
+namespace mindspore {
+namespace kernel {
+class UniqueCPUKernel : public CPUKernel {
+ public:
+  UniqueCPUKernel() = default;
+  ~UniqueCPUKernel() override = default;
+
+  void InitKernel(const CNodePtr &kernel_node) override;
+
+  bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
+              const std::vector<AddressPtr> &outputs) override;
+
+  template <typename T>
+  void LaunchKernel(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &outputs);
+
+ private:
+  void CheckParam(const CNodePtr &kernel_node);
+  size_t n_;
+  TypeId dtype_;
+};
+
+MS_REG_CPU_KERNEL(
+  Unique, KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32),
+  UniqueCPUKernel);
+
+MS_REG_CPU_KERNEL(
+  Unique, KernelAttr().AddInputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeInt32),
+  UniqueCPUKernel);
+
+MS_REG_CPU_KERNEL(
+  Unique,
+  KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeInt32),
+  UniqueCPUKernel);
+}  // namespace kernel
+}  // namespace mindspore
+
+#endif  // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_UNIQUE_CPU_KERNEL_H_

tests/st/ops/cpu/test_unique_op.py

@@ -0,0 +1,47 @@
+# Copyright 2020 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.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+import numpy as np
+
+import mindspore.context as context
+import mindspore.nn as nn
+from mindspore import Tensor
+from mindspore.common import dtype as mstype
+from mindspore.ops import operations as P
+
+context.set_context(mode=context.GRAPH_MODE, device_target='CPU')
+
+
+class Net(nn.Cell):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.uniq = P.Unique()
+
+    def construct(self, x):
+        return self.uniq(x)
+
+
+def test_net():
+    x = Tensor(np.array([1, 2, 5, 2]), mstype.float32)
+    uniq = Net()
+    output = uniq(x)
+    print("x:\n", x)
+    print("y:\n", output[0])
+    print("idx:\n", output[1])
+    expect_y_result = [1., 2., 5.]
+    expect_idx_result = [0, 1, 2, 1]
+
+    assert (output[0].asnumpy() == expect_y_result).all()
+    assert (output[1].asnumpy() == expect_idx_result).all()

tests/ut/cpp/CMakeLists.txt

@@ -128,6 +128,7 @@ file(GLOB_RECURSE MINDSPORE_SRC_LIST RELATIVE ${CMAKE_CURRENT_SOURCE_DIR}
         "../../../mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_apply_ftrl_cpu_kernel.cc"
         "../../../mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_apply_lazy_adam_cpu_kernel.cc"
         "../../../mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_apply_proximal_adagrad_cpu_kernel.cc"
+        "../../../mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.cc"
         )
 
 if (CMAKE_SYSTEM_NAME MATCHES "Windows")

tests/ut/cpp/kernel/cpu/unique_cpu_kernel_test.cc

@@ -0,0 +1,76 @@
+/**
+ * Copyright 2020 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.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <vector>
+#include "common/common_test.h"
+#define private public
+#define protected public
+#include "backend/kernel_compiler/cpu/unique_cpu_kernel.h"
+#undef private
+#undef protected
+
+namespace mindspore {
+namespace kernel {
+class UniqueCpuKernelTest : public UT::Common {
+ public:
+  UniqueCpuKernelTest() : unique_(std::make_shared<UniqueCPUKernel>()) {}
+
+  void SetUp() override {
+    unique_->n_ = 9;
+    unique_->dtype_ = kNumberTypeFloat32;
+    inputs_.clear();
+    workspace_.clear();
+    outputs_.clear();
+  }
+
+  AddressPtr CreateKernelAddress(void *addr) {
+    auto kernel_addr = std::make_shared<Address>();
+    kernel_addr->addr = addr;
+    return kernel_addr;
+  }
+
+  void CreateInputAddress() { inputs_.push_back(CreateKernelAddress(x_.data())); }
+
+  void CreateOutputAddress() {
+    outputs_.push_back(CreateKernelAddress(y_.data()));
+    outputs_.push_back(CreateKernelAddress(idx_.data()));
+  }
+
+  std::vector<float> x_;
+  std::vector<float> y_;
+  std::vector<int> idx_;
+  std::vector<AddressPtr> inputs_;
+  std::vector<AddressPtr> workspace_;
+  std::vector<AddressPtr> outputs_;
+  std::shared_ptr<UniqueCPUKernel> unique_;
+};
+
+TEST_F(UniqueCpuKernelTest, compute_test) {
+  x_ = {1, 1, 2, 4, 4, 4, 7, 8, 8};
+  y_ = {1, 1, 1, 1, 1};
+  idx_ = {1, 1, 1, 1, 1, 1, 1, 1, 1};
+  CreateInputAddress();
+  CreateOutputAddress();
+  unique_->Launch(inputs_, workspace_, outputs_);
+
+  // check compute result
+  std::vector<float> expect_y{1, 2, 4, 7, 8};
+  std::vector<int> expect_idx{0, 0, 1, 2, 2, 2, 3, 4, 4};
+  EXPECT_TRUE(y_ == expect_y);
+  EXPECT_TRUE(idx_ == expect_idx);
+}
+}  // namespace kernel
+}  // namespace mindspore