Skip to main content

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;