arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/poulpy-cpu-ref/src/module.rs
poulpy-cpu-ref/src/module.rsRaw
 1use std::ptr::NonNull;
 2
 3use poulpy_hal::{
 4    layouts::{Backend, Module},
 5    oep::ModuleNewImpl,
 6    reference::fft64::reim::{ReimFFTTable, ReimIFFTTable},
 7};
 8
 9use crate::FFT64Ref;
10
11#[repr(C)]
12pub struct FFT64RefHandle {
13    table_fft: ReimFFTTable<f64>,
14    table_ifft: ReimIFFTTable<f64>,
15}
16
17impl Backend for FFT64Ref {
18    type ScalarPrep = f64;
19    type ScalarBig = i64;
20    type Handle = FFT64RefHandle;
21    unsafe fn destroy(handle: NonNull<Self::Handle>) {
22        unsafe {
23            drop(Box::from_raw(handle.as_ptr()));
24        }
25    }
26
27    fn layout_big_word_count() -> usize {
28        1
29    }
30
31    fn layout_prep_word_count() -> usize {
32        1
33    }
34}
35
36unsafe impl ModuleNewImpl<Self> for FFT64Ref {
37    fn new_impl(n: u64) -> Module<Self> {
38        let handle: FFT64RefHandle = FFT64RefHandle {
39            table_fft: ReimFFTTable::new(n as usize >> 1),
40            table_ifft: ReimIFFTTable::new(n as usize >> 1),
41        };
42        // Leak Box to get a stable NonNull pointer
43        let ptr: NonNull<FFT64RefHandle> = NonNull::from(Box::leak(Box::new(handle)));
44        unsafe { Module::from_nonnull(ptr, n) }
45    }
46}
47
48pub trait FFT64ModuleHandle {
49    fn get_fft_table(&self) -> &ReimFFTTable<f64>;
50    fn get_ifft_table(&self) -> &ReimIFFTTable<f64>;
51}
52
53impl FFT64ModuleHandle for Module<FFT64Ref> {
54    fn get_fft_table(&self) -> &ReimFFTTable<f64> {
55        let h: &FFT64RefHandle = unsafe { &*self.ptr() };
56        &h.table_fft
57    }
58    fn get_ifft_table(&self) -> &ReimIFFTTable<f64> {
59        let h: &FFT64RefHandle = unsafe { &*self.ptr() };
60        &h.table_ifft
61    }
62}