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"); + }); +}