poulpy/poulpy-hal/src/reference/vec_znx/normalize.rs
poulpy-hal/src/reference/vec_znx/normalize.rsRaw
1use std::hint::black_box;
2
3use criterion::{BenchmarkId, Criterion};
4
5use crate::{
6 api::{ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes},
7 layouts::{Backend, FillUniform, Module, ScratchOwned, VecZnx, VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxView, ZnxViewMut},
8 reference::znx::{
9 ZnxAddInplace, ZnxCopy, ZnxExtractDigitAddMul, ZnxMulPowerOfTwoInplace, ZnxNormalizeDigit, ZnxNormalizeFinalStep,
10 ZnxNormalizeFinalStepInplace, ZnxNormalizeFirstStep, ZnxNormalizeFirstStepCarryOnly, ZnxNormalizeFirstStepInplace,
11 ZnxNormalizeMiddleStep, ZnxNormalizeMiddleStepCarryOnly, ZnxNormalizeMiddleStepInplace, ZnxZero,
12 },
13 source::Source,
14};
15
16pub fn vec_znx_normalize_tmp_bytes(n: usize) -> usize {
17 3 * n * size_of::<i64>()
18}
19
20#[allow(clippy::too_many_arguments)]
21pub fn vec_znx_normalize<R, A, ZNXARI>(
22 res: &mut R,
23 res_base2k: usize,
24 res_offset: i64,
25 res_col: usize,
26 a: &A,
27 a_base2k: usize,
28 a_col: usize,
29 carry: &mut [i64],
30) where
31 R: VecZnxToMut,
32 A: VecZnxToRef,
33 ZNXARI: ZnxZero
34 + ZnxCopy
35 + ZnxAddInplace
36 + ZnxMulPowerOfTwoInplace
37 + ZnxNormalizeFirstStepCarryOnly
38 + ZnxNormalizeMiddleStepCarryOnly
39 + ZnxNormalizeMiddleStep
40 + ZnxNormalizeFinalStep
41 + ZnxNormalizeFirstStep
42 + ZnxExtractDigitAddMul
43 + ZnxNormalizeMiddleStepInplace
44 + ZnxNormalizeFinalStepInplace
45 + ZnxNormalizeDigit,
46{
47 match res_base2k == a_base2k {
48 true => vec_znx_normalize_inter_base2k::<R, A, ZNXARI>(res_base2k, res, res_offset, res_col, a, a_col, carry),
49 false => vec_znx_normalize_cross_base2k::<R, A, ZNXARI>(res, res_base2k, res_offset, res_col, a, a_base2k, a_col, carry),
50 }
51}
52
53fn vec_znx_normalize_inter_base2k<R, A, ZNXARI>(
54 base2k: usize,
55 res: &mut R,
56 res_offset: i64,
57 res_col: usize,
58 a: &A,
59 a_col: usize,
60 carry: &mut [i64],
61) where
62 R: VecZnxToMut,
63 A: VecZnxToRef,
64 ZNXARI: ZnxZero
65 + ZnxNormalizeFirstStepCarryOnly
66 + ZnxNormalizeMiddleStepCarryOnly
67 + ZnxNormalizeMiddleStep
68 + ZnxNormalizeFinalStepInplace
69 + ZnxNormalizeMiddleStepInplace,
70{
71 let mut res: VecZnx<&mut [u8]> = res.to_mut();
72 let a: VecZnx<&[u8]> = a.to_ref();
73
74 #[cfg(debug_assertions)]
75 {
76 assert!(carry.len() >= vec_znx_normalize_tmp_bytes(res.n()) / size_of::<i64>());
77 assert_eq!(res.n(), a.n());
78 }
79
80 let n: usize = res.n();
81 let res_size: usize = res.size();
82 let a_size: usize = a.size();
83
84 let (carry, _) = carry.split_at_mut(n);
85
86 let mut lsh: i64 = res_offset % base2k as i64;
87 let mut limbs_offset: i64 = res_offset / base2k as i64;
88
89 // If res_offset is negative, makes it positive
90 // and corrects by adding an additional offset
91 // on the limbs.
92 if res_offset < 0 && lsh != 0 {
93 lsh = (lsh + base2k as i64) % (base2k as i64);
94 limbs_offset -= 1;
95 }
96
97 let lsh_pos: usize = lsh as usize;
98
99 let res_end: usize = (-limbs_offset).clamp(0, res_size as i64) as usize;
100 let res_start: usize = (a_size as i64 - limbs_offset).clamp(0, res_size as i64) as usize;
101 let a_end: usize = limbs_offset.clamp(0, a_size as i64) as usize;
102 let a_start: usize = (res_size as i64 + limbs_offset).clamp(0, a_size as i64) as usize;
103
104 let a_out_range: usize = a_size.saturating_sub(a_start);
105
106 // Computes the carry over the discarded limbs of a
107 for j in 0..a_out_range {
108 if j == 0 {
109 ZNXARI::znx_normalize_first_step_carry_only(base2k, lsh_pos, a.at(a_col, a_size - j - 1), carry);
110 } else {
111 ZNXARI::znx_normalize_middle_step_carry_only(base2k, lsh_pos, a.at(a_col, a_size - j - 1), carry);
112 }
113 }
114
115 // If no limbs were discarded, initialize carry to zero
116 if a_out_range == 0 {
117 ZNXARI::znx_zero(carry);
118 }
119
120 // Zeroes bottom limbs that will not be interacted with
121 for j in res_start..res_size {
122 ZNXARI::znx_zero(res.at_mut(res_col, j));
123 }
124
125 let mid_range: usize = a_start.saturating_sub(a_end);
126
127 // Regular normalization over the overlapping limbs of res and a.
128 for j in 0..mid_range {
129 ZNXARI::znx_normalize_middle_step(
130 base2k,
131 lsh_pos,
132 res.at_mut(res_col, res_start - j - 1),
133 a.at(a_col, a_start - j - 1),
134 carry,
135 );
136 }
137
138 // Propagates the carry over the non-overlapping limbs between res and a
139 for j in 0..res_end {
140 ZNXARI::znx_zero(res.at_mut(res_col, res_end - j - 1));
141 if j == res_end - 1 {
142 ZNXARI::znx_normalize_final_step_inplace(base2k, lsh_pos, res.at_mut(res_col, res_end - j - 1), carry);
143 } else {
144 ZNXARI::znx_normalize_middle_step_inplace(base2k, lsh_pos, res.at_mut(res_col, res_end - j - 1), carry);
145 }
146 }
147}
148
149#[allow(clippy::too_many_arguments)]
150fn vec_znx_normalize_cross_base2k<R, A, ZNXARI>(
151 res: &mut R,
152 res_base2k: usize,
153 res_offset: i64,
154 res_col: usize,
155 a: &A,
156 a_base2k: usize,
157 a_col: usize,
158 carry: &mut [i64],
159) where
160 R: VecZnxToMut,
161 A: VecZnxToRef,
162 ZNXARI: ZnxZero
163 + ZnxCopy
164 + ZnxAddInplace
165 + ZnxMulPowerOfTwoInplace
166 + ZnxNormalizeFirstStepCarryOnly
167 + ZnxNormalizeMiddleStepCarryOnly
168 + ZnxNormalizeMiddleStep
169 + ZnxNormalizeFinalStep
170 + ZnxNormalizeFirstStep
171 + ZnxExtractDigitAddMul
172 + ZnxNormalizeMiddleStepInplace
173 + ZnxNormalizeFinalStepInplace
174 + ZnxNormalizeDigit,
175{
176 let mut res: VecZnx<&mut [u8]> = res.to_mut();
177 let a: VecZnx<&[u8]> = a.to_ref();
178
179 #[cfg(debug_assertions)]
180 {
181 assert!(carry.len() >= vec_znx_normalize_tmp_bytes(res.n()) / size_of::<i64>());
182 assert_eq!(res.n(), a.n());
183 }
184
185 let n: usize = res.n();
186 let res_size: usize = res.size();
187 let a_size: usize = a.size();
188
189 let (a_norm, carry) = carry.split_at_mut(n);
190 let (res_carry, a_carry) = carry[..2 * n].split_at_mut(n);
191 ZNXARI::znx_zero(res_carry);
192
193 // Total precision (in bits) that `a` and `res` can represent.
194 let a_tot_bits: usize = a_size * a_base2k;
195 let res_tot_bits: usize = res_size * res_base2k;
196
197 // Derive intra-limb shift and cross-limb offset.
198 let mut lsh: i64 = res_offset % a_base2k as i64;
199 let mut limbs_offset: i64 = res_offset / a_base2k as i64;
200
201 // If res_offset is negative, ensures it is positive
202 // and corrects by incrementing the cross-limb offset.
203 if res_offset < 0 && lsh != 0 {
204 lsh = (lsh + a_base2k as i64) % (a_base2k as i64);
205 limbs_offset -= 1;
206 }
207
208 let lsh_pos: usize = lsh as usize;
209
210 // Derive start/stop bit indexes of the overlap between `a` and `res` (after taking into account the offset)..
211 let res_end_bit: usize = (-limbs_offset * a_base2k as i64).clamp(0, res_tot_bits as i64) as usize; // Stop bit
212 let res_start_bit: usize = (a_tot_bits as i64 - limbs_offset * a_base2k as i64).clamp(0, res_tot_bits as i64) as usize; // Start bit
213 let a_end_bit: usize = (limbs_offset * a_base2k as i64).clamp(0, a_tot_bits as i64) as usize; // Stop bit
214 let a_start_bit: usize = (res_tot_bits as i64 + limbs_offset * a_base2k as i64).clamp(0, a_tot_bits as i64) as usize; // Start bit
215
216 // Convert bits to limb indexes.
217 let res_end: usize = res_end_bit / res_base2k;
218 let res_start: usize = res_start_bit.div_ceil(res_base2k);
219 let a_end: usize = a_end_bit / a_base2k;
220 let a_start: usize = a_start_bit.div_ceil(a_base2k);
221
222 // Zero all limbs of `res`. Unlike the simple case
223 // where `res_base2k` is equal to `a_base2k`, we also
224 // need to ensure that the limbs starting from `res_end`
225 // are zero.
226 for j in 0..res_size {
227 ZNXARI::znx_zero(res.at_mut(res_col, j));
228 }
229
230 // Case where offset is positive and greater or equal
231 // to the precision of a.
232 if res_start == 0 {
233 return;
234 }
235
236 // Limbs of `a` that have a greater precision than `res`.
237 let a_out_range: usize = a_size.saturating_sub(a_start);
238
239 for j in 0..a_out_range {
240 if j == 0 {
241 ZNXARI::znx_normalize_first_step_carry_only(a_base2k, lsh_pos, a.at(a_col, a_size - j - 1), a_carry);
242 } else {
243 ZNXARI::znx_normalize_middle_step_carry_only(a_base2k, lsh_pos, a.at(a_col, a_size - j - 1), a_carry);
244 }
245 }
246
247 // Zero carry if the above loop didn't trigger.
248 if a_out_range == 0 {
249 ZNXARI::znx_zero(a_carry);
250 }
251
252 // How much is left to accumulate to fill a limb of `res`.
253 let mut res_acc_left: usize = res_base2k;
254
255 // Starting limb of `res`.
256 let mut res_limb: usize = res_start - 1;
257
258 // How many limbs of `a` overlap with `res` (after taking into account the offset).
259 let mid_range: usize = a_start.saturating_sub(a_end);
260
261 // Regular normalization over the overlapping limbs of res and a.
262 'outer: for j in 0..mid_range {
263 let a_limb: usize = a_start - j - 1;
264
265 // Current res & a limbs
266 let a_slice: &[i64] = a.at(a_col, a_limb);
267
268 // Trackers: wow much of a_norm is left to
269 // be flushed on res.
270 let mut a_take_left: usize = a_base2k;
271
272 // Normalizes the j-th limb of a and store the results into `a_norm``.
273 // This step is required to avoid overflow in the next step,
274 // which assumes that |a| is bounded by 2^{a_base2k -1} (i.e. normalized).
275 ZNXARI::znx_normalize_middle_step(a_base2k, lsh_pos, a_norm, a_slice, a_carry);
276
277 // In the first iteration we need to match the precision `res` and `a`.
278 if j == 0 {
279 // Case where `a` has more precision than `res` (after taking into account the offset)
280 //
281 // For example:
282 //
283 // a: [x x x x x][x x x x x][x x x x x][x x x x x]
284 // res: [x x x x x x][x x x x x x][x x x x x x]
285 if !(a_tot_bits - a_start_bit).is_multiple_of(a_base2k) {
286 let take: usize = (a_tot_bits - a_start_bit) % a_base2k;
287 ZNXARI::znx_mul_power_of_two_inplace(-(take as i64), a_norm);
288 a_take_left -= take;
289 // Case where `res` has more precision than `a` (after taking into account the offset)
290 //
291 // For example:
292 //
293 // a: [x x x x x][x x x x x][x x x x x][x x x x x]
294 // res: [x x x x x x][x x x x x x][x x x x x x]
295 } else if !(res_tot_bits - res_start_bit).is_multiple_of(res_base2k) {
296 res_acc_left -= (res_tot_bits - res_start_bit) % res_base2k;
297 }
298 }
299
300 // Extract bits of `a_norm` and accumulates them on res[res_limb] until
301 // res_base2k bits have been accumulated or until all bits of `a` are
302 // extracted.
303 'inner: loop {
304 // Current limb of res
305 let res_slice: &mut [i64] = res.at_mut(res_col, res_limb);
306
307 // We can take at most a_base2k bits
308 // but not more than what is left on a_norm or what is left to
309 // fully populate the current limb of res.
310 let a_take: usize = a_base2k.min(a_take_left).min(res_acc_left);
311
312 if a_take != 0 {
313 // Extract `a_take` bits from a_norm and accumulates them on `res_slice`.
314 let scale: usize = res_base2k - res_acc_left;
315 ZNXARI::znx_extract_digit_addmul(a_take, scale, res_slice, a_norm);
316 a_take_left -= a_take;
317 res_acc_left -= a_take;
318 }
319
320 // If either:
321 // * At least `res_base2k` bits have been accumulated
322 // * We have reached the last limb of a
323 // Then: Flushes them onto res
324 if res_acc_left == 0 || a_limb == 0 {
325 // This case happens only if `res_offset` is negative.
326 // If `res_offset` is negative, we need to apply the offset BEFORE
327 // the normalization to ensure the `res-offset` overflowing bits of `a`
328 // are in the MSB of `res` instead of being discarded.
329 if a_limb == 0 && a_take_left == 0 {
330 // TODO: prove no overflow can happen here (should not intuitively)
331 ZNXARI::znx_add_inplace(a_carry, a_norm);
332
333 // Usual case where for example
334 // a: [ overflow ][x x x x x][x x x x x][x x x x x][x x x x x]
335 // res: [x x x x x x][x x x x x x][x x x x x x][x x x x x x]
336 //
337 // where [overflow] are the overflowing bits of `a` (note that they are not a limb, but
338 // stored in a[0] & carry from a[1]) that are moved into the MSB of `res` due to the
339 // negative offset.
340 //
341 // In this case we populate what is left of `res_acc_left` using `a_carry`
342 //
343 // TODO: see if this can be simplified (e.g. just add).
344 if res_acc_left != 0 {
345 let scale: usize = res_base2k - res_acc_left;
346 ZNXARI::znx_extract_digit_addmul(res_acc_left, scale, res_slice, a_carry);
347 }
348
349 ZNXARI::znx_normalize_middle_step_inplace(res_base2k, 0, res_slice, res_carry);
350
351 // Previous step might not consume all bits of a_carry
352 // TODO: prove no overflow can happen here
353 ZNXARI::znx_add_inplace(res_carry, a_carry);
354
355 // We are done, so breaks out of the loop (yes we are at a[0], but
356 // this avoids possible over/under flows of tracking variables)
357 break 'outer;
358 }
359
360 // If we reached the last limb of res
361 if res_limb == 0 {
362 break 'outer;
363 }
364
365 res_acc_left += res_base2k;
366 res_limb -= 1;
367 }
368
369 // If a_norm is exhausted, breaks the inner loop.
370 if a_take_left == 0 {
371 ZNXARI::znx_add_inplace(a_carry, a_norm);
372 break 'inner;
373 }
374 }
375 }
376
377 // This case will happen if offset is negative.
378 if res_end != 0 {
379 // If there are no overlapping limbs between `res` and `a`
380 // (can happen if offset is negative), then we propagate the
381 // carry of `a` on res. Note that the carry of `a` can be
382 // greater than the precision of res.
383 //
384 // For example with offset = -8:
385 // a carry a[0] a[1] a[2] a[3]
386 // a: [---------------------- ][x x x][x x x][x x x][x x x]
387 // b: [x x x x][x x x x ]
388 // res[0] res[1]
389 //
390 // If there are overlapping limbs between `res` and `a`,
391 // we can use `res_carry`, which contains the carry of propagating
392 // the shifted reconstruction of `a` in `res_base2k` along with
393 // the carry of a[0].
394 let carry_to_use = if a_start == a_end { a_carry } else { res_carry };
395
396 for j in 0..res_end {
397 if j == res_end - 1 {
398 ZNXARI::znx_normalize_final_step_inplace(res_base2k, 0, res.at_mut(res_col, res_end - j - 1), carry_to_use);
399 } else {
400 ZNXARI::znx_normalize_middle_step_inplace(res_base2k, 0, res.at_mut(res_col, res_end - j - 1), carry_to_use);
401 }
402 }
403 }
404}
405
406pub fn vec_znx_normalize_inplace<R: VecZnxToMut, ZNXARI>(base2k: usize, res: &mut R, res_col: usize, carry: &mut [i64])
407where
408 ZNXARI: ZnxNormalizeFirstStepInplace + ZnxNormalizeMiddleStepInplace + ZnxNormalizeFinalStepInplace,
409{
410 let mut res: VecZnx<&mut [u8]> = res.to_mut();
411
412 #[cfg(debug_assertions)]
413 {
414 assert!(carry.len() >= res.n());
415 }
416
417 let res_size: usize = res.size();
418
419 for j in (0..res_size).rev() {
420 if j == res_size - 1 {
421 ZNXARI::znx_normalize_first_step_inplace(base2k, 0, res.at_mut(res_col, j), carry);
422 } else if j == 0 {
423 ZNXARI::znx_normalize_final_step_inplace(base2k, 0, res.at_mut(res_col, j), carry);
424 } else {
425 ZNXARI::znx_normalize_middle_step_inplace(base2k, 0, res.at_mut(res_col, j), carry);
426 }
427 }
428}
429
430#[test]
431fn test_vec_znx_normalize_cross_base2k() {
432 let n: usize = 8;
433
434 let mut carry: Vec<i64> = vec![0i64; vec_znx_normalize_tmp_bytes(n) / size_of::<i64>()];
435
436 use crate::reference::znx::ZnxRef;
437 use rug::ops::SubAssignRound;
438 use rug::{Float, float::Round};
439
440 let prec: usize = 128;
441
442 for in_base2k in 1..=51 {
443 for out_base2k in 1..=51 {
444 for offset in [
445 -(prec as i64),
446 -(prec as i64 - 1),
447 -(prec as i64 - in_base2k as i64),
448 -(in_base2k as i64 + 1),
449 in_base2k as i64,
450 -(in_base2k as i64 - 1),
451 0,
452 (in_base2k as i64 - 1),
453 in_base2k as i64,
454 (in_base2k as i64 + 1),
455 (prec as i64 - in_base2k as i64),
456 (prec - 1) as i64,
457 prec as i64,
458 ] {
459 let mut source: Source = Source::new([1u8; 32]);
460
461 let in_size: usize = prec.div_ceil(in_base2k);
462 let in_prec: u32 = (in_size * in_base2k) as u32;
463
464 // Ensures no loss of precision (mostly for testing purpose)
465 let out_size: usize = (in_prec as usize).div_ceil(out_base2k);
466
467 let out_prec: u32 = (out_size * out_base2k) as u32;
468 let min_prec: u32 = (in_size * in_base2k).min(out_size * out_base2k) as u32;
469 let mut want: VecZnx<Vec<u8>> = VecZnx::alloc(n, 1, in_size);
470 want.fill_uniform(60, &mut source);
471
472 let mut have: VecZnx<Vec<u8>> = VecZnx::alloc(n, 1, out_size);
473 have.fill_uniform(60, &mut source);
474 vec_znx_normalize_cross_base2k::<_, _, ZnxRef>(&mut have, out_base2k, offset, 0, &want, in_base2k, 0, &mut carry);
475
476 let mut data_have: Vec<Float> = (0..n).map(|_| Float::with_val(out_prec + 60, 0)).collect();
477 let mut data_want: Vec<Float> = (0..n).map(|_| Float::with_val(in_prec + 60, 0)).collect();
478
479 have.decode_vec_float(out_base2k, 0, &mut data_have);
480 want.decode_vec_float(in_base2k, 0, &mut data_want);
481
482 let scale: Float = Float::with_val(out_prec + 60, Float::u_pow_u(2, offset.unsigned_abs() as u32));
483
484 if offset > 0 {
485 for x in &mut data_want {
486 *x *= &scale;
487 *x %= 1;
488 }
489 } else if offset < 0 {
490 for x in &mut data_want {
491 *x /= &scale;
492 *x %= 1;
493 }
494 } else {
495 for x in &mut data_want {
496 *x %= 1;
497 }
498 }
499
500 for x in &mut data_have {
501 if *x >= 0.5 {
502 *x -= 1;
503 } else if *x < -0.5 {
504 *x += 1;
505 }
506 }
507
508 for x in &mut data_want {
509 if *x >= 0.5 {
510 *x -= 1;
511 } else if *x < -0.5 {
512 *x += 1;
513 }
514 }
515
516 for i in 0..n {
517 //println!("i:{i:02} {} {}", data_want[i], data_have[i]);
518
519 let mut err: Float = data_have[i].clone();
520 err.sub_assign_round(&data_want[i], Round::Nearest);
521 err = err.abs();
522
523 let err_log2: f64 = err.clone().max(&Float::with_val(prec as u32, 1e-60)).log2().to_f64();
524
525 assert!(err_log2 <= -(min_prec as f64) + 1.0, "{} {}", err_log2, -(min_prec as f64))
526 }
527 }
528 }
529 }
530}
531
532#[test]
533fn test_vec_znx_normalize_inter_base2k() {
534 let n: usize = 8;
535
536 let mut carry: Vec<i64> = vec![0i64; vec_znx_normalize_tmp_bytes(n) / size_of::<i64>()];
537
538 use crate::reference::znx::ZnxRef;
539 use rug::ops::SubAssignRound;
540 use rug::{Float, float::Round};
541
542 let mut source: Source = Source::new([1u8; 32]);
543
544 let prec: usize = 128;
545 let offset_range: i64 = prec as i64;
546
547 for base2k in 1..=51 {
548 for offset in (-offset_range..=offset_range).step_by(base2k + 1) {
549 let size: usize = prec.div_ceil(base2k);
550 let out_prec: u32 = (size * base2k) as u32;
551
552 // Fills "want" with uniform values
553 let mut want: VecZnx<Vec<u8>> = VecZnx::alloc(n, 1, size);
554 want.fill_uniform(60, &mut source);
555
556 // Fills "have" with the shifted normalization of "want"
557 let mut have: VecZnx<Vec<u8>> = VecZnx::alloc(n, 1, size);
558 have.fill_uniform(60, &mut source);
559 vec_znx_normalize_inter_base2k::<_, _, ZnxRef>(base2k, &mut have, offset, 0, &want, 0, &mut carry);
560
561 let mut data_have: Vec<Float> = (0..n).map(|_| Float::with_val(out_prec + 60, 0)).collect();
562 let mut data_want: Vec<Float> = (0..n).map(|_| Float::with_val(out_prec + 60, 0)).collect();
563
564 have.decode_vec_float(base2k, 0, &mut data_have);
565 want.decode_vec_float(base2k, 0, &mut data_want);
566
567 let scale: Float = Float::with_val(out_prec + 60, Float::u_pow_u(2, offset.unsigned_abs() as u32));
568
569 if offset > 0 {
570 for x in &mut data_want {
571 *x *= &scale;
572 *x %= 1;
573 }
574 } else if offset < 0 {
575 for x in &mut data_want {
576 *x /= &scale;
577 *x %= 1;
578 }
579 } else {
580 for x in &mut data_want {
581 *x %= 1;
582 }
583 }
584
585 for x in &mut data_have {
586 if *x >= 0.5 {
587 *x -= 1;
588 } else if *x < -0.5 {
589 *x += 1;
590 }
591 }
592
593 for x in &mut data_want {
594 if *x >= 0.5 {
595 *x -= 1;
596 } else if *x < -0.5 {
597 *x += 1;
598 }
599 }
600
601 for i in 0..n {
602 //println!("i:{i:02} {} {}", data_want[i], data_have[i]);
603
604 let mut err: Float = data_have[i].clone();
605 err.sub_assign_round(&data_want[i], Round::Nearest);
606 err = err.abs();
607
608 let err_log2: f64 = err.clone().max(&Float::with_val(prec as u32, 1e-60)).log2().to_f64();
609
610 assert!(err_log2 <= -(out_prec as f64), "{} {}", err_log2, -(out_prec as f64))
611 }
612 }
613 }
614}
615pub fn bench_vec_znx_normalize<B: Backend>(c: &mut Criterion, label: &str)
616where
617 Module<B>: VecZnxNormalize<B> + ModuleNew<B> + VecZnxNormalizeTmpBytes,
618 ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
619{
620 let group_name: String = format!("vec_znx_normalize::{label}");
621
622 let mut group = c.benchmark_group(group_name);
623
624 fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
625 where
626 Module<B>: VecZnxNormalize<B> + ModuleNew<B> + VecZnxNormalizeTmpBytes,
627 ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
628 {
629 let n: usize = 1 << params[0];
630 let cols: usize = params[1];
631 let size: usize = params[2];
632
633 let module: Module<B> = Module::<B>::new(n as u64);
634
635 let base2k: usize = 50;
636
637 let mut source: Source = Source::new([0u8; 32]);
638
639 let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
640 let mut res: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
641
642 // Fill a with random i64
643 a.fill_uniform(50, &mut source);
644 res.fill_uniform(50, &mut source);
645
646 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(module.vec_znx_normalize_tmp_bytes());
647 let res_offset: i64 = 0;
648 move || {
649 for i in 0..cols {
650 module.vec_znx_normalize(&mut res, base2k, res_offset, i, &a, base2k, i, scratch.borrow());
651 }
652 black_box(());
653 }
654 }
655
656 for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
657 let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
658 let mut runner = runner::<B>(params);
659 group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
660 }
661
662 group.finish();
663}
664
665pub fn bench_vec_znx_normalize_inplace<B: Backend>(c: &mut Criterion, label: &str)
666where
667 Module<B>: VecZnxNormalizeInplace<B> + ModuleNew<B> + VecZnxNormalizeTmpBytes,
668 ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
669{
670 let group_name: String = format!("vec_znx_normalize_inplace::{label}");
671
672 let mut group = c.benchmark_group(group_name);
673
674 fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
675 where
676 Module<B>: VecZnxNormalizeInplace<B> + ModuleNew<B> + VecZnxNormalizeTmpBytes,
677 ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
678 {
679 let n: usize = 1 << params[0];
680 let cols: usize = params[1];
681 let size: usize = params[2];
682
683 let module: Module<B> = Module::<B>::new(n as u64);
684
685 let base2k: usize = 50;
686
687 let mut source: Source = Source::new([0u8; 32]);
688
689 let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
690
691 // Fill a with random i64
692 a.fill_uniform(50, &mut source);
693
694 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(module.vec_znx_normalize_tmp_bytes());
695
696 move || {
697 for i in 0..cols {
698 module.vec_znx_normalize_inplace(base2k, &mut a, i, scratch.borrow());
699 }
700 black_box(());
701 }
702 }
703
704 for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
705 let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
706 let mut runner = runner::<B>(params);
707 group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
708 }
709
710 group.finish();
711}