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 &mut res_small_ref,
94 base2k,
95 0,
96 j,
97 &res_big_ref,
98 base2k,
99 j,
100 scratch_ref.borrow(),
101 );
102 module_test.vec_znx_big_normalize(
103 &mut res_small_test,
104 base2k,
105 0,
106 j,
107 &res_big_test,
108 base2k,
109 j,
110 scratch_test.borrow(),
111 );
112 }
113
114 assert_eq!(res_big_ref.digest_u64(), res_ref_digest);
115 assert_eq!(res_big_test.digest_u64(), res_test_digest);
116
117 assert_eq!(res_small_ref, res_small_test);
118 }
119 }
120 }
121 }
122}
123
124pub fn test_vmp_apply_dft_to_dft<BR: Backend, BT: Backend>(base2k: usize, module_ref: &Module<BR>, module_test: &Module<BT>)
125where
126 Module<BR>: ModuleNew<BR>
127 + VmpApplyDftToDftTmpBytes
128 + VmpApplyDftToDft<BR>
129 + VmpPMatAlloc<BR>
130 + VecZnxDftAlloc<BR>
131 + VmpPrepare<BR>
132 + VecZnxIdftApplyConsume<BR>
133 + VecZnxBigNormalize<BR>
134 + VecZnxDftApply<BR>
135 + VmpPrepareTmpBytes
136 + VecZnxBigNormalizeTmpBytes,
137 ScratchOwned<BR>: ScratchOwnedAlloc<BR> + ScratchOwnedBorrow<BR>,
138 Module<BT>: ModuleNew<BT>
139 + VmpApplyDftToDftTmpBytes
140 + VmpApplyDftToDft<BT>
141 + VmpPMatAlloc<BT>
142 + VecZnxDftAlloc<BT>
143 + VmpPrepare<BT>
144 + VecZnxIdftApplyConsume<BT>
145 + VecZnxBigNormalize<BT>
146 + VecZnxDftApply<BT>
147 + VmpPrepareTmpBytes
148 + VecZnxBigNormalizeTmpBytes,
149 ScratchOwned<BT>: ScratchOwnedAlloc<BT> + ScratchOwnedBorrow<BT>,
150{
151 assert_eq!(module_ref.n(), module_test.n());
152 let n: usize = module_ref.n();
153
154 let max_size: usize = 4;
155 let max_cols: usize = 2;
156
157 let mut source: Source = Source::new([0u8; 32]);
158
159 let mut scratch_ref: ScratchOwned<BR> = ScratchOwned::alloc(
160 module_ref
161 .vmp_apply_dft_to_dft_tmp_bytes(max_size, max_size, max_size, max_cols, max_cols, max_size)
162 .max(module_ref.vmp_prepare_tmp_bytes(max_size, max_cols, max_cols, max_size))
163 .max(module_ref.vec_znx_big_normalize_tmp_bytes()),
164 );
165 let mut scratch_test: ScratchOwned<BT> = ScratchOwned::alloc(
166 module_test
167 .vmp_apply_dft_to_dft_tmp_bytes(max_size, max_size, max_size, max_cols, max_cols, max_size)
168 .max(module_test.vmp_prepare_tmp_bytes(max_size, max_cols, max_cols, max_size))
169 .max(module_test.vec_znx_big_normalize_tmp_bytes()),
170 );
171
172 for cols_in in 1..max_cols + 1 {
173 for cols_out in 1..max_cols + 1 {
174 for size_in in 1..max_size + 1 {
175 for size_out in 1..max_size + 1 {
176 let rows: usize = size_in;
177
178 let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_in, size_in);
179 a.fill_uniform(base2k, &mut source);
180 let a_digest: u64 = a.digest_u64();
181
182 let mut a_dft_ref: VecZnxDft<Vec<u8>, BR> = module_ref.vec_znx_dft_alloc(cols_in, size_in);
183 let mut a_dft_test: VecZnxDft<Vec<u8>, BT> = module_test.vec_znx_dft_alloc(cols_in, size_in);
184
185 for j in 0..cols_in {
186 module_ref.vec_znx_dft_apply(1, 0, &mut a_dft_ref, j, &a, j);
187 module_test.vec_znx_dft_apply(1, 0, &mut a_dft_test, j, &a, j);
188 }
189
190 assert_eq!(a.digest_u64(), a_digest);
191
192 let mut mat: MatZnx<Vec<u8>> = MatZnx::alloc(n, rows, cols_in, cols_out, size_out);
193 mat.fill_uniform(base2k, &mut source);
194 let mat_digest: u64 = mat.digest_u64();
195
196 let mut pmat_ref: VmpPMat<Vec<u8>, BR> = module_ref.vmp_pmat_alloc(rows, cols_in, cols_out, size_out);
197 let mut pmat_test: VmpPMat<Vec<u8>, BT> = module_test.vmp_pmat_alloc(rows, cols_in, cols_out, size_out);
198
199 module_ref.vmp_prepare(&mut pmat_ref, &mat, scratch_ref.borrow());
200 module_test.vmp_prepare(&mut pmat_test, &mat, scratch_test.borrow());
201
202 assert_eq!(mat.digest_u64(), mat_digest);
203
204 let mut res_dft_ref: VecZnxDft<Vec<u8>, BR> = module_ref.vec_znx_dft_alloc(cols_out, size_out);
205 let mut res_dft_test: VecZnxDft<Vec<u8>, BT> = module_test.vec_znx_dft_alloc(cols_out, size_out);
206
207 source.fill_bytes(res_dft_ref.data_mut());
208 source.fill_bytes(res_dft_test.data_mut());
209
210 module_ref.vmp_apply_dft_to_dft(&mut res_dft_ref, &a_dft_ref, &pmat_ref, scratch_ref.borrow());
211 module_test.vmp_apply_dft_to_dft(&mut res_dft_test, &a_dft_test, &pmat_test, scratch_test.borrow());
212
213 let res_big_ref: VecZnxBig<Vec<u8>, BR> = module_ref.vec_znx_idft_apply_consume(res_dft_ref);
214 let res_big_test: VecZnxBig<Vec<u8>, BT> = module_test.vec_znx_idft_apply_consume(res_dft_test);
215
216 let mut res_small_ref: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_out, size_out);
217 let mut res_small_test: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_out, size_out);
218
219 let res_ref_digest: u64 = res_big_ref.digest_u64();
220 let res_test_digest: u64 = res_big_test.digest_u64();
221
222 for j in 0..cols_out {
223 module_ref.vec_znx_big_normalize(
224 &mut res_small_ref,
225 base2k,
226 0,
227 j,
228 &res_big_ref,
229 base2k,
230 j,
231 scratch_ref.borrow(),
232 );
233 module_test.vec_znx_big_normalize(
234 &mut res_small_test,
235 base2k,
236 0,
237 j,
238 &res_big_test,
239 base2k,
240 j,
241 scratch_test.borrow(),
242 );
243 }
244
245 assert_eq!(res_big_ref.digest_u64(), res_ref_digest);
246 assert_eq!(res_big_test.digest_u64(), res_test_digest);
247
248 assert_eq!(res_small_ref, res_small_test);
249 }
250 }
251 }
252 }
253}
254
255pub fn test_vmp_apply_dft_to_dft_add<BR: Backend, BT: Backend>(base2k: usize, module_ref: &Module<BR>, module_test: &Module<BT>)
256where
257 Module<BR>: ModuleNew<BR>
258 + VmpApplyDftToDftAddTmpBytes
259 + VmpApplyDftToDftAdd<BR>
260 + VmpPMatAlloc<BR>
261 + VecZnxDftAlloc<BR>
262 + VmpPrepare<BR>
263 + VecZnxIdftApplyConsume<BR>
264 + VecZnxBigNormalize<BR>
265 + VecZnxDftApply<BR>
266 + VmpPrepareTmpBytes
267 + VecZnxBigNormalizeTmpBytes,
268 ScratchOwned<BR>: ScratchOwnedAlloc<BR> + ScratchOwnedBorrow<BR>,
269 Module<BT>: ModuleNew<BT>
270 + VmpApplyDftToDftAddTmpBytes
271 + VmpApplyDftToDftAdd<BT>
272 + VmpPMatAlloc<BT>
273 + VecZnxDftAlloc<BT>
274 + VmpPrepare<BT>
275 + VecZnxIdftApplyConsume<BT>
276 + VecZnxBigNormalize<BT>
277 + VecZnxDftApply<BT>
278 + VmpPrepareTmpBytes
279 + VecZnxBigNormalizeTmpBytes,
280 ScratchOwned<BT>: ScratchOwnedAlloc<BT> + ScratchOwnedBorrow<BT>,
281{
282 assert_eq!(module_ref.n(), module_test.n());
283 let n: usize = module_ref.n();
284
285 let max_size: usize = 4;
286 let max_cols: usize = 2;
287
288 let mut source: Source = Source::new([0u8; 32]);
289
290 let mut scratch_ref: ScratchOwned<BR> = ScratchOwned::alloc(
291 module_ref
292 .vmp_apply_dft_to_dft_add_tmp_bytes(max_size, max_size, max_size, max_cols, max_cols, max_size)
293 .max(module_ref.vmp_prepare_tmp_bytes(max_size, max_cols, max_cols, max_size))
294 .max(module_ref.vec_znx_big_normalize_tmp_bytes()),
295 );
296 let mut scratch_test: ScratchOwned<BT> = ScratchOwned::alloc(
297 module_test
298 .vmp_apply_dft_to_dft_add_tmp_bytes(max_size, max_size, max_size, max_cols, max_cols, max_size)
299 .max(module_test.vmp_prepare_tmp_bytes(max_size, max_cols, max_cols, max_size))
300 .max(module_test.vec_znx_big_normalize_tmp_bytes()),
301 );
302
303 for cols_in in 1..max_cols + 1 {
304 for cols_out in 1..max_cols + 1 {
305 for size_in in 1..max_size + 1 {
306 for size_out in 1..max_size + 1 {
307 let rows: usize = size_in;
308
309 let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_in, size_in);
310 a.fill_uniform(base2k, &mut source);
311 let a_digest: u64 = a.digest_u64();
312
313 let mut a_dft_ref: VecZnxDft<Vec<u8>, BR> = module_ref.vec_znx_dft_alloc(cols_in, size_in);
314 let mut a_dft_test: VecZnxDft<Vec<u8>, BT> = module_test.vec_znx_dft_alloc(cols_in, size_in);
315
316 for j in 0..cols_in {
317 module_ref.vec_znx_dft_apply(1, 0, &mut a_dft_ref, j, &a, j);
318 module_test.vec_znx_dft_apply(1, 0, &mut a_dft_test, j, &a, j);
319 }
320
321 assert_eq!(a.digest_u64(), a_digest);
322
323 let mut mat: MatZnx<Vec<u8>> = MatZnx::alloc(n, rows, cols_in, cols_out, size_out);
324 mat.fill_uniform(base2k, &mut source);
325 let mat_digest: u64 = mat.digest_u64();
326
327 let mut pmat_ref: VmpPMat<Vec<u8>, BR> = module_ref.vmp_pmat_alloc(rows, cols_in, cols_out, size_out);
328 let mut pmat_test: VmpPMat<Vec<u8>, BT> = module_test.vmp_pmat_alloc(rows, cols_in, cols_out, size_out);
329
330 module_ref.vmp_prepare(&mut pmat_ref, &mat, scratch_ref.borrow());
331 module_test.vmp_prepare(&mut pmat_test, &mat, scratch_test.borrow());
332
333 assert_eq!(mat.digest_u64(), mat_digest);
334
335 for limb_offset in 0..size_out {
336 let mut res: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_out, size_out);
337 res.fill_uniform(base2k, &mut source);
338 let res_digest: u64 = res.digest_u64();
339
340 let mut res_dft_ref: VecZnxDft<Vec<u8>, BR> = module_ref.vec_znx_dft_alloc(cols_out, size_out);
341 let mut res_dft_test: VecZnxDft<Vec<u8>, BT> = module_test.vec_znx_dft_alloc(cols_out, size_out);
342
343 for j in 0..cols_out {
344 module_ref.vec_znx_dft_apply(1, 0, &mut res_dft_ref, j, &res, j);
345 module_test.vec_znx_dft_apply(1, 0, &mut res_dft_test, j, &res, j);
346 }
347
348 assert_eq!(res.digest_u64(), res_digest);
349
350 module_ref.vmp_apply_dft_to_dft_add(
351 &mut res_dft_ref,
352 &a_dft_ref,
353 &pmat_ref,
354 limb_offset * cols_out,
355 scratch_ref.borrow(),
356 );
357 module_test.vmp_apply_dft_to_dft_add(
358 &mut res_dft_test,
359 &a_dft_test,
360 &pmat_test,
361 limb_offset * cols_out,
362 scratch_test.borrow(),
363 );
364
365 let res_big_ref: VecZnxBig<Vec<u8>, BR> = module_ref.vec_znx_idft_apply_consume(res_dft_ref);
366 let res_big_test: VecZnxBig<Vec<u8>, BT> = module_test.vec_znx_idft_apply_consume(res_dft_test);
367
368 let mut res_small_ref: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_out, size_out);
369 let mut res_small_test: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols_out, size_out);
370
371 let res_ref_digest: u64 = res_big_ref.digest_u64();
372 let res_test_digest: u64 = res_big_test.digest_u64();
373
374 for j in 0..cols_out {
375 module_ref.vec_znx_big_normalize(
376 &mut res_small_ref,
377 base2k,
378 0,
379 j,
380 &res_big_ref,
381 base2k,
382 j,
383 scratch_ref.borrow(),
384 );
385 module_test.vec_znx_big_normalize(
386 &mut res_small_test,
387 base2k,
388 0,
389 j,
390 &res_big_test,
391 base2k,
392 j,
393 scratch_test.borrow(),
394 );
395 }
396
397 assert_eq!(res_big_ref.digest_u64(), res_ref_digest);
398 assert_eq!(res_big_test.digest_u64(), res_test_digest);
399
400 assert_eq!(res_small_ref, res_small_test);
401 }
402 }
403 }
404 }
405 }
406}