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 <cmath>
18 #include "src/litert/kernel_registry.h"
19 #include "nnacl/softmax_parameter.h"
20 #include "nnacl/fp32/softmax_fp32.h"
21 #include "nnacl/fp32_grad/softmax_cross_entropy_with_logits.h"
22 #include "src/litert/kernel/cpu/fp32_grad/softmax_cross_entropy_with_logits.h"
23 #include "include/errorcode.h"
24
25 using mindspore::lite::KernelRegistrar;
26 using mindspore::lite::RET_ERROR;
27 using mindspore::lite::RET_OK;
28 using mindspore::schema::PrimitiveType_SoftmaxCrossEntropyWithLogits;
29
30 namespace mindspore::kernel {
Prepare()31 int SoftmaxCrossEntropyWithLogitsCPUKernel::Prepare() { return ReSize(); }
32
DoExecute(int task_id)33 int SoftmaxCrossEntropyWithLogitsCPUKernel::DoExecute(int task_id) {
34 auto ins = reinterpret_cast<float *>(in_tensors_.at(0)->data());
35 CHECK_NULL_RETURN(ins);
36 auto labels = reinterpret_cast<float *>(in_tensors_.at(1)->data());
37 CHECK_NULL_RETURN(labels);
38 float *out = reinterpret_cast<float *>(out_tensors_.at(0)->data());
39 CHECK_NULL_RETURN(out);
40 float *grads = nullptr;
41 if (IsTrain() && out_tensors_.size() > 1) {
42 grads = reinterpret_cast<float *>(out_tensors_.at(1)->data());
43 }
44 size_t data_size = in_tensors_.at(0)->ElementsNum();
45
46 float *losses_ = static_cast<float *>(workspace());
47 float *sum_data_ = losses_ + data_size;
48 std::fill(losses_, losses_ + data_size, 0);
49 std::fill(sum_data_, sum_data_ + input_shape_[0], 0);
50 Softmax(ins, losses_, sum_data_, sm_params_.axis_, n_dim_, input_shape_);
51 ForwardPostExecute(labels, losses_, grads, out, param_->number_of_classes_, param_->batch_size_);
52 return RET_OK;
53 }
54
SoftmaxCrossEntropyWithLogitsRun(void * cdata,int task_id,float lhs_scale,float rhs_scale)55 int SoftmaxCrossEntropyWithLogitsRun(void *cdata, int task_id, float lhs_scale, float rhs_scale) {
56 CHECK_NULL_RETURN(cdata);
57 auto softmax_kernel = reinterpret_cast<SoftmaxCrossEntropyWithLogitsCPUKernel *>(cdata);
58 auto error_code = softmax_kernel->DoExecute(task_id);
59 if (error_code != RET_OK) {
60 MS_LOG(ERROR) << "SoftmaxCrossEntropy error task_id[" << task_id << "] error_code[" << error_code << "]";
61 return RET_ERROR;
62 }
63 return RET_OK;
64 }
65
Run()66 int SoftmaxCrossEntropyWithLogitsCPUKernel::Run() {
67 int error_code = ParallelLaunch(this->ms_context_, SoftmaxCrossEntropyWithLogitsRun, this, 1);
68 if (error_code != RET_OK) {
69 MS_LOG(ERROR) << "SoftmaxCrossEntropy function error error_code[" << error_code << "]";
70 return RET_ERROR;
71 }
72 return RET_OK;
73 }
74
ReSize()75 int SoftmaxCrossEntropyWithLogitsCPUKernel::ReSize() {
76 CHECK_LESS_RETURN(in_tensors_.size(), DIMENSION_2D);
77 CHECK_LESS_RETURN(out_tensors_.size(), DIMENSION_2D);
78 CHECK_NULL_RETURN(param_);
79 CHECK_NULL_RETURN(in_tensors_.at(0));
80 CHECK_NULL_RETURN(in_tensors_.at(1));
81 CHECK_NULL_RETURN(out_tensors_.at(0));
82 auto dims = in_tensors_.at(0)->shape();
83 param_->n_dim_ = 2;
84 CHECK_LESS_RETURN(dims.size(), DIMENSION_2D);
85 param_->number_of_classes_ = dims.at(1);
86 param_->batch_size_ = dims.at(0);
87 for (unsigned int i = 0; i < dims.size(); i++) param_->input_shape_[i] = dims.at(i);
88 if (this->in_tensors_.size() != 2) {
89 MS_LOG(ERROR) << "softmax entropy loss should have two inputs";
90 return RET_ERROR;
91 }
92 auto *in0 = in_tensors_.front();
93 if (in0 == nullptr) {
94 MS_LOG(ERROR) << "softmax etropy loss in0 have no data";
95 return RET_ERROR;
96 }
97
98 size_t data_size = in_tensors_.at(0)->ElementsNum();
99 set_workspace_size((data_size + static_cast<size_t>(dims.at(0))) * sizeof(float));
100 n_dim_ = Num2;
101 element_size_ = data_size;
102 sm_params_.axis_ = 1;
103 for (size_t i = 0; i < dims.size(); i++) input_shape_[i] = dims.at(i);
104
105 return RET_OK;
106 }
107
CpuSoftmaxCrossEntropyFp32KernelCreator(const std::vector<lite::Tensor * > & inputs,const std::vector<lite::Tensor * > & outputs,OpParameter * opParameter,const lite::InnerContext * ctx,const kernel::KernelKey & desc)108 kernel::LiteKernel *CpuSoftmaxCrossEntropyFp32KernelCreator(const std::vector<lite::Tensor *> &inputs,
109 const std::vector<lite::Tensor *> &outputs,
110 OpParameter *opParameter, const lite::InnerContext *ctx,
111 const kernel::KernelKey &desc) {
112 MS_ASSERT(opParameter != nullptr);
113 MS_ASSERT(desc.type == schema::PrimitiveType_SoftmaxCrossEntropyWithLogits);
114 auto *kernel = new (std::nothrow) SoftmaxCrossEntropyWithLogitsCPUKernel(opParameter, inputs, outputs, ctx);
115 if (kernel == nullptr) {
116 MS_LOG(ERROR) << "new SoftmaxCrossEntropyWithLogitsCPUKernel failed";
117 free(opParameter);
118 return nullptr;
119 }
120 return kernel;
121 }
122
123 REG_KERNEL(kCPU, kNumberTypeFloat32, PrimitiveType_SoftmaxCrossEntropyWithLogits,
124 CpuSoftmaxCrossEntropyFp32KernelCreator)
125 } // namespace mindspore::kernel
126