fpm_rs/backend/
cpu.rs

1use std::sync::Arc;
2
3use num_complex::Complex64;
4use rustfft::{Fft, FftPlanner};
5
6use crate::{Result, error::Error};
7
8use super::{
9    Backend, BackendCapabilities, ComplexBuffer, FftDirection, MemoryLocation, RealBuffer,
10    ResidentBackend,
11};
12
13#[derive(Debug)]
14struct CpuComplexBuffer {
15    values: Vec<Complex64>,
16}
17
18impl ComplexBuffer for CpuComplexBuffer {
19    fn len(&self) -> usize {
20        self.values.len()
21    }
22
23    fn location(&self) -> MemoryLocation {
24        MemoryLocation::Host
25    }
26
27    fn as_any(&self) -> &dyn std::any::Any {
28        self
29    }
30
31    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
32        self
33    }
34}
35
36#[derive(Debug)]
37struct CpuRealBuffer {
38    values: Vec<f64>,
39}
40
41impl RealBuffer for CpuRealBuffer {
42    fn len(&self) -> usize {
43        self.values.len()
44    }
45
46    fn location(&self) -> MemoryLocation {
47        MemoryLocation::Host
48    }
49
50    fn as_any(&self) -> &dyn std::any::Any {
51        self
52    }
53
54    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
55        self
56    }
57}
58
59#[derive(Clone)]
60struct CpuFftPlan {
61    shape: (usize, usize),
62    row_forward: Arc<dyn Fft<f64>>,
63    row_inverse: Arc<dyn Fft<f64>>,
64    column_forward: Arc<dyn Fft<f64>>,
65    column_inverse: Arc<dyn Fft<f64>>,
66}
67
68impl CpuFftPlan {
69    fn new(shape: (usize, usize)) -> Self {
70        let mut planner = FftPlanner::new();
71        Self {
72            shape,
73            row_forward: planner.plan_fft_forward(shape.1),
74            row_inverse: planner.plan_fft_inverse(shape.1),
75            column_forward: planner.plan_fft_forward(shape.0),
76            column_inverse: planner.plan_fft_inverse(shape.0),
77        }
78    }
79}
80
81/// rustfft-based CPU backend with plans cached for low- and high-resolution grids.
82#[derive(Clone)]
83pub struct CpuBackend {
84    low: CpuFftPlan,
85    high: CpuFftPlan,
86}
87
88impl CpuBackend {
89    pub fn new(low_shape: (usize, usize), high_shape: (usize, usize)) -> Result<Self> {
90        if low_shape.0 == 0 || low_shape.1 == 0 || high_shape.0 == 0 || high_shape.1 == 0 {
91            return Err(Error::InvalidShape(
92                "FFT dimensions must be non-zero".into(),
93            ));
94        }
95        Ok(Self {
96            low: CpuFftPlan::new(low_shape),
97            high: CpuFftPlan::new(high_shape),
98        })
99    }
100
101    fn plan(&self, shape: (usize, usize)) -> Result<&CpuFftPlan> {
102        if shape == self.low.shape {
103            Ok(&self.low)
104        } else if shape == self.high.shape {
105            Ok(&self.high)
106        } else {
107            Err(Error::InvalidShape(format!(
108                "FFT shape {shape:?} is neither configured shape {:?} nor {:?}",
109                self.low.shape, self.high.shape
110            )))
111        }
112    }
113}
114
115impl Backend for CpuBackend {
116    fn capabilities(&self) -> BackendCapabilities {
117        BackendCapabilities {
118            preferred_memory: MemoryLocation::Host,
119            resident_buffers: true,
120        }
121    }
122
123    fn resident_backend(&self) -> Option<&dyn ResidentBackend> {
124        Some(self)
125    }
126
127    fn fft2(
128        &self,
129        values: &mut [Complex64],
130        shape: (usize, usize),
131        direction: FftDirection,
132        column_scratch: &mut [Complex64],
133    ) -> Result<()> {
134        let expected = shape.0 * shape.1;
135        if values.len() != expected || column_scratch.len() < shape.0 {
136            return Err(Error::InvalidShape(format!(
137                "FFT {:?} requires {expected} values and {} column scratch values",
138                shape, shape.0
139            )));
140        }
141        let plan = self.plan(shape)?;
142        let (row_fft, column_fft) = match direction {
143            FftDirection::Forward => (&plan.row_forward, &plan.column_forward),
144            FftDirection::Inverse => (&plan.row_inverse, &plan.column_inverse),
145        };
146        for row in values.chunks_exact_mut(shape.1) {
147            row_fft.process(row);
148        }
149        let column = &mut column_scratch[..shape.0];
150        for column_index in 0..shape.1 {
151            for row_index in 0..shape.0 {
152                column[row_index] = values[row_index * shape.1 + column_index];
153            }
154            column_fft.process(column);
155            for row_index in 0..shape.0 {
156                values[row_index * shape.1 + column_index] = column[row_index];
157            }
158        }
159        // The forward transform is normalized by 1/N and the inverse is
160        // unnormalized. This pair is exactly invertible and, unlike the common
161        // inverse-normalized convention, preserves the amplitude of a constant
162        // object when a high-resolution spectrum is cropped and inverse-
163        // transformed on a smaller grid.
164        if direction == FftDirection::Forward {
165            let normalization = expected as f64;
166            for value in values {
167                *value /= normalization;
168            }
169        }
170        Ok(())
171    }
172}
173
174impl ResidentBackend for CpuBackend {
175    fn allocate_complex(&self, len: usize) -> Result<Box<dyn ComplexBuffer>> {
176        Ok(Box::new(CpuComplexBuffer {
177            values: vec![Complex64::default(); len],
178        }))
179    }
180
181    fn allocate_real(&self, len: usize) -> Result<Box<dyn RealBuffer>> {
182        Ok(Box::new(CpuRealBuffer {
183            values: vec![0.0; len],
184        }))
185    }
186
187    fn upload_complex(
188        &self,
189        destination: &mut dyn ComplexBuffer,
190        source: &[Complex64],
191    ) -> Result<()> {
192        let destination = cpu_complex_mut(destination)?;
193        validate_transfer_length(destination.values.len(), source.len())?;
194        destination.values.copy_from_slice(source);
195        Ok(())
196    }
197
198    fn download_complex(
199        &self,
200        source: &dyn ComplexBuffer,
201        destination: &mut [Complex64],
202    ) -> Result<()> {
203        let source = cpu_complex(source)?;
204        validate_transfer_length(destination.len(), source.values.len())?;
205        destination.copy_from_slice(&source.values);
206        Ok(())
207    }
208
209    fn upload_real(&self, destination: &mut dyn RealBuffer, source: &[f64]) -> Result<()> {
210        let destination = cpu_real_mut(destination)?;
211        validate_transfer_length(destination.values.len(), source.len())?;
212        destination.values.copy_from_slice(source);
213        Ok(())
214    }
215
216    fn download_real(&self, source: &dyn RealBuffer, destination: &mut [f64]) -> Result<()> {
217        let source = cpu_real(source)?;
218        validate_transfer_length(destination.len(), source.values.len())?;
219        destination.copy_from_slice(&source.values);
220        Ok(())
221    }
222
223    fn fft2_resident(
224        &self,
225        values: &mut dyn ComplexBuffer,
226        shape: (usize, usize),
227        direction: FftDirection,
228    ) -> Result<()> {
229        let values = cpu_complex_mut(values)?;
230        let mut column = vec![Complex64::default(); shape.0];
231        self.fft2(&mut values.values, shape, direction, &mut column)
232    }
233}
234
235fn cpu_complex(buffer: &dyn ComplexBuffer) -> Result<&CpuComplexBuffer> {
236    buffer
237        .as_any()
238        .downcast_ref()
239        .ok_or_else(|| Error::InvalidParameter {
240            name: "complex buffer",
241            reason: "buffer was not allocated by CpuBackend".into(),
242        })
243}
244
245fn cpu_complex_mut(buffer: &mut dyn ComplexBuffer) -> Result<&mut CpuComplexBuffer> {
246    buffer
247        .as_any_mut()
248        .downcast_mut()
249        .ok_or_else(|| Error::InvalidParameter {
250            name: "complex buffer",
251            reason: "buffer was not allocated by CpuBackend".into(),
252        })
253}
254
255fn cpu_real(buffer: &dyn RealBuffer) -> Result<&CpuRealBuffer> {
256    buffer
257        .as_any()
258        .downcast_ref()
259        .ok_or_else(|| Error::InvalidParameter {
260            name: "real buffer",
261            reason: "buffer was not allocated by CpuBackend".into(),
262        })
263}
264
265fn cpu_real_mut(buffer: &mut dyn RealBuffer) -> Result<&mut CpuRealBuffer> {
266    buffer
267        .as_any_mut()
268        .downcast_mut()
269        .ok_or_else(|| Error::InvalidParameter {
270            name: "real buffer",
271            reason: "buffer was not allocated by CpuBackend".into(),
272        })
273}
274
275fn validate_transfer_length(destination: usize, source: usize) -> Result<()> {
276    if destination != source {
277        return Err(Error::InvalidShape(format!(
278            "buffer transfer length {source} does not match destination length {destination}"
279        )));
280    }
281    Ok(())
282}
fpm_rs/backend/cpu.rs

fpm_rs/backend/
cpu.rs
