From 09ece2d9993678bb669f04ba93c11e303d1f5261 Mon Sep 17 00:00:00 2001 From: Robert Deaton Date: Tue, 22 Sep 2026 18:04:05 -0700 Subject: [PATCH] fix(cutile-rs): keep every launch-site specialization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A generated launcher's `LaunchSite` held only its most recent resolution, so a kernel launched alternately at two specializations (the same kernel at two shapes in one model forward, or a scalar whose divisibility changes) missed on every launch and resolved again through key construction and the kernel cache. In the `host_overhead` bench (RTX 5080, mean of two alternating runs against main) that miss path, `launch_site_miss`, goes from 3.84 to 1.72 µs a launch; `launch` with one specialization is unchanged (1.65 / 1.66 µs). The new `launch_alternating` and `launch_12_specializations` cases cost 1.66 and 1.67 µs a launch. The site now keeps one entry per probe it resolved. A lookup compares against the most recently stored entry first, then looks the probe up by its FxHash. Because a site now holds several functions, evictions have to reach it. Every eviction that removes a kernel bumps the cache epoch and then empties every launch site in the process, so an eviction of any key makes every site resolve again on its next launch. The launcher reads a `CacheEpoch` before resolving, and `LaunchSite::store` discards a resolution whose epoch went stale meanwhile. This also releases modules that main kept loaded: a site that had resolved an evicted kernel held its function until it stored another resolution. `LaunchSite`, `SiteResolution` and the new `CacheEpoch` are now `#[doc(hidden)]` launcher support for generated code; `SiteResolution::new` takes the epoch, and `LaunchSite::store` takes `&'static self` so the site can register itself. A poisoned site lock is still a miss, as on main; only an eviction recovers it, to empty the site. The tests evict only their own kernel, or bump the epoch through a new test-only hook that leaves the kernel cache alone, so they don't unload modules other tests in the binary are using. Signed-off-by: Robert Deaton --- cutile-rs/CHANGELOG.md | 15 + cutile-rs/Cargo.lock | 1 + .../benches/host_overhead.rs | 91 +++++- .../src/kernel_launcher_generator.rs | 13 +- cutile-rs/cutile/Cargo.toml | 1 + cutile-rs/cutile/src/tile_kernel.rs | 291 +++++++++++++++--- cutile-rs/cutile/tests/gpu/launch_site.rs | 289 ++++++++++++++++- 7 files changed, 643 insertions(+), 58 deletions(-) diff --git a/cutile-rs/CHANGELOG.md b/cutile-rs/CHANGELOG.md index 4b7ec6eff9..e2ea314bb0 100644 --- a/cutile-rs/CHANGELOG.md +++ b/cutile-rs/CHANGELOG.md @@ -16,6 +16,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/). [nvidia.github.io/cuda-rust/cutile](https://nvidia.github.io/cuda-rust/cutile/) and release tags are namespaced `cutile-rs/vX.Y.Z`. Crate names and versions are unchanged. +- A generated launcher's launch site keeps every specialization it resolved + instead of only the most recent one, so a kernel launched at alternating + shapes no longer resolves through the kernel cache on every launch. +- Every eviction that removes a kernel (`clear_kernel_cache`, `evict_kernel`, + `retain_kernels`) empties every launch site in the process, so the next + launch of each specialization resolves again. `LaunchSite`, + `SiteResolution` and the new `CacheEpoch` are `#[doc(hidden)]` launcher + support for generated code: `SiteResolution::new` takes the `CacheEpoch` + read before resolving, and `LaunchSite::store` takes `&'static self`. + +### Fixed + +- An eviction releases the modules that launch sites held. Before, a launch + site that had resolved an evicted kernel kept its function, and the + module, loaded until that site stored another resolution. ## [0.4.0] - 2026-09-25 diff --git a/cutile-rs/Cargo.lock b/cutile-rs/Cargo.lock index d95ed3effd..e3f2aac630 100644 --- a/cutile-rs/Cargo.lock +++ b/cutile-rs/Cargo.lock @@ -484,6 +484,7 @@ dependencies = [ "once_cell", "proc-macro2", "quote", + "rustc-hash 2.1.2", "serde", "serde_json", "sha2", diff --git a/cutile-rs/cutile-benchmarks/benches/host_overhead.rs b/cutile-rs/cutile-benchmarks/benches/host_overhead.rs index ceae70d232..44bcab5218 100644 --- a/cutile-rs/cutile-benchmarks/benches/host_overhead.rs +++ b/cutile-rs/cutile-benchmarks/benches/host_overhead.rs @@ -7,18 +7,26 @@ //! //! Times only the enqueue (`async_on`), never the wait: a batch of launches //! is timed, then the stream is drained outside the timer so the launch -//! queue never fills. Two shapes that matter for decode-style serving: +//! queue never fills. Cases that matter for decode-style serving include: //! //! - `launch`: one tiny 4-tensor kernel launch (argument retention, launch //! validation, `cuLaunchKernel`). //! - `graph_replay`: one replay of a graph recording 32 such launches over //! 8 buffers (per-replay resource reacquisition plus `cuGraphLaunch`). +//! - `launch_alternating`: two launches per iteration whose inputs +//! alternate between two specializations (extent divisible by 16, and by +//! 8 only), through the generated launcher. +//! - `launch_12_specializations`: twelve launches per iteration, each at a +//! different specialization, through the generated launcher: a kernel +//! launched at many shapes must find each one in its launch site instead +//! of resolving it again. use criterion::{criterion_group, criterion_main, Criterion}; use cuda_async::cuda_graph::CudaGraph; -use cuda_core::Device; +use cuda_core::{Device, Stream}; use cutile::prelude::*; use std::future::IntoFuture; +use std::sync::Arc; use std::time::{Duration, Instant}; #[cutile::module] @@ -174,6 +182,67 @@ fn host_overhead(c: &mut Criterion) { }); }); + // ── Alternating specializations ───────────────────────────────────── + // Inputs of 128 elements (divisible by 16) and 136 (by 8 only): the + // same kernel at two specializations, as a model forward runs it. + let mut zt = api::zeros::(&[N]).sync_on(&stream).expect("zt"); + let odd = || -> Arc> { + api::ones::(&[N + 8]) + .sync_on(&stream) + .expect("odd") + .into() + }; + let (a2, b2, c2) = (odd(), odd(), odd()); + kernels::add3((&mut zt).partition([N]), &*a2, &*b2, &*c2) + .sync_on(&stream) + .expect("warm second specialization"); + + group.bench_function("launch_alternating", |bench| { + bench.iter_custom(|iters| { + enqueue_time(iters, &stream, || { + for (x, y, w) in [(&a, &b, &c3), (&a2, &b2, &c2)] { + unsafe { + kernels::add3((&mut zt).partition([N]), &**x, &**y, &**w).async_on(&stream) + } + .expect("launch"); + } + }) + }); + }); + + // Extents divisible by 16, 8, 4, 2 and 1: `a` and `b` at independent + // divisibilities give twelve specializations of one kernel. + let extents = [N, N + 8, N + 4, N + 2, N + 1]; + let input = |len: usize| -> Arc> { + api::ones::(&[len]) + .sync_on(&stream) + .expect("input") + .into() + }; + let many: Vec<(Arc>, Arc>)> = extents + .iter() + .flat_map(|&la| extents[..3].iter().map(move |&lb| (la, lb))) + .take(12) + .map(|(la, lb)| (input(la), input(lb))) + .collect(); + for (x, y) in &many { + kernels::add3((&mut zt).partition([N]), &**x, &**y, &*c3) + .sync_on(&stream) + .expect("warm specialization"); + } + group.bench_function("launch_12_specializations", |bench| { + bench.iter_custom(|iters| { + enqueue_time(iters, &stream, || { + for (x, y) in &many { + unsafe { + kernels::add3((&mut zt).partition([N]), &**x, &**y, &*c3).async_on(&stream) + } + .expect("launch"); + } + }) + }); + }); + // A graph of 32 launches over 8 buffers, replayed. let mut bufs: Vec> = (0..8) .map(|_| api::zeros::(&[N]).sync_on(&stream).expect("buf")) @@ -213,6 +282,24 @@ fn host_overhead(c: &mut Criterion) { group.finish(); } +/// Host time to enqueue `iters` calls of `enqueue`, draining the stream every +/// [`BATCH`] calls outside the timer. +fn enqueue_time(iters: u64, stream: &Stream, mut enqueue: impl FnMut()) -> Duration { + let mut elapsed = Duration::ZERO; + let mut done = 0; + while done < iters { + let batch = BATCH.min(iters - done); + let start = Instant::now(); + for _ in 0..batch { + enqueue(); + } + elapsed += start.elapsed(); + unsafe { stream.synchronize() }.expect("drain"); + done += batch; + } + elapsed +} + fn bench_config() -> Criterion { if cfg!(feature = "smoke-test") { Criterion::default() diff --git a/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs b/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs index 27dac6ab6c..8143789c95 100644 --- a/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs +++ b/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs @@ -926,10 +926,11 @@ pub fn generate_kernel_launcher( .stmts .push(scalar_hints_stmt.clone()); - // Launch-site probe: on the steady state the launcher skips key - // construction, toolchain fingerprinting, hashing, and the global cache - // entirely — the volatile specialization inputs are compared against the - // site's last resolution by value, borrowed from the arguments. + // Launch-site probe: on a hit the launcher skips key construction, + // toolchain fingerprinting and the global cache. The volatile + // specialization inputs, borrowed from the arguments, are compared by + // value against the site's newest entry, then against the entries under + // their probe hash. // Fixed-size arrays: the probe borrows them, so a steady-state launch // allocates nothing here. launcher_method.block.stmts.push(parse_stmt(format!( @@ -967,6 +968,9 @@ pub fn generate_kernel_launcher( spec_args.iter().map(|(_, s)| s.clone()).collect(); let __hints_snapshot = __hint_vals.to_vec(); let __options_snapshot = __compile_options.clone(); + // Read before resolving, so the site can discard the resolution + // if an eviction removes its function meanwhile (see `CacheEpoch`). + let __epoch = {root}::tile_kernel::CacheEpoch::now(); let (function, validator) = self.jit_compile( ctx, __module_ast_self, module_name, function_name, function_entry, @@ -976,6 +980,7 @@ pub fn generate_kernel_launcher( _SOURCE_HASH, )?; __SITE.store({root}::tile_kernel::SiteResolution::new( + __epoch, ctx.get_device_id(), __generics_snapshot, __specs_snapshot, diff --git a/cutile-rs/cutile/Cargo.toml b/cutile-rs/cutile/Cargo.toml index 2bcf776c43..0ce54a3154 100644 --- a/cutile-rs/cutile/Cargo.toml +++ b/cutile-rs/cutile/Cargo.toml @@ -32,6 +32,7 @@ linkme = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } dashmap = { workspace = true } +rustc-hash = { workspace = true } once_cell = { workspace = true } [dev-dependencies] diff --git a/cutile-rs/cutile/src/tile_kernel.rs b/cutile-rs/cutile/src/tile_kernel.rs index 503c826e95..ad7c664d75 100644 --- a/cutile-rs/cutile/src/tile_kernel.rs +++ b/cutile-rs/cutile/src/tile_kernel.rs @@ -378,17 +378,21 @@ pub fn _specialization_from_context Module>( } } -/// One macro-generated launch site's most recent resolution: the volatile -/// parts of the specialization it was resolved for, plus the resolved -/// function and validator. +/// One macro-generated launch site's resolutions: for each, the +/// volatile parts of the specialization it was resolved for, plus the +/// resolved function and validator. Launcher support for generated code, +/// not public API. /// /// The generated launcher probes this before constructing a -/// [`TileFunctionKey`]: on the steady state the probe is a handful of -/// integer/`String` comparisons and two `Arc` clones, with no allocation and -/// no toolchain-fingerprint work. Probe equality implies key equality: the -/// process-constant key fields (names, source hash, compiler version) cannot -/// differ at one site, `gpu_name` is a function of `device_id`, and stride -/// hints are derived from the compared [`SpecializationBits`]. +/// [`TileFunctionKey`]. The probe first compares the launch's specialization +/// inputs field by field against the newest entry (the most recently +/// stored); on a mismatch it hashes them (FxHash) and compares against the +/// entries in that hash's bucket. A hit clones two `Arc`s and allocates +/// nothing, with no toolchain-fingerprint work. Probe equality implies key +/// equality: the process-constant key fields (names, source hash, compiler +/// version) cannot differ at one site, `gpu_name` is a function of +/// `device_id`, and stride hints are derived from the compared +/// [`SpecializationBits`]. /// /// Toolchain env semantics: the `CUTILE_TILEIRAS_PATH` / toolkit env values /// are snapshotted when the site fills, not re-read per launch — a @@ -396,16 +400,88 @@ pub fn _specialization_from_context Module>( /// global cache still keys by fingerprint) but does not invalidate an /// already-hot launch site. Reading the env on every launch costs /// allocations on the very path this cache exists to strip. +/// +/// Size: a site keeps one entry per probe it resolved. Its functions come +/// from the in-memory kernel cache, and every eviction that removes a +/// kernel ([`clear_kernel_cache`], [`evict_kernel`], [`retain_kernels`]) +/// empties every site, so a site never holds more entries than the kernel +/// cache holds specializations of its kernel. An eviction of any key +/// empties every site in the process; each re-resolves on its next launch +/// of each specialization. +#[doc(hidden)] pub struct LaunchSite { - inner: std::sync::RwLock>>, + entries: std::sync::RwLock, + /// Puts this site in [`LAUNCH_SITES`] on its first store, so evictions + /// reach it. + registered: std::sync::Once, +} + +#[derive(Default)] +struct SiteEntries { + /// The most recently stored resolution, also in `by_hash`. + newest: Option>, + /// Every resolution, by [`probe_hash`]; a bucket holds the rare probes + /// whose hashes collide. + by_hash: std::collections::HashMap>, rustc_hash::FxBuildHasher>, +} + +/// Hashes a launch-site probe. `get` and `store` hold the specializations +/// differently, so they pass them as an iterator. +fn probe_hash<'a>( + device_id: usize, + generics: &[String], + specs: impl ExactSizeIterator, + scalar_hints: &[DivHint], + const_grid: Option<(u32, u32, u32)>, + compile_options: &CompileOptions, +) -> u64 { + use std::hash::{Hash, Hasher}; + let mut hasher = rustc_hash::FxHasher::default(); + device_id.hash(&mut hasher); + generics.hash(&mut hasher); + hasher.write_usize(specs.len()); + for spec in specs { + spec.hash(&mut hasher); + } + scalar_hints.hash(&mut hasher); + const_grid.hash(&mut hasher); + compile_options.hash(&mut hasher); + hasher.finish() +} + +/// Every launch site that has stored a resolution; evictions empty them all. +static LAUNCH_SITES: std::sync::Mutex> = std::sync::Mutex::new(Vec::new()); + +/// Empties every launch site. The functions are dropped after the locks are +/// released: dropping the last reference unloads a module. +fn clear_launch_sites() { + let sites = LAUNCH_SITES + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let drained: Vec = sites + .iter() + .map(|site| { + // Lookups treat a poisoned site as empty, but its entries still + // hold modules: empty it anyway, so the eviction releases them. + std::mem::take( + &mut *site + .entries + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ) + }) + .collect(); + drop(sites); + drop(drained); } /// The snapshot a [`LaunchSite`] holds. Constructed by the generated /// launcher on a probe miss, after the normal cache path resolved. +/// Launcher support for generated code, not public API. +#[doc(hidden)] pub struct SiteResolution { - /// Cache generation at fill time: any eviction from the global cache - /// invalidates every launch site, so no site outlives a quiesced clear. - epoch: u64, + /// Cache generation read before resolving; see [`CacheEpoch`]. + epoch: CacheEpoch, device_id: usize, #[allow(dead_code)] toolchain: ToolchainEnvSnapshot, @@ -419,8 +495,13 @@ pub struct SiteResolution { } impl SiteResolution { + /// `epoch` must be read before the function is resolved. Read later, it + /// could postdate an eviction of that function, and the site would keep + /// the evicted module loaded until the next eviction. The function would + /// still be valid to launch. #[allow(clippy::too_many_arguments)] pub fn new( + epoch: CacheEpoch, device_id: usize, generics: Vec, specs: Vec, @@ -431,7 +512,7 @@ impl SiteResolution { validator: Arc, ) -> Self { Self { - epoch: kernel_cache_epoch(), + epoch, device_id, toolchain: toolchain_env_snapshot(), generics, @@ -445,14 +526,40 @@ impl SiteResolution { } } +impl SiteResolution { + fn same_probe(&self, other: &SiteResolution) -> bool { + self.device_id == other.device_id + && self.const_grid == other.const_grid + && self.compile_options == other.compile_options + && self.generics == other.generics + && self.specs == other.specs + && self.scalar_hints == other.scalar_hints + } + + fn probe_hash(&self) -> u64 { + probe_hash( + self.device_id, + &self.generics, + self.specs.iter(), + &self.scalar_hints, + self.const_grid, + &self.compile_options, + ) + } +} + impl LaunchSite { pub const fn new() -> Self { Self { - inner: std::sync::RwLock::new(None), + entries: std::sync::RwLock::new(SiteEntries { + newest: None, + by_hash: std::collections::HashMap::with_hasher(rustc_hash::FxBuildHasher), + }), + registered: std::sync::Once::new(), } } - /// The cached resolution, if the probe matches it exactly. + /// The cached resolution whose probe matches exactly, if any. pub fn get( &self, device_id: usize, @@ -462,28 +569,102 @@ impl LaunchSite { const_grid: Option<(u32, u32, u32)>, compile_options: &CompileOptions, ) -> Option<(Arc, Arc)> { - let guard = self.inner.read().ok()?; - let r = guard.as_ref()?; - if r.epoch == kernel_cache_epoch() - && r.device_id == device_id - && r.const_grid == const_grid - && &r.compile_options == compile_options - && r.generics.as_slice() == generics - && r.specs.len() == specs.len() - && r.specs.iter().zip(specs).all(|(a, b)| a == *b) - && r.scalar_hints.as_slice() == scalar_hints - { - Some((Arc::clone(&r.function), Arc::clone(&r.validator))) - } else { - None - } + let entries = self.entries.read().ok()?; + // An eviction removes kernels from the cache, bumps the epoch, and + // then empties the sites one by one. Until it reaches this site, + // entries from before the bump are still here; skip them, so a + // launch that sees the bump resolves again through the kernel cache. + // A launch that read the epoch just before the bump can still get + // one: the eviction APIs' safety contract rules that race out, and + // the entry's `Arc` keeps its function loaded either way. + let epoch = CacheEpoch::now(); + let matches = |r: &&Arc| { + r.epoch == epoch + && r.device_id == device_id + && r.const_grid == const_grid + && &r.compile_options == compile_options + && r.generics.as_slice() == generics + && r.specs.len() == specs.len() + && r.specs.iter().zip(specs).all(|(a, b)| a == *b) + && r.scalar_hints.as_slice() == scalar_hints + }; + let r = match entries.newest.as_ref().filter(matches) { + Some(r) => r, + None => { + let hash = probe_hash( + device_id, + generics, + specs.iter().copied(), + scalar_hints, + const_grid, + compile_options, + ); + entries.by_hash.get(&hash)?.iter().find(matches)? + } + }; + Some((Arc::clone(&r.function), Arc::clone(&r.validator))) } - /// Replaces the cached resolution (single entry: last one wins). - pub fn store(&self, resolution: SiteResolution) { - if let Ok(mut guard) = self.inner.write() { - *guard = Some(std::sync::Arc::new(resolution)); + /// Adds `resolution`, replacing one a concurrent miss stored for the same + /// probe. Discards `resolution` if the kernel cache was evicted since its + /// epoch was read: its function may be one the eviction removed. Leaves a + /// poisoned site as it is (lookups miss; evictions still empty it). + /// `'static` because the site registers itself so evictions can empty it. + /// + /// Replacing is sound: either of two current resolutions of one probe is + /// a valid resolution of it. + pub fn store(&'static self, resolution: SiteResolution) { + // Invariant: every entry is current until an eviction empties the + // site. An eviction bumps the epoch, then, holding the site list, + // empties each site under its lock. This site is in the list before + // the lock below is taken (`call_once` returns only after the push, + // on every thread), and the epoch is checked under that lock. So + // either the bump is visible here and the resolution is discarded, + // or the eviction takes this site's lock after us and removes the + // entry. + self.registered.call_once(|| { + LAUNCH_SITES + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(self); + }); + let hash = resolution.probe_hash(); + let Ok(mut entries) = self.entries.write() else { + return; + }; + if resolution.epoch != CacheEpoch::now() { + // Unlock before `resolution` drops: its function may be the last + // reference to an evicted module. + drop(entries); + return; } + let resolution = Arc::new(resolution); + let bucket = entries.by_hash.entry(hash).or_default(); + let replaced = bucket + .iter() + .position(|r| r.same_probe(&resolution)) + .map(|i| bucket.swap_remove(i)); + bucket.push(Arc::clone(&resolution)); + let previous = entries.newest.replace(resolution); + drop(entries); + drop((replaced, previous)); + } +} + +impl LaunchSite { + /// The epoch of every entry this site holds (none if its lock is + /// poisoned). Test-support only, `#[doc(hidden)]`. + #[doc(hidden)] + pub fn entry_epochs_for_tests(&self) -> Vec { + let Ok(entries) = self.entries.read() else { + return Vec::new(); + }; + entries + .by_hash + .values() + .flatten() + .map(|r| r.epoch) + .collect() } } @@ -523,6 +704,19 @@ pub(crate) fn get_kernel_cache() -> &'static DashMap`) re-resolve -/// instead of keeping evicted modules alive past a quiesced clear. +/// Generation counter for the in-memory cache, bumped by every eviction that +/// removes a kernel. Launch sites hold their own `Arc`s, so an +/// eviction bumps this counter and then empties every site. Once an eviction +/// returns, no site holds a function it removed. static KERNEL_CACHE_EPOCH: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); fn kernel_cache_epoch() -> u64 { KERNEL_CACHE_EPOCH.load(std::sync::atomic::Ordering::Acquire) } +/// A generation of the in-memory kernel cache. The generated launcher reads +/// one *before* resolving a function and stores the resolution with it; +/// [`LaunchSite::store`] discards the resolution if an eviction happened +/// since. Read after resolving, it could postdate an eviction of the +/// function being stored. Launcher support for generated code, not public +/// API. +#[doc(hidden)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CacheEpoch(u64); + +impl CacheEpoch { + pub fn now() -> Self { + CacheEpoch(kernel_cache_epoch()) + } +} + fn bump_kernel_cache_epoch() { KERNEL_CACHE_EPOCH.fetch_add(1, std::sync::atomic::Ordering::AcqRel); + clear_launch_sites(); } /// Removes every kernel from the process-global in-memory cache. @@ -565,6 +777,9 @@ fn bump_kernel_cache_epoch() { /// `Arc` to its single-flight slot and completes into it; the next /// request for that key recompiles (or is served by the disk cache). /// +/// Also empties every generated launcher's [`LaunchSite`], which would +/// otherwise keep evicted modules loaded. +/// /// Returns the number of entries removed. Freed device bytes are not /// tracked host-side; per-module sizes are not observable through the /// driver's module API. diff --git a/cutile-rs/cutile/tests/gpu/launch_site.rs b/cutile-rs/cutile/tests/gpu/launch_site.rs index 82f06a61d5..a1ceebf9b7 100644 --- a/cutile-rs/cutile/tests/gpu/launch_site.rs +++ b/cutile-rs/cutile/tests/gpu/launch_site.rs @@ -3,13 +3,26 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! The per-launch-site resolution cache: steady-state launches skip key -//! construction entirely, so these tests pin the two ways a site must NOT -//! serve a stale resolution — a changed specialization (alternating -//! generics through one call site) and an evicted global cache. +//! The per-launch-site resolution cache. A site keeps every specialization +//! it resolved, so a kernel launched at alternating shapes hits on every +//! launch, and an eviction empties every site, so no site serves or keeps +//! loaded a function the eviction removed. +//! +//! Every test here takes `cache_test_lock` and evicts only `site_module`'s +//! kernel (or only bumps the epoch), so no other test's module is unloaded. +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use cuda_async::device_context::Validator; +use cuda_async::device_operation::ExecutionContext; +use cuda_core::{Device, Function}; use cutile::prelude::*; -use cutile::tile_kernel::contains_cuda_function; +use cutile::tile_kernel::{ + compile_from_context, contains_cuda_function, evict_kernel, retain_kernels, CacheEpoch, + CompileOptions, LaunchSite, SiteResolution, TileFunctionKey, +}; use crate::common; @@ -39,12 +52,38 @@ fn run(len: usize, tile: usize) -> Vec { .expect("scale kernel") } +/// Evicts every specialization of `site_module`'s kernel, and nothing else, +/// through `retain_kernels`; returns how many were removed. +/// +/// # Safety +/// No launch of `site_module`'s kernel may be in flight. +unsafe fn evict_site_module() -> usize { + unsafe { retain_kernels(|key| key.module_name() != "site_module") } +} + +/// The kernel-cache keys of `site_module`'s kernel, read without removing +/// anything. +fn site_module_keys() -> Vec { + let mut keys = Vec::new(); + // SAFETY: the predicate keeps every key, so nothing is evicted. + unsafe { + retain_kernels(|key| { + if key.module_name() == "site_module" { + keys.push(key.clone()); + } + true + }) + }; + keys +} + #[test] -fn alternating_specializations_through_one_site_stay_correct() { +fn alternating_specializations_through_one_site_write_correct_results() { common::with_test_stack(|| { let _guard = common::cache_test_lock(); // Two tile sizes = two specializations through the same call site. - // The single-entry site cache thrashes; results must stay correct. + // This checks results only; `a_site_serves_every_specialization_it_stored` + // checks that both specializations hit. for _ in 0..3 { for tile in [4usize, 8] { let host = run(32, tile); @@ -72,20 +111,242 @@ fn cache_eviction_invalidates_hot_launch_sites() { .expect("key"); run(32, 16); // fill the site (and the global cache) assert!(contains_cuda_function(&key), "filled"); - // Quiesced (nothing of ours in flight after sync): evicting must - // force the hot site to re-resolve, not serve its stale - // Arc — the epoch check. - unsafe { - cutile::tile_kernel::clear_kernel_cache_for_tests(); - } + // Quiesced (nothing of ours in flight after sync). Evicting this key + // empties the launcher's site, so the next launch must resolve + // through the global cache again instead of reusing the evicted + // function. + assert!(unsafe { evict_kernel(&key) }, "evicted"); assert!(!contains_cuda_function(&key), "evicted"); let host = run(32, 16); assert!( contains_cuda_function(&key), - "a launch after the clear re-resolved through the global cache" + "a launch after the eviction re-resolved through the global cache" ); for (i, v) in host.iter().enumerate() { assert_eq!(*v, 2.0 * i as f32, "index {i}"); } }); } + +/// Resolves `scale` for `B = b` as the generated launcher does: the epoch +/// is read before the function is. +fn resolve(ctx: &ExecutionContext, b: usize) -> (CacheEpoch, Arc, Arc) { + let epoch = CacheEpoch::now(); + let (function, validator) = compile_from_context( + ctx, + site_module::__module_ast_self, + "site_module", + "scale", + "scale_entry", + vec![b.to_string()], + vec![("z".to_string(), vec![1]), ("x".to_string(), vec![1])], + vec![], + vec![], + None, + CompileOptions::default(), + site_module::_SOURCE_HASH, + ) + .expect("compile"); + (epoch, function, validator) +} + +/// The site entry for a [`resolve`] result for `B = b`. +fn resolution( + epoch: CacheEpoch, + b: usize, + function: &Arc, + validator: &Arc, +) -> SiteResolution { + SiteResolution::new( + epoch, + 0, + vec![b.to_string()], + vec![], + vec![], + None, + CompileOptions::default(), + Arc::clone(function), + Arc::clone(validator), + ) +} + +/// The probe of a [`resolution`] for `B = b`, as `LaunchSite::get` takes it. +fn probe(site: &LaunchSite, b: usize) -> Option<(Arc, Arc)> { + site.get( + 0, + &[b.to_string()], + &[], + &[], + None, + &CompileOptions::default(), + ) +} + +fn context() -> ExecutionContext { + let stream = Device::new(0) + .expect("device") + .new_stream() + .expect("stream"); + ExecutionContext::new(stream) +} + +#[test] +fn a_site_serves_every_specialization_it_stored() { + common::with_test_stack(|| { + let _guard = common::cache_test_lock(); + static SITE: LaunchSite = LaunchSite::new(); + let ctx = context(); + let (epoch16, f16, v16) = resolve(&ctx, 16); + let (epoch8, f8, v8) = resolve(&ctx, 8); + assert!(!Arc::ptr_eq(&f16, &f8), "two specializations"); + let hit = |b: usize, want: &Arc, path: &str| { + let (f, _) = probe(&SITE, b).unwrap_or_else(|| panic!("B = {b} misses ({path})")); + assert!( + Arc::ptr_eq(&f, want), + "B = {b} serves its own function ({path})" + ); + }; + + SITE.store(resolution(epoch16, 16, &f16, &v16)); + SITE.store(resolution(epoch8, 8, &f8, &v8)); + hit(8, &f8, "newest entry"); + hit(16, &f16, "hash lookup"); + assert!( + probe(&SITE, 4).is_none(), + "an unstored specialization misses" + ); + + // Storing B = 16 again makes it the newest; B = 8 now takes the hash + // lookup. The replaced entry is dropped, not kept beside the new one. + SITE.store(resolution(epoch16, 16, &f16, &v16)); + hit(16, &f16, "newest entry"); + hit(8, &f8, "hash lookup"); + assert_eq!( + SITE.entry_epochs_for_tests().len(), + 2, + "one entry per probe" + ); + + // SAFETY: nothing has launched these resolutions. + unsafe { evict_site_module() }; + assert!( + SITE.entry_epochs_for_tests().is_empty(), + "emptied by the eviction" + ); + }); +} + +#[test] +fn launch_sites_never_keep_an_evicted_function() { + common::with_test_stack(|| { + let _guard = common::cache_test_lock(); + static SITE: LaunchSite = LaunchSite::new(); + let ctx = context(); + + // An eviction between resolving and storing: the function it removed + // must not be stored, or the site would keep its module loaded. + let (epoch, function, validator) = resolve(&ctx, 16); + // SAFETY: nothing has launched this kernel. + assert!(unsafe { evict_site_module() } > 0, "evicted"); + assert_ne!(epoch, CacheEpoch::now()); + SITE.store(resolution(epoch, 16, &function, &validator)); + assert_eq!( + Arc::strong_count(&function), + 1, + "stale resolution discarded" + ); + + // Stored while current, then evicted through `retain_kernels`: the + // eviction empties the site. + let (epoch, function, validator) = resolve(&ctx, 16); + SITE.store(resolution(epoch, 16, &function, &validator)); + assert_eq!( + Arc::strong_count(&function), + 3, + "held by the site and the cache" + ); + // SAFETY: as above. + unsafe { evict_site_module() }; + assert_eq!( + Arc::strong_count(&function), + 1, + "site emptied by retain_kernels" + ); + + // The same through `evict_kernel`, one key at a time. + let (epoch, function, validator) = resolve(&ctx, 16); + SITE.store(resolution(epoch, 16, &function, &validator)); + assert_eq!(Arc::strong_count(&function), 3); + for key in site_module_keys() { + // SAFETY: as above. + unsafe { evict_kernel(&key) }; + } + assert_eq!( + Arc::strong_count(&function), + 1, + "site emptied by evict_kernel" + ); + }); +} + +#[test] +fn launch_site_stores_racing_evictions_keep_only_current_entries() { + common::with_test_stack(|| { + let _guard = common::cache_test_lock(); + static SITE: LaunchSite = LaunchSite::new(); + let ctx = context(); + let (_, function, validator) = resolve(&ctx, 16); + // This test's handle and the kernel cache's. The evictions below only + // bump the epoch and empty sites, so the cache keeps its share. + let baseline = Arc::strong_count(&function); + + for round in 0..20 { + let storing = AtomicUsize::new(4); + std::thread::scope(|scope| { + for _ in 0..4 { + scope.spawn(|| { + for _ in 0..20 { + // Stamped as a launcher stamps it; the pause + // stands in for resolving, so an eviction often + // lands between the stamp and the store. + let epoch = CacheEpoch::now(); + std::thread::sleep(Duration::from_micros(50)); + SITE.store(resolution(epoch, 16, &function, &validator)); + } + storing.fetch_sub(1, Ordering::Release); + }); + } + // Evict until every store is in, so the last stores race an + // eviction with no later one to clean up after them. + scope.spawn(|| { + while storing.load(Ordering::Acquire) > 0 { + // SAFETY: every site keeps only functions the kernel + // cache also holds (this one included), so emptying + // them unloads no module; the cache is not touched. + unsafe { cutile::tile_kernel::bump_kernel_cache_epoch_for_tests() }; + std::thread::sleep(Duration::from_micros(10)); + } + }); + }); + // Checked on the entries themselves, not through `get`, which + // skips stale ones: every entry left must be current, and the + // site's share of the function must be exactly its entries. + let epochs = SITE.entry_epochs_for_tests(); + let now = CacheEpoch::now(); + assert!( + epochs.iter().all(|&e| e == now), + "round {round}: a stale entry survived the evictions" + ); + assert!(epochs.len() <= 1, "round {round}: one entry per probe"); + assert_eq!( + Arc::strong_count(&function), + baseline + epochs.len(), + "round {round}: the site holds only its entries" + ); + assert_eq!(probe(&SITE, 16).is_some(), !epochs.is_empty()); + } + // SAFETY: as above. + unsafe { cutile::tile_kernel::bump_kernel_cache_epoch_for_tests() }; + assert_eq!(Arc::strong_count(&function), baseline, "emptied"); + }); +}