Skip to main content

fpm_rs/backend/
traits.rs

1use std::any::Any;
2
3use num_complex::Complex64;
4
5use crate::Result;
6
7#[derive(Clone, Copy, Debug, PartialEq, Eq)]
8pub enum FftDirection {
9    Forward,
10    Inverse,
11}
12
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14pub enum MemoryLocation {
15    Host,
16    Device,
17}
18
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
20pub struct BackendCapabilities {
21    pub preferred_memory: MemoryLocation,
22    pub resident_buffers: bool,
23}
24
25impl Default for BackendCapabilities {
26    fn default() -> Self {
27        Self {
28            preferred_memory: MemoryLocation::Host,
29            resident_buffers: false,
30        }
31    }
32}
33
34pub trait ComplexBuffer: Send + Sync {
35    fn len(&self) -> usize;
36    fn is_empty(&self) -> bool {
37        self.len() == 0
38    }
39    fn location(&self) -> MemoryLocation;
40    fn as_any(&self) -> &dyn Any;
41    fn as_any_mut(&mut self) -> &mut dyn Any;
42}
43
44pub trait RealBuffer: Send + Sync {
45    fn len(&self) -> usize;
46    fn is_empty(&self) -> bool {
47        self.len() == 0
48    }
49    fn location(&self) -> MemoryLocation;
50    fn as_any(&self) -> &dyn Any;
51    fn as_any_mut(&mut self) -> &mut dyn Any;
52}
53
54/// Optional typed-buffer operations for backends that can keep reconstruction
55/// state resident in their preferred memory domain.
56pub trait ResidentBackend: Send + Sync {
57    fn allocate_complex(&self, len: usize) -> Result<Box<dyn ComplexBuffer>>;
58    fn allocate_real(&self, len: usize) -> Result<Box<dyn RealBuffer>>;
59    fn upload_complex(
60        &self,
61        destination: &mut dyn ComplexBuffer,
62        source: &[Complex64],
63    ) -> Result<()>;
64    fn download_complex(
65        &self,
66        source: &dyn ComplexBuffer,
67        destination: &mut [Complex64],
68    ) -> Result<()>;
69    fn upload_real(&self, destination: &mut dyn RealBuffer, source: &[f64]) -> Result<()>;
70    fn download_real(&self, source: &dyn RealBuffer, destination: &mut [f64]) -> Result<()>;
71    fn fft2_resident(
72        &self,
73        values: &mut dyn ComplexBuffer,
74        shape: (usize, usize),
75        direction: FftDirection,
76    ) -> Result<()>;
77}
78
79/// Minimal backend boundary. Buffers stay owned by reconstruction state.
80pub trait Backend: Send + Sync {
81    fn capabilities(&self) -> BackendCapabilities {
82        BackendCapabilities::default()
83    }
84
85    fn resident_backend(&self) -> Option<&dyn ResidentBackend> {
86        None
87    }
88
89    fn fft2(
90        &self,
91        values: &mut [Complex64],
92        shape: (usize, usize),
93        direction: FftDirection,
94        column_scratch: &mut [Complex64],
95    ) -> Result<()>;
96}