poulpy/core/src/blind_rotation/ccgi.rs
core/src/blind_rotation/ccgi.rsRaw
1use std::time::Instant;
2
3use backend::{
4 FFT64, MatZnxDftOps, MatZnxDftScratch, Module, ScalarZnxAlloc, ScalarZnxDftAlloc, ScalarZnxDftOps, Scratch, VecZnxDftOps,
5 VecZnxOps, ZnxView, ZnxViewMut, ZnxZero,
6};
7use itertools::izip;
8
9use crate::{
10 FourierGLWECiphertext, GGSWCiphertext, GLWECiphertext, GLWECiphertextToMut, Infos, LWECiphertext, 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 basek: usize,
18 k_lut: usize,
19 k_brk: usize,
20 rows: usize,
21 rank: usize,
22) -> usize {
23 let size = k_brk.div_ceil(basek);
24 GGSWCiphertext::<Vec<u8>, FFT64>::bytes_of(module, basek, k_brk, rows, 1, rank)
25 + (module.mat_znx_dft_mul_x_pow_minus_one_scratch_space(size, rank + 1)
26 | GLWECiphertext::external_product_inplace_scratch_space(module, basek, k_lut, k_brk, 1, rank))
27}
28
29pub fn cggi_blind_rotate<DataRes, DataIn>(
30 module: &Module<FFT64>,
31 res: &mut GLWECiphertext<DataRes>,
32 lwe: &LWECiphertext<DataIn>,
33 lut: &LookUpTable,
34 brk: &BlindRotationKeyCGGI<FFT64>,
35 scratch: &mut Scratch,
36) where
37 DataRes: AsRef<[u8]> + AsMut<[u8]>,
38 DataIn: AsRef<[u8]>,
39{
40 if lut.data.len() > 1 {
41 cggi_blind_rotate_block_binary_exnteded(module, res, lwe, lut, brk, scratch);
42 } else if brk.block_size() > 1 {
43 cggi_blind_rotate_block_binary(module, res, lwe, lut, brk, scratch);
44 } else {
45 todo!("implement this case")
46 }
47}
48
49pub(crate) fn cggi_blind_rotate_block_binary_exnteded<DataRes, DataIn>(
50 module: &Module<FFT64>,
51 res: &mut GLWECiphertext<DataRes>,
52 lwe: &LWECiphertext<DataIn>,
53 lut: &LookUpTable,
54 brk: &BlindRotationKeyCGGI<FFT64>,
55 scratch: &mut Scratch,
56) where
57 DataRes: AsRef<[u8]> + AsMut<[u8]>,
58 DataIn: AsRef<[u8]>,
59{
60 let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
61 let mut out_mut: GLWECiphertext<&mut [u8]> = res.to_mut();
62 let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
63 let basek: usize = out_mut.basek();
64
65 let cols: usize = out_mut.rank() + 1;
66
67 mod_switch_2n(
68 2 * module.n() * lut.extension_factor(),
69 &mut lwe_2n,
70 &lwe_ref,
71 );
72
73 let extension_factor: i64 = lut.extension_factor() as i64;
74
75 let mut acc: Vec<GLWECiphertext<Vec<u8>>> = Vec::with_capacity(lut.extension_factor());
76
77 for _ in 0..extension_factor {
78 acc.push(GLWECiphertext::alloc(
79 module,
80 basek,
81 out_mut.k(),
82 out_mut.rank(),
83 ));
84 }
85
86 let a: &[i64] = &lwe_2n[1..];
87 let b: i64 = lwe_2n[0];
88
89 let b_inner: i64 = b / extension_factor;
90 let b_outer: i64 = b % extension_factor;
91
92 for (i, j) in (0..b_outer).zip(extension_factor - b_outer..extension_factor) {
93 module.vec_znx_rotate(
94 b_inner + 1,
95 &mut acc[j as usize].data,
96 0,
97 &lut.data[i as usize],
98 0,
99 );
100 }
101 for (i, j) in (b_outer..extension_factor).zip(0..extension_factor - b_outer) {
102 module.vec_znx_rotate(
103 b_inner,
104 &mut acc[j as usize].data,
105 0,
106 &lut.data[i as usize],
107 0,
108 );
109 }
110
111 let block_size: usize = brk.block_size();
112
113 let mut acc_dft: Vec<FourierGLWECiphertext<Vec<u8>, FFT64>> = Vec::with_capacity(lut.extension_factor());
114
115 for _ in 0..extension_factor {
116 acc_dft.push(FourierGLWECiphertext::alloc(
117 module,
118 basek,
119 out_mut.k(),
120 out_mut.rank(),
121 ));
122 }
123
124 let mut vmp_res: Vec<FourierGLWECiphertext<Vec<u8>, FFT64>> = Vec::with_capacity(lut.extension_factor());
125
126 for _ in 0..extension_factor {
127 vmp_res.push(FourierGLWECiphertext::alloc(
128 module,
129 basek,
130 out_mut.k(),
131 out_mut.rank(),
132 ));
133 }
134
135 let mut acc_add_dft: Vec<FourierGLWECiphertext<Vec<u8>, FFT64>> = Vec::with_capacity(lut.extension_factor());
136
137 for _ in 0..extension_factor {
138 acc_add_dft.push(FourierGLWECiphertext::alloc(
139 module,
140 basek,
141 out_mut.k(),
142 out_mut.rank(),
143 ));
144 }
145
146 let mut xai_minus_one: backend::ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
147 let mut xai_minus_one_dft: backend::ScalarZnxDft<Vec<u8>, FFT64> = module.new_scalar_znx_dft(1);
148
149 izip!(
150 a.chunks_exact(block_size),
151 brk.data.chunks_exact(block_size)
152 )
153 .enumerate()
154 .for_each(|(i, (ai, ski))| {
155 (0..lut.extension_factor()).for_each(|i| {
156 acc[i].dft(module, &mut acc_dft[i]);
157 acc_add_dft[i].data.zero();
158 });
159
160 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
161 let aii_inner: i64 = aii / extension_factor;
162 let aii_outer: i64 = aii % extension_factor;
163
164 // vmp_res = DFT(acc) * BRK[i]
165 (0..lut.extension_factor()).for_each(|i| {
166 module.vmp_apply(&mut vmp_res[i].data, &acc_dft[i].data, &skii.data, scratch);
167 });
168
169 if aii_outer == 0 {
170 xai_minus_one.zero();
171 xai_minus_one.at_mut(0, 0)[0] = 1;
172 module.vec_znx_rotate_inplace(aii_inner, &mut xai_minus_one, 0);
173 xai_minus_one.at_mut(0, 0)[0] -= 1;
174 module.svp_prepare(&mut xai_minus_one_dft, 0, &xai_minus_one, 0);
175
176 (0..lut.extension_factor()).for_each(|j| {
177 (0..cols).for_each(|i| {
178 module.svp_apply_inplace(&mut vmp_res[j].data, i, &xai_minus_one_dft, 0);
179 module.vec_znx_dft_add_inplace(&mut acc_add_dft[j].data, i, &vmp_res[j].data, i);
180 });
181 })
182 } else {
183 xai_minus_one.zero();
184 xai_minus_one.at_mut(0, 0)[0] = 1;
185 module.vec_znx_rotate_inplace(aii_inner + 1, &mut xai_minus_one, 0);
186 xai_minus_one.at_mut(0, 0)[0] -= 1;
187 module.svp_prepare(&mut xai_minus_one_dft, 0, &xai_minus_one, 0);
188
189 for (i, j) in (0..aii_outer).zip(extension_factor - aii_outer..extension_factor) {
190 module.vec_znx_rotate(
191 b_inner + 1,
192 &mut acc[j as usize].data,
193 0,
194 &lut.data[i as usize],
195 0,
196 );
197
198 (0..cols).for_each(|k| {
199 module.svp_apply_inplace(&mut vmp_res[i as usize].data, k, &xai_minus_one_dft, 0);
200 module.vec_znx_dft_add_inplace(
201 &mut acc_add_dft[j as usize].data,
202 k,
203 &vmp_res[i as usize].data,
204 k,
205 );
206 });
207 }
208
209 xai_minus_one.zero();
210 xai_minus_one.at_mut(0, 0)[0] = 1;
211 module.vec_znx_rotate_inplace(aii_inner, &mut xai_minus_one, 0);
212 xai_minus_one.at_mut(0, 0)[0] -= 1;
213 module.svp_prepare(&mut xai_minus_one_dft, 0, &xai_minus_one, 0);
214
215 for (i, j) in (aii_outer..extension_factor).zip(0..extension_factor - aii_outer) {
216 module.vec_znx_rotate(
217 b_inner,
218 &mut acc[j as usize].data,
219 0,
220 &lut.data[i as usize],
221 0,
222 );
223
224 (0..cols).for_each(|k| {
225 module.svp_apply_inplace(&mut vmp_res[i as usize].data, k, &xai_minus_one_dft, 0);
226 module.vec_znx_dft_add_inplace(
227 &mut acc_add_dft[j as usize].data,
228 k,
229 &vmp_res[i as usize].data,
230 k,
231 );
232 });
233 }
234 }
235 });
236
237 if i == lwe.n() - block_size {
238 (0..cols).for_each(|i| {
239 module.vec_znx_dft_add_inplace(&mut acc_dft[0].data, i, &acc_add_dft[0].data, i);
240 });
241 acc_dft[0].idft(module, &mut out_mut, scratch);
242 } else {
243 (0..lut.extension_factor()).for_each(|j| {
244 (0..cols).for_each(|i| {
245 module.vec_znx_dft_add_inplace(&mut acc_dft[j].data, i, &acc_add_dft[j].data, i);
246 });
247
248 acc_dft[j].idft(module, &mut acc[j], scratch);
249 })
250 }
251 });
252}
253
254pub(crate) fn cggi_blind_rotate_block_binary<DataRes, DataIn>(
255 module: &Module<FFT64>,
256 res: &mut GLWECiphertext<DataRes>,
257 lwe: &LWECiphertext<DataIn>,
258 lut: &LookUpTable,
259 brk: &BlindRotationKeyCGGI<FFT64>,
260 scratch: &mut Scratch,
261) where
262 DataRes: AsRef<[u8]> + AsMut<[u8]>,
263 DataIn: AsRef<[u8]>,
264{
265 let basek: usize = res.basek();
266
267 let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
268 let mut out_mut: GLWECiphertext<&mut [u8]> = res.to_mut();
269 let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
270
271 let cols: usize = out_mut.rank() + 1;
272
273 mod_switch_2n(2 * module.n(), &mut lwe_2n, &lwe_ref);
274
275 let a: &[i64] = &lwe_2n[1..];
276 let b: i64 = lwe_2n[0];
277
278 out_mut.data.zero();
279
280 // Initialize out to X^{b} * LUT(X)
281 module.vec_znx_rotate(b, &mut out_mut.data, 0, &lut.data[0], 0);
282
283 let block_size: usize = brk.block_size();
284
285 // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
286
287 let (mut acc_dft, scratch1) = scratch.tmp_glwe_fourier(module, brk.basek(), out_mut.k(), out_mut.rank());
288 let (mut acc_add_dft, scratch2) = scratch1.tmp_glwe_fourier(module, brk.basek(), out_mut.k(), out_mut.rank());
289 let (mut vmp_res, scratch3) = scratch2.tmp_glwe_fourier(module, basek, out_mut.k(), out_mut.rank());
290 let (mut xai_minus_one, scratch4) = scratch3.tmp_scalar_znx(module, 1);
291 let (mut xai_minus_one_dft, scratch5) = scratch4.tmp_scalar_znx_dft(module, 1);
292
293 let start: Instant = Instant::now();
294 izip!(
295 a.chunks_exact(block_size),
296 brk.data.chunks_exact(block_size)
297 )
298 .for_each(|(ai, ski)| {
299 out_mut.dft(module, &mut acc_dft);
300 acc_add_dft.data.zero();
301
302 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
303 // vmp_res = DFT(acc) * BRK[i]
304 module.vmp_apply(&mut vmp_res.data, &acc_dft.data, &skii.data, scratch5);
305
306 // DFT(X^ai -1)
307 xai_minus_one.zero();
308 xai_minus_one.at_mut(0, 0)[0] = 1;
309 module.vec_znx_rotate_inplace(*aii, &mut xai_minus_one, 0);
310 xai_minus_one.at_mut(0, 0)[0] -= 1;
311 module.svp_prepare(&mut xai_minus_one_dft, 0, &xai_minus_one, 0);
312
313 // DFT(X^ai -1) * (DFT(acc) * BRK[i])
314 (0..cols).for_each(|i| {
315 module.svp_apply_inplace(&mut vmp_res.data, i, &xai_minus_one_dft, 0);
316 module.vec_znx_dft_add_inplace(&mut acc_add_dft.data, i, &vmp_res.data, i);
317 });
318 });
319
320 (0..cols).for_each(|i| {
321 module.vec_znx_dft_add_inplace(&mut acc_dft.data, i, &acc_add_dft.data, i);
322 });
323
324 acc_dft.idft(module, &mut out_mut, scratch5);
325 });
326 let duration: std::time::Duration = start.elapsed();
327 println!("external products: {} us", duration.as_micros());
328}
329
330pub(crate) fn mod_switch_2n(n: usize, res: &mut [i64], lwe: &LWECiphertext<&[u8]>) {
331 let basek: usize = lwe.basek();
332
333 let log2n: usize = usize::BITS as usize - (n - 1).leading_zeros() as usize + 1;
334
335 res.copy_from_slice(&lwe.data.at(0, 0));
336
337 if basek > log2n {
338 let diff: usize = basek - log2n;
339 res.iter_mut().for_each(|x| {
340 *x = div_signed_by_pow2(x, diff);
341 })
342 } else {
343 let rem: usize = basek - (log2n % basek);
344 let size: usize = log2n.div_ceil(basek);
345 (1..size).for_each(|i| {
346 if i == size - 1 && rem != basek {
347 let k_rem: usize = basek - rem;
348 izip!(lwe.data.at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
349 *y = (*y << k_rem) + (x >> rem);
350 });
351 } else {
352 izip!(lwe.data.at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
353 *y = (*y << basek) + x;
354 });
355 }
356 })
357 }
358}
359
360#[inline(always)]
361fn div_signed_by_pow2(x: &i64, k: usize) -> i64 {
362 let bias: i64 = (1 << k) - 1;
363 (x + ((x >> 63) & bias)) >> k
364}