提交 8f70af73 编写于 作者: H huanghui

add unique_with_pad cpu kernel

上级 8c377fd1
/**
* 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
/**
* 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")
......
/**
* 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
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册