arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/backend/src/implementation/cpu_spqlios/vec_znx_big_fft64.rs
backend/src/implementation/cpu_spqlios/vec_znx_big_fft64.rsRaw
  1use std::fmt;
  2
  3use rand_distr::{Distribution, Normal};
  4use sampling::source::Source;
  5
  6use crate::{
  7    hal::{
  8        api::{
  9            TakeSlice, VecZnxBigAddDistF64, VecZnxBigFillDistF64, VecZnxBigNormalizeTmpBytes, ZnxInfos, ZnxSliceSize, ZnxView,
 10            ZnxViewMut,
 11        },
 12        layouts::{
 13            Data, DataRef, Module, Scratch, VecZnx, VecZnxBig, VecZnxBigBytesOf, VecZnxBigOwned, VecZnxBigToMut, VecZnxBigToRef,
 14            VecZnxToMut, VecZnxToRef,
 15        },
 16        oep::{
 17            VecZnxBigAddDistF64Impl, VecZnxBigAddImpl, VecZnxBigAddInplaceImpl, VecZnxBigAddNormalImpl, VecZnxBigAddSmallImpl,
 18            VecZnxBigAddSmallInplaceImpl, VecZnxBigAllocBytesImpl, VecZnxBigAllocImpl, VecZnxBigAutomorphismImpl,
 19            VecZnxBigAutomorphismInplaceImpl, VecZnxBigFillDistF64Impl, VecZnxBigFillNormalImpl, VecZnxBigFromBytesImpl,
 20            VecZnxBigNegateInplaceImpl, VecZnxBigNormalizeImpl, VecZnxBigNormalizeTmpBytesImpl, VecZnxBigSubABInplaceImpl,
 21            VecZnxBigSubBAInplaceImpl, VecZnxBigSubImpl, VecZnxBigSubSmallAImpl, VecZnxBigSubSmallAInplaceImpl,
 22            VecZnxBigSubSmallBImpl, VecZnxBigSubSmallBInplaceImpl,
 23        },
 24    },
 25    implementation::cpu_spqlios::{ffi::vec_znx, module_fft64::FFT64},
 26};
 27
 28const VEC_ZNX_BIG_FFT64_WORDSIZE: usize = 1;
 29
 30impl<D: DataRef> ZnxView for VecZnxBig<D, FFT64> {
 31    type Scalar = i64;
 32}
 33
 34impl<D: Data> VecZnxBigBytesOf for VecZnxBig<D, FFT64> {
 35    fn bytes_of(n: usize, cols: usize, size: usize) -> usize {
 36        VEC_ZNX_BIG_FFT64_WORDSIZE * n * cols * size * size_of::<f64>()
 37    }
 38}
 39
 40impl<D: Data> ZnxSliceSize for VecZnxBig<D, FFT64> {
 41    fn sl(&self) -> usize {
 42        VEC_ZNX_BIG_FFT64_WORDSIZE * self.n() * self.cols()
 43    }
 44}
 45
 46unsafe impl VecZnxBigAllocImpl<FFT64> for FFT64 {
 47    fn vec_znx_big_alloc_impl(n: usize, cols: usize, size: usize) -> VecZnxBigOwned<FFT64> {
 48        VecZnxBig::<Vec<u8>, FFT64>::new(n, cols, size)
 49    }
 50}
 51
 52unsafe impl VecZnxBigFromBytesImpl<FFT64> for FFT64 {
 53    fn vec_znx_big_from_bytes_impl(n: usize, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBigOwned<FFT64> {
 54        VecZnxBig::<Vec<u8>, FFT64>::new_from_bytes(n, cols, size, bytes)
 55    }
 56}
 57
 58unsafe impl VecZnxBigAllocBytesImpl<FFT64> for FFT64 {
 59    fn vec_znx_big_alloc_bytes_impl(n: usize, cols: usize, size: usize) -> usize {
 60        VecZnxBig::<Vec<u8>, FFT64>::bytes_of(n, cols, size)
 61    }
 62}
 63
 64unsafe impl VecZnxBigAddDistF64Impl<FFT64> for FFT64 {
 65    fn add_dist_f64_impl<R: VecZnxBigToMut<FFT64>, D: Distribution<f64>>(
 66        _module: &Module<FFT64>,
 67        basek: usize,
 68        res: &mut R,
 69        res_col: usize,
 70        k: usize,
 71        source: &mut Source,
 72        dist: D,
 73        bound: f64,
 74    ) {
 75        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
 76        assert!(
 77            (bound.log2().ceil() as i64) < 64,
 78            "invalid bound: ceil(log2(bound))={} > 63",
 79            (bound.log2().ceil() as i64)
 80        );
 81
 82        let limb: usize = k.div_ceil(basek) - 1;
 83        let basek_rem: usize = (limb + 1) * basek - k;
 84
 85        if basek_rem != 0 {
 86            res.at_mut(res_col, limb).iter_mut().for_each(|x| {
 87                let mut dist_f64: f64 = dist.sample(source);
 88                while dist_f64.abs() > bound {
 89                    dist_f64 = dist.sample(source)
 90                }
 91                *x += (dist_f64.round() as i64) << basek_rem;
 92            });
 93        } else {
 94            res.at_mut(res_col, limb).iter_mut().for_each(|x| {
 95                let mut dist_f64: f64 = dist.sample(source);
 96                while dist_f64.abs() > bound {
 97                    dist_f64 = dist.sample(source)
 98                }
 99                *x += dist_f64.round() as i64
100            });
101        }
102    }
103}
104
105unsafe impl VecZnxBigAddNormalImpl<FFT64> for FFT64 {
106    fn add_normal_impl<R: VecZnxBigToMut<FFT64>>(
107        module: &Module<FFT64>,
108        basek: usize,
109        res: &mut R,
110        res_col: usize,
111        k: usize,
112        source: &mut Source,
113        sigma: f64,
114        bound: f64,
115    ) {
116        module.vec_znx_big_add_dist_f64(
117            basek,
118            res,
119            res_col,
120            k,
121            source,
122            Normal::new(0.0, sigma).unwrap(),
123            bound,
124        );
125    }
126}
127
128unsafe impl VecZnxBigFillDistF64Impl<FFT64> for FFT64 {
129    fn fill_dist_f64_impl<R: VecZnxBigToMut<FFT64>, D: Distribution<f64>>(
130        _module: &Module<FFT64>,
131        basek: usize,
132        res: &mut R,
133        res_col: usize,
134        k: usize,
135        source: &mut Source,
136        dist: D,
137        bound: f64,
138    ) {
139        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
140        assert!(
141            (bound.log2().ceil() as i64) < 64,
142            "invalid bound: ceil(log2(bound))={} > 63",
143            (bound.log2().ceil() as i64)
144        );
145
146        let limb: usize = k.div_ceil(basek) - 1;
147        let basek_rem: usize = (limb + 1) * basek - k;
148
149        if basek_rem != 0 {
150            res.at_mut(res_col, limb).iter_mut().for_each(|x| {
151                let mut dist_f64: f64 = dist.sample(source);
152                while dist_f64.abs() > bound {
153                    dist_f64 = dist.sample(source)
154                }
155                *x = (dist_f64.round() as i64) << basek_rem;
156            });
157        } else {
158            res.at_mut(res_col, limb).iter_mut().for_each(|x| {
159                let mut dist_f64: f64 = dist.sample(source);
160                while dist_f64.abs() > bound {
161                    dist_f64 = dist.sample(source)
162                }
163                *x = dist_f64.round() as i64
164            });
165        }
166    }
167}
168
169unsafe impl VecZnxBigFillNormalImpl<FFT64> for FFT64 {
170    fn fill_normal_impl<R: VecZnxBigToMut<FFT64>>(
171        module: &Module<FFT64>,
172        basek: usize,
173        res: &mut R,
174        res_col: usize,
175        k: usize,
176        source: &mut Source,
177        sigma: f64,
178        bound: f64,
179    ) {
180        module.vec_znx_big_fill_dist_f64(
181            basek,
182            res,
183            res_col,
184            k,
185            source,
186            Normal::new(0.0, sigma).unwrap(),
187            bound,
188        );
189    }
190}
191
192unsafe impl VecZnxBigAddImpl<FFT64> for FFT64 {
193    /// Adds `a` to `b` and stores the result on `c`.
194    fn vec_znx_big_add_impl<R, A, B>(
195        module: &Module<FFT64>,
196        res: &mut R,
197        res_col: usize,
198        a: &A,
199        a_col: usize,
200        b: &B,
201        b_col: usize,
202    ) where
203        R: VecZnxBigToMut<FFT64>,
204        A: VecZnxBigToRef<FFT64>,
205        B: VecZnxBigToRef<FFT64>,
206    {
207        let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
208        let b: VecZnxBig<&[u8], FFT64> = b.to_ref();
209        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
210
211        #[cfg(debug_assertions)]
212        {
213            assert_eq!(a.n(), module.n());
214            assert_eq!(b.n(), module.n());
215            assert_eq!(res.n(), module.n());
216            assert_ne!(a.as_ptr(), b.as_ptr());
217        }
218        unsafe {
219            vec_znx::vec_znx_add(
220                module.ptr(),
221                res.at_mut_ptr(res_col, 0),
222                res.size() as u64,
223                res.sl() as u64,
224                a.at_ptr(a_col, 0),
225                a.size() as u64,
226                a.sl() as u64,
227                b.at_ptr(b_col, 0),
228                b.size() as u64,
229                b.sl() as u64,
230            )
231        }
232    }
233}
234
235unsafe impl VecZnxBigAddInplaceImpl<FFT64> for FFT64 {
236    /// Adds `a` to `b` and stores the result on `b`.
237    fn vec_znx_big_add_inplace_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
238    where
239        R: VecZnxBigToMut<FFT64>,
240        A: VecZnxBigToRef<FFT64>,
241    {
242        let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
243        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
244
245        #[cfg(debug_assertions)]
246        {
247            assert_eq!(a.n(), module.n());
248            assert_eq!(res.n(), module.n());
249        }
250        unsafe {
251            vec_znx::vec_znx_add(
252                module.ptr(),
253                res.at_mut_ptr(res_col, 0),
254                res.size() as u64,
255                res.sl() as u64,
256                a.at_ptr(a_col, 0),
257                a.size() as u64,
258                a.sl() as u64,
259                res.at_ptr(res_col, 0),
260                res.size() as u64,
261                res.sl() as u64,
262            )
263        }
264    }
265}
266
267unsafe impl VecZnxBigAddSmallImpl<FFT64> for FFT64 {
268    /// Adds `a` to `b` and stores the result on `c`.
269    fn vec_znx_big_add_small_impl<R, A, B>(
270        module: &Module<FFT64>,
271        res: &mut R,
272        res_col: usize,
273        a: &A,
274        a_col: usize,
275        b: &B,
276        b_col: usize,
277    ) where
278        R: VecZnxBigToMut<FFT64>,
279        A: VecZnxBigToRef<FFT64>,
280        B: VecZnxToRef,
281    {
282        let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
283        let b: VecZnx<&[u8]> = b.to_ref();
284        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
285
286        #[cfg(debug_assertions)]
287        {
288            assert_eq!(a.n(), module.n());
289            assert_eq!(b.n(), module.n());
290            assert_eq!(res.n(), module.n());
291            assert_ne!(a.as_ptr(), b.as_ptr());
292        }
293        unsafe {
294            vec_znx::vec_znx_add(
295                module.ptr(),
296                res.at_mut_ptr(res_col, 0),
297                res.size() as u64,
298                res.sl() as u64,
299                a.at_ptr(a_col, 0),
300                a.size() as u64,
301                a.sl() as u64,
302                b.at_ptr(b_col, 0),
303                b.size() as u64,
304                b.sl() as u64,
305            )
306        }
307    }
308}
309
310unsafe impl VecZnxBigAddSmallInplaceImpl<FFT64> for FFT64 {
311    /// Adds `a` to `b` and stores the result on `b`.
312    fn vec_znx_big_add_small_inplace_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
313    where
314        R: VecZnxBigToMut<FFT64>,
315        A: VecZnxToRef,
316    {
317        let a: VecZnx<&[u8]> = a.to_ref();
318        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
319
320        #[cfg(debug_assertions)]
321        {
322            assert_eq!(a.n(), module.n());
323            assert_eq!(res.n(), module.n());
324        }
325        unsafe {
326            vec_znx::vec_znx_add(
327                module.ptr(),
328                res.at_mut_ptr(res_col, 0),
329                res.size() as u64,
330                res.sl() as u64,
331                res.at_ptr(res_col, 0),
332                res.size() as u64,
333                res.sl() as u64,
334                a.at_ptr(a_col, 0),
335                a.size() as u64,
336                a.sl() as u64,
337            )
338        }
339    }
340}
341
342unsafe impl VecZnxBigSubImpl<FFT64> for FFT64 {
343    /// Subtracts `a` to `b` and stores the result on `c`.
344    fn vec_znx_big_sub_impl<R, A, B>(
345        module: &Module<FFT64>,
346        res: &mut R,
347        res_col: usize,
348        a: &A,
349        a_col: usize,
350        b: &B,
351        b_col: usize,
352    ) where
353        R: VecZnxBigToMut<FFT64>,
354        A: VecZnxBigToRef<FFT64>,
355        B: VecZnxBigToRef<FFT64>,
356    {
357        let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
358        let b: VecZnxBig<&[u8], FFT64> = b.to_ref();
359        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
360
361        #[cfg(debug_assertions)]
362        {
363            assert_eq!(a.n(), module.n());
364            assert_eq!(b.n(), module.n());
365            assert_eq!(res.n(), module.n());
366            assert_ne!(a.as_ptr(), b.as_ptr());
367        }
368        unsafe {
369            vec_znx::vec_znx_sub(
370                module.ptr(),
371                res.at_mut_ptr(res_col, 0),
372                res.size() as u64,
373                res.sl() as u64,
374                a.at_ptr(a_col, 0),
375                a.size() as u64,
376                a.sl() as u64,
377                b.at_ptr(b_col, 0),
378                b.size() as u64,
379                b.sl() as u64,
380            )
381        }
382    }
383}
384
385unsafe impl VecZnxBigSubABInplaceImpl<FFT64> for FFT64 {
386    /// Subtracts `a` from `b` and stores the result on `b`.
387    fn vec_znx_big_sub_ab_inplace_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
388    where
389        R: VecZnxBigToMut<FFT64>,
390        A: VecZnxBigToRef<FFT64>,
391    {
392        let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
393        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
394
395        #[cfg(debug_assertions)]
396        {
397            assert_eq!(a.n(), module.n());
398            assert_eq!(res.n(), module.n());
399        }
400        unsafe {
401            vec_znx::vec_znx_sub(
402                module.ptr(),
403                res.at_mut_ptr(res_col, 0),
404                res.size() as u64,
405                res.sl() as u64,
406                res.at_ptr(res_col, 0),
407                res.size() as u64,
408                res.sl() as u64,
409                a.at_ptr(a_col, 0),
410                a.size() as u64,
411                a.sl() as u64,
412            )
413        }
414    }
415}
416
417unsafe impl VecZnxBigSubBAInplaceImpl<FFT64> for FFT64 {
418    /// Subtracts `b` from `a` and stores the result on `b`.
419    fn vec_znx_big_sub_ba_inplace_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
420    where
421        R: VecZnxBigToMut<FFT64>,
422        A: VecZnxBigToRef<FFT64>,
423    {
424        let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
425        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
426
427        #[cfg(debug_assertions)]
428        {
429            assert_eq!(a.n(), module.n());
430            assert_eq!(res.n(), module.n());
431        }
432        unsafe {
433            vec_znx::vec_znx_sub(
434                module.ptr(),
435                res.at_mut_ptr(res_col, 0),
436                res.size() as u64,
437                res.sl() as u64,
438                a.at_ptr(a_col, 0),
439                a.size() as u64,
440                a.sl() as u64,
441                res.at_ptr(res_col, 0),
442                res.size() as u64,
443                res.sl() as u64,
444            )
445        }
446    }
447}
448
449unsafe impl VecZnxBigSubSmallAImpl<FFT64> for FFT64 {
450    /// Subtracts `b` from `a` and stores the result on `c`.
451    fn vec_znx_big_sub_small_a_impl<R, A, B>(
452        module: &Module<FFT64>,
453        res: &mut R,
454        res_col: usize,
455        a: &A,
456        a_col: usize,
457        b: &B,
458        b_col: usize,
459    ) where
460        R: VecZnxBigToMut<FFT64>,
461        A: VecZnxToRef,
462        B: VecZnxBigToRef<FFT64>,
463    {
464        let a: VecZnx<&[u8]> = a.to_ref();
465        let b: VecZnxBig<&[u8], FFT64> = b.to_ref();
466        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
467
468        #[cfg(debug_assertions)]
469        {
470            assert_eq!(a.n(), module.n());
471            assert_eq!(b.n(), module.n());
472            assert_eq!(res.n(), module.n());
473            assert_ne!(a.as_ptr(), b.as_ptr());
474        }
475        unsafe {
476            vec_znx::vec_znx_sub(
477                module.ptr(),
478                res.at_mut_ptr(res_col, 0),
479                res.size() as u64,
480                res.sl() as u64,
481                a.at_ptr(a_col, 0),
482                a.size() as u64,
483                a.sl() as u64,
484                b.at_ptr(b_col, 0),
485                b.size() as u64,
486                b.sl() as u64,
487            )
488        }
489    }
490}
491
492unsafe impl VecZnxBigSubSmallAInplaceImpl<FFT64> for FFT64 {
493    /// Subtracts `a` from `res` and stores the result on `res`.
494    fn vec_znx_big_sub_small_a_inplace_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
495    where
496        R: VecZnxBigToMut<FFT64>,
497        A: VecZnxToRef,
498    {
499        let a: VecZnx<&[u8]> = a.to_ref();
500        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
501
502        #[cfg(debug_assertions)]
503        {
504            assert_eq!(a.n(), module.n());
505            assert_eq!(res.n(), module.n());
506        }
507        unsafe {
508            vec_znx::vec_znx_sub(
509                module.ptr(),
510                res.at_mut_ptr(res_col, 0),
511                res.size() as u64,
512                res.sl() as u64,
513                res.at_ptr(res_col, 0),
514                res.size() as u64,
515                res.sl() as u64,
516                a.at_ptr(a_col, 0),
517                a.size() as u64,
518                a.sl() as u64,
519            )
520        }
521    }
522}
523
524unsafe impl VecZnxBigSubSmallBImpl<FFT64> for FFT64 {
525    /// Subtracts `b` from `a` and stores the result on `c`.
526    fn vec_znx_big_sub_small_b_impl<R, A, B>(
527        module: &Module<FFT64>,
528        res: &mut R,
529        res_col: usize,
530        a: &A,
531        a_col: usize,
532        b: &B,
533        b_col: usize,
534    ) where
535        R: VecZnxBigToMut<FFT64>,
536        A: VecZnxBigToRef<FFT64>,
537        B: VecZnxToRef,
538    {
539        let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
540        let b: VecZnx<&[u8]> = b.to_ref();
541        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
542
543        #[cfg(debug_assertions)]
544        {
545            assert_eq!(a.n(), module.n());
546            assert_eq!(b.n(), module.n());
547            assert_eq!(res.n(), module.n());
548            assert_ne!(a.as_ptr(), b.as_ptr());
549        }
550        unsafe {
551            vec_znx::vec_znx_sub(
552                module.ptr(),
553                res.at_mut_ptr(res_col, 0),
554                res.size() as u64,
555                res.sl() as u64,
556                a.at_ptr(a_col, 0),
557                a.size() as u64,
558                a.sl() as u64,
559                b.at_ptr(b_col, 0),
560                b.size() as u64,
561                b.sl() as u64,
562            )
563        }
564    }
565}
566
567unsafe impl VecZnxBigSubSmallBInplaceImpl<FFT64> for FFT64 {
568    /// Subtracts `res` from `a` and stores the result on `res`.
569    fn vec_znx_big_sub_small_b_inplace_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
570    where
571        R: VecZnxBigToMut<FFT64>,
572        A: VecZnxToRef,
573    {
574        let a: VecZnx<&[u8]> = a.to_ref();
575        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
576
577        #[cfg(debug_assertions)]
578        {
579            assert_eq!(a.n(), module.n());
580            assert_eq!(res.n(), module.n());
581        }
582        unsafe {
583            vec_znx::vec_znx_sub(
584                module.ptr(),
585                res.at_mut_ptr(res_col, 0),
586                res.size() as u64,
587                res.sl() as u64,
588                a.at_ptr(a_col, 0),
589                a.size() as u64,
590                a.sl() as u64,
591                res.at_ptr(res_col, 0),
592                res.size() as u64,
593                res.sl() as u64,
594            )
595        }
596    }
597}
598
599unsafe impl VecZnxBigNegateInplaceImpl<FFT64> for FFT64 {
600    fn vec_znx_big_negate_inplace_impl<A>(module: &Module<FFT64>, a: &mut A, a_col: usize)
601    where
602        A: VecZnxBigToMut<FFT64>,
603    {
604        let mut a: VecZnxBig<&mut [u8], FFT64> = a.to_mut();
605        #[cfg(debug_assertions)]
606        {
607            assert_eq!(a.n(), module.n());
608        }
609        unsafe {
610            vec_znx::vec_znx_negate(
611                module.ptr(),
612                a.at_mut_ptr(a_col, 0),
613                a.size() as u64,
614                a.sl() as u64,
615                a.at_ptr(a_col, 0),
616                a.size() as u64,
617                a.sl() as u64,
618            )
619        }
620    }
621}
622
623unsafe impl VecZnxBigNormalizeTmpBytesImpl<FFT64> for FFT64 {
624    fn vec_znx_big_normalize_tmp_bytes_impl(module: &Module<FFT64>, n: usize) -> usize {
625        unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(module.ptr(), n as u64) as usize }
626    }
627}
628
629unsafe impl VecZnxBigNormalizeImpl<FFT64> for FFT64 {
630    fn vec_znx_big_normalize_impl<R, A>(
631        module: &Module<FFT64>,
632        basek: usize,
633        res: &mut R,
634        res_col: usize,
635        a: &A,
636        a_col: usize,
637        scratch: &mut Scratch<FFT64>,
638    ) where
639        R: VecZnxToMut,
640        A: VecZnxBigToRef<FFT64>,
641    {
642        let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
643        let mut res: VecZnx<&mut [u8]> = res.to_mut();
644
645        #[cfg(debug_assertions)]
646        {
647            assert_eq!(res.n(), a.n());
648        }
649
650        let (tmp_bytes, _) = scratch.take_slice(module.vec_znx_big_normalize_tmp_bytes(a.n()));
651        unsafe {
652            vec_znx::vec_znx_normalize_base2k(
653                module.ptr(),
654                a.n() as u64,
655                basek as u64,
656                res.at_mut_ptr(res_col, 0),
657                res.size() as u64,
658                res.sl() as u64,
659                a.at_ptr(a_col, 0),
660                a.size() as u64,
661                a.sl() as u64,
662                tmp_bytes.as_mut_ptr(),
663            );
664        }
665    }
666}
667
668unsafe impl VecZnxBigAutomorphismImpl<FFT64> for FFT64 {
669    /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
670    fn vec_znx_big_automorphism_impl<R, A>(module: &Module<FFT64>, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
671    where
672        R: VecZnxBigToMut<FFT64>,
673        A: VecZnxBigToRef<FFT64>,
674    {
675        let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
676        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
677
678        #[cfg(debug_assertions)]
679        {
680            assert_eq!(a.n(), module.n());
681            assert_eq!(res.n(), module.n());
682        }
683        unsafe {
684            vec_znx::vec_znx_automorphism(
685                module.ptr(),
686                k,
687                res.at_mut_ptr(res_col, 0),
688                res.size() as u64,
689                res.sl() as u64,
690                a.at_ptr(a_col, 0),
691                a.size() as u64,
692                a.sl() as u64,
693            )
694        }
695    }
696}
697
698unsafe impl VecZnxBigAutomorphismInplaceImpl<FFT64> for FFT64 {
699    /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
700    fn vec_znx_big_automorphism_inplace_impl<A>(module: &Module<FFT64>, k: i64, a: &mut A, a_col: usize)
701    where
702        A: VecZnxBigToMut<FFT64>,
703    {
704        let mut a: VecZnxBig<&mut [u8], FFT64> = a.to_mut();
705
706        #[cfg(debug_assertions)]
707        {
708            assert_eq!(a.n(), module.n());
709        }
710        unsafe {
711            vec_znx::vec_znx_automorphism(
712                module.ptr(),
713                k,
714                a.at_mut_ptr(a_col, 0),
715                a.size() as u64,
716                a.sl() as u64,
717                a.at_ptr(a_col, 0),
718                a.size() as u64,
719                a.sl() as u64,
720            )
721        }
722    }
723}
724
725impl<D: DataRef> fmt::Display for VecZnxBig<D, FFT64> {
726    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
727        writeln!(
728            f,
729            "VecZnxBig(n={}, cols={}, size={})",
730            self.n, self.cols, self.size
731        )?;
732
733        for col in 0..self.cols {
734            writeln!(f, "Column {}:", col)?;
735            for size in 0..self.size {
736                let coeffs = self.at(col, size);
737                write!(f, "  Size {}: [", size)?;
738
739                let max_show = 100;
740                let show_count = coeffs.len().min(max_show);
741
742                for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
743                    if i > 0 {
744                        write!(f, ", ")?;
745                    }
746                    write!(f, "{}", coeff)?;
747                }
748
749                if coeffs.len() > max_show {
750                    write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
751                }
752
753                writeln!(f, "]")?;
754            }
755        }
756        Ok(())
757    }
758}