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