poulpy/bin_fhe/blind_rotation/cggi.rs
bin_fhe/blind_rotation/cggi.rsRaw
1use backend::hal::{
2 api::{
3 ScratchAvailable, SvpApply, SvpPPolAllocBytes, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, TakeVecZnxDftSlice,
4 TakeVecZnxSlice, VecZnxAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAllocBytes, VecZnxBigNormalizeTmpBytes, VecZnxCopy,
5 VecZnxDftAdd, VecZnxDftAddInplace, VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftSubABInplace, VecZnxDftToVecZnxBig,
6 VecZnxDftToVecZnxBigTmpBytes, VecZnxDftZero, VecZnxMulXpMinusOneInplace, VecZnxNormalize, VecZnxNormalizeInplace,
7 VecZnxRotate, VecZnxSubABInplace, VmpApplyTmpBytes, ZnxView, ZnxZero,
8 },
9 layouts::{Backend, DataMut, DataRef, Module, Scratch, SvpPPol, VecZnx},
10};
11use itertools::izip;
12
13use crate::{
14 GLWEOps, Infos, LookUpTableRotationDirection, TakeGLWECt,
15 blind_rotation::{key::BlindRotationKeyCGGIExec, lut::LookUpTable},
16 dist::Distribution,
17 layouts::{GLWECiphertext, GLWECiphertextToMut, LWECiphertext, LWECiphertextToRef},
18 trait_families::GLWEExternalProductFamily,
19};
20
21pub trait CCGIBlindRotationFamily<B: Backend> = VecZnxBigAllocBytes
22 + VecZnxDftAllocBytes
23 + SvpPPolAllocBytes
24 + VmpApplyTmpBytes
25 + VecZnxBigNormalizeTmpBytes
26 + VecZnxDftToVecZnxBigTmpBytes
27 + VecZnxDftToVecZnxBig<B>
28 + VecZnxDftAdd<B>
29 + VecZnxDftAddInplace<B>
30 + VecZnxDftFromVecZnx<B>
31 + VecZnxDftZero<B>
32 + SvpApply<B>
33 + VecZnxDftSubABInplace<B>
34 + VecZnxBigAddSmallInplace<B>
35 + GLWEExternalProductFamily<B>
36 + VecZnxRotate
37 + VecZnxAddInplace
38 + VecZnxSubABInplace
39 + VecZnxNormalize<B>
40 + VecZnxNormalizeInplace<B>
41 + VecZnxCopy
42 + VecZnxMulXpMinusOneInplace;
43
44pub fn cggi_blind_rotate_scratch_space<B: Backend>(
45 module: &Module<B>,
46 n: usize,
47 block_size: usize,
48 extension_factor: usize,
49 basek: usize,
50 k_res: usize,
51 k_brk: usize,
52 rows: usize,
53 rank: usize,
54) -> usize
55where
56 Module<B>: CCGIBlindRotationFamily<B>,
57{
58 let brk_size: usize = k_brk.div_ceil(basek);
59
60 if block_size > 1 {
61 let cols: usize = rank + 1;
62 let acc_dft: usize = module.vec_znx_dft_alloc_bytes(n, cols, rows) * extension_factor;
63 let acc_big: usize = module.vec_znx_big_alloc_bytes(n, 1, brk_size);
64 let vmp_res: usize = module.vec_znx_dft_alloc_bytes(n, cols, brk_size) * extension_factor;
65 let vmp_xai: usize = module.vec_znx_dft_alloc_bytes(n, 1, brk_size);
66 let acc_dft_add: usize = vmp_res;
67 let vmp: usize = module.vmp_apply_tmp_bytes(n, brk_size, rows, rows, 2, 2, brk_size); // GGSW product: (1 x 2) x (2 x 2)
68
69 let acc: usize;
70 if extension_factor > 1 {
71 acc = VecZnx::alloc_bytes(n, cols, k_res.div_ceil(basek)) * extension_factor;
72 } else {
73 acc = 0;
74 }
75
76 return acc
77 + acc_dft
78 + acc_dft_add
79 + vmp_res
80 + vmp_xai
81 + (vmp | (acc_big + (module.vec_znx_big_normalize_tmp_bytes(n) | module.vec_znx_dft_to_vec_znx_big_tmp_bytes(n))));
82 } else {
83 GLWECiphertext::bytes_of(n, basek, k_res, rank)
84 + GLWECiphertext::external_product_scratch_space(module, n, 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>: TakeVecZnxDftSlice<B> + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnx + ScratchAvailable + TakeVecZnxSlice,
101{
102 match brk.dist {
103 Distribution::BinaryBlock(_) | Distribution::BinaryFixed(_) | Distribution::BinaryProb(_) | Distribution::ZERO => {
104 if lut.extension_factor() > 1 {
105 cggi_blind_rotate_block_binary_extended(module, res, lwe, lut, brk, scratch);
106 } else if brk.block_size() > 1 {
107 cggi_blind_rotate_block_binary(module, res, lwe, lut, brk, scratch);
108 } else {
109 cggi_blind_rotate_binary_standard(module, res, lwe, lut, brk, scratch);
110 }
111 }
112 // TODO: ternary distribution ?
113 _ => panic!(
114 "invalid BlindRotationKeyCGGI distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)"
115 ),
116 }
117}
118
119pub(crate) fn cggi_blind_rotate_block_binary_extended<DataRes, DataIn, DataBrk, B: Backend>(
120 module: &Module<B>,
121 res: &mut GLWECiphertext<DataRes>,
122 lwe: &LWECiphertext<DataIn>,
123 lut: &LookUpTable,
124 brk: &BlindRotationKeyCGGIExec<DataBrk, B>,
125 scratch: &mut Scratch<B>,
126) where
127 DataRes: DataMut,
128 DataIn: DataRef,
129 DataBrk: DataRef,
130 Module<B>: CCGIBlindRotationFamily<B>,
131 Scratch<B>: TakeVecZnxDftSlice<B> + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnxSlice,
132{
133 let n_glwe: usize = brk.n();
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, n_glwe, cols, res.size());
140 let (mut acc_dft, scratch2) = scratch1.take_vec_znx_dft_slice(extension_factor, n_glwe, cols, rows);
141 let (mut vmp_res, scratch3) = scratch2.take_vec_znx_dft_slice(extension_factor, n_glwe, cols, brk.size());
142 let (mut acc_add_dft, scratch4) = scratch3.take_vec_znx_dft_slice(extension_factor, n_glwe, cols, brk.size());
143 let (mut vmp_xai, scratch5) = scratch4.take_vec_znx_dft(n_glwe, 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 * n_glwe;
160 let two_n_ext: usize = 2 * lut.domain_size();
161
162 mod_switch_2n(two_n_ext, &mut lwe_2n, &lwe_ref, lut.rotation_direction());
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(n_glwe, 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 n_glwe: usize = brk.n();
279 let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
280 let mut out_mut: GLWECiphertext<&mut [u8]> = res.to_mut();
281 let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
282 let two_n: usize = n_glwe << 1;
283 let basek: usize = brk.basek();
284 let rows: usize = brk.rows();
285
286 let cols: usize = out_mut.rank() + 1;
287
288 mod_switch_2n(
289 2 * lut.domain_size(),
290 &mut lwe_2n,
291 &lwe_ref,
292 lut.rotation_direction(),
293 );
294
295 let a: &[i64] = &lwe_2n[1..];
296 let b: i64 = lwe_2n[0];
297
298 out_mut.data.zero();
299
300 // Initialize out to X^{b} * LUT(X)
301 module.vec_znx_rotate(b, &mut out_mut.data, 0, &lut.data[0], 0);
302
303 let block_size: usize = brk.block_size();
304
305 // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
306
307 let (mut acc_dft, scratch1) = scratch.take_vec_znx_dft(n_glwe, cols, rows);
308 let (mut vmp_res, scratch2) = scratch1.take_vec_znx_dft(n_glwe, cols, brk.size());
309 let (mut acc_add_dft, scratch3) = scratch2.take_vec_znx_dft(n_glwe, cols, brk.size());
310 let (mut vmp_xai, scratch4) = scratch3.take_vec_znx_dft(n_glwe, 1, brk.size());
311
312 let x_pow_a: &Vec<SvpPPol<Vec<u8>, B>>;
313 if let Some(b) = &brk.x_pow_a {
314 x_pow_a = b
315 } else {
316 panic!("invalid key: x_pow_a has not been initialized")
317 }
318
319 izip!(
320 a.chunks_exact(block_size),
321 brk.data.chunks_exact(block_size)
322 )
323 .for_each(|(ai, ski)| {
324 (0..cols).for_each(|j| {
325 module.vec_znx_dft_from_vec_znx(1, 0, &mut acc_dft, j, &out_mut.data, j);
326 });
327
328 module.vec_znx_dft_zero(&mut acc_add_dft);
329
330 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
331 let ai_pos: usize = ((aii + two_n as i64) & (two_n - 1) as i64) as usize;
332
333 // vmp_res = DFT(acc) * BRK[i]
334 module.vmp_apply(&mut vmp_res, &acc_dft, &skii.data, scratch4);
335
336 // DFT(X^ai -1) * (DFT(acc) * BRK[i])
337 (0..cols).for_each(|i| {
338 module.svp_apply(&mut vmp_xai, 0, &x_pow_a[ai_pos], 0, &vmp_res, i);
339 module.vec_znx_dft_add_inplace(&mut acc_add_dft, i, &vmp_xai, 0);
340 module.vec_znx_dft_sub_ab_inplace(&mut acc_add_dft, i, &vmp_res, i);
341 });
342 });
343
344 {
345 let (mut acc_add_big, scratch5) = scratch4.take_vec_znx_big(n_glwe, 1, brk.size());
346
347 (0..cols).for_each(|i| {
348 module.vec_znx_dft_to_vec_znx_big(&mut acc_add_big, 0, &acc_add_dft, i, scratch5);
349 module.vec_znx_big_add_small_inplace(&mut acc_add_big, 0, &out_mut.data, i);
350 module.vec_znx_big_normalize(basek, &mut out_mut.data, i, &acc_add_big, 0, scratch5);
351 });
352 }
353 });
354}
355
356pub(crate) fn cggi_blind_rotate_binary_standard<DataRes, DataIn, DataBrk, B: Backend>(
357 module: &Module<B>,
358 res: &mut GLWECiphertext<DataRes>,
359 lwe: &LWECiphertext<DataIn>,
360 lut: &LookUpTable,
361 brk: &BlindRotationKeyCGGIExec<DataBrk, B>,
362 scratch: &mut Scratch<B>,
363) where
364 DataRes: DataMut,
365 DataIn: DataRef,
366 DataBrk: DataRef,
367 Module<B>: CCGIBlindRotationFamily<B>,
368 Scratch<B>: TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnx + ScratchAvailable,
369{
370 #[cfg(debug_assertions)]
371 {
372 assert_eq!(
373 res.n(),
374 brk.n(),
375 "res.n(): {} != brk.n(): {}",
376 res.n(),
377 brk.n()
378 );
379 assert_eq!(
380 lut.domain_size(),
381 brk.n(),
382 "lut.n(): {} != brk.n(): {}",
383 lut.domain_size(),
384 brk.n()
385 );
386 assert_eq!(
387 res.rank(),
388 brk.rank(),
389 "res.rank(): {} != brk.rank(): {}",
390 res.rank(),
391 brk.rank()
392 );
393 assert_eq!(
394 lwe.n(),
395 brk.data.len(),
396 "lwe.n(): {} != brk.data.len(): {}",
397 lwe.n(),
398 brk.data.len()
399 );
400 }
401
402 let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
403 let mut out_mut: GLWECiphertext<&mut [u8]> = res.to_mut();
404 let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
405 let basek: usize = brk.basek();
406
407 mod_switch_2n(
408 2 * lut.domain_size(),
409 &mut lwe_2n,
410 &lwe_ref,
411 lut.rotation_direction(),
412 );
413
414 let a: &[i64] = &lwe_2n[1..];
415 let b: i64 = lwe_2n[0];
416
417 out_mut.data.zero();
418
419 // Initialize out to X^{b} * LUT(X)
420 module.vec_znx_rotate(b, &mut out_mut.data, 0, &lut.data[0], 0);
421
422 // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
423 let (mut acc_tmp, scratch1) = scratch.take_glwe_ct(out_mut.n(), basek, out_mut.k(), out_mut.rank());
424
425 // TODO: see if faster by skipping normalization in external product and keeping acc in big coeffs
426 // TODO: first iteration can be optimized to be a gglwe product
427 izip!(a.iter(), brk.data.iter()).for_each(|(ai, ski)| {
428 // acc_tmp = sk[i] * acc
429 acc_tmp.external_product(module, &out_mut, ski, scratch1);
430
431 // acc_tmp = (sk[i] * acc) * (X^{ai} - 1)
432 acc_tmp.mul_xp_minus_one_inplace(module, *ai);
433
434 // acc = acc + (sk[i] * acc) * (X^{ai} - 1)
435 out_mut.add_inplace(module, &acc_tmp);
436 });
437
438 // We can normalize only at the end because we add normalized values in [-2^{basek-1}, 2^{basek-1}]
439 // on top of each others, thus ~ 2^{63-basek} additions are supported before overflow.
440 out_mut.normalize_inplace(module, scratch1);
441}
442
443pub(crate) fn mod_switch_2n(n: usize, res: &mut [i64], lwe: &LWECiphertext<&[u8]>, rot_dir: LookUpTableRotationDirection) {
444 let basek: usize = lwe.basek();
445
446 let log2n: usize = usize::BITS as usize - (n - 1).leading_zeros() as usize + 1;
447
448 res.copy_from_slice(&lwe.data.at(0, 0));
449
450 match rot_dir {
451 LookUpTableRotationDirection::Left => {
452 res.iter_mut().for_each(|x| *x = -*x);
453 }
454 LookUpTableRotationDirection::Right => {}
455 }
456
457 if basek > log2n {
458 let diff: usize = basek - log2n;
459 res.iter_mut().for_each(|x| {
460 *x = div_round_by_pow2(x, diff);
461 })
462 } else {
463 let rem: usize = basek - (log2n % basek);
464 let size: usize = log2n.div_ceil(basek);
465 (1..size).for_each(|i| {
466 if i == size - 1 && rem != basek {
467 let k_rem: usize = basek - rem;
468 izip!(lwe.data.at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
469 *y = (*y << k_rem) + (x >> rem);
470 });
471 } else {
472 izip!(lwe.data.at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
473 *y = (*y << basek) + x;
474 });
475 }
476 })
477 }
478}
479
480#[inline(always)]
481fn div_round_by_pow2(x: &i64, k: usize) -> i64 {
482 (x + (1 << (k - 1))) >> k
483}