This commit is contained in:
yefeng 2021-03-20 14:40:19 +08:00
parent c9f1be1244
commit d2777500ad
11 changed files with 126 additions and 3 deletions

View File

@ -14,7 +14,15 @@
* limitations under the License.
*/
#include "nnacl/infer/infer_register.h"
#ifdef MS_COMPILE_IOS
extern void _ReducePrimType_ReduceFusion();
extern void _ReshapePrimType_Reshape();
void RegisterInfer() {
_ReducePrimType_ReduceFusion();
_ReshapePrimType_Reshape();
}
#endif
InferShape *g_infer_func;
__attribute__((constructor(101))) void InitInferFuncBuf() {
@ -22,6 +30,9 @@ __attribute__((constructor(101))) void InitInferFuncBuf() {
if (g_infer_func != NULL) {
memset(g_infer_func, 0, PrimType_MAX * sizeof(InferShape));
}
#ifdef MS_COMPILE_IOS
RegisterInfer();
#endif
}
InferShape GetInferFunc(int prim_type) {

View File

@ -219,8 +219,13 @@ enum PrimType {
void RegInfer(int prim_type, InferShape func);
#ifdef MS_COMPILE_IOS
#define REG_INFER(op, type, func) \
void _##op##type() { RegInfer(type, func); }
#else
#define REG_INFER(op, type, func) \
__attribute__((constructor(102))) void Reg##op##Infer() { RegInfer(type, func); }
#endif
#ifdef __cplusplus
}
#endif

View File

@ -65,6 +65,9 @@ set(LITE_SRC
${CMAKE_CURRENT_SOURCE_DIR}/dequant.cc
${CMAKE_CURRENT_SOURCE_DIR}/huffman_decode.cc
)
if(DEFINED ARCHS)
list(APPEND LITE_SRC ${CMAKE_CURRENT_SOURCE_DIR}/reg_ops.cc)
endif()
if(SUPPORT_GPU STREQUAL opencl)
file(GLOB_RECURSE OPENCL_RUNTIME_SRC

View File

@ -19,6 +19,7 @@
#include "src/common/version_manager.h"
#include "src/common/prim_util.h"
#include "nnacl/pooling_parameter.h"
#include "src/reg_kernels.h"
#ifdef ENABLE_ARM64
#if defined(__ANDROID__)
#include <asm/hwcap.h>
@ -135,4 +136,11 @@ int KernelRegistry::GetKernel(const std::vector<Tensor *> &in_tensors, const std
}
return RET_NOT_SUPPORT;
}
#ifdef MS_COMPILE_IOS
void KernelRegistry::RegisterAllKernels() {
static std::once_flag flag_kernels;
std::call_once(flag_kernels, [&]() { kernel::RegisterKernels(); });
}
#endif
} // namespace mindspore::lite

View File

@ -45,6 +45,9 @@ class KernelRegistry {
int GetKernel(const std::vector<Tensor *> &in_tensors, const std::vector<Tensor *> &out_tensors,
const InnerContext *ctx, const kernel::KernelKey &key, OpParameter *op_parameter,
kernel::LiteKernel **kernel);
#ifdef MS_COMPILE_IOS
void RegisterAllKernels();
#endif
protected:
static const int device_type_length_{kKernelArch_MAX - kKernelArch_MIN + 1};
@ -70,8 +73,15 @@ class KernelRegistrar {
}
};
#ifdef MS_COMPILE_IOS
#define REG_KERNEL(arch, data_type, op_type, kernelCreater) \
void _##arch##data_type##op_type() { \
lite::KernelRegistry::GetInstance()->RegKernel(arch, data_type, op_type, kernelCreater); \
}
#else
#define REG_KERNEL(arch, data_type, op_type, kernelCreater) \
static KernelRegistrar g_##arch##data_type##op_type##kernelReg(arch, data_type, op_type, kernelCreater);
#endif
} // namespace mindspore::lite
#endif // MINDSPORE_LITE_SRC_KERNEL_REGISTRY_H_

