/ source / mimalloc / src / bitmap.c
bitmap.c
  1  /* ----------------------------------------------------------------------------
  2  Copyright (c) 2019-2023 Microsoft Research, Daan Leijen
  3  This is free software; you can redistribute it and/or modify it under the
  4  terms of the MIT license. A copy of the license can be found in the file
  5  "LICENSE" at the root of this distribution.
  6  -----------------------------------------------------------------------------*/
  7  
  8  /* ----------------------------------------------------------------------------
  9  Concurrent bitmap that can set/reset sequences of bits atomically,
 10  represented as an array of fields where each field is a machine word (`size_t`)
 11  
 12  There are two api's; the standard one cannot have sequences that cross
 13  between the bitmap fields (and a sequence must be <= MI_BITMAP_FIELD_BITS).
 14  
 15  The `_across` postfixed functions do allow sequences that can cross over
 16  between the fields. (This is used in arena allocation)
 17  ---------------------------------------------------------------------------- */
 18  
 19  #include "mimalloc.h"
 20  #include "mimalloc/internal.h"
 21  #include "bitmap.h"
 22  
 23  /* -----------------------------------------------------------
 24    Bitmap definition
 25  ----------------------------------------------------------- */
 26  
 27  // The bit mask for a given number of blocks at a specified bit index.
 28  static inline size_t mi_bitmap_mask_(size_t count, size_t bitidx) {
 29    mi_assert_internal(count + bitidx <= MI_BITMAP_FIELD_BITS);
 30    mi_assert_internal(count > 0);
 31    if (count >= MI_BITMAP_FIELD_BITS) return MI_BITMAP_FIELD_FULL;
 32    if (count == 0) return 0;
 33    return ((((size_t)1 << count) - 1) << bitidx);
 34  }
 35  
 36  
 37  
 38  /* -----------------------------------------------------------
 39    Claim a bit sequence atomically
 40  ----------------------------------------------------------- */
 41  
 42  // Try to atomically claim a sequence of `count` bits in a single
 43  // field at `idx` in `bitmap`. Returns `true` on success.
 44  bool _mi_bitmap_try_find_claim_field(mi_bitmap_t bitmap, size_t idx, const size_t count, mi_bitmap_index_t* bitmap_idx)
 45  {
 46    mi_assert_internal(bitmap_idx != NULL);
 47    mi_assert_internal(count <= MI_BITMAP_FIELD_BITS);
 48    mi_bitmap_field_t* field = &bitmap[idx];
 49    size_t map  = mi_atomic_load_relaxed(field);
 50    if (map==MI_BITMAP_FIELD_FULL) return false; // short cut
 51  
 52    // search for 0-bit sequence of length count
 53    const size_t mask = mi_bitmap_mask_(count, 0);
 54    const size_t bitidx_max = MI_BITMAP_FIELD_BITS - count;
 55  
 56  #ifdef MI_HAVE_FAST_BITSCAN
 57    size_t bitidx = mi_ctz(~map);    // quickly find the first zero bit if possible
 58  #else
 59    size_t bitidx = 0;               // otherwise start at 0
 60  #endif
 61    size_t m = (mask << bitidx);     // invariant: m == mask shifted by bitidx
 62  
 63    // scan linearly for a free range of zero bits
 64    while (bitidx <= bitidx_max) {
 65      const size_t mapm = (map & m);
 66      if (mapm == 0) {  // are the mask bits free at bitidx?
 67        mi_assert_internal((m >> bitidx) == mask); // no overflow?
 68        const size_t newmap = (map | m);
 69        mi_assert_internal((newmap^map) >> bitidx == mask);
 70        if (!mi_atomic_cas_strong_acq_rel(field, &map, newmap)) {  // TODO: use weak cas here?
 71          // no success, another thread claimed concurrently.. keep going (with updated `map`)
 72          continue;
 73        }
 74        else {
 75          // success, we claimed the bits!
 76          *bitmap_idx = mi_bitmap_index_create(idx, bitidx);
 77          return true;
 78        }
 79      }
 80      else {
 81        // on to the next bit range
 82  #ifdef MI_HAVE_FAST_BITSCAN
 83        mi_assert_internal(mapm != 0);
 84        const size_t shift = (count == 1 ? 1 : (MI_INTPTR_BITS - mi_clz(mapm) - bitidx));
 85        mi_assert_internal(shift > 0 && shift <= count);
 86  #else
 87        const size_t shift = 1;
 88  #endif
 89        bitidx += shift;
 90        m <<= shift;
 91      }
 92    }
 93    // no bits found
 94    return false;
 95  }
 96  
 97  
 98  // Starts at idx, and wraps around to search in all `bitmap_fields` fields.
 99  // For now, `count` can be at most MI_BITMAP_FIELD_BITS and will never cross fields.
