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