1 /**
2 * Copyright 2020 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 "src/runtime/kernel/arm/fp16/log_softmax_fp16.h"
18 #include <cstring>
19 #include <vector>
20 #include "src/runtime/kernel/arm/fp16/common_fp16.h"
21 #include "nnacl/fp16/log_softmax_fp16.h"
22 #include "nnacl/fp16/cast_fp16.h"
23 #include "schema/model_generated.h"
24 #include "src/kernel_registry.h"
25 #include "include/errorcode.h"
26
27 using mindspore::kernel::KERNEL_ARCH;
28 using mindspore::lite::KernelRegistrar;
29 using mindspore::lite::RET_ERROR;
30 using mindspore::lite::RET_OK;
31 using mindspore::schema::PrimitiveType_LogSoftmax;
32
33 namespace mindspore::kernel {
~LogSoftmaxFp16CPUKernel()34 LogSoftmaxFp16CPUKernel::~LogSoftmaxFp16CPUKernel() {
35 if (tmp_data_ != nullptr) {
36 free(tmp_data_);
37 tmp_data_ = nullptr;
38 }
39 }
40
Init()41 int LogSoftmaxFp16CPUKernel::Init() {
42 auto ret = SoftmaxBaseCPUKernel::Init();
43 if (ret != RET_OK) {
44 return ret;
45 }
46
47 if (!InferShapeDone()) {
48 return RET_OK;
49 }
50 return ReSize();
51 }
52
ReSize()53 int LogSoftmaxFp16CPUKernel::ReSize() {
54 auto ret = SoftmaxBaseCPUKernel::ReSize();
55 if (ret != RET_OK) {
56 return ret;
57 }
58 auto n_dim = softmax_param_->n_dim_;
59 auto axis = softmax_param_->axis_;
60 auto in_shape = in_tensors_.front()->shape();
61 out_plane_size_ = 1;
62 for (int i = 0; i < axis; ++i) {
63 out_plane_size_ *= in_shape[i];
64 }
65 in_plane_size_ = 1;
66 for (int i = axis + 1; i < n_dim; i++) {
67 in_plane_size_ *= in_shape[i];
68 }
69 auto tmp_data_size =
70 in_plane_size_ == 1 ? out_plane_size_ * in_plane_size_ * in_shape.at(axis) : out_plane_size_ * in_plane_size_;
71 if (tmp_data_ != nullptr) {
72 free(tmp_data_);
73 }
74 tmp_data_ = reinterpret_cast<float16_t *>(malloc(tmp_data_size * sizeof(float16_t)));
75 if (tmp_data_ == nullptr) {
76 MS_LOG(ERROR) << "malloc data for softmax fail!";
77 return RET_ERROR;
78 }
79 return RET_OK;
80 }
81
DoLogSoftmaxLastAxis(int task_id)82 int LogSoftmaxFp16CPUKernel::DoLogSoftmaxLastAxis(int task_id) {
83 MS_CHECK_FALSE(op_parameter_->thread_num_ == 0, RET_ERROR);
84 int unit = UP_DIV(out_plane_size_, op_parameter_->thread_num_);
85 int begin = task_id * unit;
86 int end = MSMIN(begin + unit, out_plane_size_);
87 int channel = softmax_param_->input_shape_[softmax_param_->axis_];
88 int offset = begin * channel;
89 auto input_ptr = reinterpret_cast<float16_t *>(in_tensors_.at(kInputIndex)->data());
90 CHECK_NULL_RETURN(input_ptr);
91 auto output_ptr = reinterpret_cast<float16_t *>(out_tensors_.at(kOutputIndex)->data());
92 CHECK_NULL_RETURN(output_ptr);
93 LogSoftmaxLastAxisFp16(input_ptr + offset, output_ptr + offset, tmp_data_ + offset, end - begin, channel);
94 return RET_OK;
95 }
96
LogSoftmaxLastAxisFp16Run(void * cdata,int task_id,float lhs_scale,float rhs_scale)97 int LogSoftmaxLastAxisFp16Run(void *cdata, int task_id, float lhs_scale, float rhs_scale) {
98 auto kernel = reinterpret_cast<LogSoftmaxFp16CPUKernel *>(cdata);
99 CHECK_NULL_RETURN(kernel);
100 auto ret = kernel->DoLogSoftmaxLastAxis(task_id);
101 if (ret != RET_OK) {
102 MS_LOG(ERROR) << "DoLogSoftmaxLastAxisFp16 error task_id: " << task_id << ", ret: " << ret;
103 }
104 return ret;
105 }
106
Run()107 int LogSoftmaxFp16CPUKernel::Run() {
108 if (in_plane_size_ == 1) {
109 auto ret = ParallelLaunch(this->ms_context_, LogSoftmaxLastAxisFp16Run, this, op_parameter_->thread_num_);
110 if (ret != RET_OK) {
111 MS_LOG(ERROR) << "LogSoftmaxFp16CPUKernel ParallelLaunch failed, ret: " << ret;
112 }
113 return ret;
114 } else {
115 input_fp16_ = reinterpret_cast<float16_t *>(in_tensors_.at(0)->data());
116 CHECK_NULL_RETURN(input_fp16_);
117 output_fp16_ = reinterpret_cast<float16_t *>(out_tensors_.at(0)->data());
118 CHECK_NULL_RETURN(output_fp16_);
119 CHECK_NULL_RETURN(tmp_data_);
120 LogSoftmaxFp16(input_fp16_, output_fp16_, tmp_data_, softmax_param_);
121 }
122 return RET_OK;
123 }
124
125 REG_KERNEL(kCPU, kNumberTypeFloat16, PrimitiveType_LogSoftmax, LiteKernelCreator<LogSoftmaxFp16CPUKernel>)
126 } // namespace mindspore::kernel
127