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