poulpy/core/src/blind_rotation/ccgi.rs
core/src/blind_rotation/ccgi.rsRaw
1use std::time::Instant;
2
3use backend::{
4 MatZnxDftOps, MatZnxDftScratch, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps, Scratch, VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDftOps, VecZnxOps, ZnxInfos, ZnxView, ZnxViewMut, ZnxZero, FFT64
5};
6use itertools::izip;
7
8use crate::{
9 FourierGLWECiphertext, FourierGLWESecret, GLWECiphertext, GLWECiphertextToMut, GLWEPlaintext, Infos, LWECiphertext,
10 ScratchCore,
11 blind_rotation::{key::BlindRotationKeyCGGI, lut::LookUpTable},
12 lwe::ciphertext::LWECiphertextToRef,
13};
14
15pub fn cggi_blind_rotate_scratch_space(
16 module: &Module<FFT64>,
17 extension_factor: usize,
18 basek: usize,
19 k_lut: usize,
20 k_brk: usize,
21 rows: usize,
22 rank: usize,
23) -> usize {
24 let lut_size: usize = k_lut.div_ceil(basek);
25 let brk_size: usize = k_brk.div_ceil(basek);
26
27 let acc_dft: usize = FourierGLWECiphertext::bytes_of(module, basek, k_brk, rank) * extension_factor;
28 let acc_big: usize = module.bytes_of_vec_znx_big(rank + 1, brk_size);
29 let acc_dft_add: usize = FourierGLWECiphertext::bytes_of(module, basek, k_brk, rank) * extension_factor;
30 let vmp_res: usize = FourierGLWECiphertext::bytes_of(module, basek, k_brk, rank) * extension_factor;
31 let xai_plus_y: usize = module.bytes_of_scalar_znx(1);
32 let xai_plus_y_dft: usize = module.bytes_of_scalar_znx_dft(1);
33 let vmp: usize = module.vmp_apply_tmp_bytes(lut_size, lut_size, rows, 2, 2, brk_size); // GGSW product: (1 x 2) x (2 x 2)
34
35 let acc: usize;
36 if extension_factor > 1 {
37 acc = GLWECiphertext::bytes_of(module, basek, k_lut, rank) * extension_factor;
38 } else {
39 acc = 0;
40 }
41
42 return acc + acc_big + acc_dft + acc_dft_add + vmp_res + xai_plus_y + xai_plus_y_dft + (vmp | module.vec_znx_big_normalize_tmp_bytes());
43}
44
45pub fn cggi_blind_rotate<DataRes, DataIn>(
46 module: &Module<FFT64>,
47 res: &mut GLWECiphertext<DataRes>,
48 lwe: &LWECiphertext<DataIn>,
49 lut: &LookUpTable,
50 brk: &BlindRotationKeyCGGI<FFT64>,
51 scratch: &mut Scratch,
52) where
53 DataRes: AsRef<[u8]> + AsMut<[u8]>,
54 DataIn: AsRef<[u8]>,
55{
56 if lut.extension_factor() > 1 {
57 cggi_blind_rotate_block_binary_extended(module, res, lwe, lut, brk, scratch);
58 } else if brk.block_size() > 1 {
59 cggi_blind_rotate_block_binary(module, res, lwe, lut, brk, scratch);
60 } else {
61 todo!("implement this case")
62 }
63}
64
65pub(crate) fn cggi_blind_rotate_block_binary_extended<DataRes, DataIn>(
66 module: &Module<FFT64>,
67 res: &mut GLWECiphertext<DataRes>,
68 lwe: &LWECiphertext<DataIn>,
69 lut: &LookUpTable,
70 brk: &BlindRotationKeyCGGI<FFT64>,
71 scratch: &mut Scratch,
72) where
73 DataRes: AsRef<[u8]> + AsMut<[u8]>,
74 DataIn: AsRef<[u8]>,
75{
76 let extension_factor: usize = lut.extension_factor();
77 let basek: usize = res.basek();
78
79 let (mut acc, scratch1) = scratch.tmp_vec_glwe_ct(extension_factor, module, basek, res.k(), res.rank());
80 let (mut acc_dft, scratch2) = scratch1.tmp_vec_fourier_glwe_ct(extension_factor, module, basek, brk.k(), res.rank());
81 let (mut vmp_res, scratch3) = scratch2.tmp_vec_fourier_glwe_ct(extension_factor, module, basek, brk.k(), res.rank());
82 let (mut acc_add_dft, scratch4) = scratch3.tmp_vec_fourier_glwe_ct(extension_factor, module, basek, brk.k(), res.rank());
83
84 (0..extension_factor).for_each(|i| {
85 acc[i].data.zero();
86 });
87
88 let (mut xai_plus_y, scratch5) = scratch4.tmp_scalar_znx(module, 1);
89 let (mut xai_plus_y_dft, scratch6) = scratch5.tmp_scalar_znx_dft(module, 1);
90 let (mut acc_add_big, scratch7) = scratch6.tmp_vec_znx_big(module, 1, brk.size());
91
92 let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
93 let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
94
95 let two_n_ext: usize = 2 * lut.domain_size();
96
97 let cols: usize = res.rank() + 1;
98
99 negate_and_mod_switch_2n(two_n_ext, &mut lwe_2n, &lwe_ref);
100
101 let a: &[i64] = &lwe_2n[1..];
102 let b_pos: usize = ((lwe_2n[0] + two_n_ext as i64) % two_n_ext as i64) as usize;
103
104 let b_hi: usize = b_pos / extension_factor;
105 let b_lo: usize = b_pos % extension_factor;
106
107 for (i, j) in (0..b_lo).zip(extension_factor - b_lo..extension_factor) {
108 module.vec_znx_rotate(b_hi as i64 + 1, &mut acc[i].data, 0, &lut.data[j], 0);
109 }
110 for (i, j) in (b_lo..extension_factor).zip(0..extension_factor - b_lo) {
111 module.vec_znx_rotate(b_hi as i64, &mut acc[i].data, 0, &lut.data[j], 0);
112 }
113
114 let block_size: usize = brk.block_size();
115
116 izip!(
117 a.chunks_exact(block_size),
118 brk.data.chunks_exact(block_size)
119 )
120 .enumerate()
121 .for_each(|(i, (ai, ski))| {
122 (0..extension_factor).for_each(|i| {
123 (0..cols).for_each(|j| {
124 module.vec_znx_dft(1, 0, &mut acc_dft[i].data, j, &acc[i].data, j);
125 });
126 acc_add_dft[i].data.zero();
127 });
128
129 // TODO: first & last iterations can be optimized
130 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
131 let ai_pos: usize = ((aii + two_n_ext as i64) % two_n_ext as i64) as usize;
132 let ai_hi: usize = ai_pos / extension_factor;
133 let ai_lo: usize = ai_pos % extension_factor;
134
135 // vmp_res = DFT(acc) * BRK[i]
136 (0..extension_factor).for_each(|i| {
137 module.vmp_apply(&mut vmp_res[i].data, &acc_dft[i].data, &skii.data, scratch7);
138 });
139
140 // Trivial case: no rotation between polynomials, we can directly multiply with (X^{-ai} - 1)
141 if ai_lo == 0 {
142 // DFT X^{-ai}
143 set_xai_plus_y(
144 module,
145 ai_hi as i64,
146 -1,
147 &mut xai_plus_y_dft,
148 &mut xai_plus_y,
149 );
150
151 // Sets acc_add_dft[i] = (acc[i] * sk) * (X^{-ai} - 1)
152 (0..extension_factor).for_each(|j| {
153 (0..cols).for_each(|i| {
154 module.svp_apply_inplace(&mut vmp_res[j].data, i, &xai_plus_y_dft, 0);
155 module.vec_znx_dft_add_inplace(&mut acc_add_dft[j].data, i, &vmp_res[j].data, i);
156 });
157 });
158 // Non trivial case: rotation between polynomials
159 // In this case we can't directly multiply with (X^{-ai} - 1) because of the
160 // ring homomorphism R^{N} -> prod R^{N/extension_factor}, so we split the
161 // computation in two steps: acc_add_dft = (acc * sk) * (-1) + (acc * sk) * X^{-ai}
162 } else {
163 // Sets acc_add_dft[i] = acc[i] * sk
164 (0..extension_factor).for_each(|i| {
165 (0..cols).for_each(|k| {
166 module.vec_znx_dft_sub_ab_inplace(&mut acc_add_dft[i].data, k, &vmp_res[i].data, k);
167 })
168 });
169
170 // DFT X^{-ai+1}
171 set_xai_plus_y(
172 module,
173 ai_hi as i64 + 1,
174 0,
175 &mut xai_plus_y_dft,
176 &mut xai_plus_y,
177 );
178
179 // Sets acc_add_dft[0..ai_lo] += (acc[extension_factor - ai_lo..extension_factor] * sk) * X^{-ai+1}
180 for (i, j) in (0..ai_lo).zip(extension_factor - ai_lo..extension_factor) {
181 module.vec_znx_rotate(b_hi as i64 + 1, &mut acc[i].data, 0, &lut.data[j], 0);
182 (0..cols).for_each(|k| {
183 module.svp_apply_inplace(&mut vmp_res[j].data, k, &xai_plus_y_dft, 0);
184 module.vec_znx_dft_add_inplace(&mut acc_add_dft[i].data, k, &vmp_res[j].data, k);
185 });
186 }
187
188 // DFT X^{-ai}
189 set_xai_plus_y(
190 module,
191 ai_hi as i64,
192 0,
193 &mut xai_plus_y_dft,
194 &mut xai_plus_y,
195 );
196
197 // Sets acc_add_dft[ai_lo..extension_factor] += (acc[0..extension_factor - ai_lo] * sk) * X^{-ai}
198 for (i, j) in (ai_lo..extension_factor).zip(0..extension_factor - ai_lo) {
199 module.vec_znx_rotate(b_hi as i64, &mut acc[i].data, 0, &lut.data[j], 0);
200 (0..cols).for_each(|k| {
201 module.svp_apply_inplace(&mut vmp_res[j].data, k, &xai_plus_y_dft, 0);
202 module.vec_znx_dft_add_inplace(&mut acc_add_dft[i].data, k, &vmp_res[j].data, k);
203 });
204 }
205 }
206 });
207
208 (0..extension_factor).for_each(|j| {
209 (0..cols).for_each(|i| {
210 module.vec_znx_dft_add_inplace(&mut acc_dft[j].data, i, &acc_add_dft[j].data, i);
211 module.vec_znx_idft(&mut acc_add_big, 0, &acc_dft[j].data, i, scratch7);
212 module.vec_znx_big_normalize(basek, &mut acc[j].data, i, &acc_add_big, 0, scratch7);
213 });
214 });
215 });
216
217 (0..cols).for_each(|i| {
218 module.vec_znx_copy(&mut res.data, i, &acc[0].data, i);
219 });
220}
221
222fn set_xai_plus_y(
223 module: &Module<FFT64>,
224 k: i64,
225 y: i64,
226 res: &mut ScalarZnxDft<&mut [u8], FFT64>,
227 buf: &mut ScalarZnx<&mut [u8]>,
228) {
229 buf.zero();
230 buf.at_mut(0, 0)[0] = 1;
231 module.vec_znx_rotate_inplace(k, buf, 0);
232 buf.at_mut(0, 0)[0] += y;
233 module.svp_prepare(res, 0, buf, 0);
234}
235
236pub(crate) fn cggi_blind_rotate_block_binary<DataRes, DataIn>(
237 module: &Module<FFT64>,
238 res: &mut GLWECiphertext<DataRes>,
239 lwe: &LWECiphertext<DataIn>,
240 lut: &LookUpTable,
241 brk: &BlindRotationKeyCGGI<FFT64>,
242 scratch: &mut Scratch,
243) where
244 DataRes: AsRef<[u8]> + AsMut<[u8]>,
245 DataIn: AsRef<[u8]>,
246{
247 let basek: usize = res.basek();
248
249 let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
250 let mut out_mut: GLWECiphertext<&mut [u8]> = res.to_mut();
251 let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
252
253 let cols: usize = out_mut.rank() + 1;
254
255 negate_and_mod_switch_2n(2 * lut.domain_size(), &mut lwe_2n, &lwe_ref);
256
257 let a: &[i64] = &lwe_2n[1..];
258 let b: i64 = lwe_2n[0];
259
260 out_mut.data.zero();
261
262 // Initialize out to X^{b} * LUT(X)
263 module.vec_znx_rotate(-b, &mut out_mut.data, 0, &lut.data[0], 0);
264
265 let block_size: usize = brk.block_size();
266
267 // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
268
269 let (mut acc_dft, scratch1) = scratch.tmp_fourier_glwe_ct(module, brk.basek(), out_mut.k(), out_mut.rank());
270 let (mut acc_add_dft, scratch2) = scratch1.tmp_fourier_glwe_ct(module, brk.basek(), out_mut.k(), out_mut.rank());
271 let (mut vmp_res, scratch3) = scratch2.tmp_fourier_glwe_ct(module, basek, out_mut.k(), out_mut.rank());
272 let (mut xai_plus_y, scratch4) = scratch3.tmp_scalar_znx(module, 1);
273 let (mut xai_plus_y_dft, scratch5) = scratch4.tmp_scalar_znx_dft(module, 1);
274
275 let start: Instant = Instant::now();
276 izip!(
277 a.chunks_exact(block_size),
278 brk.data.chunks_exact(block_size)
279 )
280 .for_each(|(ai, ski)| {
281 out_mut.dft(module, &mut acc_dft);
282 acc_add_dft.data.zero();
283
284 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
285 // vmp_res = DFT(acc) * BRK[i]
286 module.vmp_apply(&mut vmp_res.data, &acc_dft.data, &skii.data, scratch5);
287
288 // DFT(X^ai -1)
289 set_xai_plus_y(module, *aii, -1, &mut xai_plus_y_dft, &mut xai_plus_y);
290
291 // DFT(X^ai -1) * (DFT(acc) * BRK[i])
292 (0..cols).for_each(|i| {
293 module.svp_apply_inplace(&mut vmp_res.data, i, &xai_plus_y_dft, 0);
294 module.vec_znx_dft_add_inplace(&mut acc_add_dft.data, i, &vmp_res.data, i);
295 });
296 });
297
298 (0..cols).for_each(|i| {
299 module.vec_znx_dft_add_inplace(&mut acc_dft.data, i, &acc_add_dft.data, i);
300 });
301
302 acc_dft.idft(module, &mut out_mut, scratch5);
303 });
304 let duration: std::time::Duration = start.elapsed();
305}
306
307pub(crate) fn negate_and_mod_switch_2n(n: usize, res: &mut [i64], lwe: &LWECiphertext<&[u8]>) {
308 let basek: usize = lwe.basek();
309
310 let log2n: usize = usize::BITS as usize - (n - 1).leading_zeros() as usize + 1;
311
312 res.copy_from_slice(&lwe.data.at(0, 0));
313 res.iter_mut().for_each(|x| *x = -*x);
314
315 if basek > log2n {
316 let diff: usize = basek - log2n;
317 res.iter_mut().for_each(|x| {
318 *x = div_signed_by_pow2(x, diff);
319 })
320 } else {
321 let rem: usize = basek - (log2n % basek);
322 let size: usize = log2n.div_ceil(basek);
323 (1..size).for_each(|i| {
324 if i == size - 1 && rem != basek {
325 let k_rem: usize = basek - rem;
326 izip!(lwe.data.at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
327 *y = (*y << k_rem) + (x >> rem);
328 });
329 } else {
330 izip!(lwe.data.at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
331 *y = (*y << basek) + x;
332 });
333 }
334 })
335 }
336}
337
338#[inline(always)]
339fn div_signed_by_pow2(x: &i64, k: usize) -> i64 {
340 let bias: i64 = (1 << k) - 1;
341 (x + ((x >> 63) & bias)) >> k
342}