diff --git a/CLAUDE.md b/CLAUDE.md index 0750d2e..1745f2c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -148,7 +148,7 @@ Performance optimization of cuCascade's disk I/O backends (GDS and kvikIO) to ap #include #include // cuDF #include // RMM -#include // system with dot +#include // system with dot #include // STL #include ## Namespace Usage diff --git a/include/cucascade/exec/spin_lock.hpp b/include/cucascade/exec/spin_lock.hpp new file mode 100644 index 0000000..d69048c --- /dev/null +++ b/include/cucascade/exec/spin_lock.hpp @@ -0,0 +1,234 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) Meta Platforms, Inc. and affiliates. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * --------------------------------------------------------------------------- + * ATTRIBUTION + * + * This file is a derivative work adapted from folly's MicroSpinLock, originally + * authored by Meta Platforms, Inc. and affiliates and licensed under the + * Apache License, Version 2.0. The original sources are: + * - folly/synchronization/MicroSpinLock.h (MicroSpinLock, SpinLockArray, MSLGuard) + * - folly/synchronization/detail/Sleeper.h (detail::Sleeper) + * Upstream: https://github.com/facebook/folly + * + * It has been minimally adapted to cucascade naming conventions and made + * self-contained so it carries no folly dependency (the folly-specific + * ThreadSanitizer annotations and portability shims were removed). The locking + * logic is a 1-to-1 map of the folly original. + * --------------------------------------------------------------------------- + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cucascade::exec { + +namespace detail { + +//===----------------------------------------------------------------------===// +// Portable CPU relax / pause primitive. +//===----------------------------------------------------------------------===// + +/** + * @brief Emit an architecture-specific "pause"/"yield" hint inside a spin loop. + * + * On x86 this maps to the PAUSE instruction, on AArch64 to YIELD. On unknown + * architectures it degrades to a no-op. + */ +inline void spin_cpu_relax() noexcept +{ +#if defined(__x86_64__) || defined(__i386__) + __builtin_ia32_pause(); +#elif defined(__aarch64__) + asm volatile("yield" ::: "memory"); +#else + // No architecture-specific relax hint available; fall through as a no-op. +#endif +} + +//===----------------------------------------------------------------------===// +// Sleeper +//===----------------------------------------------------------------------===// + +#if defined(__aarch64__) +inline constexpr bool k_is_arch_aarch64 = true; +#else +inline constexpr bool k_is_arch_aarch64 = false; +#endif + +/** + * @brief A helper object for the contended case. + * + * Starts off with eager spinning, and falls back to sleeping for small + * quantums. On AArch64 it additionally applies exponential back-off between + * pause hints. + */ +class sleeper { + const std::chrono::nanoseconds delta; + + static constexpr uint32_t k_max_active_spin = 4096; + static constexpr bool use_back_off = k_is_arch_aarch64; + + uint32_t spin_count = 0; + uint32_t spin_count_target = 1; + + public: + static constexpr std::chrono::nanoseconds k_min_yielding_sleep = std::chrono::microseconds(500); + + constexpr sleeper() noexcept : delta(k_min_yielding_sleep) {} + + explicit sleeper(std::chrono::nanoseconds d) noexcept : delta(d) {} + + void wait() noexcept + { + bool do_spin = + use_back_off ? spin_count_target <= k_max_active_spin : spin_count < k_max_active_spin; + if (do_spin) { + if constexpr (use_back_off) { + do { + spin_cpu_relax(); + } while (++spin_count < spin_count_target); + spin_count_target <<= 1; + } else { + ++spin_count; + spin_cpu_relax(); + } + } else { + /* sleep override */ + std::this_thread::sleep_for(delta); + } + } +}; + +} // namespace detail + +//===----------------------------------------------------------------------===// +// spin_lock +//===----------------------------------------------------------------------===// + +/** + * @brief A really, *really* small spinlock for fine-grained locking of lots of + * teeny-tiny data. + * + * Zero initializing these is guaranteed to be as good as calling init(), since + * the free state is guaranteed to be all-bits zero. + * + * This class should be kept a POD, so we can use it in other packed structs + * (gcc does not allow __attribute__((__packed__)) on structs that contain + * non-POD data). This means avoid adding a constructor, or making some members + * private, etc. + */ +struct spin_lock { + enum { FREE = 0, LOCKED = 1 }; + // lock_ can't be std::atomic<> to preserve POD-ness. + uint8_t lock_; + + // Initialize this spin_lock. It is unnecessary to call this if you + // zero-initialize the spin_lock. + void init() noexcept { payload()->store(FREE); } + + bool try_lock() noexcept { return xchg_acquire(LOCKED) == FREE; } + + void lock() noexcept + { + detail::sleeper sleeper; + while (xchg_acquire(LOCKED) != FREE) { + do { + sleeper.wait(); + } while (payload()->load(std::memory_order_relaxed) == LOCKED); + } + assert(payload()->load() == LOCKED); + } + + void unlock() noexcept + { + assert(payload()->load() == LOCKED); + payload()->store(FREE, std::memory_order_release); + } + + private: + std::atomic* payload() noexcept + { + return reinterpret_cast*>(&this->lock_); + } + + uint8_t xchg_acquire(uint8_t new_val) noexcept + { + return std::atomic_exchange_explicit(payload(), new_val, std::memory_order_acquire); + } +}; +static_assert(std::is_standard_layout::value && std::is_trivial::value, + "spin_lock must be kept a POD type."); + +//===----------------------------------------------------------------------===// +// spin_lock_array +//===----------------------------------------------------------------------===// + +/** + * @brief Array of spinlocks where each one is padded to prevent false sharing. + * + * Useful for shard-based locking implementations in environments where + * contention is unlikely. + */ +template +struct alignas(alignof(std::max_align_t)) spin_lock_array { + // Conservative cache-line estimate; kept as a fixed constant (rather than + // std::hardware_destructive_interference_size) to avoid ABI-instability + // warnings under -Werror. + static constexpr std::size_t destructive_interference_size = 64; + static constexpr std::size_t max_align = alignof(std::max_align_t); + + T& operator[](std::size_t i) noexcept { return data_[i].lock; } + + const T& operator[](std::size_t i) const noexcept { return data_[i].lock; } + + constexpr std::size_t size() const noexcept { return N; } + + private: + struct padded_spin_lock { + padded_spin_lock() : lock() {} + T lock; + char padding[destructive_interference_size - sizeof(T)]; + }; + static_assert(sizeof(padded_spin_lock) == destructive_interference_size, + "Invalid size of padded_spin_lock"); + + // Check if T can theoretically cross a cache line. + static_assert(max_align > 0 && destructive_interference_size % max_align == 0 && + sizeof(T) <= max_align, + "T can cross cache line boundaries"); + + char padding_[destructive_interference_size]; + std::array data_; +}; + +//===----------------------------------------------------------------------===// +// spin_lock_guard +//===----------------------------------------------------------------------===// + +using spin_lock_guard = std::lock_guard; + +} // namespace cucascade::exec diff --git a/include/cucascade/io/cache/config.hpp b/include/cucascade/io/cache/config.hpp index d29ea70..620e556 100644 --- a/include/cucascade/io/cache/config.hpp +++ b/include/cucascade/io/cache/config.hpp @@ -23,7 +23,10 @@ namespace cucascade::io::cache { struct config { - size_t inflight_io_chunk_budget = 2048; + // Maximum number of in-flight prefetch IO *tasks* (not chunks): the prefetch + // loop reserves one unit per dispatched read task and releases it on + // completion, so this caps how many prefetch reads are outstanding at once. + size_t inflight_io_chunk_budget = 16; double min_prefetching_budget_fraction{0.05}; double eviction_threshold_fraction{0.6}; bool dispose_after_use = false; diff --git a/include/cucascade/io/cache/prefetching_cache.hpp b/include/cucascade/io/cache/prefetching_cache.hpp index f532069..144d06a 100644 --- a/include/cucascade/io/cache/prefetching_cache.hpp +++ b/include/cucascade/io/cache/prefetching_cache.hpp @@ -183,7 +183,14 @@ class prefetching_cache { const io_object& obj, size_t offset, size_t size, uint8_t* dst, prefetching_handle* out_handle); struct file_entry { - std::vector update_and_get_chunks(std::span incoming, uint32_t ticker); + // @p desired_cache_from carries, in parallel with @p incoming, the signed + // page-aligned cache_from each incoming chunk offset should be populated + // with (see cached_chunk::cache_from / needed_cache_from). On first + // creation of a chunk the value is stored directly; on an existing chunk it + // is merged in (merge_cache_from) when the chunk's state permits. + std::vector update_and_get_chunks(std::span incoming, + std::span desired_cache_from, + uint32_t ticker); std::vector fetch_chunks(std::size_t offset, std::size_t size, diff --git a/include/cucascade/io/cache/types.hpp b/include/cucascade/io/cache/types.hpp index c2ccc82..24f73e6 100644 --- a/include/cucascade/io/cache/types.hpp +++ b/include/cucascade/io/cache/types.hpp @@ -22,6 +22,7 @@ // virtual interface (device_read_async_io_using). Extracted here to break the // circular include between io_context.hpp and prefetching_cache.hpp. +#include #include #include #include @@ -46,6 +47,25 @@ namespace cucascade::io::cache { +// --------------------------------------------------------------------------- +// Page-alignment helpers +// --------------------------------------------------------------------------- +// +// @c a must be a power of two (in practice @c io::IO_BLOCK_SIZE, the O_DIRECT +// page size). @c align_down rounds @p x down to the nearest multiple of @p a; +// @c align_up rounds up. Used to page-align cache sub-range reads so partial +// fills stay O_DIRECT-compatible — never hardcode 4096 at call sites. + +[[nodiscard]] constexpr std::size_t align_down(std::size_t x, std::size_t a) noexcept +{ + return x & ~(a - 1); +} + +[[nodiscard]] constexpr std::size_t align_up(std::size_t x, std::size_t a) noexcept +{ + return (x + a - 1) & ~(a - 1); +} + /** * @brief How the prefetching layer should behave on top of a given backend. * @@ -128,12 +148,19 @@ class buffer_pool { }; // --------------------------------------------------------------------------- -// entry_state — packed atomic state + pin_count +// entry_state — mutex-guarded state + pin_count // --------------------------------------------------------------------------- // -// Packs a 4-bit state enum and a 28-bit reader pin count into a single -// atomic uint32_t. Every transition is a single CAS, which eliminates the -// TOCTOU race between checking state and modifying pin_count. +// Holds a state enum and a reader pin count guarded by a single lock +// (@c entry_lock). Every transition takes the lock, verifies its precondition, +// and mutates plain members, which eliminates the TOCTOU race between checking +// state and modifying pin_count. @c entry_lock is a type alias — currently a +// @c cucascade::exec::spin_lock (a folly-derived micro spinlock). The critical +// sections are a handful of branch/assign instructions, so a spinlock avoids the +// syscall overhead of a blocking mutex under the fine-grained per-chunk locking +// here. There is no blocking wait: a reader that observes an in-flight `loading` +// chunk does not park on it — @c acquire_read() / @c mark_loading() simply fail +// and the reader falls back to reading the bytes itself. // // State machine — each row is the complete set of valid outbound transitions // for that state. Any other transition is rejected by the corresponding @@ -155,10 +182,8 @@ class buffer_pool { // `empty` is the only state with no inbound transitions other than from // `evicting` — once an entry leaves `empty`, it can only return through the // `evicting` reclamation path. `evicting` is a one-way transit state. -// -// `loading` is the only non-terminal state with a wait point: readers that -// observe `loading` park on wait_while_pending() until the IO settles to one -// of cached / in_use(1) / allocated. + +using entry_lock = cucascade::exec::spin_lock; class entry_state { public: @@ -176,27 +201,33 @@ class entry_state { [[nodiscard]] value get_state() const noexcept { - return unpack_state(_packed.load(std::memory_order_acquire)); + std::lock_guard lk(_mtx); + return _state; } [[nodiscard]] uint32_t get_pin_count() const noexcept { - return unpack_pins(_packed.load(std::memory_order_acquire)); + std::lock_guard lk(_mtx); + return _pins; } /// empty → queued. Returns false on precondition mismatch. [[nodiscard]] bool mark_queued() noexcept { - auto expected = pack(empty, 0); - return _packed.compare_exchange_strong(expected, pack(queued, 0), std::memory_order_acq_rel); + std::lock_guard lk(_mtx); + if (_state != empty) return false; + _state = queued; + return true; } /// queued → allocated. Returns false on precondition mismatch. Called by /// the allocator when it attaches chunks to a previously-queued entry. [[nodiscard]] bool mark_allocated() noexcept { - auto expected = pack(queued, 0); - return _packed.compare_exchange_strong(expected, pack(allocated, 0), std::memory_order_acq_rel); + std::lock_guard lk(_mtx); + if (_state != queued) return false; + _state = allocated; + return true; } /// loading → allocated (IO-failure revert). Returns false on precondition @@ -204,140 +235,106 @@ class entry_state { /// whose IO did not complete: the entry's chunks stay attached so a /// subsequent allocated-steal read can retry the load with a fresh /// request_context, instead of discarding the entry to `empty` and forcing - /// the next reader through a fresh queue/allocate roundtrip. Wakes any - /// threads parked in @c wait_while_pending(). + /// the next reader through a fresh queue/allocate roundtrip. [[nodiscard]] bool mark_load_failed() noexcept { - auto expected = pack(loading, 0); - bool ok = - _packed.compare_exchange_strong(expected, pack(allocated, 0), std::memory_order_acq_rel); - if (ok) { _packed.notify_all(); } - return ok; + std::lock_guard lk(_mtx); + if (_state != loading) return false; + _state = allocated; + return true; } /// allocated → loading. Returns false on precondition mismatch. [[nodiscard]] bool mark_loading() noexcept { - auto expected = pack(allocated, 0); - return _packed.compare_exchange_strong(expected, pack(loading, 0), std::memory_order_acq_rel); + std::lock_guard lk(_mtx); + if (_state != allocated) return false; + _state = loading; + return true; } - /// loading → cached. Returns false on precondition mismatch. Wakes any - /// threads parked in @c wait_while_pending(). + /// loading → cached. Returns false on precondition mismatch. [[nodiscard]] bool mark_cached() noexcept { - auto expected = pack(loading, 0); - bool ok = _packed.compare_exchange_strong(expected, pack(cached, 0), std::memory_order_acq_rel); - if (ok) { _packed.notify_all(); } - return ok; + std::lock_guard lk(_mtx); + if (_state != loading) return false; + _state = cached; + return true; } + /// loading → in_use(pin = 1). Returns false on precondition mismatch. [[nodiscard]] bool mark_loading_in_use() noexcept { - auto expected = pack(loading, 0); - bool ok = _packed.compare_exchange_strong(expected, pack(in_use, 1), std::memory_order_acq_rel); - if (ok) { _packed.notify_all(); } - return ok; + std::lock_guard lk(_mtx); + if (_state != loading) return false; + _state = in_use; + _pins = 1; + return true; } /// allocated → evicting, or cached → evicting. Returns false on /// precondition mismatch. Both source states have pin_count == 0 by /// invariant (allocated is set with pin==0 by mark_allocated(); cached is /// only entered from in_use via release_read() when pin → 0), so the two - /// strong-CAS attempts below cover every legal transition exactly. + /// accepted source states below cover every legal transition exactly. [[nodiscard]] bool mark_evicting() noexcept { - auto expected_allocated = pack(allocated, 0); - if (_packed.compare_exchange_strong( - expected_allocated, pack(evicting, 0), std::memory_order_acq_rel)) { - return true; - } - auto expected_cached = pack(cached, 0); - return _packed.compare_exchange_strong( - expected_cached, pack(evicting, 0), std::memory_order_acq_rel); + std::lock_guard lk(_mtx); + if ((_state != allocated && _state != cached) || _pins != 0) return false; + _state = evicting; + return true; } /// evicting → empty. Returns false on precondition mismatch. [[nodiscard]] bool mark_empty() noexcept { - auto expected = pack(evicting, 0); - return _packed.compare_exchange_strong(expected, pack(empty, 0), std::memory_order_acq_rel); - } - - /// Block while state is @c loading. Returns when the state transitions to - /// any other state (i.e. to cached, in_use(1), or allocated — the three - /// outbound transitions from `loading`). - void wait_while_pending() noexcept - { - uint32_t cur = _packed.load(std::memory_order_acquire); - while (unpack_state(cur) == loading) { - _packed.wait(cur, std::memory_order_relaxed); - cur = _packed.load(std::memory_order_acquire); - } - } - - /// Park while state is @c loading, then acquire a read pin via the normal - /// (cached | in_use) → in_use(pin++) path. Returns true if a read pin was - /// acquired (load completed successfully and the entry is now readable), - /// false if the load reverted to @c allocated (IO failure) or the entry was - /// otherwise made non-readable (evicted / drained) while we were parked. - /// Safe to call from any state: if the state is already past @c loading on - /// entry, the wait loop is skipped and we go straight to @c acquire_read(). - [[nodiscard]] bool acquire_read_after_loading() noexcept - { - wait_while_pending(); - return acquire_read(); + std::lock_guard lk(_mtx); + if (_state != evicting) return false; + _state = empty; + return true; } /// (cached | in_use) → in_use with pin_count += 1. /// Returns false if the entry is not in a readable state. [[nodiscard]] bool acquire_read() noexcept { - uint32_t cur = _packed.load(std::memory_order_acquire); - while (true) { - auto st = unpack_state(cur); - if (st != cached && st != in_use) return false; - auto pins = unpack_pins(cur); - auto next = pack(in_use, pins + 1); - if (_packed.compare_exchange_weak( - cur, next, std::memory_order_acq_rel, std::memory_order_acquire)) - return true; - } + std::lock_guard lk(_mtx); + if (_state != cached && _state != in_use) return false; + _state = in_use; + ++_pins; + return true; } /// Decrement pin_count. If it reaches 0, transition in_use → cached. /// Returns true if this was the last reader. bool release_read() noexcept { - uint32_t cur = _packed.load(std::memory_order_acquire); - assert(unpack_state(cur) == in_use && unpack_pins(cur) > 0); - while (true) { - auto pins = unpack_pins(cur); - auto new_pins = pins - 1; - auto new_state = new_pins == 0 ? cached : in_use; - auto next = pack(new_state, new_pins); - if (_packed.compare_exchange_weak( - cur, next, std::memory_order_acq_rel, std::memory_order_acquire)) - return new_pins == 0; - } + std::lock_guard lk(_mtx); + assert(_state == in_use && _pins > 0); + --_pins; + if (_pins == 0) { _state = cached; } + return _pins == 0; } - private: - static constexpr uint32_t STATE_BITS = 4; - static constexpr uint32_t STATE_MASK = (1U << STATE_BITS) - 1; - static constexpr uint32_t PIN_SHIFT = STATE_BITS; - - static constexpr uint32_t pack(value s, uint32_t pins) noexcept - { - return static_cast(s) | (pins << PIN_SHIFT); - } - static constexpr value unpack_state(uint32_t v) noexcept + /// Acquire the entry lock and hand it to the caller. Lets a caller perform a + /// multi-step read-modify-write against the entry (e.g. merging cache_from) + /// atomically with respect to the state transitions above. + [[nodiscard]] std::unique_lock get_lock() noexcept { - return static_cast(v & STATE_MASK); + return std::unique_lock(_mtx); } - static constexpr uint32_t unpack_pins(uint32_t v) noexcept { return v >> PIN_SHIFT; } - std::atomic _packed{pack(empty, 0)}; + /// Read the current state WITHOUT locking. Precondition: the caller must + /// already hold this entry's lock (via @c get_lock()); otherwise the read + /// races with concurrent transitions. + [[nodiscard]] value state_locked() const noexcept { return _state; } + + private: + value _state{empty}; + uint32_t _pins{0}; + // spin_lock is a POD with no constructor; value-initialize so it starts in the + // FREE (all-bits-zero) state rather than indeterminate. + mutable entry_lock _mtx{}; }; struct alignas(64) chunk_lifecycle { @@ -437,9 +434,101 @@ struct alignas(64) cached_chunk { uint8_t* data; int numa_node{-1}; entry_state state; + // Signed, page-aligned extent (in bytes, relative to @c offset) that this + // chunk's buffer is populated with: + // 0 -> the whole chunk is populated (a "full" chunk); + // +n -> only the left prefix [offset, offset + n) is populated; + // -n -> only the right suffix [offset + chunk_size - n, offset + chunk_size) + // is populated. + // See @c needed_cache_from / @c merge_cache_from / @c chunk_covers below. + std::atomic cache_from{0}; chunk_lifecycle lifecycle; }; +// The signed, page-aligned @c cache_from a single request over the byte range +// [@p req_lo, @p req_hi) (NOT yet clamped to the chunk) implies for the chunk at +// [@p chunk_off, @p chunk_off + @p chunk_size). Returns 0 for a full (or +// non-overlapping) chunk, +n for a left prefix, -n for a right suffix. The +// magnitudes are page-aligned (io::IO_BLOCK_SIZE) so the resulting sub-range +// reads stay O_DIRECT-compatible, matching the partial reads in prefetch_loop +// and device_read_async. +[[nodiscard]] inline int32_t needed_cache_from(size_t chunk_off, + size_t chunk_size, + size_t req_lo, + size_t req_hi) noexcept +{ + const size_t page = io::IO_BLOCK_SIZE; + const size_t lo = std::max(req_lo, chunk_off); + const size_t hi = std::min(req_hi, chunk_off + chunk_size); + if (lo >= hi) { return 0; } // no overlap -- caller should not call in this case + if (lo <= chunk_off && hi >= chunk_off + chunk_size) { return 0; } // full chunk + if (lo <= chunk_off) { + // Touches the left edge -> populate a page-aligned left prefix. + const size_t bytes = std::min(align_up(hi - chunk_off, page), chunk_size); + return static_cast(bytes); + } + // Right side -> populate a page-aligned right suffix. + const size_t right_bytes = (chunk_off + chunk_size) - align_down(lo, page); + return -static_cast(right_bytes); +} + +// Fold @p want into @p c.cache_from under the merge rule. The caller MUST hold +// @c c.state.get_lock() so the read-modify-write is atomic with respect to the +// state transitions that read cache_from. +// 0 (either side) wins -> the chunk is/becomes full; +// same sign -> keep the larger magnitude (wider coverage); +// opposite signs -> the two sides together span the chunk -> full. +inline void merge_cache_from(cached_chunk& c, int32_t want) noexcept +{ + const int32_t cur = c.cache_from.load(std::memory_order_relaxed); + if (cur == 0) { return; } // already full + if (want == 0) { + c.cache_from.store(0, std::memory_order_relaxed); + return; + } + const bool same_sign = (cur > 0) == (want > 0); + if (!same_sign) { + c.cache_from.store(0, std::memory_order_relaxed); // opposite sides cover the whole chunk + return; + } + const int32_t cur_mag = cur < 0 ? -cur : cur; + const int32_t want_mag = want < 0 ? -want : want; + c.cache_from.store(want_mag > cur_mag ? want : cur, std::memory_order_relaxed); +} + +// True iff the populated extent of @p c covers the request [@p req_lo, @p req_hi) +// (which the caller has clamped to the chunk's extent). +[[nodiscard]] inline bool chunk_covers(const cached_chunk& c, + size_t chunk_size, + size_t req_lo, + size_t req_hi) noexcept +{ + const int32_t cf = c.cache_from.load(std::memory_order_relaxed); + if (cf == 0) { return true; } // full chunk + if (cf > 0) { return req_hi <= c.offset + static_cast(cf); } + return req_lo >= c.offset + chunk_size - static_cast(-cf); +} + +// The half-open file byte span [seg_lo, seg_hi) that must be read to populate a +// chunk at @p offset (size @p chunk_size) to exactly the extent that its stored +// @p cache_from advertises. cache_from is edge-anchored, so the fill span is +// too — deriving the span FROM cache_from (rather than from the request range) +// guarantees the bytes read match the bytes @c chunk_covers will later claim. +// This is the single source of truth shared by prefetch_loop and +// device_read_async's load branch. +[[nodiscard]] inline std::pair chunk_fill_span(size_t offset, + size_t chunk_size, + int32_t cache_from, + size_t page) noexcept +{ + if (cache_from == 0) { return {offset, offset + chunk_size}; } // full chunk + if (cache_from > 0) { + return {offset, offset + std::min(chunk_size, align_up(static_cast(cache_from), page))}; + } + return {offset + align_down(chunk_size - static_cast(-cache_from), page), + offset + chunk_size}; +} + // Coverage requirement for find_entry. enum class coverage_policy { full, // return the chunks only when they fully cover [offset, offset + size); else none @@ -513,6 +602,11 @@ find_entry(const Chunks& chunks, // Coverage confirmed by the invariant: sorted + non-overlapping + fixed-size // means consecutive chunks differ by exactly chunk_size, so the intermediates // are forced once the first and last are at the expected positions. + // + // NOTE: this confirms POSITIONAL coverage only (the requested byte span maps + // onto existing chunks) — it does NOT confirm those chunks are populated over + // the requested sub-range. Populated-ness (cache_from) is enforced by the + // pin-holding caller via chunk_covers() after acquire_read() (Steps 6/7). std::vector result; result.reserve(expected_count); for (std::size_t i = 0; i < expected_count; ++i) { diff --git a/include/cucascade/io/config.hpp b/include/cucascade/io/config.hpp index 250120b..5719e98 100644 --- a/include/cucascade/io/config.hpp +++ b/include/cucascade/io/config.hpp @@ -27,6 +27,33 @@ namespace cucascade::io { +/** + * @brief Storage caching behavior selected at the io_context level. + * + * A single knob for two coupled decisions — whether local reads bypass the OS + * page cache with O_DIRECT, and whether cucascade's prefetching cache is on: + * - @c none : O_DIRECT on, prefetching cache off (direct, uncached) + * - @c os : O_DIRECT off, prefetching cache off (rely on the OS page cache) + * - @c prefetch : O_DIRECT on, prefetching cache on (cucascade prefetching cache) + */ +enum class cache_level { + none, + os, + prefetch, +}; + +/// True when @p level reads through O_DIRECT (i.e. everything but @c os). +[[nodiscard]] constexpr bool odirect_enabled(cache_level level) noexcept +{ + return level != cache_level::os; +} + +/// True when @p level enables the prefetching cache (only @c prefetch). +[[nodiscard]] constexpr bool prefetch_enabled(cache_level level) noexcept +{ + return level == cache_level::prefetch; +} + /** * @brief Top-level configuration for the cucascade::io datasource layer. * @@ -47,12 +74,15 @@ struct io_config { /// (each its own libcurl event loop + connection pool). std::size_t rest_n_reactors{2}; - /// Enable the prefetching cache on the ioctx. When false the cache is - /// constructed but unarmed (no background IO threads). - bool enable_prefetch_cache{false}; + /// Storage caching behavior (O_DIRECT + prefetching cache). Sources + /// @c local.use_odirect and whether the prefetching cache is enabled via + /// @c odirect_enabled() / @c prefetch_enabled(). Defaults to @c none + /// (O_DIRECT on, prefetching cache off). + cache_level caching{}; /// Local (uring) reactor configuration — bounce-slot size, O_DIRECT, - /// ring depth, etc. + /// ring depth, etc. @c local.use_odirect is derived from @c cache_level + /// when the ioctx is built through the datasource factory. uring::config local{}; /// REST (S3/object-store) reactor configuration — timeouts, TLS, chunking, diff --git a/src/io/cache/prefetching_cache.cpp b/src/io/cache/prefetching_cache.cpp index 647aead..678920e 100644 --- a/src/io/cache/prefetching_cache.cpp +++ b/src/io/cache/prefetching_cache.cpp @@ -173,10 +173,46 @@ prefetching_handle::prefetching_handle(std::unique_ptr prefetching_cache::file_entry::update_and_get_chunks( - std::span incoming, uint32_t ticker) + std::span incoming, std::span desired_cache_from, uint32_t ticker) { std::vector result(incoming.size()); + // Fold the desired cache_from into an already-present chunk, gated on its + // state (held stable under the entry lock). Only PRE-LOAD chunks may have + // their extent widened: their buffer has not been filled yet, so the pending + // load reads the final (widened) extent before any reader can hit it. + // empty -> an evicted-and-reset chunk: cache_from is 0, which + // merge_cache_from would misread as "already full" and + // drop the partial desire, forcing a full reload after + // every eviction. Treat it like a fresh insert and + // store the desired extent DIRECTLY; + // queued / allocated -> not yet loaded: widen via the merge rule. prefetch_loop + // and device_read take the same entry lock at mark_loading + // before reading cache_from, so the widened value is either + // read by that load or the merge is skipped (loading); + // loading / evicting -> being written by its sole loader / torn down; leave + // cache_from untouched (widening would race); + // cached / in_use -> already populated to its current extent and never + // re-read. Widening would advertise never-written bytes + // (a false hit / data corruption), so leave it unchanged; + // a request needing more correctly MISSES via chunk_covers + // and falls back to a bounce/direct read. + auto merge_into_existing = [&](cached_chunk* c, int32_t desired) { + auto lk = c->state.get_lock(); + auto const st = c->state.state_locked(); + switch (st) { + case entry_state::empty: // evicted-and-reset -> fresh insert + c->cache_from.store(desired, std::memory_order_relaxed); + return; + case entry_state::queued: + case entry_state::allocated: // not yet loaded -> pending load reads final extent + merge_cache_from(*c, desired); + return; + default: // loading/evicting: racing; cached/in_use: already populated -> leave unchanged + return; + } + }; + // Phase 1: classify under shared lock — find which offsets already exist. // Track the indices of incoming items that need to be inserted. std::vector missing_indices; // indices into `incoming`/`result` @@ -192,6 +228,7 @@ std::vector prefetching_cache::file_entry::update_and_get_chunks( if (s != s_end && (*s)->offset == off) { s->get()->lifecycle.on_request(ticker); + merge_into_existing(s->get(), desired_cache_from[i]); result[i] = s->get(); // existing } else { missing_indices.push_back(i); // mark for insertion @@ -220,10 +257,15 @@ std::vector prefetching_cache::file_entry::update_and_get_chunks( s = gallop_lower_bound(s, s_end, off); if (s != s_end && (*s)->offset == off) { + merge_into_existing(s->get(), desired_cache_from[idx]); result[idx] = s->get(); // someone else inserted it } else { auto chunk = std::make_unique(off); chunk->lifecycle.on_request(ticker); + // First creation: store the desired cache_from directly. merge_cache_from + // must NOT be used here — it treats a stored 0 as "already full", which + // would misread a freshly-created chunk (default 0) as fully populated. + chunk->cache_from.store(desired_cache_from[idx], std::memory_order_relaxed); result[idx] = chunk.get(); // capture raw ptr before move to_insert.push_back(std::move(chunk)); } @@ -350,8 +392,40 @@ prefetching_handle prefetching_cache::insert(const io_object& obj, } } - auto chunks_to_fetch = - file.update_and_get_chunks(chunk_offsets, _ticker.load(std::memory_order_relaxed)); + // Derive, in parallel with chunk_offsets, the signed page-aligned cache_from + // each chunk should be populated with. The coalesced ranges above dictate the + // set of chunks, but the ORIGINAL (pre-coalesce) ranges dictate how much of + // each boundary chunk is actually wanted — so a chunk touched only at its head + // or tail is loaded partially instead of in full. Chunks pulled in purely by + // coalescing (overlapping no original range) default to 0 == full, matching + // the coalesced read that pulled them in. + auto fold_cache_from = [](int32_t cur, int32_t next) -> int32_t { + if (cur == 0 || next == 0) { return 0; } // 0 == full wins + if ((cur > 0) != (next > 0)) { return 0; } // opposite sides span the chunk + const int32_t cur_mag = cur < 0 ? -cur : cur; + const int32_t next_mag = next < 0 ? -next : next; + return next_mag > cur_mag ? next : cur; + }; + std::vector desired_cache_from(chunk_offsets.size(), 0); + std::vector touched(chunk_offsets.size(), 0); + for (const auto& r : ranges) { + if (r.size() <= 0) { continue; } + const size_t r_lo = static_cast(r.offset()); + const size_t r_hi = r_lo + static_cast(r.size()); + const size_t first_chunk = (r_lo / chunk_bytes) * chunk_bytes; + const size_t last_chunk = ((r_hi - 1) / chunk_bytes) * chunk_bytes; + for (size_t off = first_chunk; off <= last_chunk; off += chunk_bytes) { + auto it = std::lower_bound(chunk_offsets.begin(), chunk_offsets.end(), off); + if (it == chunk_offsets.end() || *it != off) { continue; } + const auto idx = static_cast(it - chunk_offsets.begin()); + const int32_t nc = needed_cache_from(off, chunk_bytes, r_lo, r_hi); + desired_cache_from[idx] = touched[idx] ? fold_cache_from(desired_cache_from[idx], nc) : nc; + touched[idx] = 1; + } + } + + auto chunks_to_fetch = file.update_and_get_chunks( + chunk_offsets, desired_cache_from, _ticker.load(std::memory_order_relaxed)); auto work = std::make_shared(obj, _ticker.load()); work->chunks = std::move(chunks_to_fetch); @@ -386,18 +460,30 @@ bool prefetching_cache::host_read_from_cache_only( } } + auto const end_offset = offset + size; + auto const chunk_size = _chunk_size; + while (!chunks.empty()) { - auto iter = - std::ranges::find_if(chunks, [](cached_chunk* c) { return !c->state.acquire_read(); }); + // Sweep that both pins each chunk and confirms it is populated over the + // requested sub-range. A chunk we cannot pin, OR one whose populated extent + // (cache_from) does not cover the request, is treated as a failed pin: we + // undo it and bail so the caller does a full fallback IO (no partial memcpy). + auto iter = std::ranges::find_if(chunks, [&](cached_chunk* c) { + if (!c->state.acquire_read()) { return true; } + size_t const req_lo = std::max(offset, c->offset); + size_t const req_hi = std::min(end_offset, c->offset + chunk_size); + if (!chunk_covers(*c, chunk_size, req_lo, req_hi)) { + c->state.release_read(); + return true; + } + return false; + }); if (iter != chunks.end()) { std::for_each(chunks.begin(), iter, [](cached_chunk* c) { c->state.release_read(); }); break; } - auto const end_offset = offset + size; - auto const chunk_size = _chunk_size; - for (auto* chunk : chunks) { auto const chunk_begin = std::max(offset, chunk->offset); auto const chunk_end = std::min(end_offset, chunk->offset + chunk_size); @@ -497,39 +583,51 @@ exec::semi_future prefetching_cache::device_read_async(const io_obj } cached_chunk* c = (ci < chunks.size() && chunks[ci]->offset == off) ? chunks[ci] : nullptr; + // The portion of this chunk the request actually needs, clamped to both the + // request and the chunk extent. Hoisted above the branch so both the hit + // coverage gate and the load/miss span share it. + size_t const need_lo = std::max(off, offset); + size_t const need_hi = std::min(off + chunk_bytes, offset + size); + + bool hit = false; if (c != nullptr && c->state.acquire_read()) { - cached_chunks.push_back(c); // (1) hit -- a cached chunk is always fully valid - hits++; - } else { + // A read pin alone is not sufficient: a partially-populated chunk may not + // cover the requested sub-range. Confirm coverage (cache_from) before + // treating it as a hit -- THIS IS THE KEY CORRECTNESS GATE. If it does + // not cover, release the pin and fall through to the load/miss path. + if (chunk_covers(*c, chunk_bytes, need_lo, need_hi)) { + cached_chunks.push_back(c); // (1) hit + hits++; + hit = true; + } else { + c->state.release_read(); + } + } + + if (!hit) { if (!cache_while_reading_enabled) { every_chunk_is_cached = false; break; // (3) miss, but we can't do H2D IO, so fall back to direct device read } - // Stage a read through the chunk's cache buffer only when caching the - // WHOLE chunk is cheap enough -- a cached chunk must be fully valid, so - // caching a partially-requested chunk costs reading its non-overlapping - // remainder from disk (boundary over-read, the dominant cold-pass cost). - // Cache when that over-read is < 25% of the chunk (so a read covering - // >75% of the chunk still warms it); otherwise read just the needed, - // block-aligned span through an internal bounce slot (null host buffer) - // and leave the chunk uncached -- zero over-read. (Short-term: a heavily - // partial boundary chunk is re-read each pass; full partial caching is a - // larger redesign.) - size_t const need_lo = std::max(off, offset); - size_t const need_hi = std::min(off + chunk_bytes, offset + size); - size_t const overread = chunk_bytes - (need_hi - need_lo); - bool const worth_caching = overread * 4 < chunk_bytes; // over-read < 25% of chunk - if (worth_caching && c != nullptr && c->state.mark_loading()) { + if (c != nullptr && c->state.mark_loading()) { + // Loader is the sole writer during allocated -> loading -> cached, so a + // plain relaxed store of the populated extent is safe here. Derive the + // fill span FROM the stored cache_from (the single source of truth), so + // the bytes we read match EXACTLY what chunk_covers will later claim. + // cache_from is edge-anchored, so an interior read conservatively fills + // to the chunk edge (over-reads the tail/head, never the full chunk). assert(c->data != nullptr); + int32_t const cf = needed_cache_from(off, chunk_bytes, need_lo, need_hi); + c->cache_from.store(cf, std::memory_order_relaxed); + auto const [seg_lo, seg_hi] = chunk_fill_span(off, chunk_bytes, cf, io::IO_BLOCK_SIZE); io_chunks.push_back(c); // (2) host-to-device load into the cache buffer - io_segments.emplace_back(off, chunk_bytes, c->data); + io_segments.emplace_back(seg_lo, seg_hi - seg_lo, c->data + (seg_lo - c->offset)); h2d++; } else { - // (3) partial head/tail (or busy / missing chunk): read just the needed, - // block-aligned span via an internal bounce slot; do not touch the cache. - size_t const seg_lo = need_lo & ~(io::IO_BLOCK_SIZE - 1); - size_t const seg_hi = std::min( - off + chunk_bytes, (need_hi + io::IO_BLOCK_SIZE - 1) & ~(io::IO_BLOCK_SIZE - 1)); + // (3) busy / missing chunk: read just the needed, block-aligned span via + // an internal bounce slot (null host buffer); do not touch the cache. + size_t const seg_lo = align_down(need_lo, io::IO_BLOCK_SIZE); + size_t const seg_hi = std::min(off + chunk_bytes, align_up(need_hi, io::IO_BLOCK_SIZE)); io_segments.emplace_back(seg_lo, seg_hi - seg_lo, nullptr); // (3) miss misses++; } @@ -740,21 +838,31 @@ void prefetching_cache::prefetch_loop(const std::stop_token& st) std::vector segments; segments.reserve(allocated_chunks.size()); - allocated_chunks.erase(std::remove_if(allocated_chunks.begin(), - allocated_chunks.end(), - [&](cached_chunk* c) { - if (c->state.mark_loading()) { - segments.emplace_back( - c->offset, _chunk_size, c->data); - return false; - } - return true; - }), - allocated_chunks.end()); + allocated_chunks.erase( + std::remove_if(allocated_chunks.begin(), + allocated_chunks.end(), + [&](cached_chunk* c) { + if (c->state.mark_loading()) { + // Read only the page-aligned sub-range this chunk was + // requested for (cache_from), not always the whole chunk. + auto const [seg_lo, seg_hi] = + chunk_fill_span(c->offset, + _chunk_size, + c->cache_from.load(std::memory_order_relaxed), + io::IO_BLOCK_SIZE); + segments.emplace_back( + seg_lo, seg_hi - seg_lo, c->data + (seg_lo - c->offset)); + return false; + } + return true; + }), + allocated_chunks.end()); std::ignore = req->state->mark_loading(); - auto token = _rate_limiter.acquire(segments.size()); + // One unit per in-flight prefetch task (not per chunk): reserve on dispatch, + // release on completion when the token below is dropped in the continuation. + auto token = _rate_limiter.acquire(1); if (req->is_cancelled() || st.stop_requested()) { std::ranges::for_each(allocated_chunks, @@ -875,6 +983,7 @@ void prefetching_cache::evict_loop(const std::stop_token& st) if (c->state.mark_evicting()) { reclaim_by_numa[c->numa_node].push_back(reinterpret_cast(c->data)); c->data = nullptr; + c->cache_from.store(0, std::memory_order_relaxed); // reset populated extent static_cast(c->state.mark_empty()); ++reclaimed; ++er.n_evicted; diff --git a/src/io/datasource_factory.cpp b/src/io/datasource_factory.cpp index edefff4..b5006ad 100644 --- a/src/io/datasource_factory.cpp +++ b/src/io/datasource_factory.cpp @@ -108,6 +108,9 @@ factory_type make_uring_ioctx_factory( // and the pinned bounce-staging resource itself. auto uring_cfg = config.local; uring_cfg.bounce_size = host_mr->get_block_size(); + // O_DIRECT is driven by the io_context-level cache_level, not set on the + // local reactor config directly. + uring_cfg.use_odirect = odirect_enabled(config.caching); auto ctx = std::make_shared(uring_cfg, host_mr); return std::make_shared(config.uring_n_reactors, std::move(ctx)); } catch (const std::exception& e) { diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 03de372..f052508 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -38,6 +38,11 @@ if(NOT CUCASCADE_TOPOLOGY_ONLY) data/test_data_repository_manager.cpp data/test_disk_io_backend.cpp data/test_io_worker.cpp + # Exec tests + exec/test_spin_lock.cpp + # IO tests + io/test_io_config.cpp + io/test_prefetching_cache_chunking.cpp # Main test runner unittest.cpp) set_target_properties(cucascade_tests PROPERTIES CUDA_STANDARD 20 diff --git a/test/exec/test_spin_lock.cpp b/test/exec/test_spin_lock.cpp new file mode 100644 index 0000000..dbccc67 --- /dev/null +++ b/test/exec/test_spin_lock.cpp @@ -0,0 +1,234 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) Meta Platforms, Inc. and affiliates. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * --------------------------------------------------------------------------- + * ATTRIBUTION + * + * The correctness and try-lock tests below are ports of folly's SpinLockTest + * (folly/test/SpinLockTest.cpp, Meta Platforms, Inc., Apache-2.0), adapted to + * cucascade::exec::spin_lock and the Catch2 framework. Because Catch2's + * assertion macros are not thread-safe, worker threads record invariant + * violations into atomics and the main thread asserts on them after joining, + * in place of folly's in-thread gtest EXPECT_* calls. + * --------------------------------------------------------------------------- + */ + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using cucascade::exec::spin_lock; +using cucascade::exec::spin_lock_array; +using cucascade::exec::spin_lock_guard; + +namespace { + +std::size_t next_pow_two(std::size_t v) noexcept +{ + std::size_t p = 1; + while (p < v) { + p <<= 1; + } + return p; +} + +std::size_t worker_count() noexcept +{ + unsigned hw = std::thread::hardware_concurrency(); + return hw == 0 ? 4u : static_cast(hw); +} + +// --------------------------------------------------------------------------- +// Correctness: N threads mutate a shared array under the lock; while holding +// the lock every element must be identical (no torn interleaving), then the +// holder memsets the array to a fresh random byte. +// --------------------------------------------------------------------------- + +struct locked_val { + static constexpr std::size_t k_len = 1024; + int ar[k_len]; + // spin_lock is a POD with no constructor; value-initialize to the FREE state. + spin_lock lock{}; + + locked_val() { std::memset(ar, 0, sizeof(ar)); } +}; + +void spinlock_test_thread(std::size_t nthrs, locked_val* v, std::atomic* violations) +{ + std::size_t const max = (1u << 16) / next_pow_two(nthrs); + std::mt19937 rng(std::random_device{}()); + for (std::size_t i = 0; i < max; i++) { + cucascade::exec::detail::spin_cpu_relax(); + std::unique_lock g(v->lock); + + // Invariant under mutual exclusion: all elements equal ar[0]. + int const first = v->ar[0]; + for (std::size_t j = 0; j < locked_val::k_len; j++) { + if (v->ar[j] != first) { + violations->fetch_add(1, std::memory_order_relaxed); + break; + } + } + + int const byte = static_cast(rng() & 0xffu); + std::memset(v->ar, byte, sizeof(v->ar)); + } +} + +// --------------------------------------------------------------------------- +// TryLock: threads contend on lock2 via try_lock() while serializing bookkeeping +// under lock1. Exactly one thread may hold lock2 at a time (locked flips), and +// every successful acquire waits for at least one other thread to fail before +// releasing, so failures must accumulate. +// --------------------------------------------------------------------------- + +struct try_lock_state { + spin_lock lock1{}; + spin_lock lock2{}; + bool locked{false}; + std::uint64_t obtained{0}; + std::uint64_t failed{0}; +}; + +void trylock_test_thread(try_lock_state* state, std::uint64_t count, std::atomic* violations) +{ + while (true) { + cucascade::exec::detail::spin_cpu_relax(); + bool ret = state->lock2.try_lock(); + std::unique_lock g(state->lock1); + if (state->obtained >= count) { + if (ret) { state->lock2.unlock(); } + break; + } + + if (ret) { + // We got lock2 — no other thread must believe it holds it. + if (state->locked) { violations->fetch_add(1, std::memory_order_relaxed); } + ++state->obtained; + state->locked = true; + + // Release lock1 and wait until at least one other thread fails to obtain + // lock2 before continuing. + auto old_failed = state->failed; + while (state->failed == old_failed && state->obtained < count) { + state->lock1.unlock(); + cucascade::exec::detail::spin_cpu_relax(); + state->lock1.lock(); + } + + state->locked = false; + state->lock2.unlock(); + } else { + ++state->failed; + } + } +} + +} // namespace + +TEST_CASE("spin_lock basic acquire/release and try_lock", "[spin_lock][exec]") +{ + // Zero-initialized spin_lock starts FREE (documented POD contract). + spin_lock lock{}; + REQUIRE(lock.try_lock()); // acquired + REQUIRE_FALSE(lock.try_lock()); // already held -> fails + lock.unlock(); + REQUIRE(lock.try_lock()); // free again + lock.unlock(); + + SECTION("spin_lock_guard (std::lock_guard alias) is RAII") + { + { + spin_lock_guard g(lock); + REQUIRE_FALSE(lock.try_lock()); + } + REQUIRE(lock.try_lock()); // released by guard destructor + lock.unlock(); + } +} + +TEST_CASE("spin_lock is a POD / trivially usable in packed structs", "[spin_lock][exec]") +{ + STATIC_REQUIRE(std::is_standard_layout::value); + STATIC_REQUIRE(std::is_trivial::value); +} + +TEST_CASE("spin_lock Correctness under contention", "[spin_lock][exec]") +{ + std::size_t const nthrs = worker_count() * std::size_t{2}; + std::atomic violations{0}; + locked_val v; + + std::vector threads; + threads.reserve(nthrs); + for (std::size_t i = 0; i < nthrs; ++i) { + threads.emplace_back(spinlock_test_thread, nthrs, &v, &violations); + } + for (auto& t : threads) { + t.join(); + } + + REQUIRE(violations.load() == 0); +} + +TEST_CASE("spin_lock TryLock ping-pong", "[spin_lock][exec]") +{ + std::size_t const nthrs = worker_count() + std::size_t{4}; + std::uint64_t const count = 100; + std::atomic violations{0}; + try_lock_state state; + + std::vector threads; + threads.reserve(nthrs); + for (std::size_t i = 0; i < nthrs; ++i) { + threads.emplace_back(trylock_test_thread, &state, count, &violations); + } + for (auto& t : threads) { + t.join(); + } + + REQUIRE(violations.load() == 0); + REQUIRE(state.obtained == count); + // Every successful acquire waits for another thread to fail, except possibly + // the very last one when no other threads remain. + REQUIRE(state.failed + 1u >= state.obtained); +} + +TEST_CASE("spin_lock_array shards lock independently", "[spin_lock][exec]") +{ + spin_lock_array locks; + REQUIRE(locks.size() == 8); + + // Each shard is independent: locking one leaves the others free. + REQUIRE(locks[3].try_lock()); + REQUIRE(locks[4].try_lock()); + REQUIRE_FALSE(locks[3].try_lock()); + locks[3].unlock(); + REQUIRE(locks[3].try_lock()); + locks[3].unlock(); + locks[4].unlock(); +} diff --git a/test/io/test_io_config.cpp b/test/io/test_io_config.cpp new file mode 100644 index 0000000..2a567df --- /dev/null +++ b/test/io/test_io_config.cpp @@ -0,0 +1,54 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Tests for the io_config cache_level enum and its derivations: the single +// cache_level knob replaces the former use_odirect / enable_prefetch_cache +// pair, so these pin down the (O_DIRECT, prefetching-cache) mapping and the +// default. + +#include + +#include + +using cucascade::io::cache_level; +using cucascade::io::io_config; +using cucascade::io::odirect_enabled; +using cucascade::io::prefetch_enabled; + +TEST_CASE("cache_level maps to O_DIRECT and prefetching-cache flags", "[io_config]") +{ + // none : O_DIRECT on, prefetch off + STATIC_REQUIRE(odirect_enabled(cache_level::none)); + STATIC_REQUIRE_FALSE(prefetch_enabled(cache_level::none)); + + // os : O_DIRECT off, prefetch off + STATIC_REQUIRE_FALSE(odirect_enabled(cache_level::os)); + STATIC_REQUIRE_FALSE(prefetch_enabled(cache_level::os)); + + // prefetch : O_DIRECT on, prefetch on + STATIC_REQUIRE(odirect_enabled(cache_level::prefetch)); + STATIC_REQUIRE(prefetch_enabled(cache_level::prefetch)); +} + +TEST_CASE("io_config defaults to cache_level::none (direct, uncached)", "[io_config]") +{ + io_config cfg; + REQUIRE(cfg.caching == cache_level::none); + // Preserves the pre-refactor defaults: O_DIRECT on, prefetching cache off. + REQUIRE(odirect_enabled(cfg.caching)); + REQUIRE_FALSE(prefetch_enabled(cfg.caching)); +} diff --git a/test/io/test_prefetching_cache_chunking.cpp b/test/io/test_prefetching_cache_chunking.cpp new file mode 100644 index 0000000..3f93b9b --- /dev/null +++ b/test/io/test_prefetching_cache_chunking.cpp @@ -0,0 +1,472 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Unit tests for the partial-chunk-caching primitives in +// cucascade::io::cache (types.hpp): the request->cache_from mapping +// (needed_cache_from), the cache_from->fill-span mapping (chunk_fill_span), +// the populated-extent predicate (chunk_covers), the subsequent-insert merge +// rule (merge_cache_from), and the lock-based entry_state machine. +// +// The central invariant these tests guard is that, for any cache_from value a +// chunk can hold, the bytes that will actually be READ into its buffer +// (chunk_fill_span) are a SUPERSET of the bytes that chunk_covers will later +// advertise as populated. A violation of that invariant is the exact +// data-corruption class (a partially-populated chunk served as a full hit) +// that this layer must never allow. + +#include +#include + +#include + +#include +#include +#include +#include +#include + +using cucascade::io::cache::align_down; +using cucascade::io::cache::align_up; +using cucascade::io::cache::cached_chunk; +using cucascade::io::cache::chunk_covers; +using cucascade::io::cache::chunk_fill_span; +using cucascade::io::cache::entry_state; +using cucascade::io::cache::merge_cache_from; +using cucascade::io::cache::needed_cache_from; + +namespace { + +constexpr std::size_t chunk_size = 1ull << 20; // 1 MiB, the default cache chunk +constexpr std::size_t page = cucascade::io::IO_BLOCK_SIZE; // 4096 + +// Set a chunk's populated extent the way the cache does (relaxed store). +void set_cache_from(cached_chunk& c, std::int32_t cf) +{ + c.cache_from.store(cf, std::memory_order_relaxed); +} + +// The core safety invariant: everything chunk_covers() claims as populated for +// `cf` must lie inside the span chunk_fill_span() would actually read for `cf`. +// Checked page-by-page across the whole chunk. +void require_fill_superset_of_cover(std::size_t off, std::int32_t cf) +{ + cached_chunk c(off); + set_cache_from(c, cf); + auto const [seg_lo, seg_hi] = chunk_fill_span(off, chunk_size, cf, page); + REQUIRE(seg_lo >= off); + REQUIRE(seg_hi <= off + chunk_size); + REQUIRE(seg_lo <= seg_hi); + for (std::size_t p = off; p < off + chunk_size; p += page) { + if (chunk_covers(c, chunk_size, p, p + page)) { + // Any page advertised as covered must be inside the filled span. + REQUIRE(p >= seg_lo); + REQUIRE(p + page <= seg_hi); + } + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// align_up / align_down +// --------------------------------------------------------------------------- + +TEST_CASE("align helpers round to page boundaries", "[cache][chunking]") +{ + REQUIRE(align_down(0, page) == 0); + REQUIRE(align_down(1, page) == 0); + REQUIRE(align_down(page, page) == page); + REQUIRE(align_down(page + 1, page) == page); + REQUIRE(align_down(5000, page) == page); // 5000 -> 4096 + + REQUIRE(align_up(0, page) == 0); + REQUIRE(align_up(1, page) == page); + REQUIRE(align_up(page, page) == page); + REQUIRE(align_up(page + 1, page) == 2 * page); + REQUIRE(align_up(5000, page) == 2 * page); // 5000 -> 8192 +} + +// --------------------------------------------------------------------------- +// needed_cache_from — request pattern -> signed page-aligned extent +// --------------------------------------------------------------------------- + +TEST_CASE("needed_cache_from maps read patterns to cache_from", "[cache][chunking]") +{ + SECTION("request spanning the whole chunk -> full (0)") + { + REQUIRE(needed_cache_from(0, chunk_size, 0, chunk_size) == 0); + } + + SECTION("request overhanging both edges -> full (0)") + { + REQUIRE(needed_cache_from(0, chunk_size, 0, 4 * chunk_size) == 0); + } + + SECTION("non-overlapping request -> 0 (no-op sentinel)") + { + // Entirely below and entirely above the chunk both clamp to empty. + REQUIRE(needed_cache_from(chunk_size, chunk_size, 0, chunk_size) == 0); + } + + SECTION("left-anchored partial -> positive, page-aligned up") + { + REQUIRE(needed_cache_from(0, chunk_size, 0, page) == static_cast(page)); + REQUIRE(needed_cache_from(0, chunk_size, 0, 5000) == static_cast(2 * page)); + } + + SECTION("right-anchored partial -> negative, page-aligned down") + { + // [1044480, chunk_end): 1044480 is page-aligned -> suffix of exactly one page. + REQUIRE(needed_cache_from(0, chunk_size, chunk_size - page, chunk_size) == + -static_cast(page)); + // [1040000, chunk_end): floor-aligns 1040000 -> 1036288, suffix = 12288 (3 pages). + REQUIRE(needed_cache_from(0, chunk_size, 1040000, chunk_size) == + -static_cast(3 * page)); + } + + SECTION("interior read conservatively becomes a right suffix to the chunk end") + { + // [524288, 532480) touches neither edge; the edge-anchored representation + // rounds it to a right suffix running to the chunk end. + REQUIRE(needed_cache_from(0, chunk_size, 524288, 532480) == + -static_cast(chunk_size - 524288)); + } + + SECTION("second chunk of a boundary-spanning read -> left prefix") + { + // Chunk [1MiB, 2MiB), request tail [1MiB, 1MiB+100KiB-ish). + std::size_t const off = chunk_size; + REQUIRE(needed_cache_from(off, chunk_size, off, off + 102400) == + static_cast(102400)); // already a page multiple (25 pages) + } +} + +// --------------------------------------------------------------------------- +// chunk_fill_span vs chunk_covers — the safety invariant +// --------------------------------------------------------------------------- + +TEST_CASE("chunk_fill_span always covers what chunk_covers advertises", "[cache][chunking][safety]") +{ + // Full, left prefixes, right suffixes, and the interior->suffix case, at + // both a zero offset and a shifted offset. + for (std::size_t off : {std::size_t{0}, chunk_size, 7 * chunk_size}) { + require_fill_superset_of_cover(off, 0); + require_fill_superset_of_cover(off, static_cast(page)); + require_fill_superset_of_cover(off, static_cast(8 * page)); + require_fill_superset_of_cover(off, -static_cast(page)); + require_fill_superset_of_cover(off, -static_cast(3 * page)); + require_fill_superset_of_cover(off, -static_cast(chunk_size - 524288)); + } +} + +TEST_CASE("chunk_fill_span produces exact edge-anchored spans", "[cache][chunking]") +{ + SECTION("full chunk reads the whole extent") + { + auto const [lo, hi] = chunk_fill_span(0, chunk_size, 0, page); + REQUIRE(lo == 0); + REQUIRE(hi == chunk_size); + } + + SECTION("left prefix reads from offset") + { + auto const [lo, hi] = chunk_fill_span(0, chunk_size, static_cast(8 * page), page); + REQUIRE(lo == 0); + REQUIRE(hi == 8 * page); + } + + SECTION("right suffix reads to chunk end") + { + auto const [lo, hi] = + chunk_fill_span(0, chunk_size, -static_cast(3 * page), page); + REQUIRE(lo == chunk_size - 3 * page); + REQUIRE(hi == chunk_size); + } + + SECTION("round-trip: needed_cache_from -> chunk_fill_span reads the requested bytes") + { + // A boundary-spanning read [946176, 1150976) split across two 1 MiB chunks. + std::size_t const req_lo = 946176; // page-aligned + std::size_t const req_hi = 1150976; // page-aligned + + // Chunk 0 gets the right suffix. + std::int32_t const cf0 = needed_cache_from(0, chunk_size, req_lo, req_hi); + auto const [lo0, hi0] = chunk_fill_span(0, chunk_size, cf0, page); + REQUIRE(lo0 == req_lo); // exactly the requested bytes in chunk 0 + REQUIRE(hi0 == chunk_size); + + // Chunk 1 gets the left prefix. + std::size_t const off1 = chunk_size; + std::int32_t const cf1 = needed_cache_from(off1, chunk_size, req_lo, req_hi); + auto const [lo1, hi1] = chunk_fill_span(off1, chunk_size, cf1, page); + REQUIRE(lo1 == off1); + REQUIRE(hi1 == req_hi); // exactly the requested bytes in chunk 1 + } +} + +// --------------------------------------------------------------------------- +// chunk_covers — hit / miss on partially-populated chunks +// --------------------------------------------------------------------------- + +TEST_CASE("chunk_covers gates hits on the populated extent", "[cache][chunking]") +{ + SECTION("full chunk covers everything") + { + cached_chunk c(0); + set_cache_from(c, 0); + REQUIRE(chunk_covers(c, chunk_size, 0, chunk_size)); + REQUIRE(chunk_covers(c, chunk_size, 12345, 54321)); + } + + SECTION("left prefix: covers the prefix, misses beyond it") + { + cached_chunk c(0); + set_cache_from(c, static_cast(2 * page)); // [0, 8192) + REQUIRE(chunk_covers(c, chunk_size, 0, page)); + REQUIRE(chunk_covers(c, chunk_size, page, 2 * page)); + REQUIRE(chunk_covers(c, chunk_size, 0, 2 * page)); + REQUIRE_FALSE(chunk_covers(c, chunk_size, 2 * page, 3 * page)); // just past the edge + REQUIRE_FALSE(chunk_covers(c, chunk_size, 0, 3 * page)); // straddles the edge + } + + SECTION("right suffix: covers the suffix, misses before it") + { + cached_chunk c(0); + set_cache_from(c, -static_cast(2 * page)); // [chunk-8192, chunk) + std::size_t const edge = chunk_size - 2 * page; + REQUIRE(chunk_covers(c, chunk_size, edge, chunk_size)); + REQUIRE(chunk_covers(c, chunk_size, edge + page, chunk_size)); + REQUIRE_FALSE(chunk_covers(c, chunk_size, edge - page, edge)); // just before the edge + REQUIRE_FALSE(chunk_covers(c, chunk_size, edge - page, chunk_size)); // straddles the edge + } + + SECTION("regression: interior fill does not false-hit the uncovered side") + { + // Reproduces the interior-read scenario: a chunk loaded for [524288, ...) + // (stored as a right suffix to the chunk end) must MISS a request for the + // left side it never populated. + cached_chunk c(0); + set_cache_from(c, needed_cache_from(0, chunk_size, 524288, 532480)); + // The far-right request that shares the suffix is a legitimate hit and is + // genuinely inside the fill span (fill == cover for this cf). + REQUIRE(chunk_covers(c, chunk_size, 786432, 794624)); + // The left side was never read -> must miss. + REQUIRE_FALSE(chunk_covers(c, chunk_size, 0, page)); + REQUIRE_FALSE(chunk_covers(c, chunk_size, 262144, 262144 + page)); + } +} + +// --------------------------------------------------------------------------- +// merge_cache_from — subsequent-insert widening rules +// --------------------------------------------------------------------------- + +namespace { + +// Apply the merge the way the cache does: under the entry lock. +std::int32_t merged(std::int32_t cur, std::int32_t want) +{ + cached_chunk c(0); + set_cache_from(c, cur); + auto lk = c.state.get_lock(); + merge_cache_from(c, want); + return c.cache_from.load(std::memory_order_relaxed); +} + +} // namespace + +TEST_CASE("merge_cache_from widens coverage per the insert rules", "[cache][chunking]") +{ + std::int32_t const left_small = static_cast(page); + std::int32_t const left_big = static_cast(4 * page); + std::int32_t const right_small = -static_cast(page); + std::int32_t const right_big = -static_cast(4 * page); + + SECTION("already full stays full") + { + REQUIRE(merged(0, left_big) == 0); + REQUIRE(merged(0, right_big) == 0); + } + + SECTION("merging full request makes the chunk full") + { + REQUIRE(merged(left_small, 0) == 0); + REQUIRE(merged(right_small, 0) == 0); + } + + SECTION("same side: keep the wider extent") + { + REQUIRE(merged(left_small, left_big) == left_big); // grow + REQUIRE(merged(left_big, left_small) == left_big); // already covered -> ignore + REQUIRE(merged(right_small, right_big) == right_big); // grow + REQUIRE(merged(right_big, right_small) == right_big); // already covered -> ignore + } + + SECTION("opposite sides together span the chunk -> full") + { + REQUIRE(merged(left_small, right_small) == 0); + REQUIRE(merged(right_big, left_big) == 0); + } +} + +// --------------------------------------------------------------------------- +// entry_state — lock-based state machine +// --------------------------------------------------------------------------- + +TEST_CASE("entry_state follows the allocate/load/read lifecycle", "[cache][state_machine]") +{ + entry_state s; + REQUIRE(s.get_state() == entry_state::empty); + REQUIRE(s.get_pin_count() == 0); + + SECTION("happy path empty -> queued -> allocated -> loading -> cached") + { + REQUIRE(s.mark_queued()); + REQUIRE(s.get_state() == entry_state::queued); + REQUIRE(s.mark_allocated()); + REQUIRE(s.get_state() == entry_state::allocated); + REQUIRE(s.mark_loading()); + REQUIRE(s.get_state() == entry_state::loading); + REQUIRE(s.mark_cached()); + REQUIRE(s.get_state() == entry_state::cached); + } + + SECTION("preconditions reject out-of-order transitions") + { + REQUIRE_FALSE(s.mark_allocated()); // empty, not queued + REQUIRE_FALSE(s.mark_loading()); // empty, not allocated + REQUIRE_FALSE(s.mark_cached()); // empty, not loading + REQUIRE(s.mark_queued()); + REQUIRE_FALSE(s.mark_queued()); // no longer empty + } + + SECTION("read pins: acquire nests, release unwinds to cached") + { + REQUIRE(s.mark_queued()); + REQUIRE(s.mark_allocated()); + REQUIRE(s.mark_loading()); + REQUIRE(s.mark_cached()); + + REQUIRE(s.acquire_read()); + REQUIRE(s.get_state() == entry_state::in_use); + REQUIRE(s.get_pin_count() == 1); + REQUIRE(s.acquire_read()); + REQUIRE(s.get_pin_count() == 2); + + REQUIRE_FALSE(s.release_read()); // still pinned + REQUIRE(s.get_pin_count() == 1); + REQUIRE(s.release_read()); // last reader + REQUIRE(s.get_state() == entry_state::cached); + REQUIRE(s.get_pin_count() == 0); + } + + SECTION("eviction only from unpinned allocated/cached") + { + REQUIRE(s.mark_queued()); + REQUIRE(s.mark_allocated()); + REQUIRE(s.mark_loading()); + REQUIRE(s.mark_cached()); + REQUIRE(s.acquire_read()); + REQUIRE_FALSE(s.mark_evicting()); // pinned in_use -> rejected + REQUIRE(s.release_read()); + REQUIRE(s.mark_evicting()); // unpinned cached -> ok + REQUIRE(s.get_state() == entry_state::evicting); + REQUIRE(s.mark_empty()); + REQUIRE(s.get_state() == entry_state::empty); + } + + SECTION("load failure reverts loading -> allocated") + { + REQUIRE(s.mark_queued()); + REQUIRE(s.mark_allocated()); + REQUIRE(s.mark_loading()); + REQUIRE(s.mark_load_failed()); + REQUIRE(s.get_state() == entry_state::allocated); + } + + SECTION("mark_loading_in_use pins directly out of loading") + { + REQUIRE(s.mark_queued()); + REQUIRE(s.mark_allocated()); + REQUIRE(s.mark_loading()); + REQUIRE(s.mark_loading_in_use()); + REQUIRE(s.get_state() == entry_state::in_use); + REQUIRE(s.get_pin_count() == 1); + } +} + +TEST_CASE("entry_state get_lock/state_locked expose the guarded state", "[cache][state_machine]") +{ + entry_state s; + REQUIRE(s.mark_queued()); + { + auto lk = s.get_lock(); + REQUIRE(s.state_locked() == entry_state::queued); + } + REQUIRE(s.mark_allocated()); + { + auto lk = s.get_lock(); + REQUIRE(s.state_locked() == entry_state::allocated); + } +} + +TEST_CASE("entry_state read pins are consistent under concurrent acquire/release", + "[cache][state_machine]") +{ + // Drives the spin_lock under contention: many threads each acquire and + // release a read pin repeatedly. The pin count must always return to 0 and + // the entry must settle back to `cached`, with no lost or double counts. + entry_state s; + REQUIRE(s.mark_queued()); + REQUIRE(s.mark_allocated()); + REQUIRE(s.mark_loading()); + REQUIRE(s.mark_cached()); + + constexpr int n_threads = 8; + constexpr int iters_per_thread = 5000; + // Catch2's REQUIRE is not thread-safe, so workers record violations into + // atomics and the main thread asserts on them after the join. + std::atomic acquire_failures{0}; + std::atomic state_violations{0}; + + std::vector workers; + workers.reserve(n_threads); + for (int t = 0; t < n_threads; ++t) { + workers.emplace_back([&] { + for (int i = 0; i < iters_per_thread; ++i) { + if (!s.acquire_read()) { + acquire_failures.fetch_add(1, std::memory_order_relaxed); + continue; + } + // Holding a pin, the entry must be readable and pinned. + if (s.get_state() != entry_state::in_use || s.get_pin_count() < 1) { + state_violations.fetch_add(1, std::memory_order_relaxed); + } + s.release_read(); + } + }); + } + for (auto& w : workers) { + w.join(); + } + + // A readable entry never rejects acquire_read(); every acquire must succeed. + REQUIRE(acquire_failures.load() == 0); + REQUIRE(state_violations.load() == 0); + // All pins released -> back to cached with pin count 0. + REQUIRE(s.get_pin_count() == 0); + REQUIRE(s.get_state() == entry_state::cached); +}