poulpy/poulpy-hal/src/test_suite/vmp.rs
poulpy-hal/src/test_suite/vmp.rsRaw
1use crate::{
2 api::{
3 ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAlloc,
4 VecZnxDftApply, VecZnxIdftApplyConsume, VmpApplyDft, VmpApplyDftTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd,
5 VmpApplyDftToDftAddTmpBytes, VmpApplyDftToDftTmpBytes, VmpPMatAlloc, VmpPrepare, VmpPrepareTmpBytes,
6 },
7 layouts::{DataViewMut, DigestU64, FillUniform, MatZnx, Module, ScratchOwned, VecZnx, VecZnxBig},
8 source::Source,
9};
10use rand::RngCore;
11
12use crate::layouts::{Backend, VecZnxDft, VmpPMat};
13
14pub fn test_vmp_apply_dft<BR: Backend, BT: Backend>(base2k: usize, module_ref: &Module<BR>, module_test: &Module<BT>)
15where
16 Module<BR>: ModuleNew<BR>
17 + VmpApplyDftTmpBytes
18 + VmpApplyDft<BR>
19 + VmpPMatAlloc<BR>
20 + VecZnxDftAlloc<BR>
21 + VmpPrepare<BR>
22 + VecZnxDftAlloc<BR>
23 + VecZnxIdftApplyConsume<BR>
24 + VecZnxBigNormalize<BR>,
25 ScratchOwned<BR>: ScratchOwnedAlloc<BR> + ScratchOwnedBorrow<BR>,
26 Module<BT>: ModuleNew<BT>
27 + VmpApplyDftTmpBytes
28 + VmpApplyDft<BT>
29 + VmpPMatAlloc<BT>
30 + VecZnxDftAlloc<BT>
31 + VmpPrepare<BT>
32 + VecZnxDftAlloc<BT>
33 + VecZnxIdftApplyConsume<BT>
34 + VecZnxBigNormalize<BT>,
35 ScratchOwned<BT>: ScratchOwnedAlloc<BT> + ScratchOwnedBorrow<BT>,
36{
37 assert_eq!(module_ref.n(), module_test.n());
38 let n: usize = module_ref.n();
39
40 let max_size: usize = 4;
41 let max_cols: usize = 2;
42 let mut source: Source = Source::new([0u8; 32]);
43
44 let mut scratch_ref: ScratchOwned<BR> =
45 ScratchOwned::alloc(module_ref.vmp_apply_dft_tmp_bytes(max_size, max_size, max_size, max_cols, max_cols, max_size));
46 let mut scratch_test: ScratchOwned<BT> =
47 ScratchOwned::alloc(module_test.vmp_apply_dft_tmp_bytes(max_size, max_size, max_size, max_cols, max_cols, max_size));
48
49 for cols_in in 1..max_cols + 1 {
50 for cols_out in 1..max_cols + 1 {
51 for size_in in 1..max_size + 1 {
52 for size_out in 1..max_size + 1 {
53 let rows: usize = cols_in;
54
55 let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_in, size_in);
56 a.fill_uniform(base2k, &mut source);
57 let a_digest: u64 = a.digest_u64();
58
59 let mut mat: MatZnx<Vec<u8>> = MatZnx::alloc(n, rows, cols_in, cols_out, size_out);
60 mat.fill_uniform(base2k, &mut source);
61 let mat_digest: u64 = mat.digest_u64();
62
63 let mut pmat_ref: VmpPMat<Vec<u8>, BR> = module_ref.vmp_pmat_alloc(rows, cols_in, cols_out, size_out);
64 let mut pmat_test: VmpPMat<Vec<u8>, BT> = module_test.vmp_pmat_alloc(rows, cols_in, cols_out, size_out);
65
66 module_ref.vmp_prepare(&mut pmat_ref, &mat, scratch_ref.borrow());
67 module_test.vmp_prepare(&mut pmat_test, &mat, scratch_test.borrow());
68
69 assert_eq!(mat.digest_u64(), mat_digest);
70
71 let mut res_dft_ref: VecZnxDft<Vec<u8>, BR> = module_ref.vec_znx_dft_alloc(cols_out, size_out);
72 let mut res_dft_test: VecZnxDft<Vec<u8>, BT> = module_test.vec_znx_dft_alloc(cols_out, size_out);
73
74 source.fill_bytes(res_dft_ref.data_mut());
75 source.fill_bytes(res_dft_test.data_mut());
76
77 module_ref.vmp_apply_dft(&mut res_dft_ref, &a, &pmat_ref, scratch_ref.borrow());
78 module_test.vmp_apply_dft(&mut res_dft_test, &a, &pmat_test, scratch_test.borrow());
79
80 assert_eq!(a.digest_u64(), a_digest);
81
82 let res_big_ref: VecZnxBig<Vec<u8>, BR> = module_ref.vec_znx_idft_apply_consume(res_dft_ref);
83 let res_big_test: VecZnxBig<Vec<u8>, BT> = module_test.vec_znx_idft_apply_consume(res_dft_test);
84
85 let mut res_small_ref: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_out, size_out);
86 let mut res_small_test: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_out, size_out);
87
88 let res_ref_digest: u64 = res_big_ref.digest_u64();
89 let res_test_digest: u64 = res_big_test.digest_u64();
90
91 for j in 0..cols_out {
92 module_ref.vec_znx_big_normalize(
93 base2k,
94 &mut res_small_ref,
95 j,
96 base2k,
97 &res_big_ref,
98 j,
99 scratch_ref.borrow(),
100 );
101 module_test.vec_znx_big_normalize(
102 base2k,
103 &mut res_small_test,
104 j,
105 base2k,
106 &res_big_test,
107 j,
108 scratch_test.borrow(),
109 );
110 }
111
112 assert_eq!(res_big_ref.digest_u64(), res_ref_digest);
113 assert_eq!(res_big_test.digest_u64(), res_test_digest);
114
115 assert_eq!(res_small_ref, res_small_test);
116 }
117 }
118 }
119 }
120}
121
122pub fn test_vmp_apply_dft_to_dft<BR: Backend, BT: Backend>(base2k: usize, module_ref: &Module<BR>, module_test: &Module<BT>)
123where
124 Module<BR>: ModuleNew<BR>
125 + VmpApplyDftToDftTmpBytes
126 + VmpApplyDftToDft<BR>
127 + VmpPMatAlloc<BR>
128 + VecZnxDftAlloc<BR>
129 + VmpPrepare<BR>
130 + VecZnxIdftApplyConsume<BR>
131 + VecZnxBigNormalize<BR>
132 + VecZnxDftApply<BR>
133 + VmpPrepareTmpBytes
134 + VecZnxBigNormalizeTmpBytes,
135 ScratchOwned<BR>: ScratchOwnedAlloc<BR> + ScratchOwnedBorrow<BR>,
136 Module<BT>: ModuleNew<BT>
137 + VmpApplyDftToDftTmpBytes
138 + VmpApplyDftToDft<BT>
139 + VmpPMatAlloc<BT>
140 + VecZnxDftAlloc<BT>
141 + VmpPrepare<BT>
142 + VecZnxIdftApplyConsume<BT>
143 + VecZnxBigNormalize<BT>
144 + VecZnxDftApply<BT>
145 + VmpPrepareTmpBytes
146 + VecZnxBigNormalizeTmpBytes,
147 ScratchOwned<BT>: ScratchOwnedAlloc<BT> + ScratchOwnedBorrow<BT>,
148{
149 assert_eq!(module_ref.n(), module_test.n());
150 let n: usize = module_ref.n();
151
152 let max_size: usize = 4;
153 let max_cols: usize = 2;
154
155 let mut source: Source = Source::new([0u8; 32]);
156
157 let mut scratch_ref: ScratchOwned<BR> = ScratchOwned::alloc(
158 module_ref
159 .vmp_apply_dft_to_dft_tmp_bytes(max_size, max_size, max_size, max_cols, max_cols, max_size)
160 .max(module_ref.vmp_prepare_tmp_bytes(max_size, max_cols, max_cols, max_size))
161 .max(module_ref.vec_znx_big_normalize_tmp_bytes()),
162 );
163 let mut scratch_test: ScratchOwned<BT> = ScratchOwned::alloc(
164 module_test
165 .vmp_apply_dft_to_dft_tmp_bytes(max_size, max_size, max_size, max_cols, max_cols, max_size)
166 .max(module_test.vmp_prepare_tmp_bytes(max_size, max_cols, max_cols, max_size))
167 .max(module_test.vec_znx_big_normalize_tmp_bytes()),
168 );
169
170 for cols_in in 1..max_cols + 1 {
171 for cols_out in 1..max_cols + 1 {
172 for size_in in 1..max_size + 1 {
173 for size_out in 1..max_size + 1 {
174 let rows: usize = size_in;
175
176 let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_in, size_in);
177 a.fill_uniform(base2k, &mut source);
178 let a_digest: u64 = a.digest_u64();
179
180 let mut a_dft_ref: VecZnxDft<Vec<u8>, BR> = module_ref.vec_znx_dft_alloc(cols_in, size_in);
181 let mut a_dft_test: VecZnxDft<Vec<u8>, BT> = module_test.vec_znx_dft_alloc(cols_in, size_in);
182
183 for j in 0..cols_in {
184 module_ref.vec_znx_dft_apply(1, 0, &mut a_dft_ref, j, &a, j);
185 module_test.vec_znx_dft_apply(1, 0, &mut a_dft_test, j, &a, j);
186 }
187
188 assert_eq!(a.digest_u64(), a_digest);
189
190 let mut mat: MatZnx<Vec<u8>> = MatZnx::alloc(n, rows, cols_in, cols_out, size_out);
191 mat.fill_uniform(base2k, &mut source);
192 let mat_digest: u64 = mat.digest_u64();
193
194 let mut pmat_ref: VmpPMat<Vec<u8>, BR> = module_ref.vmp_pmat_alloc(rows, cols_in, cols_out, size_out);
195 let mut pmat_test: VmpPMat<Vec<u8>, BT> = module_test.vmp_pmat_alloc(rows, cols_in, cols_out, size_out);
196
197 module_ref.vmp_prepare(&mut pmat_ref, &mat, scratch_ref.borrow());
198 module_test.vmp_prepare(&mut pmat_test, &mat, scratch_test.borrow());
199
200 assert_eq!(mat.digest_u64(), mat_digest);
201
202 let mut res_dft_ref: VecZnxDft<Vec<u8>, BR> = module_ref.vec_znx_dft_alloc(cols_out, size_out);
203 let mut res_dft_test: VecZnxDft<Vec<u8>, BT> = module_test.vec_znx_dft_alloc(cols_out, size_out);
204
205 source.fill_bytes(res_dft_ref.data_mut());
206 source.fill_bytes(res_dft_test.data_mut());
207
208 module_ref.vmp_apply_dft_to_dft(
209 &mut res_dft_ref,
210 &a_dft_ref,
211 &pmat_ref,
212 scratch_ref.borrow(),
213 );
214 module_test.vmp_apply_dft_to_dft(
215 &mut res_dft_test,
216 &a_dft_test,
217 &pmat_test,
218 scratch_test.borrow(),
219 );
220
221 let res_big_ref: VecZnxBig<Vec<u8>, BR> = module_ref.vec_znx_idft_apply_consume(res_dft_ref);
222 let res_big_test: VecZnxBig<Vec<u8>, BT> = module_test.vec_znx_idft_apply_consume(res_dft_test);
223
224 let mut res_small_ref: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_out, size_out);
225 let mut res_small_test: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_out, size_out);
226
227 let res_ref_digest: u64 = res_big_ref.digest_u64();
228 let res_test_digest: u64 = res_big_test.digest_u64();
229
230 for j in 0..cols_out {
231 module_ref.vec_znx_big_normalize(
232 base2k,
233 &mut res_small_ref,
234 j,
235 base2k,
236 &res_big_ref,
237 j,
238 scratch_ref.borrow(),
239 );
240 module_test.vec_znx_big_normalize(
241 base2k,
242 &mut res_small_test,
243 j,
244 base2k,
245 &res_big_test,
246 j,
247 scratch_test.borrow(),
248 );
249 }
250
251 assert_eq!(res_big_ref.digest_u64(), res_ref_digest);
252 assert_eq!(res_big_test.digest_u64(), res_test_digest);
253
254 assert_eq!(res_small_ref, res_small_test);
255 }
256 }
257 }
258 }
259}
260
261pub fn test_vmp_apply_dft_to_dft_add<BR: Backend, BT: Backend>(base2k: usize, module_ref: &Module<BR>, module_test: &Module<BT>)
262where
263 Module<BR>: ModuleNew<BR>
264 + VmpApplyDftToDftAddTmpBytes
265 + VmpApplyDftToDftAdd<BR>
266 + VmpPMatAlloc<BR>
267 + VecZnxDftAlloc<BR>
268 + VmpPrepare<BR>
269 + VecZnxIdftApplyConsume<BR>
270 + VecZnxBigNormalize<BR>
271 + VecZnxDftApply<BR>
272 + VmpPrepareTmpBytes
273 + VecZnxBigNormalizeTmpBytes,
274 ScratchOwned<BR>: ScratchOwnedAlloc<BR> + ScratchOwnedBorrow<BR>,
275 Module<BT>: ModuleNew<BT>
276 + VmpApplyDftToDftAddTmpBytes
277 + VmpApplyDftToDftAdd<BT>
278 + VmpPMatAlloc<BT>
279 + VecZnxDftAlloc<BT>
280 + VmpPrepare<BT>
281 + VecZnxIdftApplyConsume<BT>
282 + VecZnxBigNormalize<BT>
283 + VecZnxDftApply<BT>
284 + VmpPrepareTmpBytes
285 + VecZnxBigNormalizeTmpBytes,
286 ScratchOwned<BT>: ScratchOwnedAlloc<BT> + ScratchOwnedBorrow<BT>,
287{
288 assert_eq!(module_ref.n(), module_test.n());
289 let n: usize = module_ref.n();
290
291 let max_size: usize = 4;
292 let max_cols: usize = 2;
293
294 let mut source: Source = Source::new([0u8; 32]);
295
296 let mut scratch_ref: ScratchOwned<BR> = ScratchOwned::alloc(
297 module_ref
298 .vmp_apply_dft_to_dft_add_tmp_bytes(max_size, max_size, max_size, max_cols, max_cols, max_size)
299 .max(module_ref.vmp_prepare_tmp_bytes(max_size, max_cols, max_cols, max_size))
300 .max(module_ref.vec_znx_big_normalize_tmp_bytes()),
301 );
302 let mut scratch_test: ScratchOwned<BT> = ScratchOwned::alloc(
303 module_test
304 .vmp_apply_dft_to_dft_add_tmp_bytes(max_size, max_size, max_size, max_cols, max_cols, max_size)
305 .max(module_test.vmp_prepare_tmp_bytes(max_size, max_cols, max_cols, max_size))
306 .max(module_test.vec_znx_big_normalize_tmp_bytes()),
307 );
308
309 for cols_in in 1..max_cols + 1 {
310 for cols_out in 1..max_cols + 1 {
311 for size_in in 1..max_size + 1 {
312 for size_out in 1..max_size + 1 {
313 let rows: usize = size_in;
314
315 let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_in, size_in);
316 a.fill_uniform(base2k, &mut source);
317 let a_digest: u64 = a.digest_u64();
318
319 let mut a_dft_ref: VecZnxDft<Vec<u8>, BR> = module_ref.vec_znx_dft_alloc(cols_in, size_in);
320 let mut a_dft_test: VecZnxDft<Vec<u8>, BT> = module_test.vec_znx_dft_alloc(cols_in, size_in);
321
322 for j in 0..cols_in {
323 module_ref.vec_znx_dft_apply(1, 0, &mut a_dft_ref, j, &a, j);
324 module_test.vec_znx_dft_apply(1, 0, &mut a_dft_test, j, &a, j);
325 }
326
327 assert_eq!(a.digest_u64(), a_digest);
328
329 let mut mat: MatZnx<Vec<u8>> = MatZnx::alloc(n, rows, cols_in, cols_out, size_out);
330 mat.fill_uniform(base2k, &mut source);
331 let mat_digest: u64 = mat.digest_u64();
332
333 let mut pmat_ref: VmpPMat<Vec<u8>, BR> = module_ref.vmp_pmat_alloc(rows, cols_in, cols_out, size_out);
334 let mut pmat_test: VmpPMat<Vec<u8>, BT> = module_test.vmp_pmat_alloc(rows, cols_in, cols_out, size_out);
335
336 module_ref.vmp_prepare(&mut pmat_ref, &mat, scratch_ref.borrow());
337 module_test.vmp_prepare(&mut pmat_test, &mat, scratch_test.borrow());
338
339 assert_eq!(mat.digest_u64(), mat_digest);
340
341 for limb_offset in 0..size_out {
342 let mut res: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_out, size_out);
343 res.fill_uniform(base2k, &mut source);
344 let res_digest: u64 = res.digest_u64();
345
346 let mut res_dft_ref: VecZnxDft<Vec<u8>, BR> = module_ref.vec_znx_dft_alloc(cols_out, size_out);
347 let mut res_dft_test: VecZnxDft<Vec<u8>, BT> = module_test.vec_znx_dft_alloc(cols_out, size_out);
348
349 for j in 0..cols_out {
350 module_ref.vec_znx_dft_apply(1, 0, &mut res_dft_ref, j, &res, j);
351 module_test.vec_znx_dft_apply(1, 0, &mut res_dft_test, j, &res, j);
352 }
353
354 assert_eq!(res.digest_u64(), res_digest);
355
356 module_ref.vmp_apply_dft_to_dft_add(
357 &mut res_dft_ref,
358 &a_dft_ref,
359 &pmat_ref,
360 limb_offset * cols_out,
361 scratch_ref.borrow(),
362 );
363 module_test.vmp_apply_dft_to_dft_add(
364 &mut res_dft_test,
365 &a_dft_test,
366 &pmat_test,
367 limb_offset * cols_out,
368 scratch_test.borrow(),
369 );
370
371 let res_big_ref: VecZnxBig<Vec<u8>, BR> = module_ref.vec_znx_idft_apply_consume(res_dft_ref);
372 let res_big_test: VecZnxBig<Vec<u8>, BT> = module_test.vec_znx_idft_apply_consume(res_dft_test);
373
374 let mut res_small_ref: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_out, size_out);
375 let mut res_small_test: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_out, size_out);
376
377 let res_ref_digest: u64 = res_big_ref.digest_u64();
378 let res_test_digest: u64 = res_big_test.digest_u64();
379
380 for j in 0..cols_out {
381 module_ref.vec_znx_big_normalize(
382 base2k,
383 &mut res_small_ref,
384 j,
385 base2k,
386 &res_big_ref,
387 j,
388 scratch_ref.borrow(),
389 );
390 module_test.vec_znx_big_normalize(
391 base2k,
392 &mut res_small_test,
393 j,
394 base2k,
395 &res_big_test,
396 j,
397 scratch_test.borrow(),
398 );
399 }
400
401 assert_eq!(res_big_ref.digest_u64(), res_ref_digest);
402 assert_eq!(res_big_test.digest_u64(), res_test_digest);
403
404 assert_eq!(res_small_ref, res_small_test);
405 }
406 }
407 }
408 }
409 }
410}