1 /**
2 * Copyright 2020-2021 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 #include "src/kernel_registry.h"
17 #include <utility>
18 #include <memory>
19 #include "include/errorcode.h"
20 #ifndef CUSTOM_KERNEL_REGISTRY_CLIP
21 #include "include/registry/register_kernel.h"
22 #endif
23 #include "src/ops/populate/populate_register.h"
24 #include "src/common/version_manager.h"
25 #include "nnacl/pooling_parameter.h"
26 #if defined(ENABLE_FP16) && defined(ENABLE_ARM)
27 #if defined(__ANDROID__)
28 #include <asm/hwcap.h>
29 #endif
30 #include "common/utils.h"
31 #include "src/common/log_adapter.h"
32 #include "src/common/utils.h"
33 #endif
34 #include "src/common/tensor_util.h"
35
36 using mindspore::kernel::kBuiltin;
37 using mindspore::kernel::kCPU;
38 using mindspore::kernel::KERNEL_ARCH;
39 using mindspore::kernel::KernelCreator;
40 using mindspore::kernel::KernelKey;
41 #ifndef CUSTOM_KERNEL_REGISTRY_CLIP
42 using mindspore::registry::CreateKernel;
43 using mindspore::registry::KernelDesc;
44 #endif
45
46 namespace mindspore::lite {
47 #ifndef CUSTOM_KERNEL_REGISTRY_CLIP
48 namespace {
49 const char *const kArchCPU = "CPU";
50 const char *const kArchGPU = "GPU";
KernelKeyToKernelDesc(const KernelKey & key,KernelDesc * desc)51 void KernelKeyToKernelDesc(const KernelKey &key, KernelDesc *desc) {
52 MS_ASSERT(desc != nullptr);
53 desc->data_type = static_cast<DataType>(key.data_type);
54 desc->type = key.type;
55 desc->arch = key.kernel_arch;
56 desc->provider = key.provider;
57 }
58 } // namespace
59 #endif
60
CreatorArraysInit()61 void KernelRegistry::CreatorArraysInit() {
62 std::unique_lock<std::mutex> malloc_creator_array(lock_);
63 if (creator_arrays_ == nullptr) {
64 creator_arrays_ = reinterpret_cast<KernelCreator *>(malloc(array_size_ * sizeof(KernelCreator)));
65 if (creator_arrays_ != nullptr) {
66 memset(creator_arrays_, 0, array_size_ * sizeof(KernelCreator));
67 }
68 }
69 return;
70 }
71
GetInstance()72 KernelRegistry *KernelRegistry::GetInstance() {
73 static KernelRegistry instance;
74 return &instance;
75 }
76
GetCreator(const KernelKey & desc)77 kernel::KernelCreator KernelRegistry::GetCreator(const KernelKey &desc) {
78 if (desc.provider == kBuiltin) {
79 int index = GetCreatorFuncIndex(desc);
80 if (index >= array_size_ || index < 0) {
81 MS_LOG(ERROR) << "invalid kernel key, arch " << desc.arch << ", data_type " << desc.data_type << ",op type "
82 << desc.type;
83 return nullptr;
84 }
85 if (creator_arrays_ != nullptr) {
86 return creator_arrays_[index];
87 }
88 }
89 MS_LOG(ERROR) << "Call wrong interface!provider: " << desc.provider;
90 return nullptr;
91 }
92
GetCreatorFuncIndex(const kernel::KernelKey desc)93 int KernelRegistry::GetCreatorFuncIndex(const kernel::KernelKey desc) {
94 int device_index = static_cast<int>(desc.arch) - kKernelArch_MIN;
95 int dType_index = static_cast<int>(desc.data_type) - kNumberTypeBegin;
96 int op_index = static_cast<int>(desc.type);
97 int index = device_index * data_type_length_ * op_type_length_ + dType_index * op_type_length_ + op_index;
98 return index;
99 }
100
RegKernel(const KernelKey desc,const kernel::KernelCreator creator)101 void KernelRegistry::RegKernel(const KernelKey desc, const kernel::KernelCreator creator) {
102 CreatorArraysInit();
103 int index = GetCreatorFuncIndex(desc);
104 if (index >= array_size_ || index < 0) {
105 MS_LOG(ERROR) << "invalid kernel key, arch " << desc.arch << ", data_type" << desc.data_type << ",op type "
106 << desc.type;
107 return;
108 }
109 if (creator_arrays_ != nullptr) {
110 creator_arrays_[index] = creator;
111 }
112 }
113
RegKernel(KERNEL_ARCH arch,TypeId data_type,int op_type,kernel::KernelCreator creator)114 void KernelRegistry::RegKernel(KERNEL_ARCH arch, TypeId data_type, int op_type, kernel::KernelCreator creator) {
115 CreatorArraysInit();
116 KernelKey desc = {arch, data_type, op_type};
117 int index = GetCreatorFuncIndex(desc);
118 if (index >= array_size_ || index < 0) {
119 MS_LOG(ERROR) << "invalid kernel key, arch " << desc.arch << ", data_type" << desc.data_type << ",op type "
120 << desc.type;
121 return;
122 }
123 if (creator_arrays_ != nullptr) {
124 creator_arrays_[index] = creator;
125 }
126 }
127
~KernelRegistry()128 KernelRegistry::~KernelRegistry() {
129 KernelRegistry *instance = GetInstance();
130 std::unique_lock<std::mutex> malloc_creator_array(instance->lock_);
131 if (instance->creator_arrays_ != nullptr) {
132 free(instance->creator_arrays_);
133 instance->creator_arrays_ = nullptr;
134 }
135 }
136
SupportKernel(const KernelKey & key)137 bool KernelRegistry::SupportKernel(const KernelKey &key) {
138 auto kernel_creator = GetCreator(key);
139 return kernel_creator != nullptr;
140 }
141
142 #ifndef CUSTOM_KERNEL_REGISTRY_CLIP
GetCustomKernel(const std::vector<Tensor * > & in_tensors,const std::vector<Tensor * > & out_tensors,const mindspore::Context * ms_ctx,const kernel::KernelKey & key,kernel::LiteKernel ** kernel,const void * primitive)143 int KernelRegistry::GetCustomKernel(const std::vector<Tensor *> &in_tensors, const std::vector<Tensor *> &out_tensors,
144 const mindspore::Context *ms_ctx, const kernel::KernelKey &key,
145 kernel::LiteKernel **kernel, const void *primitive) {
146 MS_ASSERT(ms_ctx != nullptr);
147 MS_ASSERT(kernel != nullptr);
148 KernelDesc desc;
149 KernelKeyToKernelDesc(key, &desc);
150 auto creator = registry::RegisterKernel::GetCreator(static_cast<const schema::Primitive *>(primitive), &desc);
151 if (creator == nullptr) {
152 return RET_NOT_SUPPORT;
153 }
154
155 auto base_kernel = creator(LiteTensorsToMSTensors(in_tensors), LiteTensorsToMSTensors(out_tensors),
156 static_cast<const schema::Primitive *>(primitive), ms_ctx);
157 if (base_kernel != nullptr) {
158 auto *lite_kernel = new (std::nothrow) kernel::LiteKernel(base_kernel);
159 if (lite_kernel != nullptr) {
160 kernel::KernelKey tmp_key = key;
161 if (desc.arch == kArchCPU) {
162 tmp_key.arch = kernel::kCPU;
163 } else if (desc.arch == kArchGPU) {
164 tmp_key.arch = kernel::kGPU;
165 } else {
166 tmp_key.arch = kernel::kCustom;
167 }
168 lite_kernel->set_desc(tmp_key);
169 *kernel = lite_kernel;
170 return RET_OK;
171 }
172 }
173 return RET_ERROR;
174 }
175 #endif
176
GetKernel(const std::vector<Tensor * > & in_tensors,const std::vector<Tensor * > & out_tensors,const InnerContext * ctx,const mindspore::Context * ms_ctx,const kernel::KernelKey & key,OpParameter * parameter,kernel::LiteKernel ** kernel,const void * primitive)177 int KernelRegistry::GetKernel(const std::vector<Tensor *> &in_tensors, const std::vector<Tensor *> &out_tensors,
178 const InnerContext *ctx, const mindspore::Context *ms_ctx, const kernel::KernelKey &key,
179 OpParameter *parameter, kernel::LiteKernel **kernel, const void *primitive) {
180 MS_ASSERT(ctx != nullptr);
181 MS_ASSERT(kernel != nullptr);
182 #ifndef CUSTOM_KERNEL_REGISTRY_CLIP
183 if (key.provider == kBuiltin) {
184 #endif
185 auto creator = GetCreator(key);
186 if (creator != nullptr) {
187 auto inner_kernel = creator(in_tensors, out_tensors, parameter, ctx, key);
188 if (inner_kernel != nullptr) {
189 inner_kernel->set_registry_data_type(key.data_type);
190 std::shared_ptr<kernel::Kernel> shared_kernel(inner_kernel);
191 auto *lite_kernel = new (std::nothrow) kernel::LiteKernel(shared_kernel);
192 if (lite_kernel != nullptr) {
193 lite_kernel->set_desc(key);
194 lite_kernel->set_context(ctx);
195 *kernel = lite_kernel;
196 return RET_OK;
197 }
198 }
199 return RET_ERROR;
200 }
201 #ifndef CUSTOM_KERNEL_REGISTRY_CLIP
202 } else {
203 auto ret = GetCustomKernel(in_tensors, out_tensors, ms_ctx, key, kernel, primitive);
204 if (ret == RET_OK) {
205 (*kernel)->set_context(ctx);
206 }
207 return ret;
208 }
209 #endif
210 return RET_NOT_SUPPORT;
211 }
212 } // namespace mindspore::lite
213