diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_with_pad_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_with_pad_cpu_kernel.cc new file mode 100644 index 0000000000..31e1e3195a --- /dev/null +++ b/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 &inputs, + const std::vector & /*workspace*/, + const std::vector &outputs) { + if (dtype_ == kNumberTypeInt32) { + LaunchKernel(inputs, outputs); + } else if (dtype_ == kNumberTypeInt64) { + LaunchKernel(inputs, outputs); + } else { + MS_LOG(EXCEPTION) << "Only unsupported int32 or int64 dtype"; + } + return true; +} + +template +void UniqueWithPadCPUKernel::LaunchKernel(const std::vector &inputs, + const std::vector &outputs) { + T *a = reinterpret_cast(inputs[0]->addr); + T pad_num = *reinterpret_cast(inputs[1]->addr); + T *out = reinterpret_cast(outputs[0]->addr); + T *idx_vec = reinterpret_cast(outputs[1]->addr); + + for (int64_t i = 0; i < n_; ++i) { + out[i] = pad_num; + } + std::unordered_map 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 diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_with_pad_cpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_with_pad_cpu_kernel.h new file mode 100644 index 0000000000..d7b7563342 --- /dev/null +++ b/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 +#include +#include +#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 &inputs, const std::vector &workspace, + const std::vector &outputs) override; + + template + void LaunchKernel(const std::vector &inputs, const std::vector &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_ diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index 42ae20c4c1..e43a8272ca 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/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") diff --git a/tests/ut/cpp/kernel/cpu/unique_with_pad_cpu_kernel_test.cc b/tests/ut/cpp/kernel/cpu/unique_with_pad_cpu_kernel_test.cc new file mode 100644 index 0000000000..d7be8e94e4 --- /dev/null +++ b/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 +#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()) {} + + 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
(); + 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 x_; + int pad_dim_; + std::vector out_; + std::vector idx_; + std::vector inputs_; + std::vector workspace_; + std::vector outputs_; + std::shared_ptr 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 expect_out{1, 5, 4, 3, 2, 8, 8, 8, 8, 8}; + std::vector 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