diff --git a/FEXCore/Source/Utils/atomic_bitset.h b/FEXCore/Source/Utils/atomic_bitset.h new file mode 100644 index 000000000..b4add119c --- /dev/null +++ b/FEXCore/Source/Utils/atomic_bitset.h @@ -0,0 +1,584 @@ +// SPDX-License-Identifier: MIT +#pragma once +#include +#include +#include + +#include +#include +#include +#include +#include + +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 +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(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(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(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 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::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(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(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(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(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(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(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(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(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(base[center_index]); + center_word_atomic.store(0); + } + }; + + auto set_tail_word = [&]() -> bool { + auto tail_word_atomic = std::atomic_ref(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(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(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(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(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(base[base_index + remaining_center_words + 1]); + tail_word_atomic.fetch_and(~tail_mask); + } + } +}; +} // namespace FEXCore::Utils diff --git a/FEXCore/include/FEXCore/Utils/MathUtils.h b/FEXCore/include/FEXCore/Utils/MathUtils.h index f62bd6a3c..f537b262e 100644 --- a/FEXCore/include/FEXCore/Utils/MathUtils.h +++ b/FEXCore/include/FEXCore/Utils/MathUtils.h @@ -18,6 +18,18 @@ constexpr uint64_t AlignDown(uint64_t value, uint64_t size) { return value - value % size; } +[[nodiscard]] +constexpr uint64_t AlignUpPowerOf2(uint64_t value, uint64_t size) { + LOGMAN_THROW_A_FMT(std::popcount(size) == 1, "Alignment needs to be power of 2"); + return (value + size - 1) & ~(size - 1); +} + +[[nodiscard]] +constexpr uint64_t AlignDownPowerOf2(uint64_t value, uint64_t size) { + LOGMAN_THROW_A_FMT(std::popcount(size) == 1, "Alignment needs to be power of 2"); + return value & ~(size - 1); +} + // Returns the ilog2 of a power-of-2 integer. // Asserts in the case that the passed in integer is not a power-of-2. template diff --git a/FEXCore/unittests/APITests/AtomicBitset.cpp b/FEXCore/unittests/APITests/AtomicBitset.cpp new file mode 100644 index 000000000..c9f8a885c --- /dev/null +++ b/FEXCore/unittests/APITests/AtomicBitset.cpp @@ -0,0 +1,657 @@ +// SPDX-License-Identifier: MIT +#include +#include + +#include "Utils/atomic_bitset.h" + +bool CheckMemoryIsZero(void* ptr, size_t size) { + REQUIRE(size % sizeof(uint64_t) == 0); + + auto ptr_u64 = reinterpret_cast(ptr); + for (size_t i = 0; i < (size / sizeof(uint64_t)); ++i) { + if (ptr_u64[i] != 0) { + return false; + } + } + + return true; +} + +bool CheckMemoryIsSet(void* ptr, size_t size) { + REQUIRE(size % sizeof(uint64_t) == 0); + + auto ptr_u64 = reinterpret_cast(ptr); + for (size_t i = 0; i < (size / sizeof(uint64_t)); ++i) { + if (ptr_u64[i] != ~0ULL) { + return false; + } + } + + return true; +} + +struct buffer { + uint8_t* ptr; + + uint8_t* ptr_base; + size_t size; +}; + +buffer AllocateProtectedBuffer(size_t size) { + buffer buf { + .size = size + 4096 * 2, + }; + buf.ptr_base = reinterpret_cast(mmap(nullptr, buf.size, PROT_NONE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0)); + REQUIRE(buf.ptr_base != nullptr); + + buf.ptr = buf.ptr_base + 4096; + + // RW only the pages requested. + mprotect(buf.ptr, size, PROT_READ | PROT_WRITE); + + return buf; +} + +void FreeProtectedBuffer(buffer buf) { + munmap(buf.ptr_base, buf.size); +} + +TEST_CASE("Single") { + constexpr size_t size = 4096; + constexpr size_t size_bits = size * 8; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + set.init(buf.ptr, size_bits); + + REQUIRE(set.size_in_bits() == size_bits); + + // Basic allocation check. + auto slot = set.allocate(1); + REQUIRE(slot != set.invalid()); + CHECK(slot == 0); + + set.free(slot, 1); + + CHECK(CheckMemoryIsZero(buf.ptr, size)); + + FreeProtectedBuffer(buf); +} + +TEST_CASE("All") { + constexpr size_t size = 4096; + constexpr size_t size_bits = size * 8; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + set.init(buf.ptr, size_bits); + + REQUIRE(set.size_in_bits() == size_bits); + + // Allocate all bits, ensuring all can be allocated. + for (size_t i = 0; i < size_bits; ++i) { + auto slot = set.allocate(1); + REQUIRE(slot != set.invalid()); + } + + CHECK(CheckMemoryIsSet(buf.ptr, size)); + + // Ensure that overallocation fails. + CHECK(set.allocate(1) == set.invalid()); + + // Free all the bits + for (size_t i = 0; i < size_bits; ++i) { + set.free(i, 1); + } + + CHECK(CheckMemoryIsZero(buf.ptr, size)); + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Large") { + constexpr size_t size = 4096; + constexpr size_t size_bits = size * 8; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + set.init(buf.ptr, size_bits); + + REQUIRE(set.size_in_bits() == size_bits); + + // Allocate a single 64-bit word. + auto slot = set.allocate(64); + REQUIRE(slot != set.invalid()); + CHECK(slot == 0); + + set.free(slot, 64); + + CHECK(CheckMemoryIsZero(buf.ptr, size)); + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Large Sparse") { + constexpr size_t size = 4096; + constexpr size_t size_bits = size * 8; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + set.init(buf.ptr, size_bits); + + REQUIRE(set.size_in_bits() == size_bits); + + // Allocate a single bit. + auto slot = set.allocate(1); + REQUIRE(slot != set.invalid()); + CHECK(slot == 0); + + // Allocate a single 64-bit contiguous region. + // Due to implementation behaviour, this should be a full word ahead of the previous. + auto slot64 = set.allocate(64); + REQUIRE(slot64 != set.invalid()); + CHECK(slot64 == 64); + + set.free(slot, 1); + set.free(slot64, 64); + + CHECK(CheckMemoryIsZero(buf.ptr, size)); + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Large Sparse - in-fill") { + constexpr size_t size = 4096; + constexpr size_t size_bits = size * 8; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + set.init(buf.ptr, size_bits); + + REQUIRE(set.size_in_bits() == size_bits); + + // Allocate a single bit. + auto slot = set.allocate(1); + REQUIRE(slot != set.invalid()); + CHECK(slot == 0); + + // Allocate a single 64-bit contiguous region. + // Due to implementation behaviour, this should be a full word ahead of the previous. + auto slot64 = set.allocate(64); + REQUIRE(slot64 != set.invalid()); + CHECK(slot64 == 64); + + std::vector sparse {}; + + for (size_t i = 1; i < 64; ++i) { + // Allocation of single elements should fill in sparsity. + auto new_slot = set.allocate(1); + REQUIRE(new_slot != set.invalid()); + CHECK(new_slot == i); + sparse.emplace_back(new_slot); + } + + for (auto it : sparse) { + set.free(it, 1); + } + + set.free(slot, 1); + set.free(slot64, 64); + + CHECK(CheckMemoryIsZero(buf.ptr, size)); + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Large Sparse - chunk") { + constexpr size_t size = 4096; + constexpr size_t size_bits = size * 8; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + set.init(buf.ptr, size_bits); + + REQUIRE(set.size_in_bits() == size_bits); + + // Allocate a single bit. + auto slot = set.allocate(1); + REQUIRE(slot != set.invalid()); + CHECK(slot == 0); + + // Allocate a single 64-bit contiguous region. + // Due to implementation behaviour, this should be a full word ahead of the previous. + auto slot64 = set.allocate(64); + REQUIRE(slot64 != set.invalid()); + CHECK(slot64 == 64); + + // A smaller allocation that fits within an empty word should still sub allocate. + auto slot32 = set.allocate(32); + REQUIRE(slot32 != set.invalid()); + CHECK(slot32 < slot64); + + set.free(slot, 1); + set.free(slot64, 64); + set.free(slot32, 32); + + CHECK(CheckMemoryIsZero(buf.ptr, size)); + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Large Sparse - chunk in-fill") { + constexpr size_t size = 4096; + constexpr size_t size_bits = size * 8; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + set.init(buf.ptr, size_bits); + + REQUIRE(set.size_in_bits() == size_bits); + + // Allocate a single bit. + auto slot = set.allocate(1); + REQUIRE(slot != set.invalid()); + CHECK(slot == 0); + + // Allocate a single 64-bit contiguous region. + // Due to implementation behaviour, this should be a full word ahead of the previous. + auto slot64 = set.allocate(64); + REQUIRE(slot64 != set.invalid()); + CHECK(slot64 == 64); + + // 63-bits should in-fill between the previous two allocations + auto slot63 = set.allocate(63); + REQUIRE(slot63 != set.invalid()); + CHECK(slot63 < slot64); + CHECK(slot63 == (slot + 1)); + + set.free(slot, 1); + set.free(slot64, 64); + set.free(slot63, 63); + + CHECK(CheckMemoryIsZero(buf.ptr, size)); + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Large to small") { + constexpr size_t size = 4096; + constexpr size_t size_bits = size * 8; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + set.init(buf.ptr, size_bits); + + REQUIRE(set.size_in_bits() == size_bits); + + // Allocate a single word. + auto slot = set.allocate(64); + REQUIRE(slot != set.invalid()); + CHECK(slot == 0); + + // Clearing the sub bits individually should work. + for (size_t i = 0; i < 64; ++i) { + set.free(slot + i, 1); + } + + CHECK(CheckMemoryIsZero(buf.ptr, size)); + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Small to large") { + constexpr size_t size = 4096; + constexpr size_t size_bits = size * 8; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + set.init(buf.ptr, size_bits); + + REQUIRE(set.size_in_bits() == size_bits); + + // Allocate a single word using single allocations. + auto slot0 = set.allocate(1); + REQUIRE(slot0 != set.invalid()); + CHECK(slot0 == 0); + + for (size_t i = 1; i < 64; ++i) { + auto slot = set.allocate(1); + REQUIRE(slot != set.invalid()); + CHECK(slot == i); + } + + // Freeing smaller continguous slots using a larger size should work. + set.free(slot0, 64); + + CHECK(CheckMemoryIsZero(buf.ptr, size)); + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Large Clear") { + constexpr size_t size = 4096 * 4; + constexpr size_t size_bits = size * 8; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + set.init(buf.ptr, size_bits); + + REQUIRE(set.size_in_bits() == size_bits); + + // Allocate the whole set + while (set.allocate(64) != set.invalid()) + ; + + CHECK(CheckMemoryIsSet(buf.ptr, size)); + + // Clearing the set should reset everything. + set.clear(); + + CHECK(CheckMemoryIsZero(buf.ptr, size)); + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Larger than word") { + constexpr size_t size = 4096; + constexpr size_t size_bits = size * 8; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + set.init(buf.ptr, size_bits); + + REQUIRE(set.size_in_bits() == size_bits); + + struct test_data { + uint64_t offset_count {}; + uint64_t total_size {}; + }; + + auto run_test = [&set](test_data data) { + if (data.offset_count) { + REQUIRE(set.allocate(data.offset_count) == 0); + } + + REQUIRE(set.allocate(data.total_size) == data.offset_count); + + for (size_t i = 0; i < data.offset_count; ++i) { + REQUIRE(set.is_set(i)); + } + + for (size_t i = 0; i < data.total_size; ++i) { + REQUIRE(set.is_set(data.offset_count + i)); + } + + // Reset the buffer. + set.clear(); + }; + + // Test matrix to hit all the code paths for larger than word allocations + // + // | Head offset | Head | Head+Center | Head+Center+Tail | + // | ----------- | ---- | ----------- | ---------------- | + // | Offset(0) | 🗹 | 🗹 | 🗹 | + // | Offset(1) | 🗹 | 🗹 | 🗹 | + // | Offset(63) | 🗹 | 🗹 | 🗹 | + // | Offset(511) | 🗹 | 🗹 | 🗹 | + + constexpr static test_data tests[] = { + {0, 64 + 1}, // Offset(0) + Head + Tail + {1, 63 + 2}, // Offset(1) + Head + Tail + {62, 2 + 63}, // Offset(62) + Head + Tail + {510, 2 + 63}, // Offset(510) + Head + Tail + {0, 64 + 64}, // Offset(0) + Head + Center + {1, 63 + 64}, // Offset(1) + Head + Center + {63, 1 + 64}, // Offset(63) + Head + Center + {511, 1 + 64}, // Offset(511) + Head + Center + {0, 64 + 64 + 1}, // Offset(0) + Head + Center + Tail + {1, 63 + 64 + 1}, // Offset(1) + Head + Center + Tail + {63, 1 + 64 + 1}, // Offset(63) + Head + Center + Tail + {511, 1 + 64 + 1}, // Offset(511) + Head + Center + Tail + }; + + for (auto test : tests) { + run_test(test); + } + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Larger than word - Free") { + constexpr size_t size = 4096; + constexpr size_t size_bits = size * 8; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + set.init(buf.ptr, size_bits); + + REQUIRE(set.size_in_bits() == size_bits); + + struct test_data { + uint64_t offset_count {}; + uint64_t total_size {}; + }; + + auto run_test = [&set, &buf](test_data data) { + if (data.offset_count) { + REQUIRE(set.allocate(data.offset_count) == 0); + } + + REQUIRE(set.allocate(data.total_size) == data.offset_count); + + for (size_t i = 0; i < data.offset_count; ++i) { + REQUIRE(set.is_set(i)); + } + + for (size_t i = 0; i < data.total_size; ++i) { + REQUIRE(set.is_set(data.offset_count + i)); + } + + if (data.offset_count) { + set.free(0, data.offset_count); + } + + set.free(data.offset_count, data.total_size); + + REQUIRE(CheckMemoryIsZero(buf.ptr, size)); + + // Reset the buffer. + set.clear(); + }; + + // Test matrix to hit all the code paths for larger than word allocations + // + // | Head offset | Head | Head+Center | Head+Center+Tail | + // | ----------- | ---- | ----------- | ---------------- | + // | Offset(0) | 🗹 | 🗹 | 🗹 | + // | Offset(1) | 🗹 | 🗹 | 🗹 | + // | Offset(63) | 🗹 | 🗹 | 🗹 | + constexpr static test_data tests[] = { + {0, 64 + 1}, // Offset(0) + Head + Tail + {1, 63 + 2}, // Offset(1) + Head + Tail + {62, 2 + 63}, // Offset(62) + Head + Tail + {0, 64 + 64}, // Offset(0) + Head + Center + {1, 63 + 64}, // Offset(1) + Head + Center + {63, 1 + 64}, // Offset(63) + Head + Center + {0, 64 + 64 + 1}, // Offset(0) + Head + Center + Tail + {1, 63 + 64 + 1}, // Offset(1) + Head + Center + Tail + {63, 1 + 64 + 1}, // Offset(63) + Head + Center + Tail + }; + + for (auto test : tests) { + run_test(test); + } + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Larger than Word - 128-bits") { + // Test to ensure on small bitset size, a larger than word allocation still fits. + constexpr size_t size = 4096; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + + // Move the set up to the edge of the page to detect overruns. + set.init(reinterpret_cast(buf.ptr) + (4096 - 16), 128); + + REQUIRE(set.size_in_bits() == 128); + + for (size_t i = 0; i < 128; ++i) { + auto slot = set.allocate(i + 1); + REQUIRE(slot != set.invalid()); + set.free(slot, i + 1); + } + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Race acquire - Single") { + // Test to ensure that allocating 1 slot should never fail unless it is actually full. + // Basic race condition check. + constexpr size_t size = 4096; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + + // Move the set up to the edge of the page to detect overruns. + set.init(reinterpret_cast(buf.ptr) + (4096 - 32), 256); + + REQUIRE(set.size_in_bits() == 256); + + // Allocate all 256-bits + for (size_t i = 0; i < 256; ++i) { + REQUIRE(set.allocate(1) != set.invalid()); + } + + // Free the first 128-bits + for (size_t i = 0; i < 128; ++i) { + set.free(i, 1); + } + + std::atomic Acquire {}; + std::atomic Waiters {}; + std::vector threads {}; + std::atomic slots[129] {}; + threads.reserve(128); + + auto acquire = [&](int idx) { + ++Waiters; + + // Spin until allowed to race. + while (!Acquire.load()) { + // Be nice to valgrind. + std::this_thread::yield(); + } + + auto slot = set.allocate(1); + + if (slot == set.invalid()) { + // Set to invalid slot. Should never occur. + slot = 128; + } + + // Increment the slot counter for the number of times this slot has allocated. + slots[slot]++; + }; + + for (size_t i = 0; i < 128; ++i) { + threads.emplace_back(acquire, i); + } + + // Wait until all threads are claimed to be ready. + while (Waiters.load() != 128) { + // Be nice to valgrind. + std::this_thread::yield(); + } + + Acquire = true; + + // Wait for threads to exit. + for (auto& t : threads) { + t.join(); + } + + // Every slot should only ever be acquired once. + for (size_t i = 0; i < 128; ++i) { + CHECK(slots[i].load() == 1); + } + + // There should be no invalid slots returned. + CHECK(slots[128].load() == 0); + + FreeProtectedBuffer(buf); +} + +TEST_CASE("Larger than word - rewind") { + constexpr size_t size = 4096; + auto buf = AllocateProtectedBuffer(size); + + FEXCore::Utils::atomic_bitset set {}; + + // Move the set up to the edge of the page to detect overruns. + set.init(reinterpret_cast(buf.ptr) + (4096 - 32), 256); + + REQUIRE(set.size_in_bits() == 256); + + // Allocate all 256-bits + for (size_t i = 0; i < 256; ++i) { + REQUIRE(set.allocate(1) != set.invalid()); + } + + // Free the first 128-bits + for (size_t i = 0; i < 128; ++i) { + set.free(i, 1); + } + + std::atomic Running {}; + std::atomic Stop {}; + std::thread t {[&]() { + LogMan::Msg::DFmt("Spinning"); + // Acquire and free 1-bit back to back + while (!Stop) { + auto slot = set.allocate(1); + + // Introduce some variability by yielding here. + std::this_thread::yield(); + + Running = true; + if (slot != set.invalid()) { + set.free(slot, 1); + } else { + LogMan::Msg::DFmt("We're full!"); + break; + } + } + }}; + + LogMan::Msg::DFmt("Waiting for thread to start!"); + while (!Running.load()) { + // Be nice to valgrind. + std::this_thread::yield(); + } + + LogMan::Msg::DFmt("Attempting to allocate 128-bit while contended"); + + // Try and acquire 128-bits while contended. + size_t attempts {}; + for (;;) { + auto slot = set.allocate(128); + if (slot == set.invalid()) { + ++attempts; + // Be nice to valgrind. + std::this_thread::yield(); + continue; + } + + // Can't fit in anything other than slot0 + REQUIRE(slot == 0); + LogMan::Msg::DFmt("We got 128-bits in slot: {} after {} attempts", slot, attempts); + set.free(slot, 128); + break; + } + Stop = true; + t.join(); + + FreeProtectedBuffer(buf); +}