FEXCore: Adds a lock-free atomic bitset that supports contiguous range allocations

This thing is a bit intense, so some requirements from the start:
- It needs to be lock-free and thread-safe
- It needs to support contiguous range allocations
- It needs to support allocations larger than a single atomic word

These requirements kind of fly in the face of most bitset allocators
where they will support some parts of these requirements, or just throw
a mutex in front of the whole thing.

Some implementation details:
- If allocating only 1-bit, trivial and always succeeds if there is space
- If allocating <= 64-bit, then always succeeds if there is at least
  those many contiguous bits within a single atomic word
  - Allocation can fail if there are cross-word contiguous bits of the
    size available
  - Introduces some sparsity
- If allocating > 64-bits then it falls down the longer scan path.
  - Searches for contiguous bits of free space between multiple atomic
    words.
  - If found, will attempt to allocate tracking which bits were allocated
  - If allocation fails, unwind bits already acquired and continue
    scanning

Some downsides to this implementation:
- Allocations can fail if sparsity builds up
- Heavily contended allocations can be worse than a lock
  - If larger than atomic word allocations are in flight.
- Unwinding larger than word allocations and continuing scanning adds
  overhead, a lock would have won at that point.
- A small bit of false sharing where an atomic word is read without
  acquire semantics for scanning can technically overlook some
  allocations that no longer exist.
- Slower than a linear allocator, but that's not unexpected.

Most of these downsides are okay for our use case, which is code buffer
allocations with the ability to do partial invalidation. If the atomic
bitset fails to fit an allocation, we can throw away the code buffer
like we currently do.

The bitmap allocator that uses this lock-free atomic bitset is still
in-flight but this is one complex container that can land independently.

FEXCore/Source/Utils/atomic_bitset.h

