tenferro_fft/backend.rs
1use std::fmt;
2
3use tenferro_runtime::ExtensionCacheStore;
4use tenferro_tensor::{Engine, Tensor, TensorRead};
5
6use crate::{FftPlanCache, FftPlanSpec};
7
8#[derive(Clone, Copy, Debug)]
9enum FftCacheOwner {
10 CallerOwned,
11 EngineOwned,
12}
13
14/// Execution-cache state supplied to an [`FftBackend`].
15///
16/// Direct repeated calls use a caller-owned [`FftPlanCache`], while traced
17/// and eager execution use the executing engine's own [`ExtensionCacheStore`]
18/// (lent by the runtime for the run). The two differ in what they may hold:
19///
20/// * a caller-owned cache may outlive the engine and be used with several
21/// engines, so it holds only engine-independent plans (the CPU's rustfft
22/// plans);
23/// * an engine-owned store dies with its engine, so a device backend keeps
24/// device-bound plans and workspaces (cuFFT plans) only there. Given a
25/// caller-owned cache, a device backend keeps its device-bound plans in its
26/// own engine instead.
27///
28/// Constructors keep the owner representation closed.
29///
30/// # Examples
31///
32/// ```
33/// use tenferro_fft::{FftExecutionCache, FftPlanCache};
34///
35/// let mut plans = FftPlanCache::default();
36/// let cache = FftExecutionCache::caller_owned(&mut plans);
37/// assert!(format!("{cache:?}").contains("CallerOwned"));
38/// ```
39pub struct FftExecutionCache<'a> {
40 owner: FftCacheOwner,
41 store: &'a mut ExtensionCacheStore,
42}
43
44impl fmt::Debug for FftExecutionCache<'_> {
45 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
46 f.debug_struct("FftExecutionCache")
47 .field("owner", &self.owner)
48 .field(
49 "stats",
50 &self
51 .store
52 .stats(tenferro_runtime::ExtensionCacheSelector::All),
53 )
54 .finish_non_exhaustive()
55 }
56}
57
58impl<'a> FftExecutionCache<'a> {
59 /// Build a context backed by a caller-owned typed FFT execution cache.
60 pub fn caller_owned(cache: &'a mut FftPlanCache) -> Self {
61 Self {
62 owner: FftCacheOwner::CallerOwned,
63 store: cache.store_mut(),
64 }
65 }
66
67 /// Build a context backed by the executing engine's own extension cache
68 /// store, which may hold device-bound plans.
69 ///
70 /// # Examples
71 ///
72 /// ```
73 /// use tenferro_fft::FftExecutionCache;
74 /// use tenferro_runtime::ExtensionCacheStore;
75 ///
76 /// let mut store = ExtensionCacheStore::default();
77 /// let cache = FftExecutionCache::engine_owned(&mut store);
78 /// assert!(cache.is_engine_owned());
79 /// ```
80 pub fn engine_owned(cache: &'a mut ExtensionCacheStore) -> Self {
81 Self {
82 owner: FftCacheOwner::EngineOwned,
83 store: cache,
84 }
85 }
86
87 /// Whether this store belongs to the executing engine, so device-bound
88 /// entries may be retained in it.
89 ///
90 /// # Examples
91 ///
92 /// ```
93 /// use tenferro_fft::{FftExecutionCache, FftPlanCache};
94 ///
95 /// let mut plans = FftPlanCache::default();
96 /// assert!(!FftExecutionCache::caller_owned(&mut plans).is_engine_owned());
97 /// ```
98 pub fn is_engine_owned(&self) -> bool {
99 matches!(self.owner, FftCacheOwner::EngineOwned)
100 }
101
102 /// Borrow the bounded typed store owned by the caller or extension runtime.
103 ///
104 /// Backend implementations should use a stable family/cache namespace and
105 /// include every plan-identity field in the key discriminator. The store
106 /// owns LRU bounds, typed retrieval, clear behavior, entry counts, and the
107 /// retained-byte estimates supplied at insertion.
108 pub fn store_mut(&mut self) -> &mut ExtensionCacheStore {
109 self.store
110 }
111}
112
113/// Explicit backend capability required by concrete and traced FFT execution.
114///
115/// Implementations must execute on the input's existing placement. Unsupported
116/// dtypes, layouts, placements, or operations return an error; they must never
117/// transfer the tensor or select a different backend.
118///
119/// # Examples
120/// Backend surface required by the FFT extension runtime and the built-in
121/// concrete FFT dispatch.
122///
123/// This is a backend SPI: third-party backend implementers implement this
124/// trait on their execution sessions to provide FFT capability. The concrete
125/// [`TensorFftExt`](crate::TensorFftExt)/[`TensorReadFftExt`](crate::TensorReadFftExt) methods and [`FftExecutor`](crate::FftExecutor) take an
126/// erased `&mut dyn Engine` and dispatch **internally** to the
127/// built-in CPU/CUDA/WebGPU execution sessions; a session without this
128/// capability returns a typed `Error::Unsupported` instead of being accepted
129/// as an arbitrary third-party backend (issue #1680 Phase 3).
130///
131/// # Examples
132///
133/// ```
134/// use tenferro_cpu::CpuEngine;
135/// use tenferro_fft::FftBackend;
136///
137/// fn accepts_fft_backend<B: FftBackend>() {}
138///
139/// accepts_fft_backend::<CpuEngine>();
140/// ```
141///
142/// A generic engine does not expose the FFT capability at the SPI level:
143///
144/// ```compile_fail
145/// use tenferro_tensor::Engine;
146/// use tenferro_fft::FftBackend;
147///
148/// fn requires_fft_backend<B: Engine>(engine: &mut B) {
149/// let _: &mut dyn FftBackend = engine;
150/// }
151/// ```
152pub trait FftBackend: Engine {
153 /// Validate a borrowed runtime input before the extension read path
154 /// materializes it into a compact tensor.
155 ///
156 /// Backends that can materialize their own placement may keep the default
157 /// no-op. CPU overrides this to report device inputs as unsupported FFT
158 /// placement errors instead of leaking a generic host-materialization error.
159 ///
160 /// # Errors
161 ///
162 /// Returns [`tenferro_tensor::Error::Unsupported`] when this backend cannot
163 /// consume the borrowed input placement without an explicit transfer.
164 fn validate_fft_read_input(
165 &self,
166 _op: &'static str,
167 _input: &TensorRead<'_>,
168 ) -> tenferro_tensor::Result<()> {
169 Ok(())
170 }
171
172 /// Execute a validated FFT from borrowed storage.
173 ///
174 /// The default explicitly canonicalizes view inputs. Backends supporting
175 /// compact borrowed storage override this method to avoid input copies.
176 ///
177 /// # Examples
178 /// The borrowed public surface dispatches through this backend hook:
179 /// ```
180 /// use num_complex::Complex64;
181 /// use tenferro_cpu::CpuEngine;
182 /// use tenferro_fft::{FftNorm, TensorReadFftExt};
183 /// use tenferro_tensor::{Tensor, TensorRead};
184 /// let input = Tensor::from_vec_col_major([2], vec![Complex64::new(1., 0.); 2])?;
185 /// let mut engine = CpuEngine::serial();
186 /// let output =
187 /// TensorRead::Tensor(&input).fft_read(None, 0, FftNorm::Backward, &mut engine)?;
188 /// assert_eq!(output.as_slice::<Complex64>()?, &[Complex64::new(2., 0.), Complex64::new(0., 0.)]);
189 /// # Ok::<(), Box<dyn std::error::Error>>(())
190 /// ```
191 ///
192 /// # Errors
193 /// Returns [`tenferro_tensor::Error::Unsupported`] for unsupported dtype,
194 /// layout, or placement, and [`tenferro_tensor::Error::Validation`] when
195 /// input metadata disagrees with the validated spec. Same-placement
196 /// canonicalization, plan creation, and execution failures preserve the
197 /// selected backend's typed error source; no implicit transfer is performed.
198 fn execute_fft_read(
199 &mut self,
200 input: TensorRead<'_>,
201 spec: &FftPlanSpec,
202 cache: FftExecutionCache<'_>,
203 ) -> tenferro_tensor::Result<Tensor> {
204 match input {
205 TensorRead::Tensor(input) => self.execute_fft(input, spec, cache),
206 input => {
207 let owned = self.to_contiguous_read(input)?;
208 self.execute_fft(&owned, spec, cache)
209 }
210 }
211 }
212
213 /// Execute one validated FFT request on `input`'s existing placement.
214 ///
215 /// # Errors
216 ///
217 /// Returns [`tenferro_tensor::Error::Unsupported`] for an unsupported
218 /// operation, dtype, layout, or placement;
219 /// [`tenferro_tensor::Error::Validation`] for inconsistent input/spec
220 /// metadata or checked shape arithmetic; and a typed backend or runtime
221 /// source when plan creation, cache access, or execution fails.
222 fn execute_fft(
223 &mut self,
224 input: &Tensor,
225 spec: &FftPlanSpec,
226 cache: FftExecutionCache<'_>,
227 ) -> tenferro_tensor::Result<Tensor>;
228}
229
230#[cfg(test)]
231mod tests;