poulpy/core/src/glwe/keyswitch.rs
core/src/glwe/keyswitch.rsRaw
1use backend::{
2 FFT64, MatZnxDftOps, MatZnxDftScratch, Module, Scratch, VecZnxBig, VecZnxBigOps, VecZnxBigScratch, VecZnxDftAlloc,
3 VecZnxDftOps, ZnxZero,
4};
5
6use crate::{FourierGLWECiphertext, GLWECiphertext, GLWESwitchingKey, Infos};
7
8impl GLWECiphertext<Vec<u8>> {
9 pub fn keyswitch_scratch_space(
10 module: &Module<FFT64>,
11 basek: usize,
12 k_out: usize,
13 k_in: usize,
14 k_ksk: usize,
15 digits: usize,
16 rank_in: usize,
17 rank_out: usize,
18 ) -> usize {
19 let res_dft: usize = FourierGLWECiphertext::bytes_of(module, basek, k_out, rank_out + 1);
20 let in_size: usize = k_in.div_ceil(basek).div_ceil(digits);
21 let out_size: usize = k_out.div_ceil(basek);
22 let ksk_size: usize = k_ksk.div_ceil(basek);
23 let ai_dft: usize = module.bytes_of_vec_znx_dft(rank_in, in_size);
24 let vmp: usize = module.vmp_apply_tmp_bytes(out_size, in_size, in_size, rank_in, rank_out + 1, ksk_size)
25 + module.bytes_of_vec_znx_dft(rank_in, in_size);
26 let normalize: usize = module.vec_znx_big_normalize_tmp_bytes();
27 return res_dft + ((ai_dft + vmp) | normalize);
28 }
29
30 pub fn keyswitch_from_fourier_scratch_space(
31 module: &Module<FFT64>,
32 basek: usize,
33 k_out: usize,
34 k_in: usize,
35 k_ksk: usize,
36 digits: usize,
37 rank_in: usize,
38 rank_out: usize,
39 ) -> usize {
40 Self::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, digits, rank_in, rank_out)
41 }
42
43 pub fn keyswitch_inplace_scratch_space(
44 module: &Module<FFT64>,
45 basek: usize,
46 k_out: usize,
47 k_ksk: usize,
48 digits: usize,
49 rank: usize,
50 ) -> usize {
51 Self::keyswitch_scratch_space(module, basek, k_out, k_out, k_ksk, digits, rank, rank)
52 }
53}
54
55impl<DataSelf: AsRef<[u8]> + AsMut<[u8]>> GLWECiphertext<DataSelf> {
56 pub fn keyswitch<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
57 &mut self,
58 module: &Module<FFT64>,
59 lhs: &GLWECiphertext<DataLhs>,
60 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
61 scratch: &mut Scratch,
62 ) {
63 Self::keyswitch_private::<_, _, 0>(self, 0, module, lhs, rhs, scratch);
64 }
65
66 pub fn keyswitch_inplace<DataRhs: AsRef<[u8]>>(
67 &mut self,
68 module: &Module<FFT64>,
69 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
70 scratch: &mut Scratch,
71 ) {
72 unsafe {
73 let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
74 self.keyswitch(&module, &*self_ptr, rhs, scratch);
75 }
76 }
77
78 pub(crate) fn keyswitch_private<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>, const OP: u8>(
79 &mut self,
80 apply_auto: i64,
81 module: &Module<FFT64>,
82 lhs: &GLWECiphertext<DataLhs>,
83 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
84 scratch: &mut Scratch,
85 ) {
86 let basek: usize = self.basek();
87
88 #[cfg(debug_assertions)]
89 {
90 assert_eq!(
91 lhs.rank(),
92 rhs.rank_in(),
93 "lhs.rank(): {} != rhs.rank_in(): {}",
94 lhs.rank(),
95 rhs.rank_in()
96 );
97 assert_eq!(
98 self.rank(),
99 rhs.rank_out(),
100 "self.rank(): {} != rhs.rank_out(): {}",
101 self.rank(),
102 rhs.rank_out()
103 );
104 assert_eq!(self.basek(), basek);
105 assert_eq!(lhs.basek(), basek);
106 assert_eq!(rhs.n(), module.n());
107 assert_eq!(self.n(), module.n());
108 assert_eq!(lhs.n(), module.n());
109 assert!(
110 scratch.available()
111 >= GLWECiphertext::keyswitch_scratch_space(
112 module,
113 self.basek(),
114 self.k(),
115 lhs.k(),
116 rhs.k(),
117 rhs.digits(),
118 rhs.rank_in(),
119 rhs.rank_out(),
120 )
121 );
122 }
123
124 let cols_in: usize = rhs.rank_in();
125 let cols_out: usize = rhs.rank_out() + 1;
126 let digits: usize = rhs.digits();
127
128 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols_out, rhs.size()); // Todo optimise
129 let (mut ai_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols_in, (lhs.size() + digits - 1) / digits);
130 ai_dft.zero();
131 {
132 (0..digits).for_each(|di| {
133 ai_dft.set_size((lhs.size() + di) / digits);
134
135 // Small optimization for digits > 2
136 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
137 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(digits-1) * B}.
138 // As such we can ignore the last digits-2 limbs safely of the sum of vmp products.
139 // It is possible to further ignore the last digits-1 limbs, but this introduce
140 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
141 // noise is kept with respect to the ideal functionality.
142 res_dft.set_size(rhs.size() - ((digits - di) as isize - 2).max(0) as usize);
143
144 (0..cols_in).for_each(|col_i| {
145 module.vec_znx_dft(
146 digits,
147 digits - di - 1,
148 &mut ai_dft,
149 col_i,
150 &lhs.data,
151 col_i + 1,
152 );
153 });
154
155 if di == 0 {
156 module.vmp_apply(&mut res_dft, &ai_dft, &rhs.key.data, scratch2);
157 } else {
158 module.vmp_apply_add(&mut res_dft, &ai_dft, &rhs.key.data, di, scratch2);
159 }
160 });
161 }
162
163 let mut res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
164
165 module.vec_znx_big_add_small_inplace(&mut res_big, 0, &lhs.data, 0);
166
167 (0..cols_out).for_each(|i| {
168 if apply_auto != 0 {
169 module.vec_znx_big_automorphism_inplace(apply_auto, &mut res_big, i);
170 }
171
172 match OP {
173 1 => module.vec_znx_big_add_small_inplace(&mut res_big, i, &lhs.data, i),
174 2 => module.vec_znx_big_sub_small_a_inplace(&mut res_big, i, &lhs.data, i),
175 3 => module.vec_znx_big_sub_small_b_inplace(&mut res_big, i, &lhs.data, i),
176 _ => {}
177 }
178 module.vec_znx_big_normalize(basek, &mut self.data, i, &res_big, i, scratch1);
179 });
180 }
181
182 pub(crate) fn keyswitch_from_fourier<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
183 &mut self,
184 module: &Module<FFT64>,
185 lhs: &FourierGLWECiphertext<DataLhs, FFT64>,
186 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
187 scratch: &mut Scratch,
188 ) {
189 let basek: usize = self.basek();
190
191 #[cfg(debug_assertions)]
192 {
193 assert_eq!(lhs.rank(), rhs.rank_in());
194 assert_eq!(self.rank(), rhs.rank_out());
195 assert_eq!(self.basek(), basek);
196 assert_eq!(lhs.basek(), basek);
197 assert_eq!(rhs.n(), module.n());
198 assert_eq!(self.n(), module.n());
199 assert_eq!(lhs.n(), module.n());
200 assert!(
201 scratch.available()
202 >= GLWECiphertext::keyswitch_from_fourier_scratch_space(
203 module,
204 self.basek(),
205 self.k(),
206 lhs.k(),
207 rhs.k(),
208 rhs.digits(),
209 rhs.rank_in(),
210 rhs.rank_out(),
211 )
212 );
213 }
214
215 let cols_in: usize = rhs.rank_in();
216 let cols_out: usize = rhs.rank_out() + 1;
217
218 // Buffer of the result of VMP in DFT
219 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols_out, rhs.size()); // Todo optimise
220
221 {
222 let digits = rhs.digits();
223
224 (0..digits).for_each(|di| {
225 // (lhs.size() + di) / digits = (a - (digit - di - 1) + digit - 1) / digits
226 let (mut ai_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols_in, (lhs.size() + di) / digits);
227
228 (0..cols_in).for_each(|col_i| {
229 module.vec_znx_dft_copy(
230 digits,
231 digits - 1 - di,
232 &mut ai_dft,
233 col_i,
234 &lhs.data,
235 col_i + 1,
236 );
237 });
238
239 if di == 0 {
240 module.vmp_apply(&mut res_dft, &ai_dft, &rhs.key.data, scratch2);
241 } else {
242 module.vmp_apply_add(&mut res_dft, &ai_dft, &rhs.key.data, di, scratch2);
243 }
244 });
245 }
246
247 module.vec_znx_dft_add_inplace(&mut res_dft, 0, &lhs.data, 0);
248
249 // Switches result of VMP outside of DFT
250 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume::<&mut [u8]>(res_dft);
251
252 (0..cols_out).for_each(|i| {
253 module.vec_znx_big_normalize(basek, &mut self.data, i, &res_big, i, scratch1);
254 });
255 }
256}