View File

@ -15,7 +15,6 @@
*/
#include "src/ops/populate/populate_register.h"
#include "nnacl/matmul_parameter.h"
namespace mindspore {
namespace lite {

View File

@ -25,6 +25,9 @@
namespace mindspore {
namespace lite {
#ifdef MS_COMPILE_IOS
void RegisterOps();
#endif
typedef OpParameter *(*ParameterGen)(const void *prim);
class PopulateRegistry {
public:
@ -48,6 +51,12 @@ class PopulateRegistry {
param_creator = iter->second;
return param_creator;
}
#ifdef MS_COMPILE_IOS
void RegisterAllOps() {
static std::once_flag flag_ops;
std::call_once(flag_ops, [&]() { RegisterOps(); });
}
#endif
protected:
// key:type * 1000 + schema_version
@ -62,6 +71,16 @@ class Registry {
~Registry() = default;
};
#ifdef MS_COMPILE_IOS
#define REG_POPULATE(primitive_type, creator, version) \
void _##primitive_type##version() { \
PopulateRegistry::GetInstance()->InsertParameterMap(primitive_type, creator, version); \
}
#else
#define REG_POPULATE(primitive_type, creator, version) \
static Registry g_##primitive_type##version(primitive_type, creator, version);
#endif
} // namespace lite
} // namespace mindspore
#endif

View File

@ -16,7 +16,7 @@
#include <memory>
#include "src/ops/populate/populate_register.h"
#include "nnacl/reduce_parameter.h"
using mindspore::schema::PrimitiveType_ReduceFusion;
namespace mindspore {
namespace lite {
@ -37,7 +37,7 @@ OpParameter *PopulateReduceParameter(const void *prim) {
return reinterpret_cast<OpParameter *>(reduce_param);
}
Registry ReduceParameterRegistry(schema::PrimitiveType_ReduceFusion, PopulateReduceParameter, SCHEMA_CUR);
REG_POPULATE(PrimitiveType_ReduceFusion, PopulateReduceParameter, SCHEMA_CUR)
} // namespace lite
} // namespace mindspore

View File

@ -0,0 +1,37 @@
/**
* Copyright 2021 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef MINDSPORE_LITE_SRC_REG_KERNELS_H
#define MINDSPORE_LITE_SRC_REG_KERNELS_H
namespace mindspore {
namespace kernel {
#ifdef MS_COMPILE_IOS
// populate
extern void _kCPUkNumberTypeFloat32PrimitiveType_ReduceFusion();
extern void _kCPUkNumberTypeBoolPrimitiveType_ReduceFusion();
void RegisterKernels() {
_kCPUkNumberTypeFloat32PrimitiveType_ReduceFusion();
_kCPUkNumberTypeBoolPrimitiveType_ReduceFusion();
}
#endif
} // namespace kernel
} // namespace mindspore
#endif // MINDSPORE_LITE_SRC_REG_KERNELS_H

View File

@ -0,0 +1,27 @@
/**
* Copyright 2021 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
namespace mindspore {
namespace lite {
#ifdef MS_COMPILE_IOS
// populate
extern void _PrimitiveType_ReduceFusionSCHEMA_CUR();
void RegisterOps() { _PrimitiveType_ReduceFusionSCHEMA_CUR(); }
#endif
} // namespace lite
} // namespace mindspore

View File

@ -52,6 +52,10 @@ constexpr int kMainSubGraphIndex = 0;
} // namespace
int Scheduler::Schedule(std::vector<kernel::LiteKernel *> *dst_kernels) {
#ifdef MS_COMPILE_IOS
PopulateRegistry::GetInstance()->RegisterAllOps();
KernelRegistry::GetInstance()->RegisterAllKernels();
#endif
if (src_model_ == nullptr) {
MS_LOG(ERROR) << "Input model is nullptr";
return RET_PARAM_INVALID;