!13327 [MS][LITE][STABLE]support dynamic generate infershape function buffer

From: @jpc_chenjianping
Reviewed-by: @zhanghaibo5,@hangangqiang
Signed-off-by: @hangangqiang
pull/13327/MERGE
mindspore-ci-bot 4 years ago committed by Gitee
commit dbc24752ff

@ -15,17 +15,24 @@
*/ */
#include "nnacl/infer/infer_register.h" #include "nnacl/infer/infer_register.h"
InferShape g_infer_func[PrimType_MAX]; InferShape *g_infer_func;
__attribute__((constructor(101))) void InitInferFuncBuf() {
g_infer_func = malloc(PrimType_MAX * sizeof(InferShape));
if (g_infer_func != NULL) {
memset(g_infer_func, 0, PrimType_MAX * sizeof(InferShape));
}
}
InferShape GetInferFunc(int prim_type) { InferShape GetInferFunc(int prim_type) {
if (prim_type < PrimType_MAX) { if (g_infer_func != NULL && prim_type < PrimType_MAX) {
return g_infer_func[prim_type]; return g_infer_func[prim_type];
} }
return NULL; return NULL;
} }
void RegInfer(int prim_type, InferShape func) { void RegInfer(int prim_type, InferShape func) {
if (prim_type < PrimType_MAX) { if (g_infer_func != NULL && prim_type < PrimType_MAX) {
g_infer_func[prim_type] = func; g_infer_func[prim_type] = func;
} }
} }

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

Loading…
Cancel
Save