!4418 Add UniqueWithPad cpu kernel

Merge pull request !4418 from huanghui/unique-with-pad-cpu-kernel

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

@@ -0,0 +1,82 @@
+/**
+ * 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_with_pad_cpu_kernel.h"
+#include "runtime/device/cpu/cpu_device_address.h"
+
+namespace mindspore {
+namespace kernel {
+void UniqueWithPadCPUKernel::InitKernel(const CNodePtr &kernel_node) {
+  CheckParam(kernel_node);
+  auto input_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
+  n_ = SizeToLong(input_shape[0]);
+  dtype_ = AnfAlgo::GetPrevNodeOutputInferDataType(kernel_node, 0);
+}
+
+bool UniqueWithPadCPUKernel::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_ == kNumberTypeInt64) {
+    LaunchKernel<int64_t>(inputs, outputs);
+  } else {
+    MS_LOG(EXCEPTION) << "Only unsupported int32 or int64 dtype";
+  }
+  return true;
+}
+
+template <typename T>
+void UniqueWithPadCPUKernel::LaunchKernel(const std::vector<AddressPtr> &inputs,
+                                          const std::vector<AddressPtr> &outputs) {
+  T *a = reinterpret_cast<T *>(inputs[0]->addr);
+  T pad_num = *reinterpret_cast<T *>(inputs[1]->addr);
+  T *out = reinterpret_cast<T *>(outputs[0]->addr);
+  T *idx_vec = reinterpret_cast<T *>(outputs[1]->addr);
+
+  for (int64_t i = 0; i < n_; ++i) {
+    out[i] = pad_num;
+  }
+  std::unordered_map<T, int> uniq;
+  uniq.reserve(n_);
+  for (int64_t i = 0, j = 0; i < n_; ++i) {
+    auto it = uniq.emplace(a[i], j);
+    idx_vec[i] = it.first->second;
+    if (it.second) {
+      ++j;
+    }
+  }
+  for (const auto &it : uniq) {
+    out[it.second] = it.first;
+  }
+}
+
+void UniqueWithPadCPUKernel::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 != 2) {
+    MS_LOG(EXCEPTION) << "Input number is " << input_num << ", but UniqueCPUKernel needs 2 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_with_pad_cpu_kernel.h

@@ -0,0 +1,65 @@
+/**
+ * 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_WITH_PAD_CPU_KERNEL_H_
+#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_UNIQUE_WITH_PAD_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 UniqueWithPadCPUKernel : public CPUKernel {
+ public:
+  UniqueWithPadCPUKernel() = default;
+  ~UniqueWithPadCPUKernel() 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);
+  int64_t n_;
+  TypeId dtype_;
+};
+
+MS_REG_CPU_KERNEL(UniqueWithPad,
+                  KernelAttr()
+                    .AddInputAttr(kNumberTypeInt32)
+                    .AddInputAttr(kNumberTypeInt32)
+                    .AddOutputAttr(kNumberTypeInt32)
+                    .AddOutputAttr(kNumberTypeInt32),
+                  UniqueWithPadCPUKernel);
+
+MS_REG_CPU_KERNEL(UniqueWithPad,
+                  KernelAttr()
+                    .AddInputAttr(kNumberTypeInt64)
+                    .AddInputAttr(kNumberTypeInt64)
+                    .AddOutputAttr(kNumberTypeInt64)
+                    .AddOutputAttr(kNumberTypeInt64),
+                  UniqueWithPadCPUKernel);
+
+}  // namespace kernel
+}  // namespace mindspore
+
+#endif  // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_UNIQUE_WITH_PAD_CPU_KERNEL_H_

tests/ut/cpp/CMakeLists.txt

@@ -129,6 +129,7 @@ file(GLOB_RECURSE MINDSPORE_SRC_LIST RELATIVE ${CMAKE_CURRENT_SOURCE_DIR}
         "../../../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"
+        "../../../mindspore/ccsrc/backend/kernel_compiler/cpu/unique_with_pad_cpu_kernel.cc"
         )
 
 if (CMAKE_SYSTEM_NAME MATCHES "Windows")

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

@@ -0,0 +1,82 @@
+/**
+ * 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_with_pad_cpu_kernel.h"
+#undef private
+#undef protected
+
+namespace mindspore {
+namespace kernel {
+class UniqueWithPadCpuKernelTest : public UT::Common {
+ public:
+  UniqueWithPadCpuKernelTest() : unique_with_pad_(std::make_shared<UniqueWithPadCPUKernel>()) {}
+
+  void SetUp() override {
+    unique_with_pad_->n_ = 10;
+    unique_with_pad_->dtype_ = kNumberTypeInt32;
+    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()));
+    inputs_.push_back(CreateKernelAddress(&pad_dim_));
+    ;
+  }
+
+  void CreateOutputAddress() {
+    outputs_.push_back(CreateKernelAddress(out_.data()));
+    outputs_.push_back(CreateKernelAddress(idx_.data()));
+  }
+
+  std::vector<int> x_;
+  int pad_dim_;
+  std::vector<int> out_;
+  std::vector<int> idx_;
+  std::vector<AddressPtr> inputs_;
+  std::vector<AddressPtr> workspace_;
+  std::vector<AddressPtr> outputs_;
+  std::shared_ptr<UniqueWithPadCPUKernel> unique_with_pad_;
+};
+
+TEST_F(UniqueWithPadCpuKernelTest, compute_test) {
+  x_ = {1, 1, 5, 5, 4, 4, 3, 3, 2, 2};
+  pad_dim_ = 8;
+  out_ = {1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
+  idx_ = {1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
+  CreateInputAddress();
+  CreateOutputAddress();
+  unique_with_pad_->Launch(inputs_, workspace_, outputs_);
+
+  // check compute result
+  std::vector<int> expect_out{1, 5, 4, 3, 2, 8, 8, 8, 8, 8};
+  std::vector<int> expect_idx{0, 0, 1, 1, 2, 2, 3, 3, 4, 4};
+  EXPECT_TRUE(out_ == expect_out);
+  EXPECT_TRUE(idx_ == expect_idx);
+}
+}  // namespace kernel
+}  // namespace mindspore