tenferro_fft/
cache.rs

1use std::collections::hash_map::DefaultHasher;
2use std::fmt;
3use std::hash::{Hash, Hasher};
4use std::num::NonZeroUsize;
5use std::sync::Arc;
6
7use num_traits::{Float, FromPrimitive};
8use realfft::{ComplexToReal, RealFftPlanner, RealToComplex};
9use rustfft::{Fft, FftNum, FftPlanner};
10use tenferro_runtime::{
11    ExtensionCacheKey, ExtensionCacheLimits, ExtensionCacheSelector, ExtensionCacheStore,
12};
13use tenferro_tensor::{CacheStats, RuntimeCacheControl};
14
15use crate::FFT_EXTENSION_FAMILY_ID;
16
17/// Runtime cache namespace used for private RustFFT plans.
18pub const FFT_PLAN_CACHE_NAME: &str = "rustfft-plans";
19
20/// Default number of typed entries retained by a caller-owned [`FftPlanCache`].
21pub const DEFAULT_FFT_PLAN_CACHE_CAPACITY: usize = 64;
22
23/// Select the private CPU RustFFT plan entries in an extension runtime cache.
24///
25/// Other [`crate::FftBackend`] implementations use distinct cache names in the
26/// same FFT extension family.
27///
28/// # Examples
29///
30/// ```
31/// let selector = tenferro_fft::fft_plan_cache_selector();
32/// assert!(matches!(
33///     selector,
34///     tenferro_runtime::ExtensionCacheSelector::Cache { .. }
35/// ));
36/// ```
37pub const fn fft_plan_cache_selector() -> ExtensionCacheSelector {
38    ExtensionCacheSelector::Cache {
39        family_id: FFT_EXTENSION_FAMILY_ID,
40        cache_name: FFT_PLAN_CACHE_NAME,
41    }
42}
43
44/// Bounded, caller-owned typed cache for backend FFT plans and workspaces.
45///
46/// [`crate::FftExecutor`] owns one cache and passes its store to every
47/// [`crate::FftBackend`] through [`crate::FftExecutionCache`]. Backends retain
48/// private `Send + Sync + 'static` values under their own
49/// [`ExtensionCacheKey`] namespace. Entry limits, LRU eviction, clearing,
50/// aggregate entry counts, and retained-byte accounting apply uniformly to CPU
51/// and non-CPU entries.
52///
53/// The CPU backend stores RustFFT plans in the private `rustfft-plans`
54/// namespace. Its retained-byte estimate includes the exact plan key and the
55/// cache-owned `Arc` handle; allocations opaque to RustFFT are excluded.
56///
57/// # Examples
58///
59/// ```
60/// use std::num::NonZeroUsize;
61/// use tenferro_fft::FftPlanCache;
62///
63/// let cache = FftPlanCache::with_capacity(NonZeroUsize::new(2).unwrap());
64/// assert_eq!(cache.capacity().get(), 2);
65/// assert_eq!(cache.stats().entries, 0);
66/// ```
67pub struct FftPlanCache {
68    store: ExtensionCacheStore,
69}
70
71impl fmt::Debug for FftPlanCache {
72    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73        f.debug_struct("FftPlanCache")
74            .field("capacity", &self.capacity())
75            .field("stats", &self.stats())
76            .finish_non_exhaustive()
77    }
78}
79
80impl FftPlanCache {
81    /// Create an empty typed cache with an explicit maximum entry count.
82    pub fn with_capacity(capacity: NonZeroUsize) -> Self {
83        Self {
84            store: ExtensionCacheStore::with_limits(ExtensionCacheLimits::new(capacity)),
85        }
86    }
87
88    /// Maximum number of retained backend entries across all namespaces.
89    pub fn capacity(&self) -> NonZeroUsize {
90        self.store.limits().max_entries()
91    }
92
93    /// Return complete cache retention limits.
94    pub fn limits(&self) -> ExtensionCacheLimits {
95        self.store.limits()
96    }
97
98    /// Replace complete cache retention limits.
99    pub fn set_limits(&mut self, limits: ExtensionCacheLimits) {
100        self.store.set_limits(limits);
101    }
102
103    /// Resize the cache, evicting least-recently-used entries when necessary.
104    pub fn set_capacity(&mut self, capacity: NonZeroUsize) {
105        let mut limits = ExtensionCacheLimits::new(capacity);
106        if let Some(max_retained_bytes) = self.store.limits().max_retained_bytes() {
107            limits = limits.with_max_retained_bytes(max_retained_bytes);
108        }
109        self.store.set_limits(limits);
110    }
111
112    /// Remove every retained backend plan or workspace.
113    pub fn clear(&mut self) {
114        self.store.clear();
115    }
116
117    /// Snapshot aggregate entries and backend-reported retained bytes.
118    pub fn stats(&self) -> CacheStats {
119        self.store.stats(ExtensionCacheSelector::All)
120    }
121
122    pub(crate) fn store_mut(&mut self) -> &mut ExtensionCacheStore {
123        &mut self.store
124    }
125
126    pub(crate) fn plan_f32(&mut self, len: usize, forward: bool) -> Arc<dyn Fft<f32>> {
127        ExtensionFftPlanCache::new(&mut self.store).plan_f32(len, forward)
128    }
129
130    pub(crate) fn plan_f64(&mut self, len: usize, forward: bool) -> Arc<dyn Fft<f64>> {
131        ExtensionFftPlanCache::new(&mut self.store).plan_f64(len, forward)
132    }
133
134    pub(crate) fn real_plan_f32(&mut self, len: usize, forward: bool) -> RealPlan<f32> {
135        ExtensionFftPlanCache::new(&mut self.store).real_plan_f32(len, forward)
136    }
137
138    pub(crate) fn real_plan_f64(&mut self, len: usize, forward: bool) -> RealPlan<f64> {
139        ExtensionFftPlanCache::new(&mut self.store).real_plan_f64(len, forward)
140    }
141
142    #[cfg(test)]
143    pub(crate) fn contains_f64(&mut self, len: usize, forward: bool) -> bool {
144        let key = FftPlanKey {
145            len,
146            forward,
147            dtype: FftPlanDType::F64,
148            real: false,
149        };
150        self.store
151            .get::<ExtensionFftPlanEntry>(&extension_plan_key(key))
152            .is_some_and(|entry| entry.matches_f64(key))
153    }
154}
155
156impl Default for FftPlanCache {
157    fn default() -> Self {
158        Self::with_capacity(
159            NonZeroUsize::new(DEFAULT_FFT_PLAN_CACHE_CAPACITY).unwrap_or(NonZeroUsize::MIN),
160        )
161    }
162}
163
164impl RuntimeCacheControl for FftPlanCache {
165    fn clear(&mut self) {
166        Self::clear(self);
167    }
168
169    fn stats(&self) -> CacheStats {
170        Self::stats(self)
171    }
172}
173
174#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
175enum FftPlanDType {
176    F32,
177    F64,
178}
179
180#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
181struct FftPlanKey {
182    len: usize,
183    forward: bool,
184    dtype: FftPlanDType,
185    real: bool,
186}
187
188enum CachedFftPlan {
189    F32(Arc<dyn Fft<f32>>),
190    F64(Arc<dyn Fft<f64>>),
191    RealF32(Arc<dyn RealToComplex<f32>>),
192    RealF64(Arc<dyn RealToComplex<f64>>),
193    InverseF32(Arc<dyn ComplexToReal<f32>>),
194    InverseF64(Arc<dyn ComplexToReal<f64>>),
195}
196
197/// Cached real-FFT plan for one transform direction.
198pub(crate) enum RealPlan<T> {
199    Forward(Arc<dyn RealToComplex<T>>),
200    Inverse(Arc<dyn ComplexToReal<T>>),
201}
202
203struct ExtensionFftPlanEntry {
204    key: FftPlanKey,
205    plan: CachedFftPlan,
206}
207
208impl ExtensionFftPlanEntry {
209    #[cfg(test)]
210    fn matches_f64(&self, key: FftPlanKey) -> bool {
211        self.key == key && matches!(self.plan, CachedFftPlan::F64(_))
212    }
213}
214
215pub(crate) trait FftPlanProvider: Send {
216    fn plan_f32(&mut self, len: usize, forward: bool) -> Arc<dyn Fft<f32>>;
217    fn plan_f64(&mut self, len: usize, forward: bool) -> Arc<dyn Fft<f64>>;
218    fn real_plan_f32(&mut self, len: usize, forward: bool) -> RealPlan<f32>;
219    fn real_plan_f64(&mut self, len: usize, forward: bool) -> RealPlan<f64>;
220}
221
222impl FftPlanProvider for FftPlanCache {
223    fn plan_f32(&mut self, len: usize, forward: bool) -> Arc<dyn Fft<f32>> {
224        Self::plan_f32(self, len, forward)
225    }
226
227    fn plan_f64(&mut self, len: usize, forward: bool) -> Arc<dyn Fft<f64>> {
228        Self::plan_f64(self, len, forward)
229    }
230
231    fn real_plan_f32(&mut self, len: usize, forward: bool) -> RealPlan<f32> {
232        Self::real_plan_f32(self, len, forward)
233    }
234
235    fn real_plan_f64(&mut self, len: usize, forward: bool) -> RealPlan<f64> {
236        Self::real_plan_f64(self, len, forward)
237    }
238}
239
240pub(crate) trait CachedFftPlanScalar: FftNum + Float + FromPrimitive + 'static {
241    fn plan<P: FftPlanProvider + ?Sized>(
242        plans: &mut P,
243        len: usize,
244        forward: bool,
245    ) -> Arc<dyn Fft<Self>>;
246
247    fn real_plan<P: FftPlanProvider + ?Sized>(
248        plans: &mut P,
249        len: usize,
250        forward: bool,
251    ) -> RealPlan<Self>;
252}
253
254impl CachedFftPlanScalar for f32 {
255    fn plan<P: FftPlanProvider + ?Sized>(
256        plans: &mut P,
257        len: usize,
258        forward: bool,
259    ) -> Arc<dyn Fft<Self>> {
260        plans.plan_f32(len, forward)
261    }
262
263    fn real_plan<P: FftPlanProvider + ?Sized>(
264        plans: &mut P,
265        len: usize,
266        forward: bool,
267    ) -> RealPlan<Self> {
268        plans.real_plan_f32(len, forward)
269    }
270}
271
272impl CachedFftPlanScalar for f64 {
273    fn plan<P: FftPlanProvider + ?Sized>(
274        plans: &mut P,
275        len: usize,
276        forward: bool,
277    ) -> Arc<dyn Fft<Self>> {
278        plans.plan_f64(len, forward)
279    }
280
281    fn real_plan<P: FftPlanProvider + ?Sized>(
282        plans: &mut P,
283        len: usize,
284        forward: bool,
285    ) -> RealPlan<Self> {
286        plans.real_plan_f64(len, forward)
287    }
288}
289
290pub(crate) fn cached_fft_plan<T: CachedFftPlanScalar, P: FftPlanProvider + ?Sized>(
291    plans: &mut P,
292    len: usize,
293    forward: bool,
294) -> Arc<dyn Fft<T>> {
295    T::plan(plans, len, forward)
296}
297
298pub(crate) fn cached_real_plan<T: CachedFftPlanScalar, P: FftPlanProvider + ?Sized>(
299    plans: &mut P,
300    len: usize,
301    forward: bool,
302) -> RealPlan<T> {
303    T::real_plan(plans, len, forward)
304}
305
306pub(crate) struct ExtensionFftPlanCache<'a> {
307    entries: &'a mut ExtensionCacheStore,
308}
309
310impl<'a> ExtensionFftPlanCache<'a> {
311    pub(crate) fn new(entries: &'a mut ExtensionCacheStore) -> Self {
312        Self { entries }
313    }
314}
315
316fn extension_plan_key(key: FftPlanKey) -> ExtensionCacheKey {
317    let mut hasher = DefaultHasher::new();
318    key.hash(&mut hasher);
319    ExtensionCacheKey::new(
320        FFT_EXTENSION_FAMILY_ID,
321        FFT_PLAN_CACHE_NAME,
322        hasher.finish(),
323    )
324}
325
326impl FftPlanProvider for ExtensionFftPlanCache<'_> {
327    fn plan_f32(&mut self, len: usize, forward: bool) -> Arc<dyn Fft<f32>> {
328        let key = FftPlanKey {
329            len,
330            forward,
331            dtype: FftPlanDType::F32,
332            real: false,
333        };
334        let cache_key = extension_plan_key(key);
335        if let Some(cached) = self.entries.get::<ExtensionFftPlanEntry>(&cache_key) {
336            if cached.key == key {
337                if let CachedFftPlan::F32(plan) = &cached.plan {
338                    return Arc::clone(plan);
339                }
340            }
341        }
342        let plan = build_fft_plan::<f32>(len, forward);
343        self.entries.put(
344            cache_key,
345            ExtensionFftPlanEntry {
346                key,
347                plan: CachedFftPlan::F32(Arc::clone(&plan)),
348            },
349            fft_plan_retained_bytes(),
350        );
351        plan
352    }
353
354    fn plan_f64(&mut self, len: usize, forward: bool) -> Arc<dyn Fft<f64>> {
355        let key = FftPlanKey {
356            len,
357            forward,
358            dtype: FftPlanDType::F64,
359            real: false,
360        };
361        let cache_key = extension_plan_key(key);
362        if let Some(cached) = self.entries.get::<ExtensionFftPlanEntry>(&cache_key) {
363            if cached.key == key {
364                if let CachedFftPlan::F64(plan) = &cached.plan {
365                    return Arc::clone(plan);
366                }
367            }
368        }
369        let plan = build_fft_plan::<f64>(len, forward);
370        self.entries.put(
371            cache_key,
372            ExtensionFftPlanEntry {
373                key,
374                plan: CachedFftPlan::F64(Arc::clone(&plan)),
375            },
376            fft_plan_retained_bytes(),
377        );
378        plan
379    }
380
381    fn real_plan_f32(&mut self, len: usize, forward: bool) -> RealPlan<f32> {
382        let key = FftPlanKey {
383            len,
384            forward,
385            dtype: FftPlanDType::F32,
386            real: true,
387        };
388        let cache_key = extension_plan_key(key);
389        if let Some(cached) = self.entries.get::<ExtensionFftPlanEntry>(&cache_key) {
390            if cached.key == key {
391                match (&cached.plan, forward) {
392                    (CachedFftPlan::RealF32(plan), true) => {
393                        return RealPlan::Forward(Arc::clone(plan));
394                    }
395                    (CachedFftPlan::InverseF32(plan), false) => {
396                        return RealPlan::Inverse(Arc::clone(plan));
397                    }
398                    _ => {}
399                }
400            }
401        }
402        let plan = build_real_plan::<f32>(len, forward);
403        let cached = match &plan {
404            RealPlan::Forward(plan) => CachedFftPlan::RealF32(Arc::clone(plan)),
405            RealPlan::Inverse(plan) => CachedFftPlan::InverseF32(Arc::clone(plan)),
406        };
407        self.entries.put(
408            cache_key,
409            ExtensionFftPlanEntry { key, plan: cached },
410            fft_plan_retained_bytes(),
411        );
412        plan
413    }
414
415    fn real_plan_f64(&mut self, len: usize, forward: bool) -> RealPlan<f64> {
416        let key = FftPlanKey {
417            len,
418            forward,
419            dtype: FftPlanDType::F64,
420            real: true,
421        };
422        let cache_key = extension_plan_key(key);
423        if let Some(cached) = self.entries.get::<ExtensionFftPlanEntry>(&cache_key) {
424            if cached.key == key {
425                match (&cached.plan, forward) {
426                    (CachedFftPlan::RealF64(plan), true) => {
427                        return RealPlan::Forward(Arc::clone(plan));
428                    }
429                    (CachedFftPlan::InverseF64(plan), false) => {
430                        return RealPlan::Inverse(Arc::clone(plan));
431                    }
432                    _ => {}
433                }
434            }
435        }
436        let plan = build_real_plan::<f64>(len, forward);
437        let cached = match &plan {
438            RealPlan::Forward(plan) => CachedFftPlan::RealF64(Arc::clone(plan)),
439            RealPlan::Inverse(plan) => CachedFftPlan::InverseF64(Arc::clone(plan)),
440        };
441        self.entries.put(
442            cache_key,
443            ExtensionFftPlanEntry { key, plan: cached },
444            fft_plan_retained_bytes(),
445        );
446        plan
447    }
448}
449
450fn build_real_plan<T: FftNum + 'static>(len: usize, forward: bool) -> RealPlan<T> {
451    let mut planner = RealFftPlanner::<T>::new();
452    if forward {
453        RealPlan::Forward(planner.plan_fft_forward(len))
454    } else {
455        RealPlan::Inverse(planner.plan_fft_inverse(len))
456    }
457}
458
459fn build_fft_plan<T: FftNum + 'static>(len: usize, forward: bool) -> Arc<dyn Fft<T>> {
460    let mut planner = FftPlanner::<T>::new();
461    if forward {
462        planner.plan_fft_forward(len)
463    } else {
464        planner.plan_fft_inverse(len)
465    }
466}
467
468const fn fft_plan_retained_bytes() -> usize {
469    std::mem::size_of::<FftPlanKey>() + std::mem::size_of::<CachedFftPlan>()
470}
tenferro_fft/cache.rs

tenferro_fft/
cache.rs
