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