Browse Source

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

From: @jpc_chenjianping
Reviewed-by: @zhanghaibo5,@hangangqiang
Signed-off-by: @hangangqiang
tags/v1.2.0-rc1
mindspore-ci-bot Gitee 4 years ago
parent
commit
dbc24752ff
2 changed files with 11 additions and 4 deletions
  1. +10
    -3
      mindspore/lite/nnacl/infer/infer_register.c
  2. +1
    -1
      mindspore/lite/nnacl/infer/infer_register.h

+ 10
- 3
mindspore/lite/nnacl/infer/infer_register.c View File

@@ -15,17 +15,24 @@
*/
#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) {
if (prim_type < PrimType_MAX) {
if (g_infer_func != NULL && prim_type < PrimType_MAX) {
return g_infer_func[prim_type];
}
return NULL;
}

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;
}
}

+ 1
- 1
mindspore/lite/nnacl/infer/infer_register.h View File

@@ -219,7 +219,7 @@ enum PrimType {
void RegInfer(int prim_type, InferShape 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
}
#endif


Loading…
Cancel
Save