1 /**
2 * Copyright 2023 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 "nnacl/kernel/softmax.h"
18 #include "nnacl/nnacl_common.h"
19 #include "nnacl/fp32/softmax_fp32.h"
20 #include "nnacl/kernel/default_kernel_base.h"
21 #include "nnacl/tensor_c_utils.h"
22 #ifdef ENABLE_FP16
23 #include "nnacl/fp16/softmax_fp16.h"
24 #endif
25
SoftmaxLastAxisRun(void * cdata,int task_id,float l,float r)26 int SoftmaxLastAxisRun(void *cdata, int task_id, float l, float r) {
27 SoftmaxStruct *softmax = (SoftmaxStruct *)cdata;
28 NNACL_CHECK_NULL_RETURN_ERR(softmax);
29
30 NNACL_CHECK_ZERO_RETURN_ERR(softmax->base_.thread_nr_);
31 int unit = UP_DIV(softmax->out_plane_size_, softmax->base_.thread_nr_);
32
33 int *in_shape = softmax->base_.in_[FIRST_INPUT]->shape_;
34 NNACL_CHECK_INT_MUL_NOT_OVERFLOW(task_id, unit, NNACL_ERR);
35 int begin = task_id * unit;
36 int end = MSMIN(begin + unit, softmax->out_plane_size_);
37 int channel = in_shape[softmax->axis_];
38
39 NNACL_CHECK_INT_MUL_NOT_OVERFLOW(begin, channel, NNACL_ERR);
40 int offset = begin * channel;
41
42 void *input_ptr = softmax->base_.in_[FIRST_INPUT]->data_;
43 NNACL_CHECK_NULL_RETURN_ERR(input_ptr);
44 void *output_ptr = softmax->base_.out_[OUTPUT_INDEX]->data_;
45 NNACL_CHECK_NULL_RETURN_ERR(output_ptr);
46
47 #ifdef ENABLE_FP16
48 if (softmax->data_type_ == kNumberTypeFloat16) {
49 SoftmaxLastAxisFp16((float16_t *)input_ptr + offset, (float16_t *)output_ptr + offset, end - begin, channel);
50 return NNACL_OK;
51 }
52 #endif
53 return SoftmaxLastAxis((float *)input_ptr + offset, (float *)output_ptr + offset, end - begin, channel);
54 }
55
SoftmaxRelease(struct KernelBase * self)56 int SoftmaxRelease(struct KernelBase *self) {
57 SoftmaxStruct *softmax = (SoftmaxStruct *)self;
58 NNACL_CHECK_NULL_RETURN_ERR(softmax);
59 if (softmax->sum_data_ != NULL) {
60 self->env_->Free(self->env_->allocator_, softmax->sum_data_);
61 }
62 softmax->sum_data_ = NULL;
63 return NNACL_OK;
64 }
65
InitSoftmaxParam(SoftmaxStruct * softmax)66 int InitSoftmaxParam(SoftmaxStruct *softmax) {
67 TensorC *in_tensor = softmax->base_.in_[FIRST_INPUT];
68 NNACL_CHECK_NULL_RETURN_ERR(in_tensor);
69 int *in_shape = in_tensor->shape_;
70
71 softmax->n_dim_ = (int)in_tensor->shape_size_;
72 int origin_axis = ((SoftmaxParameter *)softmax->base_.param_)->axis_;
73 softmax->axis_ = origin_axis == -1 ? origin_axis + softmax->n_dim_ : origin_axis;
74
75 NNACL_CHECK_TRUE_RET(softmax->axis_ >= 0, NNACL_SOFTMAX_AXIS_INVALID);
76 NNACL_CHECK_TRUE_RET(softmax->axis_ < (int)in_tensor->shape_size_, NNACL_SOFTMAX_AXIS_INVALID);
77
78 int out_plane_size = 1;
79 for (int i = 0; i < softmax->axis_; ++i) {
80 out_plane_size *= in_shape[i];
81 }
82 int in_plane_size = 1;
83 for (int i = softmax->axis_ + 1; i < softmax->n_dim_; i++) {
84 in_plane_size *= in_shape[i];
85 }
86
87 ExecEnv *env = softmax->base_.env_;
88 NNACL_CHECK_NULL_RETURN_ERR(env);
89
90 softmax->in_plane_size_ = in_plane_size;
91 softmax->out_plane_size_ = out_plane_size;
92
93 (void)softmax->base_.Release(&softmax->base_);
94 if (softmax->in_plane_size_ > 1) {
95 NNACL_CHECK_INT_MUL_NOT_OVERFLOW(out_plane_size, in_plane_size, NNACL_ERR);
96 int sum_data_size = out_plane_size * in_plane_size;
97 NNACL_CHECK_INT_MUL_NOT_OVERFLOW(sum_data_size, (int)DataTypeCSize(softmax->data_type_), NNACL_ERR);
98 softmax->sum_data_ = env->Alloc(env->allocator_, sum_data_size * DataTypeCSize(softmax->data_type_));
99 NNACL_MALLOC_CHECK_NULL_RETURN_ERR(softmax->sum_data_);
100 }
101 return NNACL_OK;
102 }
103
SoftmaxResize(struct KernelBase * self)104 int SoftmaxResize(struct KernelBase *self) {
105 SoftmaxStruct *softmax = (SoftmaxStruct *)self;
106 NNACL_CHECK_NULL_RETURN_ERR(softmax);
107 InitSoftmaxParam(softmax);
108
109 TensorC *in_tensor = self->in_[FIRST_INPUT];
110 int *in_shape = in_tensor->shape_;
111
112 self->thread_nr_ = self->UpdateThread(TC_PTYPE(PrimType_Softmax), in_shape[softmax->axis_], in_shape[softmax->axis_],
113 GetElementNum(self->out_[OUTPUT_INDEX]), self->thread_nr_);
114 return NNACL_OK;
115 }
116
SoftmaxCompute(struct KernelBase * self)117 int SoftmaxCompute(struct KernelBase *self) {
118 SoftmaxStruct *softmax = (SoftmaxStruct *)self;
119 NNACL_CHECK_NULL_RETURN_ERR(softmax);
120
121 if (softmax->in_plane_size_ == 1) {
122 return self->env_->ParallelLaunch(self->env_->thread_pool_, SoftmaxLastAxisRun, softmax, self->thread_nr_);
123 }
124
125 void *input_ptr = self->in_[FIRST_INPUT]->data_;
126 NNACL_CHECK_NULL_RETURN_ERR(input_ptr);
127 void *output_ptr = self->out_[OUTPUT_INDEX]->data_;
128 NNACL_CHECK_NULL_RETURN_ERR(output_ptr);
129 NNACL_CHECK_NULL_RETURN_ERR(softmax->sum_data_);
130 #ifdef ENABLE_FP16
131 if (softmax->data_type_ == kNumberTypeFloat16) {
132 SoftmaxFp16((float16_t *)input_ptr, (float16_t *)output_ptr, (float16_t *)softmax->sum_data_, softmax->axis_,
133 softmax->n_dim_, self->in_[FIRST_INPUT]->shape_);
134 return NNACL_OK;
135 }
136 #endif
137 Softmax((float *)input_ptr, (float *)output_ptr, (float *)softmax->sum_data_, softmax->axis_, softmax->n_dim_,
138 self->in_[FIRST_INPUT]->shape_);
139 return NNACL_OK;
140 }
141
CreateSoftmax(OpParameter * param,int data_type)142 KernelBase *CreateSoftmax(OpParameter *param, int data_type) {
143 SoftmaxStruct *softmax = (SoftmaxStruct *)malloc(sizeof(SoftmaxStruct));
144 NNACL_MALLOC_CHECK_NULL_RETURN_NULL(softmax);
145 memset(softmax, 0, sizeof(SoftmaxStruct));
146
147 softmax->sum_data_ = NULL;
148 softmax->data_type_ = data_type;
149 softmax->base_.Release = SoftmaxRelease;
150 softmax->base_.Prepare = DefaultPrepare1In1Out;
151 softmax->base_.Resize = SoftmaxResize;
152 softmax->base_.Compute = SoftmaxCompute;
153 return (KernelBase *)softmax;
154 }
155
156 REG_KERNEL_CREATOR(PrimType_Softmax, kNumberTypeFloat16, CreateSoftmax)
157 REG_KERNEL_CREATOR(PrimType_Softmax, kNumberTypeFloat32, CreateSoftmax)
158