1 /**
2 * Copyright 2020-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/sparse_matrix_softmax_cpu_kernel.h"
18
19 #include <cmath>
20 #include <thread>
21 #include <vector>
22 #include <algorithm>
23 #include <map>
24 #include <string>
25 #include <functional>
26
27 #include "kernel/common_utils.h"
28
29 namespace mindspore {
30 namespace kernel {
31 namespace {
32 constexpr size_t kInputNum = 5;
33 constexpr size_t kOutputNum = 5;
34 constexpr size_t logits_dense_shape = 0;
35 constexpr size_t logits_batch_pointers = 1;
36 constexpr size_t logits_col_indices = 2;
37 constexpr size_t logits_row_pointers = 3;
38 constexpr size_t logits_values = 4;
39 constexpr char kKernelName[] = "sparse_matrix_softmax";
40 } // namespace
Init(const std::vector<KernelTensor * > & inputs,const std::vector<KernelTensor * > & outputs)41 bool SparseMatrixSoftmaxCpuKernelMod::Init(const std::vector<KernelTensor *> &inputs,
42 const std::vector<KernelTensor *> &outputs) {
43 dtype_ = inputs[logits_values]->dtype_id();
44 size_t input_num = inputs.size();
45 if (input_num != kInputNum) {
46 MS_LOG(ERROR) << "For " << kernel_name_
47 << ", input should be x_dense_shape, x_batch_pointers, x_row_pointers, x_col_indices, x_values "
48 << kInputNum << " tensors, but get " << input_num;
49 return false;
50 }
51 return true;
52 }
Resize(const std::vector<KernelTensor * > & inputs,const std::vector<KernelTensor * > & outputs)53 int SparseMatrixSoftmaxCpuKernelMod::Resize(const std::vector<KernelTensor *> &inputs,
54 const std::vector<KernelTensor *> &outputs) {
55 int ret = KernelMod::Resize(inputs, outputs);
56 if (ret != 0) {
57 return ret;
58 }
59 return 0;
60 }
61
Launch(const std::vector<kernel::KernelTensor * > & inputs,const std::vector<kernel::KernelTensor * > &,const std::vector<kernel::KernelTensor * > & outputs)62 bool SparseMatrixSoftmaxCpuKernelMod::Launch(const std::vector<kernel::KernelTensor *> &inputs,
63 const std::vector<kernel::KernelTensor *> &,
64 const std::vector<kernel::KernelTensor *> &outputs) {
65 if (dtype_ == kNumberTypeFloat32) {
66 LaunchKernel<float>(inputs, outputs);
67 } else if (dtype_ == kNumberTypeFloat64) {
68 LaunchKernel<double>(inputs, outputs);
69 } else {
70 MS_LOG(EXCEPTION) << "For '" << kernel_name_
71 << "', the dtype of 'logits_values' must be Float32 or Float64, but got "
72 << TypeIdToType(dtype_)->ToString();
73 }
74 return true;
75 }
76
77 template <typename T>
78
LaunchKernel(const std::vector<kernel::KernelTensor * > & inputs,const std::vector<kernel::KernelTensor * > &)79 void SparseMatrixSoftmaxCpuKernelMod::LaunchKernel(const std::vector<kernel::KernelTensor *> &inputs,
80 const std::vector<kernel::KernelTensor *> &) {
81 auto *input_logits_values = reinterpret_cast<T *>(inputs[logits_values]->device_ptr());
82 auto *input_logits_dense_shape = reinterpret_cast<int *>(inputs[logits_dense_shape]->device_ptr());
83 auto *input_logits_col_indices = reinterpret_cast<int *>(inputs[logits_col_indices]->device_ptr());
84 T total = 0;
85 T MAX = input_logits_values[0];
86 int row_index = input_logits_dense_shape[0];
87 int start = 0;
88 for (int i = 1; i <= row_index; i++) {
89 int single_index = (input_logits_col_indices[i] - input_logits_col_indices[i - 1]);
90 for (int k = 0; k < single_index; k++) {
91 if (input_logits_values[k + start] > MAX) {
92 MAX = input_logits_values[k + start];
93 }
94 }
95 for (int k = 0; k < single_index; k++) {
96 total = total + exp(input_logits_values[k + start] - MAX);
97 }
98 for (int k = 0; k < single_index; k++) {
99 input_logits_values[k + start] = exp(input_logits_values[k + start] - MAX) / total;
100 }
101 start = start + single_index;
102 MAX = input_logits_values[start];
103 total = 0;
104 }
105 }
106
GetOpSupport()107 std::vector<KernelAttr> SparseMatrixSoftmaxCpuKernelMod::GetOpSupport() {
108 static std::vector<KernelAttr> kernel_attr_list = {KernelAttr()
109 .AddInputAttr(kNumberTypeInt32)
110 .AddInputAttr(kNumberTypeInt32)
111 .AddInputAttr(kNumberTypeInt32)
112 .AddInputAttr(kNumberTypeInt32)
113 .AddInputAttr(kNumberTypeFloat32)
114 .AddOutputAttr(kNumberTypeInt32)
115 .AddOutputAttr(kNumberTypeInt32)
116 .AddOutputAttr(kNumberTypeInt32)
117 .AddOutputAttr(kNumberTypeInt32)
118 .AddOutputAttr(kNumberTypeFloat32)
119 .AddOutInRef(0, 0)
120 .AddOutInRef(1, 1)
121 .AddOutInRef(2, 2)
122 .AddOutInRef(3, 3)
123 .AddOutInRef(4, 4),
124 KernelAttr()
125 .AddInputAttr(kNumberTypeInt64)
126 .AddInputAttr(kNumberTypeInt64)
127 .AddInputAttr(kNumberTypeInt64)
128 .AddInputAttr(kNumberTypeInt64)
129 .AddInputAttr(kNumberTypeFloat32)
130 .AddOutputAttr(kNumberTypeInt64)
131 .AddOutputAttr(kNumberTypeInt64)
132 .AddOutputAttr(kNumberTypeInt64)
133 .AddOutputAttr(kNumberTypeInt64)
134 .AddOutputAttr(kNumberTypeFloat32)
135 .AddOutInRef(0, 0)
136 .AddOutInRef(1, 1)
137 .AddOutInRef(2, 2)
138 .AddOutInRef(3, 3)
139 .AddOutInRef(4, 4),
140 KernelAttr()
141 .AddInputAttr(kNumberTypeInt32)
142 .AddInputAttr(kNumberTypeInt32)
143 .AddInputAttr(kNumberTypeInt32)
144 .AddInputAttr(kNumberTypeInt32)
145 .AddInputAttr(kNumberTypeFloat64)
146 .AddOutputAttr(kNumberTypeInt32)
147 .AddOutputAttr(kNumberTypeInt32)
148 .AddOutputAttr(kNumberTypeInt32)
149 .AddOutputAttr(kNumberTypeInt32)
150 .AddOutputAttr(kNumberTypeFloat64)
151 .AddOutInRef(0, 0)
152 .AddOutInRef(1, 1)
153 .AddOutInRef(2, 2)
154 .AddOutInRef(3, 3)
155 .AddOutInRef(4, 4),
156 KernelAttr()
157 .AddInputAttr(kNumberTypeInt64)
158 .AddInputAttr(kNumberTypeInt64)
159 .AddInputAttr(kNumberTypeInt64)
160 .AddInputAttr(kNumberTypeInt64)
161 .AddInputAttr(kNumberTypeFloat64)
162 .AddOutputAttr(kNumberTypeInt64)
163 .AddOutputAttr(kNumberTypeInt64)
164 .AddOutputAttr(kNumberTypeInt64)
165 .AddOutputAttr(kNumberTypeInt64)
166 .AddOutputAttr(kNumberTypeFloat64)
167 .AddOutInRef(0, 0)
168 .AddOutInRef(1, 1)
169 .AddOutInRef(2, 2)
170 .AddOutInRef(3, 3)
171 .AddOutInRef(4, 4)};
172 return kernel_attr_list;
173 }
174
175 MS_KERNEL_FACTORY_REG(NativeCpuKernelMod, SparseMatrixSoftmax, SparseMatrixSoftmaxCpuKernelMod);
176 } // namespace kernel
177 } // namespace mindspore
178