poulpy/core/src/blind_rotation/cggi.rs
core/src/blind_rotation/cggi.rsRaw
1use backend::{
2 FFT64, MatZnxDftOps, MatZnxDftScratch, Module, ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps,
3 Scratch, VecZnxAlloc, VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDftAlloc, VecZnxDftOps, VecZnxOps, ZnxView,
4 ZnxViewMut, ZnxZero,
5};
6use itertools::izip;
7
8use crate::{
9 GLWECiphertext, GLWECiphertextToMut, GLWEOps, Infos, LWECiphertext, ScratchCore,
10 blind_rotation::{key::BlindRotationKeyCGGI, lut::LookUpTable},
11 dist::Distribution,
12 lwe::ciphertext::LWECiphertextToRef,
13};
14
15pub fn cggi_blind_rotate_scratch_space(
16 module: &Module<FFT64>,
17 block_size: usize,
18 extension_factor: usize,
19 basek: usize,
20 k_res: usize,
21 k_brk: usize,
22 rows: usize,
23 rank: usize,
24) -> usize {
25 let brk_size: usize = k_brk.div_ceil(basek);
26
27 if block_size > 1 {
28 let cols: usize = rank + 1;
29 let acc_dft: usize = module.bytes_of_vec_znx_dft(cols, rows) * extension_factor;
30 let acc_big: usize = module.bytes_of_vec_znx_big(1, brk_size);
31 let vmp_res: usize = module.bytes_of_vec_znx_dft(cols, brk_size) * extension_factor;
32 let acc_dft_add: usize = vmp_res;
33 let xai_plus_y: usize = module.bytes_of_scalar_znx_dft(1);
34 let xai_plus_y_dft: usize = module.bytes_of_scalar_znx_dft(1);
35 let vmp: usize = module.vmp_apply_tmp_bytes(brk_size, rows, rows, 2, 2, brk_size); // GGSW product: (1 x 2) x (2 x 2)
36
37 let acc: usize;
38 if extension_factor > 1 {
39 acc = module.bytes_of_vec_znx(cols, k_res.div_ceil(basek)) * extension_factor;
40 } else {
41 acc = 0;
42 }
43
44 return acc
45 + acc_dft
46 + acc_dft_add
47 + vmp_res
48 + xai_plus_y
49 + xai_plus_y_dft
50 + (vmp | (acc_big + (module.vec_znx_big_normalize_tmp_bytes() | module.vec_znx_idft_tmp_bytes())));
51 } else {
52 2 * GLWECiphertext::bytes_of(module, basek, k_res, rank)
53 + GLWECiphertext::external_product_scratch_space(module, basek, k_res, k_res, k_brk, 1, rank)
54 }
55}
56
57pub fn cggi_blind_rotate<DataRes, DataIn, DataBrk>(
58 module: &Module<FFT64>,
59 res: &mut GLWECiphertext<DataRes>,
60 lwe: &LWECiphertext<DataIn>,
61 lut: &LookUpTable,
62 brk: &BlindRotationKeyCGGI<DataBrk, FFT64>,
63 scratch: &mut Scratch,
64) where
65 DataRes: AsRef<[u8]> + AsMut<[u8]>,
66 DataIn: AsRef<[u8]>,
67 DataBrk: AsRef<[u8]>,
68{
69 match brk.dist {
70 Distribution::BinaryBlock(_) | Distribution::BinaryFixed(_) | Distribution::BinaryProb(_) | Distribution::ZERO => {
71 if lut.extension_factor() > 1 {
72 cggi_blind_rotate_block_binary_extended(module, res, lwe, lut, brk, scratch);
73 } else if brk.block_size() > 1 {
74 cggi_blind_rotate_block_binary(module, res, lwe, lut, brk, scratch);
75 } else {
76 cggi_blind_rotate_binary_standard(module, res, lwe, lut, brk, scratch);
77 }
78 }
79 // TODO: ternary distribution ?
80 _ => panic!(
81 "invalid BlindRotationKeyCGGI distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)"
82 ),
83 }
84}
85
86pub(crate) fn cggi_blind_rotate_block_binary_extended<DataRes, DataIn, DataBrk>(
87 module: &Module<FFT64>,
88 res: &mut GLWECiphertext<DataRes>,
89 lwe: &LWECiphertext<DataIn>,
90 lut: &LookUpTable,
91 brk: &BlindRotationKeyCGGI<DataBrk, FFT64>,
92 scratch: &mut Scratch,
93) where
94 DataRes: AsRef<[u8]> + AsMut<[u8]>,
95 DataIn: AsRef<[u8]>,
96 DataBrk: AsRef<[u8]>,
97{
98 let extension_factor: usize = lut.extension_factor();
99 let basek: usize = res.basek();
100 let rows: usize = brk.rows();
101 let cols: usize = res.rank() + 1;
102
103 let (mut acc, scratch1) = scratch.tmp_slice_vec_znx(extension_factor, module, cols, res.size());
104 let (mut acc_dft, scratch2) = scratch1.tmp_slice_vec_znx_dft(extension_factor, module, cols, rows);
105 let (mut vmp_res, scratch3) = scratch2.tmp_slice_vec_znx_dft(extension_factor, module, cols, brk.size());
106 let (mut acc_add_dft, scratch4) = scratch3.tmp_slice_vec_znx_dft(extension_factor, module, cols, brk.size());
107 let (mut minus_one, scratch5) = scratch4.tmp_scalar_znx_dft(module, 1);
108 let (mut xai_plus_y_dft, scratch6) = scratch5.tmp_scalar_znx_dft(module, 1);
109
110 minus_one.raw_mut()[..module.n() >> 1].fill(-1.0);
111
112 (0..extension_factor).for_each(|i| {
113 acc[i].zero();
114 });
115
116 let x_pow_a: &Vec<ScalarZnxDft<Vec<u8>, FFT64>>;
117 if let Some(b) = &brk.x_pow_a {
118 x_pow_a = b
119 } else {
120 panic!("invalid key: x_pow_a has not been initialized")
121 }
122
123 let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
124 let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
125
126 let two_n: usize = 2 * module.n();
127 let two_n_ext: usize = 2 * lut.domain_size();
128
129 negate_and_mod_switch_2n(two_n_ext, &mut lwe_2n, &lwe_ref);
130
131 let a: &[i64] = &lwe_2n[1..];
132 let b_pos: usize = ((lwe_2n[0] + two_n_ext as i64) & (two_n_ext - 1) as i64) as usize;
133
134 let b_hi: usize = b_pos / extension_factor;
135 let b_lo: usize = b_pos & (extension_factor - 1);
136
137 for (i, j) in (0..b_lo).zip(extension_factor - b_lo..extension_factor) {
138 module.vec_znx_rotate(b_hi as i64 + 1, &mut acc[i], 0, &lut.data[j], 0);
139 }
140 for (i, j) in (b_lo..extension_factor).zip(0..extension_factor - b_lo) {
141 module.vec_znx_rotate(b_hi as i64, &mut acc[i], 0, &lut.data[j], 0);
142 }
143
144 let block_size: usize = brk.block_size();
145
146 izip!(
147 a.chunks_exact(block_size),
148 brk.data.chunks_exact(block_size)
149 )
150 .for_each(|(ai, ski)| {
151 (0..extension_factor).for_each(|i| {
152 (0..cols).for_each(|j| {
153 module.vec_znx_dft(1, 0, &mut acc_dft[i], j, &acc[i], j);
154 });
155 acc_add_dft[i].zero();
156 });
157
158 // TODO: first & last iterations can be optimized
159 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
160 let ai_pos: usize = ((aii + two_n_ext as i64) & (two_n_ext - 1) as i64) as usize;
161 let ai_hi: usize = ai_pos / extension_factor;
162 let ai_lo: usize = ai_pos & (extension_factor - 1);
163
164 // vmp_res = DFT(acc) * BRK[i]
165 (0..extension_factor).for_each(|i| {
166 module.vmp_apply(&mut vmp_res[i], &acc_dft[i], &skii.data, scratch6);
167 });
168
169 // Trivial case: no rotation between polynomials, we can directly multiply with (X^{-ai} - 1)
170 if ai_lo == 0 {
171 // Sets acc_add_dft[i] = (acc[i] * sk) * (X^{-ai} - 1)
172 if ai_hi != 0 {
173 // DFT X^{-ai}
174 module.vec_znx_dft_add(&mut xai_plus_y_dft, 0, &x_pow_a[ai_hi], 0, &minus_one, 0);
175 (0..extension_factor).for_each(|j| {
176 (0..cols).for_each(|i| {
177 module.svp_apply_inplace(&mut vmp_res[j], i, &xai_plus_y_dft, 0);
178 module.vec_znx_dft_add_inplace(&mut acc_add_dft[j], i, &vmp_res[j], i);
179 });
180 });
181 }
182
183 // Non trivial case: rotation between polynomials
184 // In this case we can't directly multiply with (X^{-ai} - 1) because of the
185 // ring homomorphism R^{N} -> prod R^{N/extension_factor}, so we split the
186 // computation in two steps: acc_add_dft = (acc * sk) * (-1) + (acc * sk) * X^{-ai}
187 } else {
188 // Sets acc_add_dft[i] = acc[i] * sk
189
190 // Sets acc_add_dft[0..ai_lo] -= acc[..ai_lo] * sk
191 if (ai_hi + 1) & (two_n - 1) != 0 {
192 for i in 0..ai_lo {
193 (0..cols).for_each(|k| {
194 module.vec_znx_dft_sub_ab_inplace(&mut acc_add_dft[i], k, &vmp_res[i], k);
195 });
196 }
197 }
198
199 // Sets acc_add_dft[ai_lo..extension_factor] -= acc[ai_lo..extension_factor] * sk
200 if ai_hi != 0 {
201 for i in ai_lo..extension_factor {
202 (0..cols).for_each(|k: usize| {
203 module.vec_znx_dft_sub_ab_inplace(&mut acc_add_dft[i], k, &vmp_res[i], k);
204 });
205 }
206 }
207
208 // Sets acc_add_dft[0..ai_lo] += (acc[extension_factor - ai_lo..extension_factor] * sk) * X^{-ai+1}
209 if (ai_hi + 1) & (two_n - 1) != 0 {
210 for (i, j) in (0..ai_lo).zip(extension_factor - ai_lo..extension_factor) {
211 (0..cols).for_each(|k| {
212 module.svp_apply_inplace(&mut vmp_res[j], k, &x_pow_a[ai_hi + 1], 0);
213 module.vec_znx_dft_add_inplace(&mut acc_add_dft[i], k, &vmp_res[j], k);
214 });
215 }
216 }
217
218 // Sets acc_add_dft[ai_lo..extension_factor] += (acc[0..extension_factor - ai_lo] * sk) * X^{-ai}
219 if ai_hi != 0 {
220 // Sets acc_add_dft[ai_lo..extension_factor] += (acc[0..extension_factor - ai_lo] * sk) * X^{-ai}
221 for (i, j) in (ai_lo..extension_factor).zip(0..extension_factor - ai_lo) {
222 (0..cols).for_each(|k| {
223 module.svp_apply_inplace(&mut vmp_res[j], k, &x_pow_a[ai_hi], 0);
224 module.vec_znx_dft_add_inplace(&mut acc_add_dft[i], k, &vmp_res[j], k);
225 });
226 }
227 }
228 }
229 });
230
231 {
232 let (mut acc_add_big, scratch7) = scratch6.tmp_vec_znx_big(module, 1, brk.size());
233
234 (0..extension_factor).for_each(|j| {
235 (0..cols).for_each(|i| {
236 module.vec_znx_idft(&mut acc_add_big, 0, &acc_add_dft[j], i, scratch7);
237 module.vec_znx_big_add_small_inplace(&mut acc_add_big, 0, &acc[j], i);
238 module.vec_znx_big_normalize(basek, &mut acc[j], i, &acc_add_big, 0, scratch7);
239 });
240 });
241 }
242 });
243
244 (0..cols).for_each(|i| {
245 module.vec_znx_copy(&mut res.data, i, &acc[0], i);
246 });
247}
248
249pub(crate) fn cggi_blind_rotate_block_binary<DataRes, DataIn, DataBrk>(
250 module: &Module<FFT64>,
251 res: &mut GLWECiphertext<DataRes>,
252 lwe: &LWECiphertext<DataIn>,
253 lut: &LookUpTable,
254 brk: &BlindRotationKeyCGGI<DataBrk, FFT64>,
255 scratch: &mut Scratch,
256) where
257 DataRes: AsRef<[u8]> + AsMut<[u8]>,
258 DataIn: AsRef<[u8]>,
259 DataBrk: AsRef<[u8]>,
260{
261 let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
262 let mut out_mut: GLWECiphertext<&mut [u8]> = res.to_mut();
263 let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
264 let two_n: usize = module.n() << 1;
265 let basek: usize = brk.basek();
266 let rows = brk.rows();
267
268 let cols: usize = out_mut.rank() + 1;
269
270 negate_and_mod_switch_2n(2 * lut.domain_size(), &mut lwe_2n, &lwe_ref);
271
272 let a: &[i64] = &lwe_2n[1..];
273 let b: i64 = lwe_2n[0];
274
275 out_mut.data.zero();
276
277 // Initialize out to X^{b} * LUT(X)
278 module.vec_znx_rotate(b, &mut out_mut.data, 0, &lut.data[0], 0);
279
280 let block_size: usize = brk.block_size();
281
282 // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
283
284 let (mut acc_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rows);
285 let (mut vmp_res, scratch2) = scratch1.tmp_vec_znx_dft(module, cols, brk.size());
286 let (mut acc_add_dft, scratch3) = scratch2.tmp_vec_znx_dft(module, cols, brk.size());
287 let (mut minus_one, scratch4) = scratch3.tmp_scalar_znx_dft(module, 1);
288 let (mut xai_plus_y_dft, scratch5) = scratch4.tmp_scalar_znx_dft(module, 1);
289
290 minus_one.raw_mut()[..module.n() >> 1].fill(-1.0);
291
292 let x_pow_a: &Vec<ScalarZnxDft<Vec<u8>, FFT64>>;
293 if let Some(b) = &brk.x_pow_a {
294 x_pow_a = b
295 } else {
296 panic!("invalid key: x_pow_a has not been initialized")
297 }
298
299 izip!(
300 a.chunks_exact(block_size),
301 brk.data.chunks_exact(block_size)
302 )
303 .for_each(|(ai, ski)| {
304 (0..cols).for_each(|j| {
305 module.vec_znx_dft(1, 0, &mut acc_dft, j, &out_mut.data, j);
306 });
307
308 acc_add_dft.zero();
309
310 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
311 let ai_pos: usize = ((aii + two_n as i64) & (two_n - 1) as i64) as usize;
312
313 // vmp_res = DFT(acc) * BRK[i]
314 module.vmp_apply(&mut vmp_res, &acc_dft, &skii.data, scratch5);
315
316 // DFT(X^ai -1)
317 module.vec_znx_dft_add(&mut xai_plus_y_dft, 0, &x_pow_a[ai_pos], 0, &minus_one, 0);
318
319 // DFT(X^ai -1) * (DFT(acc) * BRK[i])
320 (0..cols).for_each(|i| {
321 module.svp_apply_inplace(&mut vmp_res, i, &xai_plus_y_dft, 0);
322 module.vec_znx_dft_add_inplace(&mut acc_add_dft, i, &vmp_res, i);
323 });
324 });
325
326 {
327 let (mut acc_add_big, scratch6) = scratch5.tmp_vec_znx_big(module, 1, brk.size());
328
329 (0..cols).for_each(|i| {
330 module.vec_znx_idft(&mut acc_add_big, 0, &acc_add_dft, i, scratch6);
331 module.vec_znx_big_add_small_inplace(&mut acc_add_big, 0, &out_mut.data, i);
332 module.vec_znx_big_normalize(basek, &mut out_mut.data, i, &acc_add_big, 0, scratch6);
333 });
334 }
335 });
336}
337
338pub(crate) fn cggi_blind_rotate_binary_standard<DataRes, DataIn, DataBrk>(
339 module: &Module<FFT64>,
340 res: &mut GLWECiphertext<DataRes>,
341 lwe: &LWECiphertext<DataIn>,
342 lut: &LookUpTable,
343 brk: &BlindRotationKeyCGGI<DataBrk, FFT64>,
344 scratch: &mut Scratch,
345) where
346 DataRes: AsRef<[u8]> + AsMut<[u8]>,
347 DataIn: AsRef<[u8]>,
348 DataBrk: AsRef<[u8]>,
349{
350 #[cfg(debug_assertions)]
351 {
352 assert_eq!(
353 res.n(),
354 module.n(),
355 "res.n(): {} != brk.n(): {}",
356 res.n(),
357 module.n()
358 );
359 assert_eq!(
360 lut.domain_size(),
361 module.n(),
362 "lut.n(): {} != brk.n(): {}",
363 lut.domain_size(),
364 module.n()
365 );
366 assert_eq!(
367 brk.n(),
368 module.n(),
369 "brk.n(): {} != brk.n(): {}",
370 brk.n(),
371 module.n()
372 );
373 assert_eq!(
374 res.rank(),
375 brk.rank(),
376 "res.rank(): {} != brk.rank(): {}",
377 res.rank(),
378 brk.rank()
379 );
380 assert_eq!(
381 lwe.n(),
382 brk.data.len(),
383 "lwe.n(): {} != brk.data.len(): {}",
384 lwe.n(),
385 brk.data.len()
386 );
387 }
388
389 let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
390 let mut out_mut: GLWECiphertext<&mut [u8]> = res.to_mut();
391 let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
392 let basek: usize = brk.basek();
393
394 negate_and_mod_switch_2n(2 * lut.domain_size(), &mut lwe_2n, &lwe_ref);
395
396 let a: &[i64] = &lwe_2n[1..];
397 let b: i64 = lwe_2n[0];
398
399 out_mut.data.zero();
400
401 // Initialize out to X^{b} * LUT(X)
402 module.vec_znx_rotate(b, &mut out_mut.data, 0, &lut.data[0], 0);
403
404 // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
405 let (mut acc_tmp, scratch1) = scratch.tmp_glwe_ct(module, basek, out_mut.k(), out_mut.rank());
406 let (mut acc_tmp_rot, scratch2) = scratch1.tmp_glwe_ct(module, basek, out_mut.k(), out_mut.rank());
407
408 // TODO: see if faster by skipping normalization in external product and keeping acc in big coeffs
409 // TODO: first iteration can be optimized to be a gglwe product
410 izip!(a.iter(), brk.data.iter()).for_each(|(ai, ski)| {
411 // acc_tmp = sk[i] * acc
412 acc_tmp.external_product(module, &out_mut, ski, scratch2);
413
414 // acc_tmp = (sk[i] * acc) * X^{ai}
415 acc_tmp_rot.rotate(module, *ai, &acc_tmp);
416
417 // acc = acc + (sk[i] * acc) * X^{ai}
418 out_mut.add_inplace(module, &acc_tmp_rot);
419
420 // acc = acc + (sk[i] * acc) * X^{ai} - (sk[i] * acc) = acc + (sk[i] * acc) * (X^{ai} - 1)
421 out_mut.sub_inplace_ab(module, &acc_tmp);
422 });
423
424 // We can normalize only at the end because we add normalized values in [-2^{basek-1}, 2^{basek-1}]
425 // on top of each others, thus ~ 2^{63-basek} additions are supported before overflow.
426 out_mut.normalize_inplace(module, scratch2);
427}
428
429pub(crate) fn negate_and_mod_switch_2n(n: usize, res: &mut [i64], lwe: &LWECiphertext<&[u8]>) {
430 let basek: usize = lwe.basek();
431
432 let log2n: usize = usize::BITS as usize - (n - 1).leading_zeros() as usize + 1;
433
434 res.copy_from_slice(&lwe.data.at(0, 0));
435 res.iter_mut().for_each(|x| *x = -*x);
436
437 if basek > log2n {
438 let diff: usize = basek - log2n;
439 res.iter_mut().for_each(|x| {
440 *x = div_round_by_pow2(x, diff);
441 })
442 } else {
443 let rem: usize = basek - (log2n % basek);
444 let size: usize = log2n.div_ceil(basek);
445 (1..size).for_each(|i| {
446 if i == size - 1 && rem != basek {
447 let k_rem: usize = basek - rem;
448 izip!(lwe.data.at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
449 *y = (*y << k_rem) + (x >> rem);
450 });
451 } else {
452 izip!(lwe.data.at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
453 *y = (*y << basek) + x;
454 });
455 }
456 })
457 }
458}
459
460#[inline(always)]
461fn div_round_by_pow2(x: &i64, k: usize) -> i64 {
462 (x + (1 << (k - 1))) >> k
463}