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
54pub 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
79pub 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}