100  bool _mi_bitmap_try_find_from_claim(mi_bitmap_t bitmap, const size_t bitmap_fields, const size_t start_field_idx, const size_t count, mi_bitmap_index_t* bitmap_idx) {
101    size_t idx = start_field_idx;
102    for (size_t visited = 0; visited < bitmap_fields; visited++, idx++) {
103      if (idx >= bitmap_fields) { idx = 0; } // wrap
104      if (_mi_bitmap_try_find_claim_field(bitmap, idx, count, bitmap_idx)) {
105        return true;
106      }
107    }
108    return false;
109  }
110  
111  
112  // Set `count` bits at `bitmap_idx` to 0 atomically
113  // Returns `true` if all `count` bits were 1 previously.
114  bool _mi_bitmap_unclaim(mi_bitmap_t bitmap, size_t bitmap_fields, size_t count, mi_bitmap_index_t bitmap_idx) {
115    const size_t idx = mi_bitmap_index_field(bitmap_idx);
116    const size_t bitidx = mi_bitmap_index_bit_in_field(bitmap_idx);
117    const size_t mask = mi_bitmap_mask_(count, bitidx);
118    mi_assert_internal(bitmap_fields > idx); MI_UNUSED(bitmap_fields);
119    // mi_assert_internal((bitmap[idx] & mask) == mask);
120    const size_t prev = mi_atomic_and_acq_rel(&bitmap[idx], ~mask);
121    return ((prev & mask) == mask);
122  }
123  
124  
125  // Set `count` bits at `bitmap_idx` to 1 atomically
126  // Returns `true` if all `count` bits were 0 previously. `any_zero` is `true` if there was at least one zero bit.
127  bool _mi_bitmap_claim(mi_bitmap_t bitmap, size_t bitmap_fields, size_t count, mi_bitmap_index_t bitmap_idx, bool* any_zero) {
128    const size_t idx = mi_bitmap_index_field(bitmap_idx);
129    const size_t bitidx = mi_bitmap_index_bit_in_field(bitmap_idx);
130    const size_t mask = mi_bitmap_mask_(count, bitidx);
131    mi_assert_internal(bitmap_fields > idx); MI_UNUSED(bitmap_fields);
132    //mi_assert_internal(any_zero != NULL || (bitmap[idx] & mask) == 0);
133    size_t prev = mi_atomic_or_acq_rel(&bitmap[idx], mask);
134    if (any_zero != NULL) { *any_zero = ((prev & mask) != mask); }
135    return ((prev & mask) == 0);
136  }
137  
138  // Returns `true` if all `count` bits were 1. `any_ones` is `true` if there was at least one bit set to one.
139  static bool mi_bitmap_is_claimedx(mi_bitmap_t bitmap, size_t bitmap_fields, size_t count, mi_bitmap_index_t bitmap_idx, bool* any_ones) {
140    const size_t idx = mi_bitmap_index_field(bitmap_idx);
141    const size_t bitidx = mi_bitmap_index_bit_in_field(bitmap_idx);
142    const size_t mask = mi_bitmap_mask_(count, bitidx);
143    mi_assert_internal(bitmap_fields > idx); MI_UNUSED(bitmap_fields);
144    const size_t field = mi_atomic_load_relaxed(&bitmap[idx]);
145    if (any_ones != NULL) { *any_ones = ((field & mask) != 0); }
146    return ((field & mask) == mask);
147  }
148  
149  // Try to set `count` bits at `bitmap_idx` from 0 to 1 atomically. 
150  // Returns `true` if successful when all previous `count` bits were 0.
151  bool _mi_bitmap_try_claim(mi_bitmap_t bitmap, size_t bitmap_fields, size_t count, mi_bitmap_index_t bitmap_idx) {
152    const size_t idx = mi_bitmap_index_field(bitmap_idx);
153    const size_t bitidx = mi_bitmap_index_bit_in_field(bitmap_idx);
154    const size_t mask = mi_bitmap_mask_(count, bitidx);
155    mi_assert_internal(bitmap_fields > idx); MI_UNUSED(bitmap_fields);
156    size_t expected = mi_atomic_load_relaxed(&bitmap[idx]);
157    do  {    
158      if ((expected & mask) != 0) return false;
159    } 
160    while (!mi_atomic_cas_strong_acq_rel(&bitmap[idx], &expected, expected | mask));
161    mi_assert_internal((expected & mask) == 0);
162    return true;
163  }
164  
165  
166  bool _mi_bitmap_is_claimed(mi_bitmap_t bitmap, size_t bitmap_fields, size_t count, mi_bitmap_index_t bitmap_idx) {
167    return mi_bitmap_is_claimedx(bitmap, bitmap_fields, count, bitmap_idx, NULL);
168  }
169  
170  bool _mi_bitmap_is_any_claimed(mi_bitmap_t bitmap, size_t bitmap_fields, size_t count, mi_bitmap_index_t bitmap_idx) {
171    bool any_ones;
172    mi_bitmap_is_claimedx(bitmap, bitmap_fields, count, bitmap_idx, &any_ones);
173    return any_ones;
174  }
175  
176  
177  //--------------------------------------------------------------------------
178  // the `_across` functions work on bitmaps where sequences can cross over
179  // between the fields. This is used in arena allocation
180  //--------------------------------------------------------------------------
181  
182  // Try to atomically claim a sequence of `count` bits starting from the field
183  // at `idx` in `bitmap` and crossing into subsequent fields. Returns `true` on success.
184  // Only needs to consider crossing into the next fields (see `mi_bitmap_try_find_from_claim_across`)
185  static bool mi_bitmap_try_find_claim_field_across(mi_bitmap_t bitmap, size_t bitmap_fields, size_t idx, const size_t count, const size_t retries, mi_bitmap_index_t* bitmap_idx, mi_stats_t* stats)
186  {
187    mi_assert_internal(bitmap_idx != NULL);
188  
189    // check initial trailing zeros
190    mi_bitmap_field_t* field = &bitmap[idx];
191    size_t map = mi_atomic_load_relaxed(field);
192    const size_t initial = mi_clz(map);  // count of initial zeros starting at idx
193    mi_assert_internal(initial <= MI_BITMAP_FIELD_BITS);
194    if (initial == 0)     return false;
195    if (initial >= count) return _mi_bitmap_try_find_claim_field(bitmap, idx, count, bitmap_idx);    // no need to cross fields (this case won't happen for us)
196    if (_mi_divide_up(count - initial, MI_BITMAP_FIELD_BITS) >= (bitmap_fields - idx)) return false; // not enough entries
197    
198    // scan ahead
199    size_t found = initial;
200    size_t mask = 0;     // mask bits for the final field
201    while(found < count) {
202      field++;
203      map = mi_atomic_load_relaxed(field);
204      const size_t mask_bits = (found + MI_BITMAP_FIELD_BITS <= count ? MI_BITMAP_FIELD_BITS : (count - found));
205      mi_assert_internal(mask_bits > 0 && mask_bits <= MI_BITMAP_FIELD_BITS);
206      mask = mi_bitmap_mask_(mask_bits, 0);
207      if ((map & mask) != 0) return false;  // some part is already claimed
208      found += mask_bits;
209    }
210    mi_assert_internal(field < &bitmap[bitmap_fields]);
211  
212    // we found a range of contiguous zeros up to the final field; mask contains mask in the final field
213    // now try to claim the range atomically
214    mi_bitmap_field_t* const final_field = field;
215    const size_t final_mask = mask;
216    mi_bitmap_field_t* const initial_field = &bitmap[idx];
217    const size_t initial_idx = MI_BITMAP_FIELD_BITS - initial;
218    const size_t initial_mask = mi_bitmap_mask_(initial, initial_idx);
219  
220    // initial field
221    size_t newmap;
222    field = initial_field;
223    map = mi_atomic_load_relaxed(field);
224    do {
225      newmap = (map | initial_mask);
226      if ((map & initial_mask) != 0) { goto rollback; };
227    } while (!mi_atomic_cas_strong_acq_rel(field, &map, newmap));
228  
229    // intermediate fields
230    while (++field < final_field) {
231      newmap = mi_bitmap_mask_(MI_BITMAP_FIELD_BITS, 0);
232      map = 0;
233      if (!mi_atomic_cas_strong_acq_rel(field, &map, newmap)) { goto rollback; }
234    }
235  
236    // final field
237    mi_assert_internal(field == final_field);
238    map = mi_atomic_load_relaxed(field);
239    do {
240      newmap = (map | final_mask);
241      if ((map & final_mask) != 0) { goto rollback; }
242    } while (!mi_atomic_cas_strong_acq_rel(field, &map, newmap));
243  
244    // claimed!
245    mi_stat_counter_increase(stats->arena_crossover_count,1);
246    *bitmap_idx = mi_bitmap_index_create(idx, initial_idx);
247    return true;
248  
249  rollback:
250    // roll back intermediate fields
251    // (we just failed to claim `field` so decrement first)
252    while (--field > initial_field) {
253      newmap = 0;
254      map = mi_bitmap_mask_(MI_BITMAP_FIELD_BITS, 0);
255      mi_assert_internal(mi_atomic_load_relaxed(field) == map);
256      mi_atomic_store_release(field, newmap);
257    }
258    if (field == initial_field) {               // (if we failed on the initial field, `field + 1 == initial_field`)
259      map = mi_atomic_load_relaxed(field);
260      do {
261        mi_assert_internal((map & initial_mask) == initial_mask);
262        newmap = (map & ~initial_mask);
263      } while (!mi_atomic_cas_strong_acq_rel(field, &map, newmap));
264    }
265    mi_stat_counter_increase(stats->arena_rollback_count,1);
266    // retry? (we make a recursive call instead of goto to be able to use const declarations)
267    if (retries <= 2) {
268      return mi_bitmap_try_find_claim_field_across(bitmap, bitmap_fields, idx, count, retries+1, bitmap_idx, stats);
269    }
270    else {
271      return false;
272    }
273  }
274  
275  
276  // Find `count` bits of zeros and set them to 1 atomically; returns `true` on success.
277  // Starts at idx, and wraps around to search in all `bitmap_fields` fields.
278  bool _mi_bitmap_try_find_from_claim_across(mi_bitmap_t bitmap, const size_t bitmap_fields, const size_t start_field_idx, const size_t count, mi_bitmap_index_t* bitmap_idx, mi_stats_t* stats) {
279    mi_assert_internal(count > 0);
280    if (count <= 2) {
281      // we don't bother with crossover fields for small counts
282      return _mi_bitmap_try_find_from_claim(bitmap, bitmap_fields, start_field_idx, count, bitmap_idx);
283    }
284  
285    // visit the fields
286    size_t idx = start_field_idx;
287    for (size_t visited = 0; visited < bitmap_fields; visited++, idx++) {
288      if (idx >= bitmap_fields) { idx = 0; } // wrap
289      // first try to claim inside a field
290      /*
291      if (count <= MI_BITMAP_FIELD_BITS) {
292        if (_mi_bitmap_try_find_claim_field(bitmap, idx, count, bitmap_idx)) {
293          return true;
294        }
295      }
296      */
297      // if that fails, then try to claim across fields
298      if (mi_bitmap_try_find_claim_field_across(bitmap, bitmap_fields, idx, count, 0, bitmap_idx, stats)) {
299        return true;
300      }
301    }
302    return false;
303  }
304  
305  // Helper for masks across fields; returns the mid count, post_mask may be 0
306  static size_t mi_bitmap_mask_across(mi_bitmap_index_t bitmap_idx, size_t bitmap_fields, size_t count, size_t* pre_mask, size_t* mid_mask, size_t* post_mask) {
307    MI_UNUSED(bitmap_fields);
308    const size_t bitidx = mi_bitmap_index_bit_in_field(bitmap_idx);
309    if mi_likely(bitidx + count <= MI_BITMAP_FIELD_BITS) {
310      *pre_mask = mi_bitmap_mask_(count, bitidx);
311      *mid_mask = 0;
312      *post_mask = 0;
313      mi_assert_internal(mi_bitmap_index_field(bitmap_idx) < bitmap_fields);
314      return 0;
315    }
316    else {
317      const size_t pre_bits = MI_BITMAP_FIELD_BITS - bitidx;
318      mi_assert_internal(pre_bits < count);
319      *pre_mask = mi_bitmap_mask_(pre_bits, bitidx);
320      count -= pre_bits;
321      const size_t mid_count = (count / MI_BITMAP_FIELD_BITS);
322      *mid_mask = MI_BITMAP_FIELD_FULL;
323      count %= MI_BITMAP_FIELD_BITS;
324      *post_mask = (count==0 ? 0 : mi_bitmap_mask_(count, 0));
325      mi_assert_internal(mi_bitmap_index_field(bitmap_idx) + mid_count + (count==0 ? 0 : 1) < bitmap_fields);
326      return mid_count;
327    }
328  }
329  
330  // Set `count` bits at `bitmap_idx` to 0 atomically
331  // Returns `true` if all `count` bits were 1 previously.
332  bool _mi_bitmap_unclaim_across(mi_bitmap_t bitmap, size_t bitmap_fields, size_t count, mi_bitmap_index_t bitmap_idx) {
333    size_t idx = mi_bitmap_index_field(bitmap_idx);
334    size_t pre_mask;
335    size_t mid_mask;
336    size_t post_mask;
337    size_t mid_count = mi_bitmap_mask_across(bitmap_idx, bitmap_fields, count, &pre_mask, &mid_mask, &post_mask);
338    bool all_one = true;
339    mi_bitmap_field_t* field = &bitmap[idx];
340    size_t prev = mi_atomic_and_acq_rel(field++, ~pre_mask);   // clear first part
341    if ((prev & pre_mask) != pre_mask) all_one = false;
342    while(mid_count-- > 0) {
343      prev = mi_atomic_and_acq_rel(field++, ~mid_mask);        // clear mid part
344      if ((prev & mid_mask) != mid_mask) all_one = false;
345    }
346    if (post_mask!=0) {
347      prev = mi_atomic_and_acq_rel(field, ~post_mask);         // clear end part
348      if ((prev & post_mask) != post_mask) all_one = false;
349    }
350    return all_one;
351  }
352  
353  // Set `count` bits at `bitmap_idx` to 1 atomically
354  // Returns `true` if all `count` bits were 0 previously. `any_zero` is `true` if there was at least one zero bit.
355  bool _mi_bitmap_claim_across(mi_bitmap_t bitmap, size_t bitmap_fields, size_t count, mi_bitmap_index_t bitmap_idx, bool* pany_zero) {
356    size_t idx = mi_bitmap_index_field(bitmap_idx);
357    size_t pre_mask;
358    size_t mid_mask;
359    size_t post_mask;
360    size_t mid_count = mi_bitmap_mask_across(bitmap_idx, bitmap_fields, count, &pre_mask, &mid_mask, &post_mask);
361    bool all_zero = true;
362    bool any_zero = false;
363    _Atomic(size_t)*field = &bitmap[idx];
364    size_t prev = mi_atomic_or_acq_rel(field++, pre_mask);
365    if ((prev & pre_mask) != 0) all_zero = false;
366    if ((prev & pre_mask) != pre_mask) any_zero = true;
367    while (mid_count-- > 0) {
368      prev = mi_atomic_or_acq_rel(field++, mid_mask);
369      if ((prev & mid_mask) != 0) all_zero = false;
370      if ((prev & mid_mask) != mid_mask) any_zero = true;
371    }
372    if (post_mask!=0) {
373      prev = mi_atomic_or_acq_rel(field, post_mask);
374      if ((prev & post_mask) != 0) all_zero = false;
375      if ((prev & post_mask) != post_mask) any_zero = true;
376    }
377    if (pany_zero != NULL) { *pany_zero = any_zero; }
378    return all_zero;
379  }
380  
381  
382  // Returns `true` if all `count` bits were 1.
383  // `any_ones` is `true` if there was at least one bit set to one.
384  static bool mi_bitmap_is_claimedx_across(mi_bitmap_t bitmap, size_t bitmap_fields, size_t count, mi_bitmap_index_t bitmap_idx, bool* pany_ones) {
385    size_t idx = mi_bitmap_index_field(bitmap_idx);
386    size_t pre_mask;
387    size_t mid_mask;
388    size_t post_mask;
389    size_t mid_count = mi_bitmap_mask_across(bitmap_idx, bitmap_fields, count, &pre_mask, &mid_mask, &post_mask);
390    bool all_ones = true;
391    bool any_ones = false;
392    mi_bitmap_field_t* field = &bitmap[idx];
393    size_t prev = mi_atomic_load_relaxed(field++);
394    if ((prev & pre_mask) != pre_mask) all_ones = false;
395    if ((prev & pre_mask) != 0) any_ones = true;
396    while (mid_count-- > 0) {
397      prev = mi_atomic_load_relaxed(field++);
398      if ((prev & mid_mask) != mid_mask) all_ones = false;
399      if ((prev & mid_mask) != 0) any_ones = true;
400    }
401    if (post_mask!=0) {
402      prev = mi_atomic_load_relaxed(field);
403      if ((prev & post_mask) != post_mask) all_ones = false;
404      if ((prev & post_mask) != 0) any_ones = true;
405    }
406    if (pany_ones != NULL) { *pany_ones = any_ones; }
407    return all_ones;
408  }
409  
410  bool _mi_bitmap_is_claimed_across(mi_bitmap_t bitmap, size_t bitmap_fields, size_t count, mi_bitmap_index_t bitmap_idx) {
411    return mi_bitmap_is_claimedx_across(bitmap, bitmap_fields, count, bitmap_idx, NULL);
412  }
413  
414  bool _mi_bitmap_is_any_claimed_across(mi_bitmap_t bitmap, size_t bitmap_fields, size_t count, mi_bitmap_index_t bitmap_idx) {
415    bool any_ones;
416    mi_bitmap_is_claimedx_across(bitmap, bitmap_fields, count, bitmap_idx, &any_ones);
417    return any_ones;
418  }