@@ -0,0 +1,584 @@
+// SPDX-License-Identifier: MIT
+#pragma once
+#include <FEXCore/Utils/Allocator.h>
+#include <FEXCore/Utils/LogManager.h>
+#include <FEXCore/Utils/MathUtils.h>
+
+#include <atomic>
+#include <bit>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+
+namespace FEXCore::Utils {
+/**
+ * A bitset that supports allocating contiguous ranges atomically.
+ * - Lock-free, with the caveat that on contention for > 64-bit it would be faster to acquire a lock.
+ *   - If low-contention then atomic-behaviour wins.
+ * - Three modes of allocation:
+ *   - 1-bit, single-atomic.
+ *   - <= 64-bit, single-atomic, contained within single word (introduces sparsity).
+ *   - > 64-bit, multiple-atomic, contiguous, roll-back on contiguous allocation failure.
+ * - Can return failure to allocate even if there is space in certain circumstances.
+ *   - If the allocated size crosses multiple words.
+ *   - Race to allocation caused contention.
+ * - Remembers last allocation/free for inner-word allocations to improve performance.
+ *   - Large greater than atomic-word scans always scan from the start.
+ * - Resetting the bitset with clear() is lower cost when page_size=true.
+ *   - MADV_DONTNEED replaces pages with zero-page
+ *   - When page_size=false, basic memset is also fairly quick.
+ */
+template<bool track_last_allocation = false, bool page_sized = true>
+class atomic_bitset final {
+public:
+  void init(void* ptr, size_t bits) {
+    LOGMAN_THROW_A_FMT(bits != 0, "Can't init zero");
+
+    base = reinterpret_cast<uint64_t*>(ptr);
+    bits_to_track = bits;
+    words_to_track = bits / WORD_SIZE_BITS;
+    LOGMAN_THROW_A_FMT(bits % WORD_SIZE_BITS == 0, "Bits to track must match uint64_t");
+    if constexpr (page_sized) {
+      LOGMAN_THROW_A_FMT(bits % (4096 * 8) == 0, "Bits to track must match bit count in page");
+    }
+
+    last_allocation_track.set_last_allocation(0);
+  }
+
+  // Allocate a contiguous buffer of bits.
+  // Returns initial bit offset on success, ~0ULL on failure.
+  size_t allocate(size_t count) {
+    LOGMAN_THROW_A_FMT(count != 0, "Can't allocate zero");
+    LOGMAN_THROW_A_FMT(count <= bits_to_track, "Can't allocate larger than size");
+
+    if (count == 1) [[likely]] {
+      // Common and trivial case.
+      return allocate_one(last_allocation_track.get_last_allocation(), words_to_track);
+    } else if (count <= WORD_SIZE_BITS) [[likely]] {
+      // Allocate up to a single word. Don't allow cross-word allocations
+      // Could cause some sparsity
+      return allocate_inside_word(count, last_allocation_track.get_last_allocation(), words_to_track);
+    }
+
+    // TODO: Always scans from beginning to end.
+    // Support iterative scanning.
+    return allocate_large_amount(count, 0, words_to_track);
+  }
+
+  // Frees a contiguous set of bits.
+  void free(size_t index, size_t count) {
+    LOGMAN_THROW_A_FMT(count != 0, "Can't free zero");
+    LOGMAN_THROW_A_FMT(index < bits_to_track, "Can't free beyond end");
+
+    if (count == 1) [[likely]] {
+      free_one(index);
+      return;
+    } else if ((index % WORD_SIZE_BITS + count) <= WORD_SIZE_BITS) [[likely]] {
+      free_inside_word(index, count);
+      return;
+    }
+
+    free_large_amount(index, count);
+  }
+
+  // Clears the entire bitset.
+  // Not thread safe!
+  void clear() {
+    const size_t bytes = words_to_track * sizeof(uint64_t);
+    if constexpr (page_sized) {
+      // VirtualDontNeed replaces pages with zero page.
+      FEXCore::Allocator::VirtualDontNeed(base, bytes);
+    } else {
+      memset(base, 0, bytes);
+    }
+
+    last_allocation_track.set_last_allocation(0);
+  }
+
+  // Checks if a single bit is set.
+  bool is_set(size_t index) const {
+    const size_t word_index = index / WORD_SIZE_BITS;
+    const size_t word_offset = index % WORD_SIZE_BITS;
+    auto word_atomic = std::atomic_ref<uint64_t>(base[word_index]);
+
+    const uint64_t bit_mask = 1ULL << word_offset;
+    return (word_atomic.load() & bit_mask) != 0;
+  }
+
+  size_t size_in_bits() const {
+    return bits_to_track;
+  }
+
+  constexpr static size_t invalid() {
+    return ~0ULL;
+  }
+
+  // Debug interface
+  // non-atomically returns the number of set bits in the bitset.
+  size_t popcount() const {
+    size_t count {};
+
+    // Just ensure all store are visible.
+    std::atomic_thread_fence(std::memory_order_release);
+
+    for (size_t word_index = 0; word_index < words_to_track; ++word_index) {
+      auto word_atomic = std::atomic_ref<uint64_t>(base[word_index]);
+      count += std::popcount(word_atomic.load(std::memory_order_relaxed));
+    }
+    return count;
+  }
+
+private:
+  uint64_t* base {};
+  size_t bits_to_track {};
+  size_t words_to_track {};
+  struct data_to_track_nop {
+    constexpr static size_t get_last_allocation() {
+      return 0;
+    }
+    constexpr static void set_last_allocation(size_t) {}
+  };
+
+  struct data_to_track {
+    std::atomic<size_t> last_allocation_word {};
+
+    constexpr size_t get_last_allocation() const {
+      // It's okay if this isn't up to date, full scan of the region still occurs.
+      return last_allocation_word.load(std::memory_order_relaxed);
+    }
+    constexpr void set_last_allocation(size_t word) {
+      last_allocation_word = word;
+    }
+  };
+  using data_type = typename std::conditional<track_last_allocation, data_to_track, data_to_track_nop>::type;
+  data_type last_allocation_track {};
+
+  constexpr static size_t WORD_SIZE_BITS = sizeof(uint64_t) * 8;
+
+  size_t allocate_one(size_t beginning_word_index, size_t ending_word_index) {
+    // Trivial spin.
+    for (size_t i = beginning_word_index; i < ending_word_index; ++i) {
+      auto word_atomic = std::atomic_ref<uint64_t>(base[i]);
+      auto expected_word = word_atomic.load(std::memory_order_relaxed);
+
+      if (expected_word == ~0ULL) {
+        // Won't pass.
+        continue;
+      }
+
+      // Spin on the word trying to acquire a bit.
+      // Uncontended case should immediately succeed.
+      // Contended case can spin the whole word and lose every acquire.
+
+      do {
+        const auto zero_bit = std::countr_one(expected_word);
+        const auto bit_mask = 1ULL << zero_bit;
+
+        // If mask was already set, then we raced to acquire (returned value will be 1).
+        // If mask not set, then we will have acquired (returned value will be 0).
+        expected_word = word_atomic.fetch_or(bit_mask);
+        if ((expected_word & bit_mask) == 0) {
+          // Acquired the bit, return the offset.
+          last_allocation_track.set_last_allocation(i);
+          return i * WORD_SIZE_BITS + zero_bit;
+        }
+
+        // Bit was already acquired.
+        expected_word |= bit_mask;
+      } while (expected_word != ~0ULL);
+    }
+
+    if constexpr (track_last_allocation) {
+      if (beginning_word_index) {
+        // One more chance to get an allocation.
+        // Scan before the previous allocation to see if any free slots have appeared.
+        return allocate_one(0, beginning_word_index);
+      }
+    }
+
+    // Failure to acquire here.
+    return invalid();
+  }
+
+  size_t allocate_inside_word(size_t count, size_t beginning_word_index, size_t ending_word_index) {
+    for (size_t i = beginning_word_index; i < ending_word_index; ++i) {
+      auto word_atomic = std::atomic_ref<uint64_t>(base[i]);
+      auto expected_word = word_atomic.load(std::memory_order_relaxed);
+
+      // Spin on the word trying to acquire a bit.
+      // Uncontended case should immediately succeed.
+      // Contended case can spin the whole word and lose every acquire.
+      while (expected_word != ~0ULL) {
+        auto zero_bit = std::countr_one(expected_word);
+        bool fits = false;
+        uint64_t bit_mask = count == WORD_SIZE_BITS ? ~0ULL : ((1ULL << count) - 1);
+        for (; (zero_bit + count) <= WORD_SIZE_BITS; ++zero_bit) {
+          // Check if the bits fit.
+          uint64_t tmp_bit_mask = bit_mask << zero_bit;
+          if ((expected_word & tmp_bit_mask) == 0) {
+            fits = true;
+            break;
+          }
+        }
+
+        // Could never fit, early exit.
+        if (!fits) {
+          break;
+        }
+
+        // Shift bit_mask to the desired location.
+        bit_mask <<= zero_bit;
+
+        uint64_t desired_word {};
+        bool acquired = true;
+
+        do {
+          if (expected_word & bit_mask) {
+            // Couldn't acquire this field, move to the next.
+            acquired = false;
+            break;
+          }
+
+          // We desire setting a single bit.
+          desired_word = expected_word | bit_mask;
+        } while (!word_atomic.compare_exchange_strong(expected_word, desired_word));
+
+        if (acquired) {
+          // Acquired the bit, return the offset.
+          last_allocation_track.set_last_allocation(i);
+          return i * WORD_SIZE_BITS + zero_bit;
+        }
+      }
+    }
+
+    if constexpr (track_last_allocation) {
+      if (beginning_word_index) {
+        // One more chance to get an allocation.
+        // Scan before the previous allocation to see if any free slots have appeared.
+        return allocate_inside_word(count, 0, beginning_word_index);
+      }
+    }
+
+    // Failure to acquire here.
+    return invalid();
+  }
+
+  size_t allocate_large_amount(size_t count, size_t beginning_word_index, size_t ending_word_index) {
+    // This version of the code needs to explicitly deal with large allocations that can't fit in a word.
+    // So we are always scanning minimum 2 words.
+    const size_t num_words_to_scan = FEXCore::AlignUpPowerOf2(count, WORD_SIZE_BITS) / WORD_SIZE_BITS;
+    const size_t last_word_to_scan = ending_word_index - num_words_to_scan - 1;
+
+    // Scan forward to find the first word.
+    for (size_t base_index = beginning_word_index; base_index < last_word_to_scan;) {
+      uint64_t leading_zeros {};
+
+      size_t center_word_index = 1;
+      bool has_center {};
+
+      size_t tail_bits {};
+
+      size_t remaining_bits = count;
+
+      auto check_head_fitment = [&]() -> bool {
+        auto base_word_atomic = std::atomic_ref<uint64_t>(base[base_index]);
+        auto base_expected_word = base_word_atomic.load(std::memory_order_relaxed);
+
+        // Count the leading zeros, if it is above zero then we can start here.
+        leading_zeros = std::countl_zero(base_expected_word);
+
+        if (leading_zeros == 0) {
+          // Nope.
+          ++base_index;
+          return false;
+        }
+
+        return true;
+      };
+
+      auto check_center_fitment = [&]() -> bool {
+        // Subtract the number of zeros.
+        remaining_bits -= leading_zeros;
+
+        has_center = remaining_bits >= WORD_SIZE_BITS;
+
+        while (remaining_bits >= WORD_SIZE_BITS) {
+          // All words in-between head and tail must be zero.
+          auto center_word_atomic = std::atomic_ref<uint64_t>(base[base_index + center_word_index]);
+          auto center_expected_word = center_word_atomic.load(std::memory_order_relaxed);
+
+          if (center_expected_word != 0) {
+            // Couldn't fit, won't ever fit in this range, so jump ahead to the last scanned item.
+            // Might still be able to start on the tail of this word.
+            base_index += center_word_index;
+            return false;
+          }
+
+          ++center_word_index;
+          remaining_bits -= WORD_SIZE_BITS;
+        }
+
+        return true;
+      };
+
+      auto check_tail_fitment = [&]() -> bool {
+        tail_bits = remaining_bits;
+        if (tail_bits) {
+          // Now for the tail (if it is necessary).
+          size_t tail_index = center_word_index;
+
+          // Count the trailing zeros, if it fits out remaining bits then we can try and allocate.
+          auto tail_word_atomic = std::atomic_ref<uint64_t>(base[base_index + tail_index]);
+          auto tail_expected_word = tail_word_atomic.load(std::memory_order_relaxed);
+
+          const auto trailing_zeros = std::countr_zero(tail_expected_word);
+
+          if (trailing_zeros < remaining_bits) {
+            // Couldn't fit, but also won't ever fit in this range. Jump ahead to this tail item.
+            // Might still be able to start on the tail of this word.
+            base_index += tail_index;
+            return false;
+          }
+        }
+
+        return true;
+      };
+
+      if (!check_head_fitment()) {
+        continue;
+      }
+
+      if (!check_center_fitment()) {
+        continue;
+      }
+
+      if (!check_tail_fitment()) {
+        continue;
+      }
+
+      // We can try fitting!
+      if (attempt_allocate_range_from_base(base_index, count, leading_zeros, tail_bits, has_center)) {
+        const size_t head_leading_offset = (WORD_SIZE_BITS - leading_zeros);
+        return base_index * WORD_SIZE_BITS + head_leading_offset;
+      }
+
+      // Failure to fit here means we could never fit in this full range. Jump past all the bits
+      // Rescanning the tail to avoid fragmented sparsity on the tails.
+      base_index += num_words_to_scan - 1;
+    }
+
+    return invalid();
+  }
+
+  bool attempt_allocate_range_from_base(size_t base_index, size_t count, size_t head_bits, size_t tail_bits, bool has_center) {
+    const bool has_tail = tail_bits != 0;
+
+    const uint64_t head_bit_mask = head_bits == WORD_SIZE_BITS ? ~0ULL : (((1ULL << head_bits) - 1) << (WORD_SIZE_BITS - head_bits));
+    const uint64_t tail_bit_mask = (1ULL << tail_bits) - 1;
+    const uint64_t center_words_count = (count - head_bits - tail_bits) / WORD_SIZE_BITS;
+    const uint64_t center_words_base_index = base_index + 1;
+    const uint64_t tail_word_base_index = center_words_base_index + center_words_count;
+
+    // Three distinct sections, all of which need to support rewinding.
+    // - Head: Setting the leading zeros to one
+    //   - Always exists. Can be a full word, or partial.
+    // - Center: Setting all in-between words to ~0ULL
+    //   - Might not exist if tail is smaller than a word
+    //   - Always full words if it does exist.
+    // - Tail: Set all trailing zeros up to the size to 1
+    //   - Might not exist if Center perfectly aligned to word edge.
+    //   - Always partial words, otherwise it would be considered "Center".
+
+    bool set_head {true};
+    bool set_center {true};
+    bool set_tail {true};
+    size_t num_center_set {};
+
+    // Head first.
+    auto set_head_word = [&]() -> bool {
+      auto head_word_atomic = std::atomic_ref<uint64_t>(base[base_index]);
+      auto head_expected_word = head_word_atomic.load(std::memory_order_relaxed);
+
+      uint64_t desired_word {};
+      do {
+        if (head_expected_word & head_bit_mask) {
+          // Another thread raced and allocated.
+          return false;
+        }
+
+        // Set the whole mask in one atomic operation.
+        desired_word = head_expected_word | head_bit_mask;
+      } while (!head_word_atomic.compare_exchange_strong(head_expected_word, desired_word));
+
+      return true;
+    };
+
+    auto clear_head_word = [&]() {
+      auto head_word_atomic = std::atomic_ref<uint64_t>(base[base_index]);
+      head_word_atomic.fetch_and(~head_bit_mask);
+    };
+
+    auto set_center_words = [&]() -> bool {
+      const size_t end_center_word_index = center_words_base_index + center_words_count;
+      for (size_t center_index = center_words_base_index; center_index < end_center_word_index; ++center_index) {
+        auto center_word_atomic = std::atomic_ref<uint64_t>(base[center_index]);
+        auto center_expected_word = center_word_atomic.load(std::memory_order_relaxed);
+
+        // Center words set a full ~0ULL mask.
+        const uint64_t desired_word {~0ULL};
+        do {
+          if (center_expected_word) {
+            // Another thread raced and allocated.
+            return false;
+          }
+        } while (!center_word_atomic.compare_exchange_strong(center_expected_word, desired_word));
+
+        ++num_center_set;
+      }
+
+      return true;
+    };
+
+    auto clear_center_words = [&]() {
+      const size_t end_center_word_index = center_words_base_index + num_center_set;
+      for (size_t center_index = center_words_base_index; center_index < end_center_word_index; ++center_index) {
+        auto center_word_atomic = std::atomic_ref<uint64_t>(base[center_index]);
+        center_word_atomic.store(0);
+      }
+    };
+
+    auto set_tail_word = [&]() -> bool {
+      auto tail_word_atomic = std::atomic_ref<uint64_t>(base[tail_word_base_index]);
+      auto tail_expected_word = tail_word_atomic.load(std::memory_order_relaxed);
+
+      uint64_t desired_word {};
+      do {
+        if (tail_expected_word & tail_bit_mask) {
+          // Another thread raced and allocated.
+          return false;
+        }
+
+        // Set the whole mask in one atomic operation.
+        desired_word = tail_expected_word | tail_bit_mask;
+      } while (!tail_word_atomic.compare_exchange_strong(tail_expected_word, desired_word));
+
+      return true;
+    };
+
+    set_head = set_head_word();
+
+    // Do the center if it exists.
+    if (set_head && has_center) {
+      set_center = set_center_words();
+    }
+
+    // Do the tail if it exists.
+    if (set_head && set_center && has_tail) {
+      set_tail = set_tail_word();
+    }
+
+    if (set_head && set_center && set_tail) {
+      return true;
+    }
+
+    // Some stage failed to set, rewind everything.
+    if (set_head) {
+      // Clear the head if it was set
+      clear_head_word();
+    }
+
+    if (has_center && num_center_set) {
+      // `set_center` might not be set, but it still managed to set some of the words.
+      clear_center_words();
+    }
+
+    // Tail doesn't need to rewind as it will never have been set if we got here.
+    return false;
+  }
+
+#if defined(__ARM_FEATURE_ATOMICS) && __ARM_FEATURE_ATOMICS == 1
+  // Might violate memory-ordering requirements?
+  // TODO: Verify and enable or delete depending.
+  // Provides an 11% (Cortex-X4) to 25% (AmpereOneA) performance improvement.
+  constexpr static bool use_stclr {};
+  static inline void stclr(uint64_t value, uint64_t* addr) {
+    asm volatile("stclrl %[Val], [%[addr]];" ::[Val] "r"(value), [addr] "r"(addr) : "memory");
+  }
+#endif
+
+  void free_one(size_t index) {
+    const size_t word_index = index / WORD_SIZE_BITS;
+    const size_t word_offset = index % WORD_SIZE_BITS;
+    last_allocation_track.set_last_allocation(word_index);
+
+    const uint64_t bic_bit_mask = 1ULL << word_offset;
+#if defined(__ARM_FEATURE_ATOMICS) && __ARM_FEATURE_ATOMICS == 1
+    if constexpr (use_stclr) {
+      stclr(bic_bit_mask, &base[word_index]);
+      return;
+    }
+#endif
+
+    auto word_atomic = std::atomic_ref<uint64_t>(base[word_index]);
+    word_atomic.fetch_and(~bic_bit_mask);
+  }
+
+  void free_inside_word(size_t index, size_t count) {
+    const size_t word_index = index / WORD_SIZE_BITS;
+    const size_t word_offset = index % WORD_SIZE_BITS;
+    auto word_atomic = std::atomic_ref<uint64_t>(base[word_index]);
+    last_allocation_track.set_last_allocation(word_index);
+
+    if (count == WORD_SIZE_BITS) {
+      word_atomic.store(0);
+      return;
+    }
+
+    const uint64_t bic_bit_mask = ((1ULL << count) - 1) << word_offset;
+#if defined(__ARM_FEATURE_ATOMICS) && __ARM_FEATURE_ATOMICS == 1
+    if constexpr (use_stclr) {
+      stclr(bic_bit_mask, &base[word_index]);
+      return;
+    }
+#endif
+    word_atomic.fetch_and(~bic_bit_mask);
+  }
+
+  void free_large_amount(size_t index, size_t count) {
+    // Incoming count can be less than WORD_SIZE_BITS if it is unaligned and crossing multiple words.
+    // Needs to always handle at minimum a head plus center and/or tail arrangement.
+    const size_t base_index = index / WORD_SIZE_BITS;
+    const size_t last_index = FEXCore::AlignUpPowerOf2(index + count, WORD_SIZE_BITS) / WORD_SIZE_BITS;
+    const size_t num_words_to_scan = last_index - base_index;
+    LOGMAN_THROW_A_FMT(num_words_to_scan > 1, "Needs to be larger than 1 ({}, {})", index, count);
+
+    size_t remaining_bits = count;
+
+    const uint64_t head_offset_start = index % WORD_SIZE_BITS;
+    const uint64_t head_bits = WORD_SIZE_BITS - head_offset_start;
+
+    uint64_t head_mask = head_bits == WORD_SIZE_BITS ? ~0ULL : (((1ULL << head_bits) - 1) << head_offset_start);
+
+    remaining_bits -= head_bits;
+    auto head_word_atomic = std::atomic_ref<uint64_t>(base[base_index]);
+    head_word_atomic.fetch_and(~head_mask);
+
+    const size_t remaining_center_words = remaining_bits / WORD_SIZE_BITS;
+    for (size_t i = 0; i < remaining_center_words; ++i) {
+      // Handle centers if they exist.
+      auto center_word_atomic = std::atomic_ref<uint64_t>(base[base_index + i + 1]);
+      center_word_atomic.store(0);
+      remaining_bits -= WORD_SIZE_BITS;
+    }
+
+    if (remaining_bits) {
+      // Handle tail if they exist, must always be less than WORD_SIZE_BITS.
+      LOGMAN_THROW_A_FMT(remaining_bits < WORD_SIZE_BITS, "Too large");
+      const uint64_t tail_mask = (1ULL << remaining_bits) - 1;
+
+      auto tail_word_atomic = std::atomic_ref<uint64_t>(base[base_index + remaining_center_words + 1]);
+      tail_word_atomic.fetch_and(~tail_mask);
+    }
+  }
+};
+} // namespace FEXCore::Utils