mimalloc: changes to support separate heaps

These are changes to support separate heaps for Python objects,
Python objects with GC header, and non Python objects.

Author: colesbury

Include/mimalloc/mimalloc-internal.h

@@ -81,6 +81,8 @@ void       _mi_segment_page_free(mi_page_t* page, bool force, mi_segments_tld_t*
 void       _mi_segment_page_abandon(mi_page_t* page, mi_segments_tld_t* tld);
 uint8_t*   _mi_segment_page_start(const mi_segment_t* segment, const mi_page_t* page, size_t block_size, size_t* page_size, size_t* pre_size); // page start for any page
 void       _mi_segment_huge_page_free(mi_segment_t* segment, mi_page_t* page, mi_block_t* block);
+mi_segment_t* _mi_segment_abandoned(void);
+mi_segment_t* _mi_segment_abandoned_visited(void);
 
 void       _mi_segment_thread_collect(mi_segments_tld_t* tld);
 void       _mi_abandoned_reclaim_all(mi_heap_t* heap, mi_segments_tld_t* tld);

Include/mimalloc/mimalloc-types.h

@@ -58,7 +58,9 @@ terms of the MIT license. A copy of the license can be found in the file
 // Encoded free lists allow detection of corrupted free lists
 // and can detect buffer overflows, modify after free, and double `free`s.
 #if (MI_SECURE>=3 || MI_DEBUG>=1 || defined(MI_PADDING))
-#define MI_ENCODE_FREELIST  1
+// TODO(sgross): Don't encode free-list because it breaks the constraint that
+// freed blocks do not have the LSB of the first word set.
+//#define MI_ENCODE_FREELIST  1
 #endif
 
 
@@ -207,6 +209,7 @@ typedef struct mi_page_s {
   uint8_t               is_reset:1;        // `true` if the page memory was reset
   uint8_t               is_committed:1;    // `true` if the page virtual memory is committed
   uint8_t               is_zero_init:1;    // `true` if the page was zero initialized
+  unsigned int          tag:4;
 
   // layout like this to optimize access in `mi_malloc` and `mi_free`
   uint16_t              capacity;          // number of blocks committed, must be the first field, see `segment.c:page_clear`
@@ -331,6 +334,8 @@ struct mi_heap_s {
   size_t                page_count;                          // total number of pages in the `pages` queues.
   mi_heap_t*            next;                                // list of heaps per thread
   bool                  no_reclaim;                          // `true` if this heap should not reclaim abandoned pages
+  unsigned char         tag;
+  bool                  visited;                                     // used by gcmodule.c
 };
 
 
@@ -471,6 +476,8 @@ struct mi_tld_s {
   bool                recurse;       // true if deferred was called; used to prevent infinite recursion.
   mi_heap_t*          heap_backing;  // backing heap of this thread (cannot be deleted)
   mi_heap_t*          heaps;         // list of heaps in this thread (so we can abandon all when the thread terminates)
+  mi_heap_t*          heap_obj;
+  mi_heap_t*          heap_gc;
   mi_segments_tld_t   segments;      // segment tld
   mi_os_tld_t         os;            // os tld
   mi_stats_t          stats;         // statistics

Include/mimalloc/mimalloc.h

@@ -178,6 +178,8 @@ mi_decl_export void       mi_heap_destroy(mi_heap_t* heap);
 mi_decl_export mi_heap_t* mi_heap_set_default(mi_heap_t* heap);
 mi_decl_export mi_heap_t* mi_heap_get_default(void);
 mi_decl_export mi_heap_t* mi_heap_get_backing(void);
+mi_decl_export mi_heap_t* mi_heap_get_obj(void);
+mi_decl_export mi_heap_t* mi_heap_get_gc(void);
 mi_decl_export void       mi_heap_collect(mi_heap_t* heap, bool force) mi_attr_noexcept;
 
 mi_decl_nodiscard mi_decl_export mi_decl_restrict void* mi_heap_malloc(mi_heap_t* heap, size_t size) mi_attr_noexcept mi_attr_malloc mi_attr_alloc_size(2);
@@ -228,6 +230,13 @@ mi_decl_nodiscard mi_decl_export void* mi_heap_recalloc_aligned(mi_heap_t* heap,
 mi_decl_nodiscard mi_decl_export void* mi_heap_recalloc_aligned_at(mi_heap_t* heap, void* p, size_t newcount, size_t size, size_t alignment, size_t offset) mi_attr_noexcept mi_attr_alloc_size2(3,4);
 
 
+typedef enum mi_heap_tag_e {
+  mi_heap_tag_default,
+  mi_heap_tag_obj,
+  mi_heap_tag_gc,
+  _mi_heap_tag_last
+} mi_heap_tag_t;
+
 // ------------------------------------------------------
 // Analysis
 // ------------------------------------------------------

Objects/mimalloc/heap.c

@@ -189,6 +189,16 @@ mi_heap_t* mi_heap_get_backing(void) {
   return bheap;
 }
 
+mi_heap_t* mi_heap_get_obj(void) {
+  mi_heap_t* def = mi_heap_get_default();
+  return def->tld->heap_obj;
+}
+
+mi_heap_t* mi_heap_get_gc(void) {
+  mi_heap_t* def = mi_heap_get_default();
+  return def->tld->heap_gc;
+}
+
 mi_heap_t* mi_heap_new(void) {
   mi_heap_t* bheap = mi_heap_get_backing();
   mi_heap_t* heap = mi_heap_malloc_tp(bheap, mi_heap_t);  // todo: OS allocate in secure mode?
@@ -329,7 +339,7 @@ void mi_heap_destroy(mi_heap_t* heap) {
 ----------------------------------------------------------- */
 
 // Tranfer the pages from one heap to the other
-static void mi_heap_absorb(mi_heap_t* heap, mi_heap_t* from) {
+void mi_heap_absorb(mi_heap_t* heap, mi_heap_t* from) {
   mi_assert_internal(heap!=NULL);
   if (from==NULL || from->page_count == 0) return;
 

Objects/mimalloc/init.c

@@ -13,6 +13,7 @@ terms of the MIT license. A copy of the license can be found in the file
 // Empty page used to initialize the small free pages array
 const mi_page_t _mi_page_empty = {
   0, false, false, false, false,
+  0,       // tag
   0,       // capacity
   0,       // reserved capacity
   { 0 },   // flags
@@ -98,6 +99,8 @@ const mi_heap_t _mi_heap_empty = {
   { {0}, {0}, 0 },
   0,                // page count
   NULL,             // next
+  false,
+  0,
   false
 };
 
@@ -110,9 +113,15 @@ mi_decl_thread mi_heap_t* _mi_heap_default = (mi_heap_t*)&_mi_heap_empty;
 
 extern mi_heap_t _mi_heap_main;
 
+// Initialized in mi_heap_init
+static mi_heap_t _mi_heap_main_obj;
+static mi_heap_t _mi_heap_main_gc;
+
 static mi_tld_t tld_main = {
   0, false,
   &_mi_heap_main, &_mi_heap_main,
+  &_mi_heap_main_obj,
+  &_mi_heap_main_gc,
   { { NULL, NULL }, {NULL ,NULL}, {NULL ,NULL, 0},
     0, 0, 0, 0, 0, 0, NULL,
     tld_main_stats, tld_main_os
@@ -163,9 +172,22 @@ mi_heap_t* _mi_heap_main_get(void) {
 // note: in x64 in release build `sizeof(mi_thread_data_t)` is under 4KiB (= OS page size).
 typedef struct mi_thread_data_s {
   mi_heap_t  heap;  // must come first due to cast in `_mi_heap_done`
+  mi_heap_t  heap_obj;
+  mi_heap_t  heap_gc;
   mi_tld_t   tld;
 } mi_thread_data_t;
 
+static void _mi_heap_init_ex(mi_heap_t* heap, mi_tld_t* tld, int tag) {
+    memcpy(heap, &_mi_heap_empty, sizeof(*heap));
+    heap->thread_id = _mi_thread_id();
+    _mi_random_init(&heap->random);
+    heap->cookie  = _mi_heap_random_next(heap) | 1;
+    heap->keys[0] = _mi_heap_random_next(heap);
+    heap->keys[1] = _mi_heap_random_next(heap);
+    heap->tld = tld;
+    heap->tag = tag;
+}
+
 // Initialize the thread local default heap, called from `mi_thread_init`
 static bool _mi_heap_init(void) {
   if (mi_heap_is_initialized(mi_get_default_heap())) return true;
@@ -185,24 +207,23 @@ static bool _mi_heap_init(void) {
     }
     // OS allocated so already zero initialized
     mi_tld_t*  tld = &td->tld;
-    mi_heap_t* heap = &td->heap;
-    memcpy(heap, &_mi_heap_empty, sizeof(*heap));
-    heap->thread_id = _mi_thread_id();
-    _mi_random_init(&heap->random);
-    heap->cookie  = _mi_heap_random_next(heap) | 1;
-    heap->keys[0] = _mi_heap_random_next(heap);
-    heap->keys[1] = _mi_heap_random_next(heap);
-    heap->tld = tld;
-    tld->heap_backing = heap;
-    tld->heaps = heap;
+    _mi_heap_init_ex(&td->heap, &td->tld, mi_heap_tag_default);
+    _mi_heap_init_ex(&td->heap_obj, &td->tld, mi_heap_tag_obj);
+    _mi_heap_init_ex(&td->heap_gc, &td->tld, mi_heap_tag_gc);
+    tld->heap_backing = &td->heap;
+    tld->heaps = &td->heap;
+    tld->heap_obj = &td->heap_obj;
+    tld->heap_gc = &td->heap_gc;
     tld->segments.stats = &tld->stats;
     tld->segments.os = &tld->os;
     tld->os.stats = &tld->stats;
-    _mi_heap_set_default_direct(heap);
+    _mi_heap_set_default_direct(&td->heap);
   }
   return false;
 }
 
+void mi_heap_absorb(mi_heap_t* heap, mi_heap_t* from);
+
 // Free the thread local default heap (called from `mi_thread_done`)
 static bool _mi_heap_done(mi_heap_t* heap) {
   if (!mi_heap_is_initialized(heap)) return true;
@@ -230,6 +251,8 @@ static bool _mi_heap_done(mi_heap_t* heap) {
 
   // collect if not the main thread
   if (heap != &_mi_heap_main) {
+    mi_heap_absorb(heap, heap->tld->heap_obj);
+    mi_heap_absorb(heap, heap->tld->heap_gc);
     _mi_heap_collect_abandon(heap);
   }
 
@@ -461,6 +484,8 @@ void mi_process_init(void) mi_attr_noexcept {
   _mi_verbose_message("process init: 0x%zx\n", _mi_thread_id());
   _mi_os_init();
   mi_heap_main_init();
+  _mi_heap_init_ex(&_mi_heap_main_obj, &tld_main, mi_heap_tag_obj);
+  _mi_heap_init_ex(&_mi_heap_main_gc, &tld_main, mi_heap_tag_gc);
   #if (MI_DEBUG)
   _mi_verbose_message("debug level : %d\n", MI_DEBUG);
   #endif

Objects/mimalloc/page.c

@@ -234,6 +234,7 @@ void _mi_page_reclaim(mi_heap_t* heap, mi_page_t* page) {
   mi_assert_internal(mi_page_thread_free_flag(page) != MI_NEVER_DELAYED_FREE);
   mi_assert_internal(_mi_page_segment(page)->page_kind != MI_PAGE_HUGE);
   mi_assert_internal(!page->is_reset);
+  mi_assert_internal(page->tag == heap->tag);
   // TODO: push on full queue immediately if it is full?
   mi_page_queue_t* pq = mi_page_queue(heap, mi_page_block_size(page));
   mi_page_queue_push(heap, pq, page);
@@ -613,6 +614,7 @@ static void mi_page_init(mi_heap_t* heap, mi_page_t* page, size_t block_size, mi
   page->keys[1] = _mi_heap_random_next(heap);
   #endif
   page->is_zero = page->is_zero_init;
+  page->tag = heap->tag;
 
   mi_assert_internal(page->capacity == 0);
   mi_assert_internal(page->free == NULL);

Objects/mimalloc/segment.c

@@ -989,6 +989,16 @@ static mi_segment_t* mi_abandoned_pop(void) {
    Abandon segment/page
 ----------------------------------------------------------- */
 
+extern mi_segment_t* _mi_segment_abandoned(void) {
+  mi_tagged_segment_t ts = mi_atomic_read(&abandoned);
+  mi_segment_t *segment = mi_tagged_segment_ptr(ts);
+  return segment;
+}
+
+extern mi_segment_t* _mi_segment_abandoned_visited(void) {
+  return mi_atomic_read_ptr(mi_segment_t, &abandoned_visited);
+}
+
 static void mi_segment_abandon(mi_segment_t* segment, mi_segments_tld_t* tld) {
   mi_assert_internal(segment->used == segment->abandoned);
   mi_assert_internal(segment->used > 0);
@@ -1066,6 +1076,21 @@ static bool mi_segment_check_free(mi_segment_t* segment, size_t block_size, bool
   return has_page;
 }
 
+static mi_heap_t* mi_heap_from_tag(mi_heap_t* base, unsigned int tag)
+{
+  if (tag == base->tag) {
+    return base;
+  } else if (tag == mi_heap_tag_default) {
+    return base->tld->heap_backing;
+  } else if (tag == mi_heap_tag_obj) {
+    return base->tld->heap_obj;
+  } else if (tag == mi_heap_tag_gc) {
+    return base->tld->heap_gc;
+  } else {
+    _mi_error_message(EINVAL, "unknown page tag: %u\n", tag);
+    return NULL;
+  }
+}
 
 // Reclaim a segment; returns NULL if the segment was freed
 // set `right_page_reclaimed` to `true` if it reclaimed a page of the right `block_size` that was not full.
@@ -1083,6 +1108,7 @@ static mi_segment_t* mi_segment_reclaim(mi_segment_t* segment, mi_heap_t* heap,
   for (size_t i = 0; i < segment->capacity; i++) {
     mi_page_t* page = &segment->pages[i];
     if (page->segment_in_use) {
+      mi_heap_t* target_heap = mi_heap_from_tag(heap, page->tag);
       mi_assert_internal(!page->is_reset);
       mi_assert_internal(page->is_committed);
       mi_assert_internal(mi_page_not_in_queue(page, tld));
@@ -1092,7 +1118,7 @@ static mi_segment_t* mi_segment_reclaim(mi_segment_t* segment, mi_heap_t* heap,
       mi_assert(page->next == NULL);
       _mi_stat_decrease(&tld->stats->pages_abandoned, 1);
       // set the heap again and allow heap thread delayed free again.
-      mi_page_set_heap(page, heap);
+      mi_page_set_heap(page, target_heap);
       _mi_page_use_delayed_free(page, MI_USE_DELAYED_FREE, true); // override never (after heap is set)
       // TODO: should we not collect again given that we just collected in `check_free`?
       _mi_page_free_collect(page, false); // ensure used count is up to date
@@ -1102,8 +1128,9 @@ static mi_segment_t* mi_segment_reclaim(mi_segment_t* segment, mi_heap_t* heap,
       }
       else {
         // otherwise reclaim it into the heap
-        _mi_page_reclaim(heap, page);
-        if (requested_block_size == page->xblock_size && mi_page_has_any_available(page)) {
+        _mi_page_reclaim(target_heap, page);
+        if (heap == target_heap &&
+            requested_block_size == page->xblock_size && mi_page_has_any_available(page)) {
           if (right_page_reclaimed != NULL) { *right_page_reclaimed = true; }
         }
       }