scirs2_fft/
fftw_plan_cache.rs

1//! FFTW Plan Cache for performance optimization
2//!
3//! Caches FFTW plans to avoid expensive plan creation on repeated FFT calls.
4//! Plans are stored in global caches keyed by size (1D) or dimensions (2D).
5//!
6//! Note: Plans require &mut self for execution, so we use Mutex to ensure
7//! exclusive access. This is fine for Python bindings where the GIL serializes calls.
8
9use fftw::plan::*;
10use fftw::types::*;
11use std::collections::HashMap;
12use std::sync::Mutex;
13
14use crate::error::{FFTError, FFTResult};
15
16// ========================================
17// 1D PLAN CACHES
18// ========================================
19
20/// Global cache for R2C (real-to-complex) plans, keyed by input size
21static R2C_CACHE: Mutex<Option<HashMap<usize, R2CPlan64>>> = Mutex::new(None);
22
23/// Global cache for C2C forward plans, keyed by size
24static C2C_FWD_CACHE: Mutex<Option<HashMap<usize, C2CPlan64>>> = Mutex::new(None);
25
26/// Global cache for C2C backward plans, keyed by size
27static C2C_BWD_CACHE: Mutex<Option<HashMap<usize, C2CPlan64>>> = Mutex::new(None);
28
29/// Global cache for C2R (complex-to-real) plans, keyed by output size
30static C2R_CACHE: Mutex<Option<HashMap<usize, C2RPlan64>>> = Mutex::new(None);
31
32/// Global cache for R2R DCT-II plans
33static DCT2_CACHE: Mutex<Option<HashMap<usize, R2RPlan64>>> = Mutex::new(None);
34
35/// Global cache for R2R IDCT-II (DCT-III) plans
36static IDCT2_CACHE: Mutex<Option<HashMap<usize, R2RPlan64>>> = Mutex::new(None);
37
38/// Global cache for R2R DST-II plans
39static DST2_CACHE: Mutex<Option<HashMap<usize, R2RPlan64>>> = Mutex::new(None);
40
41/// Global cache for R2R IDST-II (DST-III) plans
42static IDST2_CACHE: Mutex<Option<HashMap<usize, R2RPlan64>>> = Mutex::new(None);
43
44// ========================================
45// 2D PLAN CACHES
46// ========================================
47
48/// Global cache for 2D R2C plans, keyed by (rows, cols)
49static R2C_2D_CACHE: Mutex<Option<HashMap<(usize, usize), R2CPlan64>>> = Mutex::new(None);
50
51/// Global cache for 2D C2C forward plans
52static C2C_2D_FWD_CACHE: Mutex<Option<HashMap<(usize, usize), C2CPlan64>>> = Mutex::new(None);
53
54/// Global cache for 2D C2C backward plans
55static C2C_2D_BWD_CACHE: Mutex<Option<HashMap<(usize, usize), C2CPlan64>>> = Mutex::new(None);
56
57/// Global cache for 2D C2R plans
58static C2R_2D_CACHE: Mutex<Option<HashMap<(usize, usize), C2RPlan64>>> = Mutex::new(None);
59
60// ========================================
61// 1D PLAN EXECUTION FUNCTIONS
62// ========================================
63
64/// Execute R2C FFT with cached plan
65pub fn execute_r2c(input: &mut [f64], output: &mut [c64]) -> FFTResult<()> {
66    let n = input.len();
67
68    let mut cache = R2C_CACHE
69        .lock()
70        .map_err(|e| FFTError::ComputationError(format!("Failed to lock R2C cache: {}", e)))?;
71
72    // Initialize cache if needed
73    if cache.is_none() {
74        *cache = Some(HashMap::new());
75    }
76
77    let cache_map = cache.as_mut().unwrap();
78
79    // Create plan if not cached
80    if !cache_map.contains_key(&n) {
81        let plan = R2CPlan64::aligned(&[n], Flag::ESTIMATE).map_err(|e| {
82            FFTError::ComputationError(format!("Failed to create R2C plan: {:?}", e))
83        })?;
84        cache_map.insert(n, plan);
85    }
86
87    // Execute with cached plan
88    let plan = cache_map.get_mut(&n).unwrap();
89    plan.r2c(input, output)
90        .map_err(|e| FFTError::ComputationError(format!("R2C execution failed: {:?}", e)))
91}
92
93/// Execute C2C forward FFT with cached plan
94pub fn execute_c2c_forward(input: &mut [c64], output: &mut [c64]) -> FFTResult<()> {
95    let n = input.len();
96
97    let mut cache = C2C_FWD_CACHE.lock().map_err(|e| {
98        FFTError::ComputationError(format!("Failed to lock C2C forward cache: {}", e))
99    })?;
100
101    if cache.is_none() {
102        *cache = Some(HashMap::new());
103    }
104
105    let cache_map = cache.as_mut().unwrap();
106
107    if !cache_map.contains_key(&n) {
108        let plan = C2CPlan64::aligned(&[n], Sign::Forward, Flag::ESTIMATE).map_err(|e| {
109            FFTError::ComputationError(format!("Failed to create C2C forward plan: {:?}", e))
110        })?;
111        cache_map.insert(n, plan);
112    }
113
114    let plan = cache_map.get_mut(&n).unwrap();
115    plan.c2c(input, output)
116        .map_err(|e| FFTError::ComputationError(format!("C2C forward execution failed: {:?}", e)))
117}
118
119/// Execute C2C backward (inverse) FFT with cached plan
120pub fn execute_c2c_backward(input: &mut [c64], output: &mut [c64]) -> FFTResult<()> {
121    let n = input.len();
122
123    let mut cache = C2C_BWD_CACHE.lock().map_err(|e| {
124        FFTError::ComputationError(format!("Failed to lock C2C backward cache: {}", e))
125    })?;
126
127    if cache.is_none() {
128        *cache = Some(HashMap::new());
129    }
130
131    let cache_map = cache.as_mut().unwrap();
132
133    if !cache_map.contains_key(&n) {
134        let plan = C2CPlan64::aligned(&[n], Sign::Backward, Flag::ESTIMATE).map_err(|e| {
135            FFTError::ComputationError(format!("Failed to create C2C backward plan: {:?}", e))
136        })?;
137        cache_map.insert(n, plan);
138    }
139
140    let plan = cache_map.get_mut(&n).unwrap();
141    plan.c2c(input, output)
142        .map_err(|e| FFTError::ComputationError(format!("C2C backward execution failed: {:?}", e)))
143}
144
145/// Execute C2R (inverse real) FFT with cached plan
146pub fn execute_c2r(input: &mut [c64], output: &mut [f64], n: usize) -> FFTResult<()> {
147    let mut cache = C2R_CACHE
148        .lock()
149        .map_err(|e| FFTError::ComputationError(format!("Failed to lock C2R cache: {}", e)))?;
150
151    if cache.is_none() {
152        *cache = Some(HashMap::new());
153    }
154
155    let cache_map = cache.as_mut().unwrap();
156
157    if !cache_map.contains_key(&n) {
158        let plan = C2RPlan64::aligned(&[n], Flag::ESTIMATE).map_err(|e| {
159            FFTError::ComputationError(format!("Failed to create C2R plan: {:?}", e))
160        })?;
161        cache_map.insert(n, plan);
162    }
163
164    let plan = cache_map.get_mut(&n).unwrap();
165    plan.c2r(input, output)
166        .map_err(|e| FFTError::ComputationError(format!("C2R execution failed: {:?}", e)))
167}
168
169/// Execute DCT-II with cached plan
170pub fn execute_dct2(input: &mut [f64], output: &mut [f64]) -> FFTResult<()> {
171    let n = input.len();
172
173    let mut cache = DCT2_CACHE
174        .lock()
175        .map_err(|e| FFTError::ComputationError(format!("Failed to lock DCT2 cache: {}", e)))?;
176
177    if cache.is_none() {
178        *cache = Some(HashMap::new());
179    }
180
181    let cache_map = cache.as_mut().unwrap();
182
183    if !cache_map.contains_key(&n) {
184        let plan =
185            R2RPlan64::aligned(&[n], R2RKind::FFTW_REDFT10, Flag::ESTIMATE).map_err(|e| {
186                FFTError::ComputationError(format!("Failed to create DCT2 plan: {:?}", e))
187            })?;
188        cache_map.insert(n, plan);
189    }
190
191    let plan = cache_map.get_mut(&n).unwrap();
192    plan.r2r(input, output)
193        .map_err(|e| FFTError::ComputationError(format!("DCT2 execution failed: {:?}", e)))
194}
195
196/// Execute IDCT-II (DCT-III) with cached plan
197pub fn execute_idct2(input: &mut [f64], output: &mut [f64]) -> FFTResult<()> {
198    let n = input.len();
199
200    let mut cache = IDCT2_CACHE
201        .lock()
202        .map_err(|e| FFTError::ComputationError(format!("Failed to lock IDCT2 cache: {}", e)))?;
203
204    if cache.is_none() {
205        *cache = Some(HashMap::new());
206    }
207
208    let cache_map = cache.as_mut().unwrap();
209
210    if !cache_map.contains_key(&n) {
211        let plan =
212            R2RPlan64::aligned(&[n], R2RKind::FFTW_REDFT01, Flag::ESTIMATE).map_err(|e| {
213                FFTError::ComputationError(format!("Failed to create IDCT2 plan: {:?}", e))
214            })?;
215        cache_map.insert(n, plan);
216    }
217
218    let plan = cache_map.get_mut(&n).unwrap();
219    plan.r2r(input, output)
220        .map_err(|e| FFTError::ComputationError(format!("IDCT2 execution failed: {:?}", e)))
221}
222
223/// Execute DST-II with cached plan
224pub fn execute_dst2(input: &mut [f64], output: &mut [f64]) -> FFTResult<()> {
225    let n = input.len();
226
227    let mut cache = DST2_CACHE
228        .lock()
229        .map_err(|e| FFTError::ComputationError(format!("Failed to lock DST2 cache: {}", e)))?;
230
231    if cache.is_none() {
232        *cache = Some(HashMap::new());
233    }
234
235    let cache_map = cache.as_mut().unwrap();
236
237    if !cache_map.contains_key(&n) {
238        let plan =
239            R2RPlan64::aligned(&[n], R2RKind::FFTW_RODFT10, Flag::ESTIMATE).map_err(|e| {
240                FFTError::ComputationError(format!("Failed to create DST2 plan: {:?}", e))
241            })?;
242        cache_map.insert(n, plan);
243    }
244
245    let plan = cache_map.get_mut(&n).unwrap();
246    plan.r2r(input, output)
247        .map_err(|e| FFTError::ComputationError(format!("DST2 execution failed: {:?}", e)))
248}
249
250/// Execute IDST-II (DST-III) with cached plan
251pub fn execute_idst2(input: &mut [f64], output: &mut [f64]) -> FFTResult<()> {
252    let n = input.len();
253
254    let mut cache = IDST2_CACHE
255        .lock()
256        .map_err(|e| FFTError::ComputationError(format!("Failed to lock IDST2 cache: {}", e)))?;
257
258    if cache.is_none() {
259        *cache = Some(HashMap::new());
260    }
261
262    let cache_map = cache.as_mut().unwrap();
263
264    if !cache_map.contains_key(&n) {
265        let plan =
266            R2RPlan64::aligned(&[n], R2RKind::FFTW_RODFT01, Flag::ESTIMATE).map_err(|e| {
267                FFTError::ComputationError(format!("Failed to create IDST2 plan: {:?}", e))
268            })?;
269        cache_map.insert(n, plan);
270    }
271
272    let plan = cache_map.get_mut(&n).unwrap();
273    plan.r2r(input, output)
274        .map_err(|e| FFTError::ComputationError(format!("IDST2 execution failed: {:?}", e)))
275}
276
277// ========================================
278// 2D PLAN EXECUTION FUNCTIONS
279// ========================================
280
281/// Execute 2D R2C FFT with cached plan
282pub fn execute_r2c_2d(
283    input: &mut [f64],
284    output: &mut [c64],
285    rows: usize,
286    cols: usize,
287) -> FFTResult<()> {
288    let key = (rows, cols);
289
290    let mut cache = R2C_2D_CACHE
291        .lock()
292        .map_err(|e| FFTError::ComputationError(format!("Failed to lock 2D R2C cache: {}", e)))?;
293
294    if cache.is_none() {
295        *cache = Some(HashMap::new());
296    }
297
298    let cache_map = cache.as_mut().unwrap();
299
300    if !cache_map.contains_key(&key) {
301        let plan = R2CPlan64::aligned(&[rows, cols], Flag::ESTIMATE).map_err(|e| {
302            FFTError::ComputationError(format!("Failed to create 2D R2C plan: {:?}", e))
303        })?;
304        cache_map.insert(key, plan);
305    }
306
307    let plan = cache_map.get_mut(&key).unwrap();
308    plan.r2c(input, output)
309        .map_err(|e| FFTError::ComputationError(format!("2D R2C execution failed: {:?}", e)))
310}
311
312/// Execute 2D C2C forward FFT with cached plan
313pub fn execute_c2c_2d_forward(
314    input: &mut [c64],
315    output: &mut [c64],
316    rows: usize,
317    cols: usize,
318) -> FFTResult<()> {
319    let key = (rows, cols);
320
321    let mut cache = C2C_2D_FWD_CACHE.lock().map_err(|e| {
322        FFTError::ComputationError(format!("Failed to lock 2D C2C forward cache: {}", e))
323    })?;
324
325    if cache.is_none() {
326        *cache = Some(HashMap::new());
327    }
328
329    let cache_map = cache.as_mut().unwrap();
330
331    if !cache_map.contains_key(&key) {
332        let plan =
333            C2CPlan64::aligned(&[rows, cols], Sign::Forward, Flag::ESTIMATE).map_err(|e| {
334                FFTError::ComputationError(format!("Failed to create 2D C2C forward plan: {:?}", e))
335            })?;
336        cache_map.insert(key, plan);
337    }
338
339    let plan = cache_map.get_mut(&key).unwrap();
340    plan.c2c(input, output).map_err(|e| {
341        FFTError::ComputationError(format!("2D C2C forward execution failed: {:?}", e))
342    })
343}
344
345/// Execute 2D C2C backward FFT with cached plan
346pub fn execute_c2c_2d_backward(
347    input: &mut [c64],
348    output: &mut [c64],
349    rows: usize,
350    cols: usize,
351) -> FFTResult<()> {
352    let key = (rows, cols);
353
354    let mut cache = C2C_2D_BWD_CACHE.lock().map_err(|e| {
355        FFTError::ComputationError(format!("Failed to lock 2D C2C backward cache: {}", e))
356    })?;
357
358    if cache.is_none() {
359        *cache = Some(HashMap::new());
360    }
361
362    let cache_map = cache.as_mut().unwrap();
363
364    if !cache_map.contains_key(&key) {
365        let plan =
366            C2CPlan64::aligned(&[rows, cols], Sign::Backward, Flag::ESTIMATE).map_err(|e| {
367                FFTError::ComputationError(format!(
368                    "Failed to create 2D C2C backward plan: {:?}",
369                    e
370                ))
371            })?;
372        cache_map.insert(key, plan);
373    }
374
375    let plan = cache_map.get_mut(&key).unwrap();
376    plan.c2c(input, output).map_err(|e| {
377        FFTError::ComputationError(format!("2D C2C backward execution failed: {:?}", e))
378    })
379}
380
381/// Execute 2D C2R FFT with cached plan
382pub fn execute_c2r_2d(
383    input: &mut [c64],
384    output: &mut [f64],
385    rows: usize,
386    cols: usize,
387) -> FFTResult<()> {
388    let key = (rows, cols);
389
390    let mut cache = C2R_2D_CACHE
391        .lock()
392        .map_err(|e| FFTError::ComputationError(format!("Failed to lock 2D C2R cache: {}", e)))?;
393
394    if cache.is_none() {
395        *cache = Some(HashMap::new());
396    }
397
398    let cache_map = cache.as_mut().unwrap();
399
400    if !cache_map.contains_key(&key) {
401        let plan = C2RPlan64::aligned(&[rows, cols], Flag::ESTIMATE).map_err(|e| {
402            FFTError::ComputationError(format!("Failed to create 2D C2R plan: {:?}", e))
403        })?;
404        cache_map.insert(key, plan);
405    }
406
407    let plan = cache_map.get_mut(&key).unwrap();
408    plan.c2r(input, output)
409        .map_err(|e| FFTError::ComputationError(format!("2D C2R execution failed: {:?}", e)))
410}
411
412// ========================================
413// CACHE MANAGEMENT
414// ========================================
415
416/// Clear all cached plans (useful for testing or memory management)
417pub fn clear_all_caches() {
418    if let Ok(mut cache) = R2C_CACHE.lock() {
419        *cache = None;
420    }
421    if let Ok(mut cache) = C2C_FWD_CACHE.lock() {
422        *cache = None;
423    }
424    if let Ok(mut cache) = C2C_BWD_CACHE.lock() {
425        *cache = None;
426    }
427    if let Ok(mut cache) = C2R_CACHE.lock() {
428        *cache = None;
429    }
430    if let Ok(mut cache) = DCT2_CACHE.lock() {
431        *cache = None;
432    }
433    if let Ok(mut cache) = IDCT2_CACHE.lock() {
434        *cache = None;
435    }
436    if let Ok(mut cache) = DST2_CACHE.lock() {
437        *cache = None;
438    }
439    if let Ok(mut cache) = IDST2_CACHE.lock() {
440        *cache = None;
441    }
442    if let Ok(mut cache) = R2C_2D_CACHE.lock() {
443        *cache = None;
444    }
445    if let Ok(mut cache) = C2C_2D_FWD_CACHE.lock() {
446        *cache = None;
447    }
448    if let Ok(mut cache) = C2C_2D_BWD_CACHE.lock() {
449        *cache = None;
450    }
451    if let Ok(mut cache) = C2R_2D_CACHE.lock() {
452        *cache = None;
453    }
454}
scirs2_fft/fftw_plan_cache.rs

scirs2_fft/
fftw_plan_cache.rs
