bricks/media/pb_type_image: Use umm_malloc for image data.

dlech · dlech · commit 2a15f7204215 · 2025-07-18T17:49:32.000-05:00
Add support for a new memory allocator (umm_malloc) that is better
for allocating large blocks of memory. Since MicroPython's garbage
collector has to scan every single byte of memory it manages, to check
for pointers, it is not efficient for large blocks of memory that
are not expected to contain pointers. This is especially true for
image data.

So to avoid slow garbage collection/allocation times and filling up
the MicroPython heap, we can use a separate allocator with a separate
heap for image data.
diff --git a/.github/build-each-commit.py b/.github/build-each-commit.py
@@ -77,6 +77,8 @@
         pybricks.git.submodule(
             "update", "--init", "--checkout", "lib/STM32_USB_Device_Library"
         )
+    if args.hub == "ev3":
+        pybricks.git.submodule("update", "--init", "--checkout", "lib/umm_malloc")
 
     # build the firmware
     subprocess.check_call(
diff --git a/.github/build-missing-commits.py b/.github/build-missing-commits.py
@@ -100,6 +100,7 @@
     pybricks.git.submodule(
         "update", "--init", "--checkout", "lib/STM32_USB_Device_Library"
     )
+    pybricks.git.submodule("update", "--init", "--checkout", "lib/umm_malloc")
 
     # Make mpy-cross once
     print("Building mpy-cross")
diff --git a/.gitmodules b/.gitmodules
@@ -10,3 +10,7 @@
 	path = lib/STM32_USB_Device_Library
 	url = https://github.com/STMicroelectronics/stm32_mw_usb_device.git
 	update = none
+[submodule "lib/umm_malloc"]
+	path = lib/umm_malloc
+	url = https://github.com/rhempel/umm_malloc.git
+	update = none
diff --git a/.vscode/c_cpp_properties.json b/.vscode/c_cpp_properties.json
@@ -73,6 +73,7 @@
                 "${workspaceFolder}/lib/tiam1808/tiam1808/armv5",
                 "${workspaceFolder}/lib/tiam1808/tiam1808/armv5/am1808",
                 "${workspaceFolder}/lib/tiam1808/tiam1808/usblib",
+                "${workspaceFolder}/lib/umm_malloc/src"
             ],
             "defines": [
                 "am1808",
diff --git a/bricks/_common/arm_none_eabi.mk b/bricks/_common/arm_none_eabi.mk
@@ -70,6 +70,15 @@ $(error failed)
 endif
 endif
 endif
+ifeq ($(PB_LIB_UMM_MALLOC),1)
+ifeq ("$(wildcard $(PBTOP)/lib/umm_malloc/README.md)","")
+$(info GIT cloning umm_malloc submodule)
+$(info $(shell cd $(PBTOP) && git submodule update --checkout --init lib/umm_malloc))
+ifeq ("$(wildcard $(PBTOP)/lib/umm_malloc/README.md)","")
+$(error failed)
+endif
+endif
+endif
 
 # lets micropython make files work with external files
 USER_C_MODULES = $(PBTOP)
@@ -479,6 +488,16 @@ SRC_STM32_USB_DEV += $(addprefix lib/pbio/drv/usb/stm32_usbd/,\
 	usbd_pybricks.c \
 	)
 
+# umm_malloc library
+
+SRC_UMM_MALLOC = lib/umm_malloc/src/umm_malloc.c
+
+ifeq ($(PB_LIB_UMM_MALLOC),1)
+CFLAGS += -I$(PBTOP)/lib/umm_malloc/src
+endif
+
+# NXT OS
+
 NXOS_SRC_C = $(addprefix lib/pbio/platform/nxt/nxos/,\
 	_abort.c \
 	assert.c \
@@ -553,6 +572,10 @@ OBJ += $(addprefix $(BUILD)/, $(SRC_STM32_USB_DEV:.c=.o))
 $(BUILD)/lib/STM32_USB_Device_Library/%.o: CFLAGS += -Wno-sign-compare
 endif
 
+ifeq ($(PB_LIB_UMM_MALLOC),1)
+OBJ += $(addprefix $(BUILD)/, $(SRC_UMM_MALLOC:.c=.o))
+endif
+
 ifeq ($(PBIO_PLATFORM),nxt)
 OBJ += $(addprefix $(BUILD)/, $(NXOS_SRC_C:.c=.o))
 OBJ += $(addprefix $(BUILD)/, $(NXOS_SRC_S:.s=.o))
diff --git a/bricks/ev3/Makefile b/bricks/ev3/Makefile
@@ -1,4 +1,6 @@
 PBIO_PLATFORM = ev3
 PB_MCU_FAMILY = TIAM1808
 
+PB_LIB_UMM_MALLOC = 1
+
 include ../_common/arm_none_eabi.mk
diff --git a/lib/pbio/platform/ev3/platform.c b/lib/pbio/platform/ev3/platform.c
@@ -51,6 +51,8 @@
 #include <tiam1808/psc.h>
 #include <tiam1808/uart.h>
 
+#include <umm_malloc.h>
+
 #include <pbdrv/ioport.h>
 #include <pbio/port_interface.h>
 
@@ -720,6 +722,11 @@ void SystemInit(void) {
     } else {
         memcpy(pbdrv_ev3_bluetooth_mac_address, &i2c_buf[0], 6);
     }
+
+    // Separate heap for large allocations - defined in linker script.
+    extern char pb_umm_heap_start;
+    extern char pb_umm_heap_end;
+    umm_init_heap(&pb_umm_heap_start, &pb_umm_heap_end - &pb_umm_heap_start);
 }
 
 
diff --git a/lib/pbio/platform/ev3/platform.ld b/lib/pbio/platform/ev3/platform.ld
@@ -14,6 +14,8 @@ MEMORY
 
 _minimal_stack_size = 4M;
 pbdrv_stack_end = ORIGIN(DDR) + LENGTH(DDR) - 4;
+/* Extra heap for large allocations (images, etc). */
+pb_umm_heap_size = 2M;
 
 SECTIONS
 {
@@ -68,6 +70,11 @@ SECTIONS
     {
         . = ALIGN(4);
         *(.noinit)
+
+        . = ALIGN(4);
+        pb_umm_heap_start = .;
+        . = . + pb_umm_heap_size;
+        pb_umm_heap_end = .;
     } > DDR
 
     /* Exception vector table */
diff --git a/lib/pbio/platform/ev3/umm_malloc_cfgport.h b/lib/pbio/platform/ev3/umm_malloc_cfgport.h
@@ -0,0 +1,9 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2025 The Pybricks Authors
+//
+// Configuration for umm_malloc memory allocator for EV3 platform.
+// See lib/umm_malloc/src/umm_malloc_cfg.h for details.
+
+// We only use this for large allocations, like images and we have a lot of RAM
+// on the EV3, so we can use a larger block size to reduce overhead.
+#define UMM_BLOCK_BODY_SIZE 256
diff --git a/lib/umm_malloc b/lib/umm_malloc
@@ -0,0 +1 @@
+Subproject commit ed68d0b5f40d83b34fd14cbbd99c4fc1e8aad77d
diff --git a/pybricks/media/pb_type_image.c b/pybricks/media/pb_type_image.c
@@ -5,6 +5,12 @@
 
 #if PYBRICKS_PY_MEDIA && PYBRICKS_PY_MEDIA_IMAGE
 
+#if !MICROPY_ENABLE_FINALISER
+#error "MICROPY_ENABLE_FINALISER must be enabled."
+#endif
+
+#include <umm_malloc.h>
+
 #include "py/mphal.h"
 #include "py/obj.h"
 
@@ -23,9 +29,9 @@ extern const mp_obj_type_t pb_type_Image;
 // pybricks.media.Image class object
 typedef struct _pb_type_Image_obj_t {
     mp_obj_base_t base;
-    // TODO: avoid using MicroPython heap for large allocation, use a separate
-    // memory allocator and use a python object to handle memory release.
-    uint8_t *buffer;  // m_malloc'd buffer, may be shared, NULL for display.
+    // For images with shared memory, we need to keep a reference to the object
+    // that owns the memory.
+    mp_obj_t owner;
     pbio_image_t image;
     bool is_display;
 } pb_type_Image_obj_t;
@@ -45,7 +51,7 @@ static int get_color(mp_obj_t obj) {
 
 mp_obj_t pb_type_Image_display_obj_new(void) {
     pb_type_Image_obj_t *self = mp_obj_malloc(pb_type_Image_obj_t, &pb_type_Image);
-    self->buffer = NULL;
+    self->owner = MP_OBJ_NULL;
     self->is_display = true;
     self->image = *pbdrv_display_get_image();
 
@@ -62,34 +68,42 @@ static mp_obj_t pb_type_Image_make_new(const mp_obj_type_t *type,
         PB_ARG_DEFAULT_NONE(x2),
         PB_ARG_DEFAULT_NONE(y2));
 
-    pb_type_Image_obj_t *self = mp_obj_malloc(pb_type_Image_obj_t, type);
+    pb_type_Image_obj_t *self;
 
     if (mp_obj_is_type(source_in, &pb_type_Image)) {
         pb_type_Image_obj_t *source = MP_OBJ_TO_PTR(source_in);
         if (!mp_obj_is_true(sub_in)) {
             // Copy.
             int width = source->image.width;
             int height = source->image.height;
-            // TODO: avoid using MicroPython heap for large allocation.
-            self->buffer = m_malloc(width * height * sizeof(uint8_t));
+
+            void *buf = umm_malloc(width * height * sizeof(uint8_t));
+            if (!buf) {
+                mp_raise_type(&mp_type_MemoryError);
+            }
+
+            self = mp_obj_malloc_with_finaliser(pb_type_Image_obj_t, &pb_type_Image);
+            self->owner = MP_OBJ_NULL;
             self->is_display = false;
-            pbio_image_init(&self->image, self->buffer, width, height, width);
+            pbio_image_init(&self->image, buf, width, height, width);
             pbio_image_draw_image(&self->image, &source->image, 0, 0);
         } else {
             // Sub-image.
             mp_int_t x1 = pb_obj_get_int(x1_in);
             mp_int_t y1 = pb_obj_get_int(y1_in);
             mp_int_t x2 = x2_in == mp_const_none ? source->image.width - 1 : pb_obj_get_int(x2_in);
             mp_int_t y2 = y2_in == mp_const_none ? source->image.height - 1 : pb_obj_get_int(y2_in);
-            self->buffer = source->buffer;
+            self = mp_obj_malloc(pb_type_Image_obj_t, &pb_type_Image);
+            self->owner = source_in;
             self->is_display = false;
             int width = x2 - x1 + 1;
             int height = y2 - y1 + 1;
             pbio_image_init_sub(&self->image, &source->image, x1, y1, width, height);
         }
     } else if (mp_obj_equal(source_in, MP_OBJ_NEW_QSTR(MP_QSTR__screen_))) {
         // Screen.
-        self->buffer = NULL;
+        self = mp_obj_malloc(pb_type_Image_obj_t, &pb_type_Image);
+        self->owner = MP_OBJ_NULL;
         self->is_display = true;
         self->image = *pbdrv_display_get_image();
     } else if (mp_obj_is_str(source_in)) {
@@ -102,6 +116,17 @@ static mp_obj_t pb_type_Image_make_new(const mp_obj_type_t *type,
     return MP_OBJ_FROM_PTR(self);
 }
 
+static mp_obj_t pb_type_Image_close(mp_obj_t self_in) {
+    pb_type_Image_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    // If we own the memory, free it.
+    if (self->owner == MP_OBJ_NULL && !self->is_display && self->image.pixels) {
+        umm_free(self->image.pixels);
+        self->image.pixels = NULL;
+    }
+    return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_1(pb_type_Image_close_obj, pb_type_Image_close);
+
 static mp_obj_t pb_type_Image_empty(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
     PB_PARSE_ARGS_FUNCTION(n_args, pos_args, kw_args,
         PB_ARG_DEFAULT_NONE(width),
@@ -116,12 +141,15 @@ static mp_obj_t pb_type_Image_empty(size_t n_args, const mp_obj_t *pos_args, mp_
         mp_raise_ValueError(MP_ERROR_TEXT("Image width or height is less than 1"));
     }
 
-    pb_type_Image_obj_t *self = mp_obj_malloc(pb_type_Image_obj_t, &pb_type_Image);
+    void *buf = umm_malloc(width * height * sizeof(uint8_t));
+    if (!buf) {
+        mp_raise_type(&mp_type_MemoryError);
+    }
 
-    // TODO: avoid using MicroPython heap for large allocation.
-    self->buffer = m_malloc0(width * height * sizeof(uint8_t));
+    pb_type_Image_obj_t *self = mp_obj_malloc_with_finaliser(pb_type_Image_obj_t, &pb_type_Image);
+    self->owner = MP_OBJ_NULL;
     self->is_display = false;
-    pbio_image_init(&self->image, self->buffer, width, height, width);
+    pbio_image_init(&self->image, buf, width, height, width);
 
     return MP_OBJ_FROM_PTR(self);
 }
@@ -328,6 +356,9 @@ static MP_DEFINE_CONST_FUN_OBJ_KW(pb_type_Image_draw_circle_obj, 1, pb_type_Imag
 
 // dir(pybricks.media.Image)
 static const mp_rom_map_elem_t pb_type_Image_locals_dict_table[] = {
+    // REVISIT: consider close() method and __enter__/__exit__ for context manager
+    // to deterministically free memory if needed.
+    { MP_ROM_QSTR(MP_QSTR___del__), MP_ROM_PTR(&pb_type_Image_close_obj) },
     { MP_ROM_QSTR(MP_QSTR_empty), MP_ROM_PTR(&pb_type_Image_empty_obj) },
     { MP_ROM_QSTR(MP_QSTR_clear), MP_ROM_PTR(&pb_type_Image_clear_obj) },
     { MP_ROM_QSTR(MP_QSTR_load_image), MP_ROM_PTR(&pb_type_Image_load_image_obj) },

Original file line number	Diff line number	Diff line change
`@@ -77,6 +77,8 @@`
`77`	`77`	`pybricks.git.submodule(`
`78`	`78`	`"update", "--init", "--checkout", "lib/STM32_USB_Device_Library"`
`79`	`79`	`)`
	`80`	`+ if args.hub == "ev3":`
	`81`	`+ pybricks.git.submodule("update", "--init", "--checkout", "lib/umm_malloc")`
`80`	`82`
`81`	`83`	`# build the firmware`
`82`	`84`	`subprocess.check_call(`
Original file line number	Diff line number	Diff line change
`@@ -100,6 +100,7 @@`
`100`	`100`	`pybricks.git.submodule(`
`101`	`101`	`"update", "--init", "--checkout", "lib/STM32_USB_Device_Library"`
`102`	`102`	`)`
	`103`	`+ pybricks.git.submodule("update", "--init", "--checkout", "lib/umm_malloc")`
`103`	`104`
`104`	`105`	`# Make mpy-cross once`
`105`	`106`	`print("Building mpy-cross")`