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            ScratchAvailableImpl, ScratchFromBytesImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, SvpPPolAllocBytesImpl,
 10            TakeMatZnxImpl, TakeScalarZnxImpl, TakeSliceImpl, TakeSvpPPolImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl,
 11            TakeVecZnxDftSliceImpl, TakeVecZnxImpl, TakeVecZnxSliceImpl, TakeVmpPMatImpl, VecZnxBigAllocBytesImpl,
 12            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,
 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, ScalarZnx::alloc_bytes(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,
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(&mut scratch.data, VecZnx::alloc_bytes(n, cols, size));
109        (
110            VecZnx::from_data(take_slice, n, cols, size),
111            Scratch::from_bytes(rem_slice),
112        )
113    }
114}
115
116unsafe impl<B: Backend> TakeVecZnxBigImpl<B> for B
117where
118    B: CPUAVX + VecZnxBigAllocBytesImpl<B>,
119{
120    fn take_vec_znx_big_impl(
121        scratch: &mut Scratch<B>,
122        n: usize,
123        cols: usize,
124        size: usize,
125    ) -> (VecZnxBig<&mut [u8], B>, &mut Scratch<B>) {
126        let (take_slice, rem_slice) = take_slice_aligned(
127            &mut scratch.data,
128            B::vec_znx_big_alloc_bytes_impl(n, cols, size),
129        );
130        (
131            VecZnxBig::from_data(take_slice, n, cols, size),
132            Scratch::from_bytes(rem_slice),
133        )
134    }
135}
136
137unsafe impl<B: Backend> TakeVecZnxDftImpl<B> for B
138where
139    B: CPUAVX + VecZnxDftAllocBytesImpl<B>,
140{
141    fn take_vec_znx_dft_impl(
142        scratch: &mut Scratch<B>,
143        n: usize,
144        cols: usize,
145        size: usize,
146    ) -> (VecZnxDft<&mut [u8], B>, &mut Scratch<B>) {
147        let (take_slice, rem_slice) = take_slice_aligned(
148            &mut scratch.data,
149            B::vec_znx_dft_alloc_bytes_impl(n, cols, size),
150        );
151
152        (
153            VecZnxDft::from_data(take_slice, n, cols, size),
154            Scratch::from_bytes(rem_slice),
155        )
156    }
157}
158
159unsafe impl<B: Backend> TakeVecZnxDftSliceImpl<B> for B
160where
161    B: CPUAVX + VecZnxDftAllocBytesImpl<B>,
162{
163    fn take_vec_znx_dft_slice_impl(
164        scratch: &mut Scratch<B>,
165        len: usize,
166        n: usize,
167        cols: usize,
168        size: usize,
169    ) -> (Vec<VecZnxDft<&mut [u8], B>>, &mut Scratch<B>) {
170        let mut scratch: &mut Scratch<B> = scratch;
171        let mut slice: Vec<VecZnxDft<&mut [u8], B>> = Vec::with_capacity(len);
172        for _ in 0..len {
173            let (znx, new_scratch) = B::take_vec_znx_dft_impl(scratch, n, cols, size);
174            scratch = new_scratch;
175            slice.push(znx);
176        }
177        (slice, scratch)
178    }
179}
180
181unsafe impl<B: Backend> TakeVecZnxSliceImpl<B> for B
182where
183    B: CPUAVX,
184{
185    fn take_vec_znx_slice_impl(
186        scratch: &mut Scratch<B>,
187        len: usize,
188        n: usize,
189        cols: usize,
190        size: usize,
191    ) -> (Vec<VecZnx<&mut [u8]>>, &mut Scratch<B>) {
192        let mut scratch: &mut Scratch<B> = scratch;
193        let mut slice: Vec<VecZnx<&mut [u8]>> = Vec::with_capacity(len);
194        for _ in 0..len {
195            let (znx, new_scratch) = B::take_vec_znx_impl(scratch, n, cols, size);
196            scratch = new_scratch;
197            slice.push(znx);
198        }
199        (slice, scratch)
200    }
201}
202
203unsafe impl<B: Backend> TakeVmpPMatImpl<B> for B
204where
205    B: CPUAVX + VmpPMatAllocBytesImpl<B>,
206{
207    fn take_vmp_pmat_impl(
208        scratch: &mut Scratch<B>,
209        n: usize,
210        rows: usize,
211        cols_in: usize,
212        cols_out: usize,
213        size: usize,
214    ) -> (VmpPMat<&mut [u8], B>, &mut Scratch<B>) {
215        let (take_slice, rem_slice) = take_slice_aligned(
216            &mut scratch.data,
217            B::vmp_pmat_alloc_bytes_impl(n, rows, cols_in, cols_out, size),
218        );
219        (
220            VmpPMat::from_data(take_slice, n, rows, cols_in, cols_out, size),
221            Scratch::from_bytes(rem_slice),
222        )
223    }
224}
225
226unsafe impl<B: Backend> TakeMatZnxImpl<B> for B
227where
228    B: CPUAVX,
229{
230    fn take_mat_znx_impl(
231        scratch: &mut Scratch<B>,
232        n: usize,
233        rows: usize,
234        cols_in: usize,
235        cols_out: usize,
236        size: usize,
237    ) -> (MatZnx<&mut [u8]>, &mut Scratch<B>) {
238        let (take_slice, rem_slice) = take_slice_aligned(
239            &mut scratch.data,
240            MatZnx::alloc_bytes(n, rows, cols_in, cols_out, size),
241        );
242        (
243            MatZnx::from_data(take_slice, n, rows, cols_in, cols_out, size),
244            Scratch::from_bytes(rem_slice),
245        )
246    }
247}
248
249fn take_slice_aligned(data: &mut [u8], take_len: usize) -> (&mut [u8], &mut [u8]) {
250    let ptr: *mut u8 = data.as_mut_ptr();
251    let self_len: usize = data.len();
252
253    let aligned_offset: usize = ptr.align_offset(DEFAULTALIGN);
254    let aligned_len: usize = self_len.saturating_sub(aligned_offset);
255
256    if let Some(rem_len) = aligned_len.checked_sub(take_len) {
257        unsafe {
258            let rem_ptr: *mut u8 = ptr.add(aligned_offset).add(take_len);
259            let rem_slice: &mut [u8] = &mut *std::ptr::slice_from_raw_parts_mut(rem_ptr, rem_len);
260
261            let take_slice: &mut [u8] = &mut *std::ptr::slice_from_raw_parts_mut(ptr.add(aligned_offset), take_len);
262
263            (take_slice, rem_slice)
264        }
265    } else {
266        panic!(
267            "Attempted to take {} from scratch with {} aligned bytes left",
268            take_len, aligned_len,
269        );
270    }
271}