mimalloc: faster block visitor

colesbury · DinoV · commit a4a4b1d0b3ab · 2023-09-26T13:32:53.000-07:00
diff --git a/Include/mimalloc/mimalloc/internal.h b/Include/mimalloc/mimalloc/internal.h
@@ -179,6 +179,7 @@ void       _mi_heap_set_default_direct(mi_heap_t* heap);
 bool       _mi_heap_memid_is_suitable(mi_heap_t* heap, mi_memid_t memid);
 void       _mi_heap_unsafe_destroy_all(void);
 void       _mi_heap_absorb(mi_heap_t* heap, mi_heap_t* from);
+bool       _mi_abandoned_visit_blocks(int page_tag, bool visit_blocks, mi_block_visit_fun* visitor, void* arg);
 
 // "stats.c"
 void       _mi_stats_done(mi_stats_t* stats);
@@ -456,6 +457,18 @@ static inline mi_page_t* _mi_segment_page_of(const mi_segment_t* segment, const
   return mi_slice_to_page(slice);
 }
 
+static inline const mi_slice_t* mi_segment_slices_end(const mi_segment_t* segment) {
+  return &segment->slices[segment->slice_entries];
+}
+
+static inline mi_slice_t* mi_slices_start_iterate(mi_segment_t* segment, const mi_slice_t** end) {
+  mi_slice_t* slice = &segment->slices[0];
+  *end = mi_segment_slices_end(segment);
+  mi_assert_internal(slice->slice_count>0 && slice->xblock_size>0); // segment allocated page
+  slice = slice + slice->slice_count; // skip the first segment allocated page
+  return slice;
+}
+
 // Quick page start for initialized pages
 static inline uint8_t* _mi_page_start(const mi_segment_t* segment, const mi_page_t* page, size_t* page_size) {
   return _mi_segment_page_start(segment, page, page_size);
diff --git a/Objects/mimalloc/heap.c b/Objects/mimalloc/heap.c
@@ -26,13 +26,13 @@ typedef bool (heap_page_visitor_fun)(mi_heap_t* heap, mi_page_queue_t* pq, mi_pa
 // Visit all pages in a heap; returns `false` if break was called.
 static bool mi_heap_visit_pages(mi_heap_t* heap, heap_page_visitor_fun* fn, void* arg1, void* arg2)
 {
-  if (heap==NULL || heap->page_count==0) return 0;
+  if (heap==NULL || heap->page_count==0) return true;
 
   // visit all pages
   #if MI_DEBUG>1
   size_t total = heap->page_count;
   size_t count = 0;
-  #endif  
+  #endif
 
   for (size_t i = 0; i <= MI_BIN_FULL; i++) {
     mi_page_queue_t* pq = &heap->pages[i];
@@ -120,11 +120,11 @@ static void mi_heap_collect_ex(mi_heap_t* heap, mi_collect_t collect)
 {
   if (heap==NULL || !mi_heap_is_initialized(heap)) return;
 
-  const bool force = collect >= MI_FORCE;  
+  const bool force = collect >= MI_FORCE;
   _mi_deferred_free(heap, force);
 
-  // note: never reclaim on collect but leave it to threads that need storage to reclaim 
-  const bool force_main = 
+  // note: never reclaim on collect but leave it to threads that need storage to reclaim
+  const bool force_main =
     #ifdef NDEBUG
       collect == MI_FORCE
     #else
@@ -513,6 +513,17 @@ typedef struct mi_heap_area_ex_s {
   mi_page_t*     page;
 } mi_heap_area_ex_t;
 
+static void mi_fast_divisor(size_t divisor, size_t* magic, size_t* shift) {
+  mi_assert_internal(divisor > 0 && divisor <= UINT32_MAX);
+  *shift = MI_INTPTR_BITS - mi_clz(divisor - 1);
+  *magic = (size_t)(((1ULL << 32) * ((1ULL << *shift) - divisor)) / divisor + 1);
+}
+
+static size_t mi_fast_divide(size_t n, size_t magic, size_t shift) {
+  mi_assert_internal(n <= UINT32_MAX);
+  return ((((uint64_t) n * magic) >> 32) + n) >> shift;
+}
+
 static bool mi_heap_area_visit_blocks(const mi_heap_area_ex_t* xarea, mi_block_visit_fun* visitor, void* arg) {
   mi_assert(xarea != NULL);
   if (xarea==NULL) return true;
@@ -536,10 +547,31 @@ static bool mi_heap_area_visit_blocks(const mi_heap_area_ex_t* xarea, mi_block_v
     return visitor(mi_page_heap(page), area, pstart, ubsize, arg);
   }
 
-  // create a bitmap of free blocks.
+  if (page->used == page->capacity) {
+    // optimize full pages
+    uint8_t* block = pstart;
+    for (size_t i = 0; i < page->capacity; i++) {
+        if (!visitor(mi_page_heap(page), area, block, ubsize, arg)) return false;
+        block += bsize;
+    }
+    return true;
+  }
+
+  // zero-initialize a bitmap of free blocks.
   #define MI_MAX_BLOCKS   (MI_SMALL_PAGE_SIZE / sizeof(void*))
-  uintptr_t free_map[MI_MAX_BLOCKS / sizeof(uintptr_t)];
-  memset(free_map, 0, sizeof(free_map));
+  uintptr_t free_map[MI_MAX_BLOCKS / MI_INTPTR_BITS];
+  size_t bmapsize = (page->capacity + MI_INTPTR_BITS - 1) / MI_INTPTR_BITS;
+  memset(free_map, 0, bmapsize * sizeof(uintptr_t));
+
+  if (page->capacity % MI_INTPTR_BITS != 0) {
+    size_t shift = (page->capacity % MI_INTPTR_BITS);
+    uintptr_t mask = (UINTPTR_MAX << shift);
+    free_map[bmapsize-1] = mask;
+  }
+
+  // fast repeated division by an integer
+  size_t magic, shift;
+  mi_fast_divisor(bsize, &magic, &shift);
 
   #if MI_DEBUG>1
   size_t free_count = 0;
@@ -551,10 +583,11 @@ static bool mi_heap_area_visit_blocks(const mi_heap_area_ex_t* xarea, mi_block_v
     mi_assert_internal((uint8_t*)block >= pstart && (uint8_t*)block < (pstart + psize));
     size_t offset = (uint8_t*)block - pstart;
     mi_assert_internal(offset % bsize == 0);
-    size_t blockidx = offset / bsize;  // Todo: avoid division?
-    mi_assert_internal( blockidx < MI_MAX_BLOCKS);
-    size_t bitidx = (blockidx / sizeof(uintptr_t));
-    size_t bit = blockidx - (bitidx * sizeof(uintptr_t));
+    size_t blockidx = mi_fast_divide(offset, magic, shift);
+    mi_assert_internal(blockidx == offset / bsize);
+    mi_assert_internal(blockidx < MI_MAX_BLOCKS);
+    size_t bitidx = (blockidx / MI_INTPTR_BITS);
+    size_t bit = blockidx - (bitidx * MI_INTPTR_BITS);
     free_map[bitidx] |= ((uintptr_t)1 << bit);
   }
   mi_assert_internal(page->capacity == (free_count + page->used));
@@ -563,19 +596,29 @@ static bool mi_heap_area_visit_blocks(const mi_heap_area_ex_t* xarea, mi_block_v
   #if MI_DEBUG>1
   size_t used_count = 0;
   #endif
-  for (size_t i = 0; i < page->capacity; i++) {
-    size_t bitidx = (i / sizeof(uintptr_t));
-    size_t bit = i - (bitidx * sizeof(uintptr_t));
-    uintptr_t m = free_map[bitidx];
-    if (bit == 0 && m == UINTPTR_MAX) {
-      i += (sizeof(uintptr_t) - 1); // skip a run of free blocks
+  uint8_t* block = pstart;
+  for (size_t i = 0; i < bmapsize; i++) {
+    if (free_map[i] == 0) {
+      // every block is in use
+      for (size_t j = 0; j < MI_INTPTR_BITS; j++) {
+        #if MI_DEBUG>1
+        used_count++;
+        #endif
+        if (!visitor(mi_page_heap(page), area, block, ubsize, arg)) return false;
+        block += bsize;
+      }
     }
-    else if ((m & ((uintptr_t)1 << bit)) == 0) {
-      #if MI_DEBUG>1
-      used_count++;
-      #endif
-      uint8_t* block = pstart + (i * bsize);
-      if (!visitor(mi_page_heap(page), area, block, ubsize, arg)) return false;
+    else {
+      uintptr_t m = ~free_map[i];
+      while (m) {
+        #if MI_DEBUG>1
+        used_count++;
+        #endif
+        size_t bitidx = mi_ctz(m);
+        if (!visitor(mi_page_heap(page), area, block + (bitidx * bsize), ubsize, arg)) return false;
+        m &= m - 1;
+      }
+      block += bsize * MI_INTPTR_BITS;
     }
   }
   mi_assert_internal(page->used == used_count);
@@ -584,21 +627,24 @@ static bool mi_heap_area_visit_blocks(const mi_heap_area_ex_t* xarea, mi_block_v
 
 typedef bool (mi_heap_area_visit_fun)(const mi_heap_t* heap, const mi_heap_area_ex_t* area, void* arg);
 
+static void mi_heap_xarea_init(mi_heap_area_ex_t* xarea, mi_page_t* page) {
+  const size_t bsize = mi_page_block_size(page);
+  const size_t ubsize = mi_page_usable_block_size(page);
+  xarea->page = page;
+  xarea->area.reserved = page->reserved * bsize;
+  xarea->area.committed = page->capacity * bsize;
+  xarea->area.blocks = _mi_page_start(_mi_page_segment(page), page, NULL);
+  xarea->area.used = page->used;   // number of blocks in use (#553)
+  xarea->area.block_size = ubsize;
+  xarea->area.full_block_size = bsize;
+}
 
 static bool mi_heap_visit_areas_page(mi_heap_t* heap, mi_page_queue_t* pq, mi_page_t* page, void* vfun, void* arg) {
   MI_UNUSED(heap);
   MI_UNUSED(pq);
   mi_heap_area_visit_fun* fun = (mi_heap_area_visit_fun*)vfun;
   mi_heap_area_ex_t xarea;
-  const size_t bsize = mi_page_block_size(page);
-  const size_t ubsize = mi_page_usable_block_size(page);
-  xarea.page = page;
-  xarea.area.reserved = page->reserved * bsize;
-  xarea.area.committed = page->capacity * bsize;
-  xarea.area.blocks = _mi_page_start(_mi_page_segment(page), page, NULL);
-  xarea.area.used = page->used;   // number of blocks in use (#553)
-  xarea.area.block_size = ubsize;
-  xarea.area.full_block_size = bsize;
+  mi_heap_xarea_init(&xarea, page);
   return fun(heap, &xarea, arg);
 }
 
@@ -629,5 +675,55 @@ static bool mi_heap_area_visitor(const mi_heap_t* heap, const mi_heap_area_ex_t*
 // Visit all blocks in a heap
 bool mi_heap_visit_blocks(const mi_heap_t* heap, bool visit_blocks, mi_block_visit_fun* visitor, void* arg) {
   mi_visit_blocks_args_t args = { visit_blocks, visitor, arg };
+  _mi_heap_delayed_free_partial((mi_heap_t *)heap);
   return mi_heap_visit_areas(heap, &mi_heap_area_visitor, &args);
 }
+
+typedef bool (segment_page_visitor_fun)(mi_segment_t* segment, mi_page_t* page, void* arg);
+
+// Visit all pages in a heap; returns `false` if break was called.
+static bool mi_segment_visit_pages(int page_tag, mi_segment_t *segment, segment_page_visitor_fun* fn, void* arg)
+{
+  while (segment) {
+    const mi_slice_t* end;
+    mi_slice_t* slice = mi_slices_start_iterate(segment, &end);
+    while (slice < end) {
+      if (slice->xblock_size > 0) {
+        mi_page_t* const page = mi_slice_to_page(slice);
+        if (page->tag == page_tag) {
+          if (!fn(segment, page, arg)) return false;
+        }
+      }
+      slice = slice + slice->slice_count;
+    }
+    segment = segment->abandoned_next;
+  }
+  return true;
+}
+
+// Visit all pages in a heap; returns `false` if break was called.
+static bool mi_abandoned_visit_pages(int page_tag, segment_page_visitor_fun* fn, void* arg)
+{
+  if (!mi_segment_visit_pages(page_tag, _mi_segment_abandoned(), fn, arg)) return false;
+  return mi_segment_visit_pages(page_tag, _mi_segment_abandoned_visited(), fn, arg);
+}
+
+static bool mi_segment_visitor(mi_segment_t* segment, mi_page_t* page, void* arg)
+{
+  mi_visit_blocks_args_t* args = (mi_visit_blocks_args_t*)arg;
+  mi_heap_area_ex_t xarea;
+  mi_heap_xarea_init(&xarea, page);
+  if (!args->visitor(NULL, &xarea.area, NULL, xarea.area.block_size, args->arg)) return false;
+  if (args->visit_blocks) {
+    return mi_heap_area_visit_blocks(&xarea, args->visitor, args->arg);
+  }
+  else {
+    return true;
+  }
+}
+
+// Visit all blocks in a abandoned segments
+bool _mi_abandoned_visit_blocks(int page_tag, bool visit_blocks, mi_block_visit_fun* visitor, void* arg) {
+  mi_visit_blocks_args_t args = { visit_blocks, visitor, arg };
+  return mi_abandoned_visit_pages(page_tag, &mi_segment_visitor, &args);
+}
diff --git a/Objects/mimalloc/segment.c b/Objects/mimalloc/segment.c
@@ -157,11 +157,6 @@ size_t _mi_commit_mask_next_run(const mi_commit_mask_t* cm, size_t* idx) {
    Slices
 ----------------------------------------------------------- */
 
-
-static const mi_slice_t* mi_segment_slices_end(const mi_segment_t* segment) {
-  return &segment->slices[segment->slice_entries];
-}
-
 static uint8_t* mi_slice_start(const mi_slice_t* slice) {
   mi_segment_t* segment = _mi_ptr_segment(slice);
   mi_assert_internal(slice >= segment->slices && slice < mi_segment_slices_end(segment));
@@ -1271,14 +1266,6 @@ void _mi_segment_page_abandon(mi_page_t* page, mi_segments_tld_t* tld) {
   Reclaim abandoned pages
 ----------------------------------------------------------- */
 
-static mi_slice_t* mi_slices_start_iterate(mi_segment_t* segment, const mi_slice_t** end) {
-  mi_slice_t* slice = &segment->slices[0];
-  *end = mi_segment_slices_end(segment);
-  mi_assert_internal(slice->slice_count>0 && slice->xblock_size>0); // segment allocated page
-  slice = slice + slice->slice_count; // skip the first segment allocated page
-  return slice;
-}
-
 // Possibly free pages and check if free space is available
 static bool mi_segment_check_free(mi_segment_t* segment, size_t slices_needed, size_t block_size, int tag, mi_segments_tld_t* tld)
 {
