poulpy/poulpy-hal/src/test_suite/convolution.rs
poulpy-hal/src/test_suite/convolution.rsRaw
1use rand::RngCore;
2
3use crate::{
4 api::{
5 CnvPVecAlloc, Convolution, ModuleN, ScratchOwnedAlloc, ScratchOwnedBorrow, ScratchTakeBasic, TakeSlice, VecZnxAdd,
6 VecZnxBigAlloc, VecZnxBigNormalize, VecZnxCopy, VecZnxDftAlloc, VecZnxDftApply, VecZnxIdftApplyTmpA,
7 VecZnxNormalizeInplace,
8 },
9 layouts::{
10 Backend, CnvPVecL, CnvPVecR, FillUniform, Scratch, ScratchOwned, VecZnx, VecZnxBig, VecZnxDft, VecZnxToMut, VecZnxToRef,
11 ZnxInfos, ZnxView, ZnxViewMut, ZnxZero,
12 },
13 source::Source,
14};
15
16pub fn test_convolution_by_const<M, BE: Backend>(module: &M)
17where
18 M: ModuleN + Convolution<BE> + VecZnxBigNormalize<BE> + VecZnxNormalizeInplace<BE> + VecZnxBigAlloc<BE>,
19 Scratch<BE>: ScratchTakeBasic,
20 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
21{
22 let mut source: Source = Source::new([0u8; 32]);
23
24 let base2k: usize = 12;
25
26 let a_cols: usize = 2;
27 let a_size: usize = 15;
28 let b_size: usize = 15;
29 let res_size: usize = a_size + b_size;
30
31 let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), a_cols, a_size);
32 let mut b: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, b_size);
33
34 let mut res_want: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, res_size);
35 let mut res_have: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, res_size);
36 let mut res_big: VecZnxBig<Vec<u8>, BE> = module.vec_znx_big_alloc(1, res_size);
37
38 a.fill_uniform(base2k, &mut source);
39
40 let mut b_const = vec![0i64; b_size];
41 let mask = (1 << base2k) - 1;
42 for (j, x) in b_const[..1].iter_mut().enumerate() {
43 let r = source.next_u64() & mask;
44 *x = ((r << (64 - base2k)) as i64) >> (64 - base2k);
45 b.at_mut(0, j)[0] = *x
46 }
47
48 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(module.cnv_by_const_apply_tmp_bytes(res_size, 0, a_size, b_size));
49
50 for a_col in 0..a.cols() {
51 for offset in 0..res_size {
52 module.cnv_by_const_apply(&mut res_big, offset, 0, &a, a_col, &b_const, scratch.borrow());
53 module.vec_znx_big_normalize(&mut res_have, base2k, 0, 0, &res_big, base2k, 0, scratch.borrow());
54
55 bivariate_convolution_naive(
56 module,
57 base2k,
58 (offset + 1) as i64,
59 &mut res_want,
60 0,
61 &a,
62 a_col,
63 &b,
64 0,
65 scratch.borrow(),
66 );
67
68 assert_eq!(res_want, res_have);
69 }
70 }
71}
72
73pub fn test_convolution<M, BE: Backend>(module: &M)
74where
75 M: ModuleN
76 + Convolution<BE>
77 + CnvPVecAlloc<BE>
78 + VecZnxDftAlloc<BE>
79 + VecZnxDftApply<BE>
80 + VecZnxIdftApplyTmpA<BE>
81 + VecZnxBigNormalize<BE>
82 + VecZnxNormalizeInplace<BE>
83 + VecZnxBigAlloc<BE>,
84 Scratch<BE>: ScratchTakeBasic,
85 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
86{
87 let mut source: Source = Source::new([0u8; 32]);
88
89 let base2k: usize = 12;
90
91 let a_cols: usize = 2;
92 let b_cols: usize = 2;
93 let a_size: usize = 15;
94 let b_size: usize = 15;
95 let res_size: usize = a_size + b_size;
96
97 let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), a_cols, a_size);
98 let mut b: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), b_cols, b_size);
99
100 let mut res_want: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, res_size);
101 let mut res_have: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, res_size);
102 let mut res_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(1, res_size);
103 let mut res_big: VecZnxBig<Vec<u8>, BE> = module.vec_znx_big_alloc(1, res_size);
104
105 a.fill_uniform(base2k, &mut source);
106 b.fill_uniform(base2k, &mut source);
107
108 let mut a_prep: CnvPVecL<Vec<u8>, BE> = module.cnv_pvec_left_alloc(a_cols, a_size);
109 let mut b_prep: CnvPVecR<Vec<u8>, BE> = module.cnv_pvec_right_alloc(b_cols, b_size);
110
111 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
112 module
113 .cnv_apply_dft_tmp_bytes(res_size, 0, a_size, b_size)
114 .max(module.cnv_prepare_left_tmp_bytes(res_size, a_size))
115 .max(module.cnv_prepare_right_tmp_bytes(res_size, b_size)),
116 );
117
118 module.cnv_prepare_left(&mut a_prep, &a, scratch.borrow());
119 module.cnv_prepare_right(&mut b_prep, &b, scratch.borrow());
120
121 for a_col in 0..a.cols() {
122 for b_col in 0..b.cols() {
123 for offset in 0..res_size {
124 module.cnv_apply_dft(&mut res_dft, offset, 0, &a_prep, a_col, &b_prep, b_col, scratch.borrow());
125
126 module.vec_znx_idft_apply_tmpa(&mut res_big, 0, &mut res_dft, 0);
127 module.vec_znx_big_normalize(&mut res_have, base2k, 0, 0, &res_big, base2k, 0, scratch.borrow());
128
129 bivariate_convolution_naive(
130 module,
131 base2k,
132 (offset + 1) as i64,
133 &mut res_want,
134 0,
135 &a,
136 a_col,
137 &b,
138 b_col,
139 scratch.borrow(),
140 );
141
142 assert_eq!(res_want, res_have);
143 }
144 }
145 }
146}
147
148pub fn test_convolution_pairwise<M, BE: Backend>(module: &M)
149where
150 M: ModuleN
151 + Convolution<BE>
152 + CnvPVecAlloc<BE>
153 + VecZnxDftAlloc<BE>
154 + VecZnxDftApply<BE>
155 + VecZnxIdftApplyTmpA<BE>
156 + VecZnxBigNormalize<BE>
157 + VecZnxNormalizeInplace<BE>
158 + VecZnxBigAlloc<BE>
159 + VecZnxAdd
160 + VecZnxCopy,
161 Scratch<BE>: ScratchTakeBasic,
162 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
163{
164 let mut source: Source = Source::new([0u8; 32]);
165
166 let base2k: usize = 12;
167
168 let cols = 2;
169 let a_size: usize = 15;
170 let b_size: usize = 15;
171 let res_size: usize = a_size + b_size;
172
173 let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), cols, a_size);
174 let mut b: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), cols, b_size);
175 let mut tmp_a: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, a_size);
176 let mut tmp_b: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, b_size);
177
178 let mut res_want: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, res_size);
179 let mut res_have: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, res_size);
180 let mut res_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(1, res_size);
181 let mut res_big: VecZnxBig<Vec<u8>, BE> = module.vec_znx_big_alloc(1, res_size);
182
183 a.fill_uniform(base2k, &mut source);
184 b.fill_uniform(base2k, &mut source);
185
186 let mut a_prep: CnvPVecL<Vec<u8>, BE> = module.cnv_pvec_left_alloc(cols, a_size);
187 let mut b_prep: CnvPVecR<Vec<u8>, BE> = module.cnv_pvec_right_alloc(cols, b_size);
188
189 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
190 module
191 .cnv_pairwise_apply_dft_tmp_bytes(res_size, 0, a_size, b_size)
192 .max(module.cnv_prepare_left_tmp_bytes(res_size, a_size))
193 .max(module.cnv_prepare_right_tmp_bytes(res_size, b_size)),
194 );
195
196 module.cnv_prepare_left(&mut a_prep, &a, scratch.borrow());
197 module.cnv_prepare_right(&mut b_prep, &b, scratch.borrow());
198
199 for col_i in 0..cols {
200 for col_j in 0..cols {
201 for offset in 0..res_size {
202 module.cnv_pairwise_apply_dft(&mut res_dft, offset, 0, &a_prep, &b_prep, col_i, col_j, scratch.borrow());
203
204 module.vec_znx_idft_apply_tmpa(&mut res_big, 0, &mut res_dft, 0);
205 module.vec_znx_big_normalize(&mut res_have, base2k, 0, 0, &res_big, base2k, 0, scratch.borrow());
206
207 if col_i != col_j {
208 module.vec_znx_add(&mut tmp_a, 0, &a, col_i, &a, col_j);
209 module.vec_znx_add(&mut tmp_b, 0, &b, col_i, &b, col_j);
210 } else {
211 module.vec_znx_copy(&mut tmp_a, 0, &a, col_i);
212 module.vec_znx_copy(&mut tmp_b, 0, &b, col_j);
213 }
214
215 bivariate_convolution_naive(
216 module,
217 base2k,
218 (offset + 1) as i64,
219 &mut res_want,
220 0,
221 &tmp_a,
222 0,
223 &tmp_b,
224 0,
225 scratch.borrow(),
226 );
227
228 assert_eq!(res_want, res_have);
229 }
230 }
231 }
232}
233
234#[allow(clippy::too_many_arguments)]
235pub fn bivariate_convolution_naive<R, A, B, M, BE: Backend>(
236 module: &M,
237 base2k: usize,
238 k: i64,
239 res: &mut R,
240 res_col: usize,
241 a: &A,
242 a_col: usize,
243 b: &B,
244 b_col: usize,
245 scratch: &mut Scratch<BE>,
246) where
247 R: VecZnxToMut,
248 A: VecZnxToRef,
249 B: VecZnxToRef,
250 M: VecZnxNormalizeInplace<BE>,
251 Scratch<BE>: TakeSlice,
252{
253 let res: &mut VecZnx<&mut [u8]> = &mut res.to_mut();
254 let a: &VecZnx<&[u8]> = &a.to_ref();
255 let b: &VecZnx<&[u8]> = &b.to_ref();
256
257 for j in 0..res.size() {
258 res.zero_at(res_col, j);
259 }
260
261 for a_limb in 0..a.size() {
262 for b_limb in 0..b.size() {
263 let res_scale_abs = k.unsigned_abs() as usize;
264
265 let mut res_limb: usize = a_limb + b_limb + 1;
266
267 if k <= 0 {
268 res_limb += res_scale_abs;
269
270 if res_limb < res.size() {
271 negacyclic_convolution_naive_add(res.at_mut(res_col, res_limb), a.at(a_col, a_limb), b.at(b_col, b_limb));
272 }
273 } else if res_limb >= res_scale_abs {
274 res_limb -= res_scale_abs;
275
276 if res_limb < res.size() {
277 negacyclic_convolution_naive_add(res.at_mut(res_col, res_limb), a.at(a_col, a_limb), b.at(b_col, b_limb));
278 }
279 }
280 }
281 }
282
283 module.vec_znx_normalize_inplace(base2k, res, res_col, scratch);
284}
285
286fn bivariate_tensoring_naive<R, A, B, M, BE: Backend>(
287 module: &M,
288 base2k: usize,
289 k: i64,
290 res: &mut R,
291 a: &A,
292 b: &B,
293 scratch: &mut Scratch<BE>,
294) where
295 R: VecZnxToMut,
296 A: VecZnxToRef,
297 B: VecZnxToRef,
298 M: VecZnxNormalizeInplace<BE>,
299 Scratch<BE>: TakeSlice,
300{
301 let res: &mut VecZnx<&mut [u8]> = &mut res.to_mut();
302 let a: &VecZnx<&[u8]> = &a.to_ref();
303 let b: &VecZnx<&[u8]> = &b.to_ref();
304
305 assert!(res.cols() >= a.cols() + b.cols() - 1);
306
307 res.zero();
308
309 for a_col in 0..a.cols() {
310 for a_limb in 0..a.size() {
311 for b_col in 0..b.cols() {
312 for b_limb in 0..b.size() {
313 let res_scale_abs = k.unsigned_abs() as usize;
314
315 let mut res_limb: usize = a_limb + b_limb + 1;
316
317 if k <= 0 {
318 res_limb += res_scale_abs;
319
320 if res_limb < res.size() {
321 negacyclic_convolution_naive_add(
322 res.at_mut(a_col + b_col, res_limb),
323 a.at(a_col, a_limb),
324 b.at(b_col, b_limb),
325 );
326 }
327 } else if res_limb >= res_scale_abs {
328 res_limb -= res_scale_abs;
329
330 if res_limb < res.size() {
331 negacyclic_convolution_naive_add(
332 res.at_mut(a_col + b_col, res_limb),
333 a.at(a_col, a_limb),
334 b.at(b_col, b_limb),
335 );
336 }
337 }
338 }
339 }
340 }
341 }
342
343 for i in 0..res.cols() {
344 module.vec_znx_normalize_inplace(base2k, res, i, scratch);
345 }
346}
347
348fn negacyclic_convolution_naive_add(res: &mut [i64], a: &[i64], b: &[i64]) {
349 let n: usize = res.len();
350 for i in 0..n {
351 let ai: i64 = a[i];
352 let lim: usize = n - i;
353 for j in 0..lim {
354 res[i + j] += ai * b[j];
355 }
356 for j in lim..n {
357 res[i + j - n] -= ai * b[j];
358 }
359 }
360}
361
362fn negacyclic_convolution_naive(res: &mut [i64], a: &[i64], b: &[i64]) {
363 let n: usize = res.len();
364 res.fill(0);
365 for i in 0..n {
366 let ai: i64 = a[i];
367 let lim: usize = n - i;
368 for j in 0..lim {
369 res[i + j] += ai * b[j];
370 }
371 for j in lim..n {
372 res[i + j - n] -= ai * b[j];
373 }
374 }
375}