modules/py_ml: Clean up fb_alloc path.

- Call fb_alloc_free_till_mark only if memory was actually allocated.

Signed-off-by: iabdalkader <i.abdalkader@gmail.com>
This commit is contained in:
iabdalkader 2024-12-13 07:11:39 +01:00
parent 6dde96953d
commit abc528241d

View File

@ -335,14 +335,12 @@ mp_obj_t py_ml_model_make_new(const mp_obj_type_t *type, size_t n_args, size_t n
mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)]; mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
mp_arg_parse_all_kw_array(n_args, n_kw, all_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args); mp_arg_parse_all_kw_array(n_args, n_kw, all_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
fb_alloc_mark();
const char *path = mp_obj_str_get_str(args[ARG_path].u_obj); const char *path = mp_obj_str_get_str(args[ARG_path].u_obj);
(void) path; (void) path;
py_ml_model_obj_t *model = mp_obj_malloc_with_finaliser(py_ml_model_obj_t, &py_ml_model_type); py_ml_model_obj_t *model = mp_obj_malloc_with_finaliser(py_ml_model_obj_t, &py_ml_model_type);
model->data = NULL; model->data = NULL;
model->fb_alloc = args[ARG_load_to_fb].u_int; model->fb_alloc = false;
model->labels = mp_const_none; model->labels = mp_const_none;
#if MICROPY_PY_ML_TFLM #if MICROPY_PY_ML_TFLM
@ -374,8 +372,12 @@ mp_obj_t py_ml_model_make_new(const mp_obj_type_t *type, size_t n_args, size_t n
FIL fp; FIL fp;
file_open(&fp, path, false, FA_READ | FA_OPEN_EXISTING); file_open(&fp, path, false, FA_READ | FA_OPEN_EXISTING);
model->size = f_size(&fp); model->size = f_size(&fp);
model->fb_alloc = args[ARG_load_to_fb].u_int;
if (model->fb_alloc) { if (model->fb_alloc) {
fb_alloc_mark();
model->data = fb_alloc(model->size, FB_ALLOC_PREFER_SPEED | FB_ALLOC_CACHE_ALIGN); model->data = fb_alloc(model->size, FB_ALLOC_PREFER_SPEED | FB_ALLOC_CACHE_ALIGN);
// The model's data will Not be free'd on exceptions.
fb_alloc_mark_permanent();
} else { } else {
// Keeps a reference to the GC block. // Keeps a reference to the GC block.
model->_raw = xalloc(model->size + IMLIB_ML_MODEL_ALIGN); model->_raw = xalloc(model->size + IMLIB_ML_MODEL_ALIGN);
@ -388,13 +390,6 @@ mp_obj_t py_ml_model_make_new(const mp_obj_type_t *type, size_t n_args, size_t n
#endif #endif
} }
if (model->fb_alloc) {
// The model's data will Not be free'd on exceptions.
fb_alloc_mark_permanent();
} else {
fb_alloc_free_till_mark();
}
ml_backend_init_model(model); ml_backend_init_model(model);
return MP_OBJ_FROM_PTR(model); return MP_OBJ_FROM_PTR(model);
} }