1use std::marker::PhantomData;
 2
 3use poulpy_hal::{
 4    DEFAULTALIGN, alloc_aligned,
 5    api::ScratchFromBytes,
 6    layouts::{Backend, Scratch, ScratchOwned},
 7    oep::{ScratchAvailableImpl, ScratchFromBytesImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeSliceImpl},
 8};
 9
10use crate::cpu_fft64_avx::FFT64Avx;
11
12unsafe impl<B: Backend> ScratchOwnedAllocImpl<B> for FFT64Avx {
13    fn scratch_owned_alloc_impl(size: usize) -> ScratchOwned<B> {
14        let data: Vec<u8> = alloc_aligned(size);
15        ScratchOwned {
16            data,
17            _phantom: PhantomData,
18        }
19    }
20}
21
22unsafe impl<B: Backend> ScratchOwnedBorrowImpl<B> for FFT64Avx
23where
24    B: ScratchFromBytesImpl<B>,
25{
26    fn scratch_owned_borrow_impl(scratch: &mut ScratchOwned<B>) -> &mut Scratch<B> {
27        Scratch::from_bytes(&mut scratch.data)
28    }
29}
30
31unsafe impl<B: Backend> ScratchFromBytesImpl<B> for FFT64Avx {
32    fn scratch_from_bytes_impl(data: &mut [u8]) -> &mut Scratch<B> {
33        unsafe { &mut *(data as *mut [u8] as *mut Scratch<B>) }
34    }
35}
36
37unsafe impl<B: Backend> ScratchAvailableImpl<B> for FFT64Avx {
38    fn scratch_available_impl(scratch: &Scratch<B>) -> usize {
39        let ptr: *const u8 = scratch.data.as_ptr();
40        let self_len: usize = scratch.data.len();
41        let aligned_offset: usize = ptr.align_offset(DEFAULTALIGN);
42        self_len.saturating_sub(aligned_offset)
43    }
44}
45
46unsafe impl<B: Backend> TakeSliceImpl<B> for FFT64Avx
47where
48    B: ScratchFromBytesImpl<B>,
49{
50    fn take_slice_impl<T>(scratch: &mut Scratch<B>, len: usize) -> (&mut [T], &mut Scratch<B>) {
51        let (take_slice, rem_slice) = take_slice_aligned(&mut scratch.data, len * std::mem::size_of::<T>());
52
53        unsafe {
54            (
55                &mut *(std::ptr::slice_from_raw_parts_mut(take_slice.as_mut_ptr() as *mut T, len)),
56                Scratch::from_bytes(rem_slice),
57            )
58        }
59    }
60}
61
62fn take_slice_aligned(data: &mut [u8], take_len: usize) -> (&mut [u8], &mut [u8]) {
63    let ptr: *mut u8 = data.as_mut_ptr();
64    let self_len: usize = data.len();
65
66    let aligned_offset: usize = ptr.align_offset(DEFAULTALIGN);
67    let aligned_len: usize = self_len.saturating_sub(aligned_offset);
68
69    if let Some(rem_len) = aligned_len.checked_sub(take_len) {
70        unsafe {
71            let rem_ptr: *mut u8 = ptr.add(aligned_offset).add(take_len);
72            let rem_slice: &mut [u8] = &mut *std::ptr::slice_from_raw_parts_mut(rem_ptr, rem_len);
73
74            let take_slice: &mut [u8] = &mut *std::ptr::slice_from_raw_parts_mut(ptr.add(aligned_offset), take_len);
75
76            (take_slice, rem_slice)
77        }
78    } else {
79        panic!("Attempted to take {take_len} from scratch with {aligned_len} aligned bytes left");
80    }
81}