1 /**
2 * Copyright 2020-2021 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "backend/kernel_compiler/cpu/unique_with_pad_cpu_kernel.h"
18 #include "runtime/device/cpu/cpu_device_address.h"
19
20 namespace mindspore {
21 namespace kernel {
22 namespace {
23 constexpr size_t kUniqueWithPadInputsNum = 2;
24 constexpr size_t kUniqueWithPadOutputsNum = 2;
25 } // namespace
26
Launch(const std::vector<kernel::AddressPtr> & inputs,const std::vector<kernel::AddressPtr> & workspace,const std::vector<kernel::AddressPtr> & outputs)27 bool UniqueWithPadCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
28 const std::vector<kernel::AddressPtr> &workspace,
29 const std::vector<kernel::AddressPtr> &outputs) {
30 CHECK_KERNEL_INPUTS_NUM(inputs.size(), kUniqueWithPadInputsNum, kernel_name_);
31 CHECK_KERNEL_OUTPUTS_NUM(outputs.size(), kUniqueWithPadOutputsNum, kernel_name_);
32 if (dtype_ == kNumberTypeInt32) {
33 UniqueCPUKernel::LaunchKernel<int, int>(inputs, workspace, outputs);
34 PadOutput<int>(inputs, outputs);
35 } else if (dtype_ == kNumberTypeInt64) {
36 UniqueCPUKernel::LaunchKernel<int64_t, int64_t>(inputs, workspace, outputs);
37 PadOutput<int64_t>(inputs, outputs);
38 } else if (dtype_ == kNumberTypeFloat32 || dtype_ == kNumberTypeFloat16) {
39 UniqueCPUKernel::LaunchKernel<float, int>(inputs, workspace, outputs);
40 PadOutput<float>(inputs, outputs);
41 } else {
42 MS_LOG(EXCEPTION) << "Unsupported input data type: " << dtype_;
43 }
44 return true;
45 }
46
47 template <typename T>
PadOutput(const std::vector<AddressPtr> & inputs,const std::vector<AddressPtr> & outputs) const48 void UniqueWithPadCPUKernel::PadOutput(const std::vector<AddressPtr> &inputs,
49 const std::vector<AddressPtr> &outputs) const {
50 auto pad_num = *reinterpret_cast<T *>(inputs[1]->addr);
51 auto *out = reinterpret_cast<T *>(outputs[0]->addr);
52 for (size_t i = output_size_; i < input_size_; ++i) {
53 out[i] = pad_num;
54 }
55 }
56 } // namespace kernel
57 } // namespace mindspore
58