poulpy_hal/test_suite/
convolution.rs

1use super::{download_vec_znx, upload_vec_znx, vec_znx_backend_mut, vec_znx_backend_ref};
2use crate::layouts::CnvPVecLToBackendMut;
3use crate::layouts::CnvPVecLToBackendRef;
4use crate::layouts::CnvPVecRToBackendMut;
5use crate::layouts::CnvPVecRToBackendRef;
6use crate::layouts::VecZnxBigToBackendMut;
7use crate::layouts::VecZnxBigToBackendRef;
8use crate::layouts::VecZnxDftToBackendMut;
9use crate::layouts::VecZnxDftToBackendRef;
10use rand::Rng;
11
12use crate::{
13    api::{
14        CnvPVecAlloc, Convolution, ModuleN, ScratchOwnedAlloc, VecZnxAddIntoBackend, VecZnxBigAlloc, VecZnxBigNormalize,
15        VecZnxBigNormalizeTmpBytes, VecZnxCopyBackend, VecZnxDftAddAssign, VecZnxDftAlloc, VecZnxDftApply, VecZnxIdftApplyTmpA,
16        VecZnxNormalizeAssignBackend,
17    },
18    layouts::{DataView, FillUniform, ScratchArena, ScratchOwned, VecZnx, VecZnxOwned, ZnxView, ZnxViewMut, ZnxZero},
19    source::Source,
20};
21
22use crate::layouts::VecZnxDftOwned;
23use crate::layouts::{CnvPVecLOwned, CnvPVecROwned, VecZnxBigOwned};
24
25pub fn test_convolution_by_const<M, BE: crate::test_suite::TestBackend>(module: &M, base2k: usize)
26where
27    M: ModuleN
28        + Convolution<BE>
29        + VecZnxBigNormalize<BE>
30        + VecZnxBigNormalizeTmpBytes
31        + VecZnxNormalizeAssignBackend<BE>
32        + VecZnxBigAlloc<BE>,
33    ScratchOwned<BE>: ScratchOwnedAlloc<BE>,
34{
35    let mut source: Source = Source::new([0u8; 32]);
36
37    let a_cols: usize = 2;
38    let a_size: usize = 15;
39    let b_size: usize = 15;
40    let res_size: usize = a_size + b_size;
41
42    let mut a = VecZnx::alloc(module.n(), a_cols, a_size);
43    let mut b = VecZnx::alloc(module.n(), 1, b_size);
44
45    let mut res_want = VecZnx::alloc(module.n(), 1, res_size);
46    let mut res_big: VecZnxBigOwned<BE> = module.vec_znx_big_alloc(1, res_size);
47
48    a.fill_uniform(17, &mut source);
49
50    let mask = (1 << base2k) - 1;
51    for j in 0..1 {
52        let r = source.next_u64() & mask;
53        b.at_mut(0, j)[0] = ((r << (64 - 17)) as i64) >> (64 - 17);
54    }
55
56    let a_backend = upload_vec_znx::<BE>(&a);
57    let b_backend = upload_vec_znx::<BE>(&b);
58    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
59        module
60            .cnv_by_const_apply_tmp_bytes(0, res_size, a_size, b_size)
61            .max(module.vec_znx_big_normalize_tmp_bytes()),
62    );
63
64    for a_col in 0..a.cols() {
65        for cnv_offset in 0..res_size {
66            module.cnv_by_const_apply(
67                cnv_offset,
68                &mut res_big.to_backend_mut(),
69                0,
70                &vec_znx_backend_ref::<BE>(&a_backend),
71                a_col,
72                &vec_znx_backend_ref::<BE>(&b_backend),
73                0,
74                0,
75                &mut scratch.arena(),
76            );
77
78            let res_host_template = VecZnx::alloc(module.n(), 1, res_size);
79            let mut res_have_backend = upload_vec_znx::<BE>(&res_host_template);
80            module.vec_znx_big_normalize(
81                &mut vec_znx_backend_mut::<BE>(&mut res_have_backend),
82                base2k,
83                0,
84                0,
85                &res_big.to_backend_ref(),
86                base2k,
87                0,
88                &mut scratch.arena(),
89            );
90            let res_have = download_vec_znx::<BE>(&res_have_backend);
91
92            bivariate_convolution_naive(
93                module,
94                base2k,
95                (cnv_offset + 1) as i64,
96                &mut res_want,
97                0,
98                &a,
99                a_col,
100                &b,
101                0,
102                &mut scratch.arena(),
103            );
104
105            assert_eq!(res_want, res_have);
106        }
107    }
108}
109
110pub fn test_convolution<M, BE: crate::test_suite::TestBackend>(module: &M, base2k: usize)
111where
112    M: ModuleN
113        + Convolution<BE>
114        + CnvPVecAlloc<BE>
115        + VecZnxDftAlloc<BE>
116        + VecZnxDftApply<BE>
117        + VecZnxIdftApplyTmpA<BE>
118        + VecZnxBigNormalize<BE>
119        + VecZnxBigNormalizeTmpBytes
120        + VecZnxNormalizeAssignBackend<BE>
121        + VecZnxBigAlloc<BE>,
122    ScratchOwned<BE>: ScratchOwnedAlloc<BE>,
123{
124    let mut source: Source = Source::new([0u8; 32]);
125
126    let a_cols: usize = 2;
127    let b_cols: usize = 2;
128    let a_size: usize = 15;
129    let b_size: usize = 15;
130    let res_size: usize = a_size + b_size;
131
132    let mut a = VecZnx::alloc(module.n(), a_cols, a_size);
133    let mut b = VecZnx::alloc(module.n(), b_cols, b_size);
134
135    let mut res_want = VecZnx::alloc(module.n(), 1, res_size);
136    // Two-column DFT destination written at column 1: covers the
137    // column-interleaved `VecZnxDft` indexing of the backend kernels.
138    let res_dft_col: usize = 1;
139    let mut res_dft: VecZnxDftOwned<BE> = module.vec_znx_dft_alloc(2, res_size);
140    let mut res_big: VecZnxBigOwned<BE> = module.vec_znx_big_alloc(1, res_size);
141
142    a.fill_uniform(17, &mut source);
143    b.fill_uniform(17, &mut source);
144
145    let a_backend = upload_vec_znx::<BE>(&a);
146    let b_backend = upload_vec_znx::<BE>(&b);
147
148    let mut a_prep: CnvPVecLOwned<BE> = module.cnv_pvec_left_alloc(a_cols, a_size);
149    let mut b_prep: CnvPVecROwned<BE> = module.cnv_pvec_right_alloc(b_cols, b_size);
150
151    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
152        module
153            .cnv_apply_dft_tmp_bytes(0, res_size, a_size, b_size)
154            .max(module.cnv_prepare_left_tmp_bytes(res_size, a_size))
155            .max(module.cnv_prepare_right_tmp_bytes(res_size, b_size))
156            .max(module.vec_znx_big_normalize_tmp_bytes()),
157    );
158
159    {
160        let mut a_prep_backend = a_prep.to_backend_mut();
161        module.cnv_prepare_left(
162            &mut a_prep_backend,
163            &vec_znx_backend_ref::<BE>(&a_backend),
164            !0i64,
165            &mut scratch.arena(),
166        );
167    }
168    {
169        let mut b_prep_backend = b_prep.to_backend_mut();
170        module.cnv_prepare_right(
171            &mut b_prep_backend,
172            &vec_znx_backend_ref::<BE>(&b_backend),
173            !0i64,
174            &mut scratch.arena(),
175        );
176    }
177
178    for a_col in 0..a.cols() {
179        for b_col in 0..b.cols() {
180            for cnv_offset in 0..res_size {
181                module.cnv_apply_dft(
182                    cnv_offset,
183                    &mut res_dft.to_backend_mut(),
184                    res_dft_col,
185                    &a_prep.to_backend_ref(),
186                    a_col,
187                    &b_prep.to_backend_ref(),
188                    b_col,
189                    &mut scratch.arena(),
190                );
191
192                module.vec_znx_idft_apply_tmpa(&mut res_big.to_backend_mut(), 0, &mut res_dft.to_backend_mut(), res_dft_col);
193
194                let res_host_template = VecZnx::alloc(module.n(), 1, res_size);
195                let mut res_have_backend = upload_vec_znx::<BE>(&res_host_template);
196                module.vec_znx_big_normalize(
197                    &mut vec_znx_backend_mut::<BE>(&mut res_have_backend),
198                    base2k,
199                    0,
200                    0,
201                    &res_big.to_backend_ref(),
202                    base2k,
203                    0,
204                    &mut scratch.arena(),
205                );
206                let res_have = download_vec_znx::<BE>(&res_have_backend);
207
208                bivariate_convolution_naive(
209                    module,
210                    base2k,
211                    (cnv_offset + 1) as i64,
212                    &mut res_want,
213                    0,
214                    &a,
215                    a_col,
216                    &b,
217                    b_col,
218                    &mut scratch.arena(),
219                );
220
221                assert_eq!(res_want, res_have);
222            }
223        }
224    }
225}
226
227/// `cnv_apply_dft_accumulate` matches `cnv_apply_dft` followed by a DFT add,
228/// bit-for-bit on the raw prepared data.
229pub fn test_convolution_accumulate<M, BE: crate::test_suite::TestBackend>(module: &M, _base2k: usize)
230where
231    M: ModuleN + Convolution<BE> + CnvPVecAlloc<BE> + VecZnxDftAlloc<BE> + VecZnxDftAddAssign<BE>,
232    ScratchOwned<BE>: ScratchOwnedAlloc<BE>,
233{
234    let mut source: Source = Source::new([0u8; 32]);
235
236    let cols: usize = 2;
237    let a_size: usize = 15;
238    let b_size: usize = 15;
239    let res_size: usize = a_size + b_size;
240
241    let mut a = VecZnx::alloc(module.n(), cols, a_size);
242    let mut b = VecZnx::alloc(module.n(), cols, b_size);
243    a.fill_uniform(17, &mut source);
244    b.fill_uniform(17, &mut source);
245
246    let a_backend = upload_vec_znx::<BE>(&a);
247    let b_backend = upload_vec_znx::<BE>(&b);
248
249    let mut a_prep: CnvPVecLOwned<BE> = module.cnv_pvec_left_alloc(cols, a_size);
250    let mut b_prep: CnvPVecROwned<BE> = module.cnv_pvec_right_alloc(cols, b_size);
251
252    // Two-column accumulators exercised at both columns: covers the
253    // column-interleaved `VecZnxDft` indexing of the backend kernels.
254    let res_cols: usize = 2;
255    let mut res_acc: VecZnxDftOwned<BE> = module.vec_znx_dft_alloc(res_cols, res_size);
256    let mut res_ref: VecZnxDftOwned<BE> = module.vec_znx_dft_alloc(res_cols, res_size);
257    let mut tmp_dft: VecZnxDftOwned<BE> = module.vec_znx_dft_alloc(1, res_size);
258
259    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
260        module
261            .cnv_apply_dft_tmp_bytes(0, res_size, a_size, b_size)
262            .max(module.cnv_prepare_left_tmp_bytes(res_size, a_size))
263            .max(module.cnv_prepare_right_tmp_bytes(res_size, b_size)),
264    );
265
266    {
267        let mut a_prep_backend = a_prep.to_backend_mut();
268        module.cnv_prepare_left(
269            &mut a_prep_backend,
270            &vec_znx_backend_ref::<BE>(&a_backend),
271            !0i64,
272            &mut scratch.arena(),
273        );
274    }
275    {
276        let mut b_prep_backend = b_prep.to_backend_mut();
277        module.cnv_prepare_right(
278            &mut b_prep_backend,
279            &vec_znx_backend_ref::<BE>(&b_backend),
280            !0i64,
281            &mut scratch.arena(),
282        );
283    }
284
285    for res_col in 0..res_cols {
286        // Identical deterministic initial accumulator content for both paths.
287        module.cnv_apply_dft(
288            0,
289            &mut res_acc.to_backend_mut(),
290            res_col,
291            &a_prep.to_backend_ref(),
292            0,
293            &b_prep.to_backend_ref(),
294            0,
295            &mut scratch.arena(),
296        );
297        module.cnv_apply_dft(
298            0,
299            &mut res_ref.to_backend_mut(),
300            res_col,
301            &a_prep.to_backend_ref(),
302            0,
303            &b_prep.to_backend_ref(),
304            0,
305            &mut scratch.arena(),
306        );
307
308        for a_col in 0..cols {
309            for b_col in 0..cols {
310                for cnv_offset in (0..res_size).step_by(3) {
311                    module.cnv_apply_dft_accumulate(
312                        cnv_offset,
313                        &mut res_acc.to_backend_mut(),
314                        res_col,
315                        &a_prep.to_backend_ref(),
316                        a_col,
317                        &b_prep.to_backend_ref(),
318                        b_col,
319                        &mut scratch.arena(),
320                    );
321
322                    module.cnv_apply_dft(
323                        cnv_offset,
324                        &mut tmp_dft.to_backend_mut(),
325                        0,
326                        &a_prep.to_backend_ref(),
327                        a_col,
328                        &b_prep.to_backend_ref(),
329                        b_col,
330                        &mut scratch.arena(),
331                    );
332                    module.vec_znx_dft_add_assign(&mut res_ref.to_backend_mut(), res_col, &tmp_dft.to_backend_ref(), 0);
333
334                    assert!(
335                        BE::to_host_bytes(res_acc.data()) == BE::to_host_bytes(res_ref.data()),
336                        "accumulate != apply + add (res_col={res_col} a_col={a_col} b_col={b_col} cnv_offset={cnv_offset})"
337                    );
338                }
339            }
340        }
341    }
342}
343
344/// `cnv_accumulate_dft` matches the per-term `cnv_apply_dft` +
345/// `cnv_apply_dft_accumulate` sequence after normalization to the coefficient
346/// domain (the fused path reduces once per output, so the raw q120 lazy
347/// representatives may differ).
348pub fn test_convolution_accumulate_fused<M, BE: crate::test_suite::TestBackend>(module: &M, base2k: usize)
349where
350    M: ModuleN
351        + Convolution<BE>
352        + CnvPVecAlloc<BE>
353        + VecZnxDftAlloc<BE>
354        + VecZnxIdftApplyTmpA<BE>
355        + VecZnxBigNormalize<BE>
356        + VecZnxBigNormalizeTmpBytes
357        + VecZnxBigAlloc<BE>,
358    ScratchOwned<BE>: ScratchOwnedAlloc<BE>,
359{
360    use crate::layouts::CnvDftAccTerm;
361
362    let mut source: Source = Source::new([0u8; 32]);
363
364    let cols: usize = 2;
365    let a_size: usize = 15;
366    let b_size: usize = 15;
367    let res_size: usize = a_size + b_size;
368    // Two-column destination written at column 1: covers the column-interleaved
369    // `VecZnxDft` indexing of the backend kernels.
370    let res_col: usize = 1;
371
372    let mut a = VecZnx::alloc(module.n(), cols, a_size);
373    let mut b = VecZnx::alloc(module.n(), cols, b_size);
374    a.fill_uniform(17, &mut source);
375    b.fill_uniform(17, &mut source);
376
377    let a_backend = upload_vec_znx::<BE>(&a);
378    let b_backend = upload_vec_znx::<BE>(&b);
379
380    let mut a_prep: CnvPVecLOwned<BE> = module.cnv_pvec_left_alloc(cols, a_size);
381    let mut b_prep: CnvPVecROwned<BE> = module.cnv_pvec_right_alloc(cols, b_size);
382
383    let mut res_fused: VecZnxDftOwned<BE> = module.vec_znx_dft_alloc(2, res_size);
384    let mut res_ref: VecZnxDftOwned<BE> = module.vec_znx_dft_alloc(2, res_size);
385    let mut big_fused: VecZnxBigOwned<BE> = module.vec_znx_big_alloc(1, res_size);
386    let mut big_ref: VecZnxBigOwned<BE> = module.vec_znx_big_alloc(1, res_size);
387
388    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
389        module
390            .cnv_accumulate_dft_tmp_bytes(0, res_size, a_size, b_size)
391            .max(module.cnv_apply_dft_tmp_bytes(0, res_size, a_size, b_size))
392            .max(module.cnv_prepare_left_tmp_bytes(res_size, a_size))
393            .max(module.cnv_prepare_right_tmp_bytes(res_size, b_size))
394            .max(module.vec_znx_big_normalize_tmp_bytes()),
395    );
396
397    {
398        let mut a_prep_backend = a_prep.to_backend_mut();
399        module.cnv_prepare_left(
400            &mut a_prep_backend,
401            &vec_znx_backend_ref::<BE>(&a_backend),
402            !0i64,
403            &mut scratch.arena(),
404        );
405    }
406    {
407        let mut b_prep_backend = b_prep.to_backend_mut();
408        module.cnv_prepare_right(
409            &mut b_prep_backend,
410            &vec_znx_backend_ref::<BE>(&b_backend),
411            !0i64,
412            &mut scratch.arena(),
413        );
414    }
415
416    // Three terms mixing operand columns, like one BSGS giant step.
417    let term_cols: [(usize, usize); 3] = [(0, 0), (1, 1), (0, 1)];
418
419    for cnv_offset in (0..res_size).step_by(3) {
420        {
421            let terms: Vec<CnvDftAccTerm<'_, BE>> = term_cols
422                .iter()
423                .map(|&(a_col, b_col)| CnvDftAccTerm {
424                    a: a_prep.to_backend_ref(),
425                    a_col,
426                    b: b_prep.to_backend_ref(),
427                    b_col,
428                })
429                .collect();
430            module.cnv_accumulate_dft(
431                cnv_offset,
432                &mut res_fused.to_backend_mut(),
433                res_col,
434                &terms,
435                &mut scratch.arena(),
436            );
437        }
438
439        for (idx, &(a_col, b_col)) in term_cols.iter().enumerate() {
440            if idx == 0 {
441                module.cnv_apply_dft(
442                    cnv_offset,
443                    &mut res_ref.to_backend_mut(),
444                    res_col,
445                    &a_prep.to_backend_ref(),
446                    a_col,
447                    &b_prep.to_backend_ref(),
448                    b_col,
449                    &mut scratch.arena(),
450                );
451            } else {
452                module.cnv_apply_dft_accumulate(
453                    cnv_offset,
454                    &mut res_ref.to_backend_mut(),
455                    res_col,
456                    &a_prep.to_backend_ref(),
457                    a_col,
458                    &b_prep.to_backend_ref(),
459                    b_col,
460                    &mut scratch.arena(),
461                );
462            }
463        }
464
465        // Compare in the normalized coefficient domain.
466        module.vec_znx_idft_apply_tmpa(&mut big_fused.to_backend_mut(), 0, &mut res_fused.to_backend_mut(), res_col);
467        module.vec_znx_idft_apply_tmpa(&mut big_ref.to_backend_mut(), 0, &mut res_ref.to_backend_mut(), res_col);
468
469        let host_template = VecZnx::alloc(module.n(), 1, res_size);
470        let mut have_backend = upload_vec_znx::<BE>(&host_template);
471        let mut want_backend = upload_vec_znx::<BE>(&host_template);
472        module.vec_znx_big_normalize(
473            &mut vec_znx_backend_mut::<BE>(&mut have_backend),
474            base2k,
475            0,
476            0,
477            &big_fused.to_backend_ref(),
478            base2k,
479            0,
480            &mut scratch.arena(),
481        );
482        module.vec_znx_big_normalize(
483            &mut vec_znx_backend_mut::<BE>(&mut want_backend),
484            base2k,
485            0,
486            0,
487            &big_ref.to_backend_ref(),
488            base2k,
489            0,
490            &mut scratch.arena(),
491        );
492        let have = download_vec_znx::<BE>(&have_backend);
493        let want = download_vec_znx::<BE>(&want_backend);
494        assert_eq!(have, want, "fused accumulate != per-term sequence (cnv_offset={cnv_offset})");
495    }
496}
497
498pub fn test_convolution_pairwise<M, BE: crate::test_suite::TestBackend>(module: &M, base2k: usize)
499where
500    M: ModuleN
501        + Convolution<BE>
502        + CnvPVecAlloc<BE>
503        + VecZnxDftAlloc<BE>
504        + VecZnxDftApply<BE>
505        + VecZnxIdftApplyTmpA<BE>
506        + VecZnxBigNormalize<BE>
507        + VecZnxBigNormalizeTmpBytes
508        + VecZnxNormalizeAssignBackend<BE>
509        + VecZnxBigAlloc<BE>
510        + VecZnxAddIntoBackend<BE>
511        + VecZnxCopyBackend<BE>,
512    ScratchOwned<BE>: ScratchOwnedAlloc<BE>,
513{
514    let mut source: Source = Source::new([0u8; 32]);
515
516    let cols: usize = 2;
517    let a_size: usize = 15;
518    let b_size: usize = 15;
519    let res_size: usize = a_size + b_size;
520
521    let mut a = VecZnx::alloc(module.n(), cols, a_size);
522    let mut b = VecZnx::alloc(module.n(), cols, b_size);
523    let mut tmp_a = VecZnx::alloc(module.n(), 1, a_size);
524    let mut tmp_b = VecZnx::alloc(module.n(), 1, b_size);
525
526    let mut res_want = VecZnx::alloc(module.n(), 1, res_size);
527    // Two-column DFT destination written at column 1: covers the
528    // column-interleaved `VecZnxDft` indexing of the backend kernels.
529    let res_dft_col: usize = 1;
530    let mut res_dft: VecZnxDftOwned<BE> = module.vec_znx_dft_alloc(2, res_size);
531    let mut res_big: VecZnxBigOwned<BE> = module.vec_znx_big_alloc(1, res_size);
532
533    a.fill_uniform(17, &mut source);
534    b.fill_uniform(17, &mut source);
535
536    let a_backend = upload_vec_znx::<BE>(&a);
537    let b_backend = upload_vec_znx::<BE>(&b);
538
539    let mut a_prep: CnvPVecLOwned<BE> = module.cnv_pvec_left_alloc(cols, a_size);
540    let mut b_prep: CnvPVecROwned<BE> = module.cnv_pvec_right_alloc(cols, b_size);
541
542    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
543        module
544            .cnv_pairwise_apply_dft_tmp_bytes(0, res_size, a_size, b_size)
545            .max(module.cnv_prepare_left_tmp_bytes(res_size, a_size))
546            .max(module.cnv_prepare_right_tmp_bytes(res_size, b_size))
547            .max(module.vec_znx_big_normalize_tmp_bytes()),
548    );
549
550    {
551        let mut a_prep_backend = a_prep.to_backend_mut();
552        module.cnv_prepare_left(
553            &mut a_prep_backend,
554            &vec_znx_backend_ref::<BE>(&a_backend),
555            !0i64,
556            &mut scratch.arena(),
557        );
558    }
559    {
560        let mut b_prep_backend = b_prep.to_backend_mut();
561        module.cnv_prepare_right(
562            &mut b_prep_backend,
563            &vec_znx_backend_ref::<BE>(&b_backend),
564            !0i64,
565            &mut scratch.arena(),
566        );
567    }
568
569    for col_i in 0..cols {
570        for col_j in 0..cols {
571            for cnv_offset in 0..res_size {
572                module.cnv_pairwise_apply_dft(
573                    cnv_offset,
574                    &mut res_dft.to_backend_mut(),
575                    res_dft_col,
576                    &a_prep.to_backend_ref(),
577                    &b_prep.to_backend_ref(),
578                    col_i,
579                    col_j,
580                    &mut scratch.arena(),
581                );
582
583                module.vec_znx_idft_apply_tmpa(&mut res_big.to_backend_mut(), 0, &mut res_dft.to_backend_mut(), res_dft_col);
584
585                let res_host_template = VecZnx::alloc(module.n(), 1, res_size);
586                let mut res_have_backend = upload_vec_znx::<BE>(&res_host_template);
587                module.vec_znx_big_normalize(
588                    &mut vec_znx_backend_mut::<BE>(&mut res_have_backend),
589                    base2k,
590                    0,
591                    0,
592                    &res_big.to_backend_ref(),
593                    base2k,
594                    0,
595                    &mut scratch.arena(),
596                );
597                let res_have = download_vec_znx::<BE>(&res_have_backend);
598
599                let mut tmp_a_backend = upload_vec_znx::<BE>(&tmp_a);
600                let mut tmp_b_backend = upload_vec_znx::<BE>(&tmp_b);
601                if col_i != col_j {
602                    module.vec_znx_add_into_backend(
603                        &mut vec_znx_backend_mut::<BE>(&mut tmp_a_backend),
604                        0,
605                        &vec_znx_backend_ref::<BE>(&a_backend),
606                        col_i,
607                        &vec_znx_backend_ref::<BE>(&a_backend),
608                        col_j,
609                    );
610                    module.vec_znx_add_into_backend(
611                        &mut vec_znx_backend_mut::<BE>(&mut tmp_b_backend),
612                        0,
613                        &vec_znx_backend_ref::<BE>(&b_backend),
614                        col_i,
615                        &vec_znx_backend_ref::<BE>(&b_backend),
616                        col_j,
617                    );
618                } else {
619                    module.vec_znx_copy_backend(
620                        &mut vec_znx_backend_mut::<BE>(&mut tmp_a_backend),
621                        0,
622                        &vec_znx_backend_ref::<BE>(&a_backend),
623                        col_i,
624                    );
625                    module.vec_znx_copy_backend(
626                        &mut vec_znx_backend_mut::<BE>(&mut tmp_b_backend),
627                        0,
628                        &vec_znx_backend_ref::<BE>(&b_backend),
629                        col_j,
630                    );
631                }
632
633                tmp_a = download_vec_znx::<BE>(&tmp_a_backend);
634                tmp_b = download_vec_znx::<BE>(&tmp_b_backend);
635
636                bivariate_convolution_naive(
637                    module,
638                    base2k,
639                    (cnv_offset + 1) as i64,
640                    &mut res_want,
641                    0,
642                    &tmp_a,
643                    0,
644                    &tmp_b,
645                    0,
646                    &mut scratch.arena(),
647                );
648
649                assert_eq!(res_want, res_have);
650            }
651        }
652    }
653}
654
655#[allow(clippy::too_many_arguments)]
656pub fn bivariate_convolution_naive<M, BE: crate::test_suite::TestBackend>(
657    module: &M,
658    base2k: usize,
659    k: i64,
660    res: &mut VecZnxOwned<BE::ZnxWord>,
661    res_col: usize,
662    a: &VecZnxOwned<BE::ZnxWord>,
663    a_col: usize,
664    b: &VecZnxOwned<BE::ZnxWord>,
665    b_col: usize,
666    scratch: &mut ScratchArena<'_, BE>,
667) where
668    M: VecZnxNormalizeAssignBackend<BE>,
669    ScratchOwned<BE>: ScratchOwnedAlloc<BE>,
670{
671    for j in 0..res.size() {
672        res.zero_at(res_col, j);
673    }
674
675    for a_limb in 0..a.size() {
676        for b_limb in 0..b.size() {
677            let res_scale_abs = k.unsigned_abs() as usize;
678
679            let mut res_limb: usize = a_limb + b_limb + 1;
680
681            if k <= 0 {
682                res_limb += res_scale_abs;
683
684                if res_limb < res.size() {
685                    negacyclic_convolution_naive_add(res.at_mut(res_col, res_limb), a.at(a_col, a_limb), b.at(b_col, b_limb));
686                }
687            } else if res_limb >= res_scale_abs {
688                res_limb -= res_scale_abs;
689
690                if res_limb < res.size() {
691                    negacyclic_convolution_naive_add(res.at_mut(res_col, res_limb), a.at(a_col, a_limb), b.at(b_col, b_limb));
692                }
693            }
694        }
695    }
696
697    let mut res_backend = upload_vec_znx::<BE>(res);
698    module.vec_znx_normalize_assign_backend(base2k, &mut vec_znx_backend_mut::<BE>(&mut res_backend), res_col, scratch);
699    *res = download_vec_znx::<BE>(&res_backend);
700}
701
702fn bivariate_tensoring_naive<M, BE: crate::test_suite::TestBackend>(
703    module: &M,
704    base2k: usize,
705    k: i64,
706    res: &mut VecZnxOwned<BE::ZnxWord>,
707    a: &VecZnxOwned<BE::ZnxWord>,
708    b: &VecZnxOwned<BE::ZnxWord>,
709    scratch: &mut ScratchArena<'_, BE>,
710) where
711    M: VecZnxNormalizeAssignBackend<BE>,
712{
713    let cols = res.cols();
714
715    assert!(res.cols() >= a.cols() + b.cols() - 1);
716
717    res.zero();
718
719    for a_col in 0..a.cols() {
720        for a_limb in 0..a.size() {
721            for b_col in 0..b.cols() {
722                for b_limb in 0..b.size() {
723                    let res_scale_abs = k.unsigned_abs() as usize;
724
725                    let mut res_limb: usize = a_limb + b_limb + 1;
726
727                    if k <= 0 {
728                        res_limb += res_scale_abs;
729
730                        if res_limb < res.size() {
731                            negacyclic_convolution_naive_add(
732                                res.at_mut(a_col + b_col, res_limb),
733                                a.at(a_col, a_limb),
734                                b.at(b_col, b_limb),
735                            );
736                        }
737                    } else if res_limb >= res_scale_abs {
738                        res_limb -= res_scale_abs;
739
740                        if res_limb < res.size() {
741                            negacyclic_convolution_naive_add(
742                                res.at_mut(a_col + b_col, res_limb),
743                                a.at(a_col, a_limb),
744                                b.at(b_col, b_limb),
745                            );
746                        }
747                    }
748                }
749            }
750        }
751    }
752
753    let mut res_backend = upload_vec_znx::<BE>(res);
754    for i in 0..cols {
755        module.vec_znx_normalize_assign_backend(base2k, &mut vec_znx_backend_mut::<BE>(&mut res_backend), i, scratch);
756    }
757    *res = download_vec_znx::<BE>(&res_backend);
758}
759
760fn negacyclic_convolution_naive_add(res: &mut [i64], a: &[i64], b: &[i64]) {
761    let n: usize = res.len();
762    for i in 0..n {
763        let ai: i64 = a[i];
764        let lim: usize = n - i;
765        for j in 0..lim {
766            res[i + j] += ai * b[j];
767        }
768        for j in lim..n {
769            res[i + j - n] -= ai * b[j];
770        }
771    }
772}
773
774fn negacyclic_convolution_naive(res: &mut [i64], a: &[i64], b: &[i64]) {
775    let n: usize = res.len();
776    res.fill(0);
777    for i in 0..n {
778        let ai: i64 = a[i];
779        let lim: usize = n - i;
780        for j in 0..lim {
781            res[i + j] += ai * b[j];
782        }
783        for j in lim..n {
784            res[i + j - n] -= ai * b[j];
785        }
786    }
787}
poulpy_hal/test_suite/convolution.rs

poulpy_hal/test_suite/
convolution.rs
