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