forked from mindspore-Ecosystem/mindspore
!4418 Add UniqueWithPad cpu kernel
Merge pull request !4418 from huanghui/unique-with-pad-cpu-kernel
This commit is contained in:
commit
e9629f95e1
|
@ -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
|
|
@ -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_
|
|
@ -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")
|
||||
|
|
|
@ -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
|
Loading…
Reference in New Issue