1use std::marker::PhantomData;
  2
  3use poulpy_hal::{
  4    DEFAULTALIGN, alloc_aligned,
  5    api::ScratchFromBytes,
  6    layouts::{Backend, MatZnx, ScalarZnx, Scratch, ScratchOwned, SvpPPol, VecZnx, VecZnxBig, VecZnxDft, VmpPMat},
  7    oep::{
  8        ScratchAvailableImpl, ScratchFromBytesImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, SvpPPolAllocBytesImpl,
  9        TakeMatZnxImpl, TakeScalarZnxImpl, TakeSliceImpl, TakeSvpPPolImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl,
 10        TakeVecZnxDftSliceImpl, TakeVecZnxImpl, TakeVecZnxSliceImpl, TakeVmpPMatImpl, VecZnxBigAllocBytesImpl,
 11        VecZnxDftAllocBytesImpl, VmpPMatAllocBytesImpl,
 12    },
 13};
 14
 15use crate::cpu_fft64_ref::FFT64Ref;
 16
 17unsafe impl<B: Backend> ScratchOwnedAllocImpl<B> for FFT64Ref {
 18    fn scratch_owned_alloc_impl(size: usize) -> ScratchOwned<B> {
 19        let data: Vec<u8> = alloc_aligned(size);
 20        ScratchOwned {
 21            data,
 22            _phantom: PhantomData,
 23        }
 24    }
 25}
 26
 27unsafe impl<B: Backend> ScratchOwnedBorrowImpl<B> for FFT64Ref
 28where
 29    B: ScratchFromBytesImpl<B>,
 30{
 31    fn scratch_owned_borrow_impl(scratch: &mut ScratchOwned<B>) -> &mut Scratch<B> {
 32        Scratch::from_bytes(&mut scratch.data)
 33    }
 34}
 35
 36unsafe impl<B: Backend> ScratchFromBytesImpl<B> for FFT64Ref {
 37    fn scratch_from_bytes_impl(data: &mut [u8]) -> &mut Scratch<B> {
 38        unsafe { &mut *(data as *mut [u8] as *mut Scratch<B>) }
 39    }
 40}
 41
 42unsafe impl<B: Backend> ScratchAvailableImpl<B> for FFT64Ref {
 43    fn scratch_available_impl(scratch: &Scratch<B>) -> usize {
 44        let ptr: *const u8 = scratch.data.as_ptr();
 45        let self_len: usize = scratch.data.len();
 46        let aligned_offset: usize = ptr.align_offset(DEFAULTALIGN);
 47        self_len.saturating_sub(aligned_offset)
 48    }
 49}
 50
 51unsafe impl<B: Backend> TakeSliceImpl<B> for FFT64Ref
 52where
 53    B: ScratchFromBytesImpl<B>,
 54{
 55    fn take_slice_impl<T>(scratch: &mut Scratch<B>, len: usize) -> (&mut [T], &mut Scratch<B>) {
 56        let (take_slice, rem_slice) = take_slice_aligned(&mut scratch.data, len * std::mem::size_of::<T>());
 57
 58        unsafe {
 59            (
 60                &mut *(std::ptr::slice_from_raw_parts_mut(take_slice.as_mut_ptr() as *mut T, len)),
 61                Scratch::from_bytes(rem_slice),
 62            )
 63        }
 64    }
 65}
 66
 67unsafe impl<B: Backend> TakeScalarZnxImpl<B> for FFT64Ref
 68where
 69    B: ScratchFromBytesImpl<B>,
 70{
 71    fn take_scalar_znx_impl(scratch: &mut Scratch<B>, n: usize, cols: usize) -> (ScalarZnx<&mut [u8]>, &mut Scratch<B>) {
 72        let (take_slice, rem_slice) = take_slice_aligned(&mut scratch.data, ScalarZnx::alloc_bytes(n, cols));
 73        (
 74            ScalarZnx::from_data(take_slice, n, cols),
 75            Scratch::from_bytes(rem_slice),
 76        )
 77    }
 78}
 79
 80unsafe impl<B: Backend> TakeSvpPPolImpl<B> for FFT64Ref
 81where
 82    B: SvpPPolAllocBytesImpl<B> + ScratchFromBytesImpl<B>,
 83{
 84    fn take_svp_ppol_impl(scratch: &mut Scratch<B>, n: usize, cols: usize) -> (SvpPPol<&mut [u8], B>, &mut Scratch<B>) {
 85        let (take_slice, rem_slice) = take_slice_aligned(&mut scratch.data, B::svp_ppol_alloc_bytes_impl(n, cols));
 86        (
 87            SvpPPol::from_data(take_slice, n, cols),
 88            Scratch::from_bytes(rem_slice),
 89        )
 90    }
 91}
 92
 93unsafe impl<B: Backend> TakeVecZnxImpl<B> for FFT64Ref
 94where
 95    B: ScratchFromBytesImpl<B>,
 96{
 97    fn take_vec_znx_impl(scratch: &mut Scratch<B>, n: usize, cols: usize, size: usize) -> (VecZnx<&mut [u8]>, &mut Scratch<B>) {
 98        let (take_slice, rem_slice) = take_slice_aligned(&mut scratch.data, VecZnx::alloc_bytes(n, cols, size));
 99        (
100            VecZnx::from_data(take_slice, n, cols, size),
101            Scratch::from_bytes(rem_slice),
102        )
103    }
104}
105
106unsafe impl<B: Backend> TakeVecZnxBigImpl<B> for FFT64Ref
107where
108    B: VecZnxBigAllocBytesImpl<B> + ScratchFromBytesImpl<B>,
109{
110    fn take_vec_znx_big_impl(
111        scratch: &mut Scratch<B>,
112        n: usize,
113        cols: usize,
114        size: usize,
115    ) -> (VecZnxBig<&mut [u8], B>, &mut Scratch<B>) {
116        let (take_slice, rem_slice) = take_slice_aligned(
117            &mut scratch.data,
118            B::vec_znx_big_alloc_bytes_impl(n, cols, size),
119        );
120        (
121            VecZnxBig::from_data(take_slice, n, cols, size),
122            Scratch::from_bytes(rem_slice),
123        )
124    }
125}
126
127unsafe impl<B: Backend> TakeVecZnxDftImpl<B> for FFT64Ref
128where
129    B: VecZnxDftAllocBytesImpl<B> + ScratchFromBytesImpl<B>,
130{
131    fn take_vec_znx_dft_impl(
132        scratch: &mut Scratch<B>,
133        n: usize,
134        cols: usize,
135        size: usize,
136    ) -> (VecZnxDft<&mut [u8], B>, &mut Scratch<B>) {
137        let (take_slice, rem_slice) = take_slice_aligned(
138            &mut scratch.data,
139            B::vec_znx_dft_alloc_bytes_impl(n, cols, size),
140        );
141
142        (
143            VecZnxDft::from_data(take_slice, n, cols, size),
144            Scratch::from_bytes(rem_slice),
145        )
146    }
147}
148
149unsafe impl<B: Backend> TakeVecZnxDftSliceImpl<B> for FFT64Ref
150where
151    B: VecZnxDftAllocBytesImpl<B> + ScratchFromBytesImpl<B> + TakeVecZnxDftImpl<B>,
152{
153    fn take_vec_znx_dft_slice_impl(
154        scratch: &mut Scratch<B>,
155        len: usize,
156        n: usize,
157        cols: usize,
158        size: usize,
159    ) -> (Vec<VecZnxDft<&mut [u8], B>>, &mut Scratch<B>) {
160        let mut scratch: &mut Scratch<B> = scratch;
161        let mut slice: Vec<VecZnxDft<&mut [u8], B>> = Vec::with_capacity(len);
162        for _ in 0..len {
163            let (znx, new_scratch) = B::take_vec_znx_dft_impl(scratch, n, cols, size);
164            scratch = new_scratch;
165            slice.push(znx);
166        }
167        (slice, scratch)
168    }
169}
170
171unsafe impl<B: Backend> TakeVecZnxSliceImpl<B> for FFT64Ref
172where
173    B: ScratchFromBytesImpl<B> + TakeVecZnxImpl<B>,
174{
175    fn take_vec_znx_slice_impl(
176        scratch: &mut Scratch<B>,
177        len: usize,
178        n: usize,
179        cols: usize,
180        size: usize,
181    ) -> (Vec<VecZnx<&mut [u8]>>, &mut Scratch<B>) {
182        let mut scratch: &mut Scratch<B> = scratch;
183        let mut slice: Vec<VecZnx<&mut [u8]>> = Vec::with_capacity(len);
184        for _ in 0..len {
185            let (znx, new_scratch) = B::take_vec_znx_impl(scratch, n, cols, size);
186            scratch = new_scratch;
187            slice.push(znx);
188        }
189        (slice, scratch)
190    }
191}
192
193unsafe impl<B: Backend> TakeVmpPMatImpl<B> for FFT64Ref
194where
195    B: VmpPMatAllocBytesImpl<B> + ScratchFromBytesImpl<B>,
196{
197    fn take_vmp_pmat_impl(
198        scratch: &mut Scratch<B>,
199        n: usize,
200        rows: usize,
201        cols_in: usize,
202        cols_out: usize,
203        size: usize,
204    ) -> (VmpPMat<&mut [u8], B>, &mut Scratch<B>) {
205        let (take_slice, rem_slice) = take_slice_aligned(
206            &mut scratch.data,
207            B::vmp_pmat_alloc_bytes_impl(n, rows, cols_in, cols_out, size),
208        );
209        (
210            VmpPMat::from_data(take_slice, n, rows, cols_in, cols_out, size),
211            Scratch::from_bytes(rem_slice),
212        )
213    }
214}
215
216unsafe impl<B: Backend> TakeMatZnxImpl<B> for FFT64Ref
217where
218    B: ScratchFromBytesImpl<B>,
219{
220    fn take_mat_znx_impl(
221        scratch: &mut Scratch<B>,
222        n: usize,
223        rows: usize,
224        cols_in: usize,
225        cols_out: usize,
226        size: usize,
227    ) -> (MatZnx<&mut [u8]>, &mut Scratch<B>) {
228        let (take_slice, rem_slice) = take_slice_aligned(
229            &mut scratch.data,
230            MatZnx::alloc_bytes(n, rows, cols_in, cols_out, size),
231        );
232        (
233            MatZnx::from_data(take_slice, n, rows, cols_in, cols_out, size),
234            Scratch::from_bytes(rem_slice),
235        )
236    }
237}
238
239fn take_slice_aligned(data: &mut [u8], take_len: usize) -> (&mut [u8], &mut [u8]) {
240    let ptr: *mut u8 = data.as_mut_ptr();
241    let self_len: usize = data.len();
242
243    let aligned_offset: usize = ptr.align_offset(DEFAULTALIGN);
244    let aligned_len: usize = self_len.saturating_sub(aligned_offset);
245
246    if let Some(rem_len) = aligned_len.checked_sub(take_len) {
247        unsafe {
248            let rem_ptr: *mut u8 = ptr.add(aligned_offset).add(take_len);
249            let rem_slice: &mut [u8] = &mut *std::ptr::slice_from_raw_parts_mut(rem_ptr, rem_len);
250
251            let take_slice: &mut [u8] = &mut *std::ptr::slice_from_raw_parts_mut(ptr.add(aligned_offset), take_len);
252
253            (take_slice, rem_slice)
254        }
255    } else {
256        panic!(
257            "Attempted to take {} from scratch with {} aligned bytes left",
258            take_len, aligned_len,
259        );
260    }
261}