1 /**
2 * Copyright 2022 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 "plugin/device/cpu/kernel/set_size_cpu_kernel.h"
18 #include <algorithm>
19 #include <unordered_set>
20 #include "plugin/device/cpu/hal/device/cpu_device_address.h"
21 #include "mindspore/core/ops/set_size.h"
22
23 namespace mindspore {
24 namespace kernel {
Init(const std::vector<KernelTensor * > & inputs,const std::vector<KernelTensor * > & outputs)25 bool SetSizeCpuKernelMod::Init(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &outputs) {
26 constexpr size_t input_num = 3;
27 constexpr size_t output_num = 1;
28 CHECK_KERNEL_INPUTS_NUM(inputs.size(), input_num, kernel_name_);
29 CHECK_KERNEL_OUTPUTS_NUM(outputs.size(), output_num, kernel_name_);
30 validate_indices_ = GetValue<bool>(primitive_->GetAttr(ops::kValidateIndices));
31
32 return MatchKernelFunc(kernel_name_, inputs, outputs);
33 }
34
Resize(const std::vector<KernelTensor * > & inputs,const std::vector<KernelTensor * > & outputs)35 int SetSizeCpuKernelMod::Resize(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &outputs) {
36 if (auto ret = KernelMod::Resize(inputs, outputs); ret != KRET_OK) {
37 return ret;
38 }
39 auto values_shape = inputs[kIndex1]->GetShapeVector();
40 shape_ = inputs[kIndex2]->GetShapeVector();
41 values_size_ = static_cast<size_t>(values_shape[0]);
42 output_shape_ = outputs[kIndex0]->GetShapeVector();
43 dims_ = shape_[0];
44 return KRET_OK;
45 }
46
IndicesValid(int64_t n,const std::vector<kernel::KernelTensor * > & inputs) const47 bool SetSizeCpuKernelMod::IndicesValid(int64_t n, const std::vector<kernel::KernelTensor *> &inputs) const {
48 bool valid = true;
49 bool different = false;
50 bool increasing = true;
51 const auto *indices_t = static_cast<int64_t *>(inputs[kIndex0]->device_ptr());
52 const auto *shape_t = static_cast<int64_t *>(inputs[kIndex2]->device_ptr());
53 for (int64_t di = 0; di < dims_; ++di) {
54 if (indices_t[(n * dims_) + di] < 0 || indices_t[(n * dims_) + di] >= shape_t[di]) {
55 valid = false;
56 break;
57 }
58 if (n != 0) {
59 int64_t diff = indices_t[(n * dims_) + di] - indices_t[((n - 1) * dims_) + di];
60 if (diff > 0) {
61 different = true;
62 }
63 if (!different && diff < 0) {
64 increasing = false;
65 break;
66 }
67 }
68 }
69 if (n == 0) {
70 different = true;
71 }
72 if (!valid) {
73 MS_EXCEPTION(ValueError) << "For '" << kernel_name_ << "', indices is out of bounds, index=" << n << ".";
74 return false;
75 }
76 if (!increasing) {
77 MS_EXCEPTION(ValueError) << "For '" << kernel_name_ << "', indices is out of order, index=" << n << ".";
78 return false;
79 }
80 if (!different) {
81 MS_EXCEPTION(ValueError) << "For '" << kernel_name_ << "', indices is repeated, index=" << n << ".";
82 return false;
83 }
84 return true;
85 }
86
87 template <typename T>
LaunchKernel(const std::vector<kernel::KernelTensor * > & inputs,const std::vector<kernel::KernelTensor * > &,const std::vector<kernel::KernelTensor * > & outputs)88 bool SetSizeCpuKernelMod::LaunchKernel(const std::vector<kernel::KernelTensor *> &inputs,
89 const std::vector<kernel::KernelTensor *> &,
90 const std::vector<kernel::KernelTensor *> &outputs) {
91 if (outputs[0]->size() == 0) {
92 MS_LOG(WARNING) << "For '" << kernel_name_ << "', output memory size should be greater than 0, but got 0.";
93 return true;
94 }
95 auto ret = memset_s(outputs[0]->device_ptr(), outputs[0]->size(), 0, outputs[0]->size());
96 if (ret != EOK) {
97 MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', memset output failed. Error no: " << ret << ".";
98 }
99 for (unsigned int i = 0; i < values_size_ && validate_indices_; ++i) {
100 if (!IndicesValid(i, inputs)) {
101 return false;
102 }
103 }
104
105 auto output_t = static_cast<int32_t *>(outputs[0]->device_ptr());
106 auto indices_t = static_cast<int64_t *>(inputs[0]->device_ptr());
107 auto vals_t = static_cast<T *>(inputs[1]->device_ptr());
108 auto vals_num = values_size_;
109 std::vector<int64_t> strides(dims_ - 1);
110 auto shape_t = static_cast<int64_t *>(inputs[2]->device_ptr());
111 if (dims_ > 1) {
112 int t = 2;
113 strides[dims_ - t] = 1;
114 }
115 for (int32_t d = dims_ - 3; d >= 0; --d) {
116 strides[d] = strides[d + 1] * shape_t[d + 1];
117 }
118
119 int32_t output_size = 1;
120 for (int32_t d = 0; d < dims_ - 1; ++d) {
121 output_size = output_size * shape_t[d];
122 }
123 std::vector<std::unordered_set<T>> all_values(output_size);
124 for (unsigned int n = 0; n < vals_num; ++n) {
125 int64_t ix = 0;
126 for (int d = 0; d < dims_ - 1; ++d) {
127 const int64_t ix_n_d = indices_t[SizeToLong(n) * dims_ + d];
128 ix += strides[d] * ix_n_d;
129 }
130 all_values[ix].insert(*(vals_t + n));
131 }
132 for (int i = 0; i < output_size; ++i) {
133 output_t[i] = SizeToLong(all_values[i].size());
134 }
135 return true;
136 }
137
GetFuncList() const138 const std::vector<std::pair<KernelAttr, SetSizeCpuKernelMod::KernelRunFunc>> &SetSizeCpuKernelMod::GetFuncList() const {
139 static const std::vector<std::pair<KernelAttr, SetSizeCpuKernelMod::KernelRunFunc>> func_list = {
140 {KernelAttr()
141 .AddInputAttr(kNumberTypeInt64)
142 .AddInputAttr(kNumberTypeInt8)
143 .AddInputAttr(kNumberTypeInt64)
144 .AddOutputAttr(kNumberTypeInt32),
145 &SetSizeCpuKernelMod::LaunchKernel<int8_t>},
146 {KernelAttr()
147 .AddInputAttr(kNumberTypeInt64)
148 .AddInputAttr(kNumberTypeInt16)
149 .AddInputAttr(kNumberTypeInt64)
150 .AddOutputAttr(kNumberTypeInt32),
151 &SetSizeCpuKernelMod::LaunchKernel<int16_t>},
152 {KernelAttr()
153 .AddInputAttr(kNumberTypeInt64)
154 .AddInputAttr(kNumberTypeInt32)
155 .AddInputAttr(kNumberTypeInt64)
156 .AddOutputAttr(kNumberTypeInt32),
157 &SetSizeCpuKernelMod::LaunchKernel<int32_t>},
158 {KernelAttr()
159 .AddInputAttr(kNumberTypeInt64)
160 .AddInputAttr(kNumberTypeInt64)
161 .AddInputAttr(kNumberTypeInt64)
162 .AddOutputAttr(kNumberTypeInt32),
163 &SetSizeCpuKernelMod::LaunchKernel<int64_t>},
164 {KernelAttr()
165 .AddInputAttr(kNumberTypeInt64)
166 .AddInputAttr(kNumberTypeUInt8)
167 .AddInputAttr(kNumberTypeInt64)
168 .AddOutputAttr(kNumberTypeInt32),
169 &SetSizeCpuKernelMod::LaunchKernel<uint8_t>},
170 {KernelAttr()
171 .AddInputAttr(kNumberTypeInt64)
172 .AddInputAttr(kNumberTypeUInt16)
173 .AddInputAttr(kNumberTypeInt64)
174 .AddOutputAttr(kNumberTypeInt32),
175 &SetSizeCpuKernelMod::LaunchKernel<uint16_t>},
176 };
177 return func_list;
178 }
179
180 MS_KERNEL_FACTORY_REG(NativeCpuKernelMod, SetSize, SetSizeCpuKernelMod);
181 } // namespace kernel
182 } // namespace mindspore
183