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/fp32/power_fp32.h"
18 #include "schema/model_generated.h"
19 #include "src/kernel_registry.h"
20 #include "include/errorcode.h"
21
22 using mindspore::lite::KernelRegistrar;
23 using mindspore::lite::RET_ERROR;
24 using mindspore::lite::RET_OK;
25 using mindspore::schema::PrimitiveType_PowFusion;
26
27 namespace mindspore::kernel {
Init()28 int PowerCPUKernel::Init() {
29 CHECK_LESS_RETURN(in_tensors_.size(), C2NUM);
30 CHECK_LESS_RETURN(out_tensors_.size(), 1);
31 return RET_OK;
32 }
33
ReSize()34 int PowerCPUKernel::ReSize() { return RET_OK; }
35
PowerImpl(void * cdata,int task_id,float lhs_scale,float rhs_scale)36 int PowerImpl(void *cdata, int task_id, float lhs_scale, float rhs_scale) {
37 CHECK_NULL_RETURN(cdata);
38 auto kernel = reinterpret_cast<PowerCPUKernel *>(cdata);
39 CHECK_NULL_RETURN(kernel);
40 auto ret = kernel->RunImpl(task_id);
41 if (ret != RET_OK) {
42 MS_LOG(ERROR) << "PowerImpl error: " << ret;
43 return ret;
44 }
45 return RET_OK;
46 }
47
Run()48 int PowerCPUKernel::Run() {
49 auto ret = ParallelLaunch(this->ms_context_, PowerImpl, this, thread_count_);
50 if (ret != RET_OK) {
51 MS_LOG(ERROR) << "PowerCPUKernel error: " << ret;
52 return RET_ERROR;
53 }
54 return RET_OK;
55 }
56
RunImpl(int task_id)57 int PowerCPUKernel::RunImpl(int task_id) {
58 auto x_addr = reinterpret_cast<float *>(in_tensors_.at(0)->MutableData());
59 CHECK_NULL_RETURN(x_addr);
60 auto output_addr = reinterpret_cast<float *>(out_tensors_.at(0)->MutableData());
61 CHECK_NULL_RETURN(output_addr);
62 auto size = in_tensors_.at(0)->ElementsNum();
63 int stride = UP_DIV(size, thread_count_);
64 int len = MSMIN(stride, size - stride * task_id);
65 if (len <= 0) {
66 return RET_OK;
67 }
68 float *exp_addr = nullptr;
69 bool broadcast = true;
70 MS_ASSERT(in_tensors_.size() == 2);
71 exp_addr = reinterpret_cast<float *>(in_tensors_[1]->data());
72 CHECK_NULL_RETURN(exp_addr);
73 broadcast = in_tensors_[0]->shape() == in_tensors_[1]->shape() ? false : true;
74
75 float *cur_exp = nullptr;
76 if (broadcast) {
77 cur_exp = exp_addr;
78 } else {
79 cur_exp = exp_addr + stride * task_id;
80 }
81 CHECK_NULL_RETURN(cur_exp);
82 auto error_code =
83 Power(x_addr + stride * task_id, cur_exp, output_addr + stride * task_id, len, scale_, shift_, broadcast);
84 if (error_code != RET_OK) {
85 MS_LOG(ERROR) << "PowerCPUKernel RunImpl error task_id[" << task_id << "] error_code[" << error_code << "]";
86 return RET_ERROR;
87 }
88 return RET_OK;
89 }
90
91 REG_KERNEL(kCPU, kNumberTypeFloat32, PrimitiveType_PowFusion, LiteKernelCreator<PowerCPUKernel>)
92 } // namespace mindspore::kernel
93