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