1use crate::{
  2    api::{
  3        Convolution, ModuleN, ScratchOwnedAlloc, ScratchOwnedBorrow, ScratchTakeBasic, TakeSlice, VecZnxBigNormalize,
  4        VecZnxDftAlloc, VecZnxDftApply, VecZnxIdftApplyConsume, VecZnxNormalizeInplace,
  5    },
  6    layouts::{
  7        Backend, FillUniform, Scratch, ScratchOwned, VecZnx, VecZnxBig, VecZnxDft, VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxView,
  8        ZnxViewMut, ZnxZero,
  9    },
 10    source::Source,
 11};
 12
 13pub fn test_convolution<M, BE: Backend>(module: &M)
 14where
 15    M: ModuleN
 16        + Convolution<BE>
 17        + VecZnxDftAlloc<BE>
 18        + VecZnxDftApply<BE>
 19        + VecZnxIdftApplyConsume<BE>
 20        + VecZnxBigNormalize<BE>
 21        + VecZnxNormalizeInplace<BE>,
 22    Scratch<BE>: ScratchTakeBasic,
 23    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 24{
 25    let mut source: Source = Source::new([0u8; 32]);
 26
 27    let base2k: usize = 12;
 28
 29    for a_cols in 1..3 {
 30        for b_cols in 1..3 {
 31            for a_size in 1..5 {
 32                for b_size in 1..5 {
 33                    let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), a_cols, a_size);
 34                    let mut b: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), b_cols, b_size);
 35
 36                    let mut c_want: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), a_cols + b_cols - 1, b_size + a_size);
 37                    let mut c_have: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), c_want.cols(), c_want.size());
 38
 39                    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(module.convolution_tmp_bytes(c_want.size()));
 40
 41                    a.fill_uniform(base2k, &mut source);
 42                    b.fill_uniform(base2k, &mut source);
 43
 44                    let mut b_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(b.cols(), b.size());
 45
 46                    for i in 0..b.cols() {
 47                        module.vec_znx_dft_apply(1, 0, &mut b_dft, i, &b, i);
 48                    }
 49
 50                    for mut res_scale in 0..2 * c_want.size() as i64 + 1 {
 51                        res_scale -= c_want.size() as i64;
 52
 53                        let mut c_have_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(c_have.cols(), c_have.size());
 54                        module.convolution(&mut c_have_dft, res_scale, &a, &b_dft, scratch.borrow());
 55
 56                        let c_have_big: VecZnxBig<Vec<u8>, BE> = module.vec_znx_idft_apply_consume(c_have_dft);
 57
 58                        for i in 0..c_have.cols() {
 59                            module.vec_znx_big_normalize(
 60                                base2k,
 61                                &mut c_have,
 62                                i,
 63                                base2k,
 64                                &c_have_big,
 65                                i,
 66                                scratch.borrow(),
 67                            );
 68                        }
 69
 70                        convolution_naive(
 71                            module,
 72                            base2k,
 73                            &mut c_want,
 74                            res_scale,
 75                            &a,
 76                            &b,
 77                            scratch.borrow(),
 78                        );
 79
 80                        assert_eq!(c_want, c_have);
 81                    }
 82                }
 83            }
 84        }
 85    }
 86}
 87
 88fn convolution_naive<R, A, B, M, BE: Backend>(
 89    module: &M,
 90    base2k: usize,
 91    res: &mut R,
 92    res_scale: i64,
 93    a: &A,
 94    b: &B,
 95    scratch: &mut Scratch<BE>,
 96) where
 97    R: VecZnxToMut,
 98    A: VecZnxToRef,
 99    B: VecZnxToRef,
100    M: VecZnxNormalizeInplace<BE>,
101    Scratch<BE>: TakeSlice,
102{
103    let res: &mut VecZnx<&mut [u8]> = &mut res.to_mut();
104    let a: &VecZnx<&[u8]> = &a.to_ref();
105    let b: &VecZnx<&[u8]> = &b.to_ref();
106
107    assert!(res.cols() >= a.cols() + b.cols() - 1);
108
109    res.zero();
110
111    for a_col in 0..a.cols() {
112        for a_limb in 0..a.size() {
113            for b_col in 0..b.cols() {
114                for b_limb in 0..b.size() {
115                    let res_scale_abs = res_scale.unsigned_abs() as usize;
116
117                    let mut res_limb: usize = a_limb + b_limb + 1;
118
119                    if res_scale <= 0 {
120                        res_limb += res_scale_abs;
121
122                        if res_limb < res.size() {
123                            negacyclic_convolution_naive_add(
124                                res.at_mut(a_col + b_col, res_limb),
125                                a.at(a_col, a_limb),
126                                b.at(b_col, b_limb),
127                            );
128                        }
129                    } else if res_limb >= res_scale_abs {
130                        res_limb -= res_scale_abs;
131
132                        if res_limb < res.size() {
133                            negacyclic_convolution_naive_add(
134                                res.at_mut(a_col + b_col, res_limb),
135                                a.at(a_col, a_limb),
136                                b.at(b_col, b_limb),
137                            );
138                        }
139                    }
140                }
141            }
142        }
143    }
144
145    for i in 0..res.cols() {
146        module.vec_znx_normalize_inplace(base2k, res, i, scratch);
147    }
148}
149
150fn negacyclic_convolution_naive_add(res: &mut [i64], a: &[i64], b: &[i64]) {
151    let n: usize = res.len();
152    for i in 0..n {
153        let ai: i64 = a[i];
154        let lim: usize = n - i;
155        for j in 0..lim {
156            res[i + j] += ai * b[j];
157        }
158        for j in lim..n {
159            res[i + j - n] -= ai * b[j];
160        }
161    }
162}