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