poulpy-cpu-ref 0.4.4

The providing concrete implementations of poulpy-hal through its open extension points and reference cpu code
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

//! Backend handle and module initialization for [`FFT64Ref`](crate::FFT64Ref).
//!
//! This module defines:
//!
//! - [`FFT64RefHandle`]: the opaque handle stored inside a `Module<FFT64Ref>`,
//!   holding precomputed FFT and IFFT twiddle-factor tables.
//! - The [`Backend`] trait implementation, which defines scalar types and the
//!   handle destruction path.
//! - The [`ModuleNewImpl`] implementation, which allocates the handle on the heap
//!   and transfers ownership to the `Module`.
//! - The [`FFT64ModuleHandle`] trait, which provides typed access to the FFT tables
//!   from a `Module<FFT64Ref>`. This trait is also implemented by other FFT64 backends
//!   (e.g. AVX), enabling shared downstream code.

use std::ptr::NonNull;

use poulpy_hal::{
    layouts::{Backend, Module},
    oep::ModuleNewImpl,
    reference::fft64::reim::{ReimFFTTable, ReimIFFTTable},
};

use super::FFT64Ref;

/// Opaque handle for the [`FFT64Ref`](crate::FFT64Ref) backend.
///
/// Holds precomputed twiddle-factor tables for the forward FFT and inverse FFT
/// of size `m = n / 2`, where `n` is the ring dimension passed to
/// [`Module::new`](poulpy_hal::api::ModuleNew::new).
///
/// This struct is heap-allocated during module creation and freed when the
/// `Module<FFT64Ref>` is dropped (via [`Backend::destroy`]).
#[repr(C)]
pub struct FFT64RefHandle {
    table_fft: ReimFFTTable<f64>,
    table_ifft: ReimIFFTTable<f64>,
}

impl Backend for FFT64Ref {
    type ScalarPrep = f64;
    type ScalarBig = i64;
    type Handle = FFT64RefHandle;
    unsafe fn destroy(handle: NonNull<Self::Handle>) {
        unsafe {
            drop(Box::from_raw(handle.as_ptr()));
        }
    }
}

/// # Safety
///
/// The returned `Module` owns the heap-allocated `FFT64RefHandle`.
/// `n` must be a power of two >= 2 (asserted by `Module::from_nonnull`).
/// The FFT tables are built for dimension `m = n / 2`.
unsafe impl ModuleNewImpl<Self> for FFT64Ref {
    fn new_impl(n: u64) -> Module<Self> {
        let handle: FFT64RefHandle = FFT64RefHandle {
            table_fft: ReimFFTTable::new(n as usize >> 1),
            table_ifft: ReimIFFTTable::new(n as usize >> 1),
        };
        // Leak Box to get a stable NonNull pointer
        let ptr: NonNull<FFT64RefHandle> = NonNull::from(Box::leak(Box::new(handle)));
        unsafe { Module::from_nonnull(ptr, n) }
    }
}

/// Provides access to the precomputed FFT and IFFT twiddle-factor tables
/// stored inside a module handle.
///
/// This trait is implemented by `Module<FFT64Ref>` (this crate) and also by
/// `Module<FFT64Avx>` in `poulpy-cpu-avx`, allowing shared code in the
/// reference functions to work across both backends.
pub trait FFT64ModuleHandle {
    /// Returns a reference to the forward-FFT twiddle-factor table.
    fn get_fft_table(&self) -> &ReimFFTTable<f64>;
    /// Returns a reference to the inverse-FFT twiddle-factor table.
    fn get_ifft_table(&self) -> &ReimIFFTTable<f64>;
}

impl FFT64ModuleHandle for Module<FFT64Ref> {
    fn get_fft_table(&self) -> &ReimFFTTable<f64> {
        // SAFETY: `self.ptr()` returns a valid pointer to `FFT64RefHandle`,
        // which was heap-allocated in `new_impl` and is kept alive by the Module.
        let h: &FFT64RefHandle = unsafe { &*self.ptr() };
        &h.table_fft
    }
    fn get_ifft_table(&self) -> &ReimIFFTTable<f64> {
        // SAFETY: same as above.
        let h: &FFT64RefHandle = unsafe { &*self.ptr() };
        &h.table_ifft
    }
}