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 2 * n * size_of::<i64>()
18}
19
20pub fn vec_znx_normalize<R, A, ZNXARI>(
21 res_base2k: usize,
22 res: &mut R,
23 res_col: usize,
24 a_base2k: usize,
25 a: &A,
26 a_col: usize,
27 carry: &mut [i64],
28) where
29 R: VecZnxToMut,
30 A: VecZnxToRef,
31 ZNXARI: ZnxZero
32 + ZnxCopy
33 + ZnxAddInplace
34 + ZnxMulPowerOfTwoInplace
35 + ZnxNormalizeFirstStepCarryOnly
36 + ZnxNormalizeMiddleStepCarryOnly
37 + ZnxNormalizeMiddleStep
38 + ZnxNormalizeFinalStep
39 + ZnxNormalizeFirstStep
40 + ZnxExtractDigitAddMul
41 + ZnxNormalizeDigit,
42{
43 let mut res: VecZnx<&mut [u8]> = res.to_mut();
44 let a: VecZnx<&[u8]> = a.to_ref();
45
46 #[cfg(debug_assertions)]
47 {
48 assert!(carry.len() >= 2 * res.n());
49 assert_eq!(res.n(), a.n());
50 }
51
52 let n: usize = res.n();
53 let res_size: usize = res.size();
54 let a_size: usize = a.size();
55
56 if res_base2k == a_base2k {
57 if a_size > res_size {
58 for j in (res_size..a_size).rev() {
59 if j == a_size - 1 {
60 ZNXARI::znx_normalize_first_step_carry_only(res_base2k, 0, a.at(a_col, j), carry);
61 } else {
62 ZNXARI::znx_normalize_middle_step_carry_only(res_base2k, 0, a.at(a_col, j), carry);
63 }
64 }
65
66 for j in (1..res_size).rev() {
67 ZNXARI::znx_normalize_middle_step(res_base2k, 0, res.at_mut(res_col, j), a.at(a_col, j), carry);
68 }
69
70 ZNXARI::znx_normalize_final_step(res_base2k, 0, res.at_mut(res_col, 0), a.at(a_col, 0), carry);
71 } else {
72 for j in (0..a_size).rev() {
73 if j == a_size - 1 {
74 ZNXARI::znx_normalize_first_step(res_base2k, 0, res.at_mut(res_col, j), a.at(a_col, j), carry);
75 } else if j == 0 {
76 ZNXARI::znx_normalize_final_step(res_base2k, 0, res.at_mut(res_col, j), a.at(a_col, j), carry);
77 } else {
78 ZNXARI::znx_normalize_middle_step(res_base2k, 0, res.at_mut(res_col, j), a.at(a_col, j), carry);
79 }
80 }
81
82 for j in a_size..res_size {
83 ZNXARI::znx_zero(res.at_mut(res_col, j));
84 }
85 }
86 } else {
87 let (a_norm, carry) = carry.split_at_mut(n);
88
89 // Relevant limbs of res
90 let res_min_size: usize = (a_size * a_base2k).div_ceil(res_base2k).min(res_size);
91
92 // Relevant limbs of a
93 let a_min_size: usize = (res_size * res_base2k).div_ceil(a_base2k).min(a_size);
94
95 // Get carry for limbs of a that have higher precision than res
96 for j in (a_min_size..a_size).rev() {
97 if j == a_size - 1 {
98 ZNXARI::znx_normalize_first_step_carry_only(a_base2k, 0, a.at(a_col, j), carry);
99 } else {
100 ZNXARI::znx_normalize_middle_step_carry_only(a_base2k, 0, a.at(a_col, j), carry);
101 }
102 }
103
104 if a_min_size == a_size {
105 ZNXARI::znx_zero(carry);
106 }
107
108 // Maximum relevant precision of a
109 let a_prec: usize = a_min_size * a_base2k;
110
111 // Maximum relevant precision of res
112 let res_prec: usize = res_min_size * res_base2k;
113
114 // Res limb index
115 let mut res_idx: usize = res_min_size - 1;
116
117 // Trackers: wow much of res is left to be populated
118 // for the current limb.
119 let mut res_left: usize = res_base2k;
120
121 for j in 0..res_size {
122 ZNXARI::znx_zero(res.at_mut(res_col, j));
123 }
124
125 for j in (0..a_min_size).rev() {
126 // Trackers: wow much of a_norm is left to
127 // be flushed on res.
128 let mut a_left: usize = a_base2k;
129
130 // Normalizes the j-th limb of a and store the results into a_norm.
131 // This step is required to avoid overflow in the next step,
132 // which assumes that |a| is bounded by 2^{a_base2k -1}.
133 if j != 0 {
134 ZNXARI::znx_normalize_middle_step(a_base2k, 0, a_norm, a.at(a_col, j), carry);
135 } else {
136 ZNXARI::znx_normalize_final_step(a_base2k, 0, a_norm, a.at(a_col, j), carry);
137 }
138
139 // In the first iteration we need to match the precision of the input/output.
140 // If a_min_size * a_base2k > res_min_size * res_base2k
141 // then divround a_norm by the difference of precision and
142 // acts like if a_norm has already been partially consummed.
143 // Else acts like if res has been already populated
144 // by the difference.
145 if j == a_min_size - 1 {
146 if a_prec > res_prec {
147 ZNXARI::znx_mul_power_of_two_inplace(res_prec as i64 - a_prec as i64, a_norm);
148 a_left -= a_prec - res_prec;
149 } else if res_prec > a_prec {
150 res_left -= res_prec - a_prec;
151 }
152 }
153
154 // Flushes a into res
155 loop {
156 // Selects the maximum amount of a that can be flushed
157 let a_take: usize = a_base2k.min(a_left).min(res_left);
158
159 // Output limb
160 let res_slice: &mut [i64] = res.at_mut(res_col, res_idx);
161
162 // Scaling of the value to flush
163 let lsh: usize = res_base2k - res_left;
164
165 // Extract the bits to flush on the output and updates
166 // a_norm accordingly.
167 ZNXARI::znx_extract_digit_addmul(a_take, lsh, res_slice, a_norm);
168
169 // Updates the trackers
170 a_left -= a_take;
171 res_left -= a_take;
172
173 // If the current limb of res is full,
174 // then normalizes this limb and adds
175 // the carry on a_norm.
176 if res_left == 0 {
177 // Updates tracker
178 res_left += res_base2k;
179
180 // Normalizes res and propagates the carry on a.
181 ZNXARI::znx_normalize_digit(res_base2k, res_slice, a_norm);
182
183 // If we reached the last limb of res breaks,
184 // but we might rerun the above loop if the
185 // base2k of a is much smaller than the base2k
186 // of res.
187 if res_idx == 0 {
188 ZNXARI::znx_add_inplace(carry, a_norm);
189 break;
190 }
191
192 // Else updates the limb index of res.
193 res_idx -= 1
194 }
195
196 // If a_norm is exhausted, breaks the loop.
197 if a_left == 0 {
198 ZNXARI::znx_add_inplace(carry, a_norm);
199 break;
200 }
201 }
202 }
203 }
204}
205
206pub fn vec_znx_normalize_inplace<R: VecZnxToMut, ZNXARI>(base2k: usize, res: &mut R, res_col: usize, carry: &mut [i64])
207where
208 ZNXARI: ZnxNormalizeFirstStepInplace + ZnxNormalizeMiddleStepInplace + ZnxNormalizeFinalStepInplace,
209{
210 let mut res: VecZnx<&mut [u8]> = res.to_mut();
211
212 #[cfg(debug_assertions)]
213 {
214 assert!(carry.len() >= res.n());
215 }
216
217 let res_size: usize = res.size();
218
219 for j in (0..res_size).rev() {
220 if j == res_size - 1 {
221 ZNXARI::znx_normalize_first_step_inplace(base2k, 0, res.at_mut(res_col, j), carry);
222 } else if j == 0 {
223 ZNXARI::znx_normalize_final_step_inplace(base2k, 0, res.at_mut(res_col, j), carry);
224 } else {
225 ZNXARI::znx_normalize_middle_step_inplace(base2k, 0, res.at_mut(res_col, j), carry);
226 }
227 }
228}
229
230pub fn bench_vec_znx_normalize<B: Backend>(c: &mut Criterion, label: &str)
231where
232 Module<B>: VecZnxNormalize<B> + ModuleNew<B> + VecZnxNormalizeTmpBytes,
233 ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
234{
235 let group_name: String = format!("vec_znx_normalize::{label}");
236
237 let mut group = c.benchmark_group(group_name);
238
239 fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
240 where
241 Module<B>: VecZnxNormalize<B> + ModuleNew<B> + VecZnxNormalizeTmpBytes,
242 ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
243 {
244 let n: usize = 1 << params[0];
245 let cols: usize = params[1];
246 let size: usize = params[2];
247
248 let module: Module<B> = Module::<B>::new(n as u64);
249
250 let base2k: usize = 50;
251
252 let mut source: Source = Source::new([0u8; 32]);
253
254 let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
255 let mut res: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
256
257 // Fill a with random i64
258 a.fill_uniform(50, &mut source);
259 res.fill_uniform(50, &mut source);
260
261 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(module.vec_znx_normalize_tmp_bytes());
262
263 move || {
264 for i in 0..cols {
265 module.vec_znx_normalize(base2k, &mut res, i, base2k, &a, i, scratch.borrow());
266 }
267 black_box(());
268 }
269 }
270
271 for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
272 let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
273 let mut runner = runner::<B>(params);
274 group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
275 }
276
277 group.finish();
278}
279
280pub fn bench_vec_znx_normalize_inplace<B: Backend>(c: &mut Criterion, label: &str)
281where
282 Module<B>: VecZnxNormalizeInplace<B> + ModuleNew<B> + VecZnxNormalizeTmpBytes,
283 ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
284{
285 let group_name: String = format!("vec_znx_normalize_inplace::{label}");
286
287 let mut group = c.benchmark_group(group_name);
288
289 fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
290 where
291 Module<B>: VecZnxNormalizeInplace<B> + ModuleNew<B> + VecZnxNormalizeTmpBytes,
292 ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
293 {
294 let n: usize = 1 << params[0];
295 let cols: usize = params[1];
296 let size: usize = params[2];
297
298 let module: Module<B> = Module::<B>::new(n as u64);
299
300 let base2k: usize = 50;
301
302 let mut source: Source = Source::new([0u8; 32]);
303
304 let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
305
306 // Fill a with random i64
307 a.fill_uniform(50, &mut source);
308
309 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(module.vec_znx_normalize_tmp_bytes());
310
311 move || {
312 for i in 0..cols {
313 module.vec_znx_normalize_inplace(base2k, &mut a, i, scratch.borrow());
314 }
315 black_box(());
316 }
317 }
318
319 for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
320 let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
321 let mut runner = runner::<B>(params);
322 group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
323 }
324
325 group.finish();
326}
327
328#[test]
329fn test_vec_znx_normalize_conv() {
330 let n: usize = 8;
331
332 let mut carry: Vec<i64> = vec![0i64; 2 * n];
333
334 use crate::reference::znx::ZnxRef;
335 use rug::ops::SubAssignRound;
336 use rug::{Float, float::Round};
337
338 let mut source: Source = Source::new([1u8; 32]);
339
340 let prec: usize = 128;
341
342 let mut data: Vec<i128> = vec![0i128; n];
343
344 data.iter_mut().for_each(|x| *x = source.next_i128());
345
346 for start_base2k in 1..50 {
347 for end_base2k in 1..50 {
348 let end_size: usize = prec.div_ceil(end_base2k);
349
350 let mut want: VecZnx<Vec<u8>> = VecZnx::alloc(n, 1, end_size);
351 want.encode_vec_i128(end_base2k, 0, prec, &data);
352 vec_znx_normalize_inplace::<_, ZnxRef>(end_base2k, &mut want, 0, &mut carry);
353
354 // Creates a temporary poly where encoding is in start_base2k
355 let mut tmp: VecZnx<Vec<u8>> = VecZnx::alloc(n, 1, prec.div_ceil(start_base2k));
356 tmp.encode_vec_i128(start_base2k, 0, prec, &data);
357
358 vec_znx_normalize_inplace::<_, ZnxRef>(start_base2k, &mut tmp, 0, &mut carry);
359
360 let mut data_tmp: Vec<Float> = (0..n).map(|_| Float::with_val(prec as u32, 0)).collect();
361 tmp.decode_vec_float(start_base2k, 0, &mut data_tmp);
362
363 let mut have: VecZnx<Vec<u8>> = VecZnx::alloc(n, 1, end_size);
364 vec_znx_normalize::<_, _, ZnxRef>(end_base2k, &mut have, 0, start_base2k, &tmp, 0, &mut carry);
365
366 let out_prec: u32 = (end_size * end_base2k) as u32;
367
368 let mut data_want: Vec<Float> = (0..n)
369 .map(|_| Float::with_val(out_prec as u32, 0))
370 .collect();
371 let mut data_res: Vec<Float> = (0..n)
372 .map(|_| Float::with_val(out_prec as u32, 0))
373 .collect();
374
375 have.decode_vec_float(end_base2k, 0, &mut data_want);
376 want.decode_vec_float(end_base2k, 0, &mut data_res);
377
378 for i in 0..n {
379 let mut err: Float = data_want[i].clone();
380 err.sub_assign_round(&data_res[i], Round::Nearest);
381 err = err.abs();
382
383 let err_log2: f64 = err
384 .clone()
385 .max(&Float::with_val(prec as u32, 1e-60))
386 .log2()
387 .to_f64();
388
389 assert!(
390 err_log2 <= -(out_prec as f64) + 1.,
391 "{} {}",
392 err_log2,
393 -(out_prec as f64) + 1.
394 )
395 }
396 }
397 }
398}