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