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!(lhs.rank(), rhs.rank_in());
91 assert_eq!(self.rank(), rhs.rank_out());
92 assert_eq!(self.basek(), basek);
93 assert_eq!(lhs.basek(), basek);
94 assert_eq!(rhs.n(), module.n());
95 assert_eq!(self.n(), module.n());
96 assert_eq!(lhs.n(), module.n());
97 assert!(
98 scratch.available()
99 >= GLWECiphertext::keyswitch_scratch_space(
100 module,
101 self.basek(),
102 self.k(),
103 lhs.k(),
104 rhs.k(),
105 rhs.digits(),
106 rhs.rank_in(),
107 rhs.rank_out(),
108 )
109 );
110 }
111
112 let cols_in: usize = rhs.rank_in();
113 let cols_out: usize = rhs.rank_out() + 1;
114 let digits: usize = rhs.digits();
115
116 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols_out, rhs.size()); // Todo optimise
117 let (mut ai_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols_in, (lhs.size() + digits - 1) / digits);
118 ai_dft.zero();
119 {
120 (0..digits).for_each(|di| {
121 ai_dft.set_size((lhs.size() + di) / digits);
122
123 // Small optimization for digits > 2
124 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
125 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(digits-1) * B}.
126 // As such we can ignore the last digits-2 limbs safely of the sum of vmp products.
127 // It is possible to further ignore the last digits-1 limbs, but this introduce
128 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
129 // noise is kept with respect to the ideal functionality.
130 res_dft.set_size(rhs.size() - ((digits - di) as isize - 2).max(0) as usize);
131
132 (0..cols_in).for_each(|col_i| {
133 module.vec_znx_dft(
134 digits,
135 digits - di - 1,
136 &mut ai_dft,
137 col_i,
138 &lhs.data,
139 col_i + 1,
140 );
141 });
142
143 if di == 0 {
144 module.vmp_apply(&mut res_dft, &ai_dft, &rhs.0.data, scratch2);
145 } else {
146 module.vmp_apply_add(&mut res_dft, &ai_dft, &rhs.0.data, di, scratch2);
147 }
148 });
149 }
150
151 let mut res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
152
153 module.vec_znx_big_add_small_inplace(&mut res_big, 0, &lhs.data, 0);
154
155 (0..cols_out).for_each(|i| {
156 if apply_auto != 0 {
157 module.vec_znx_big_automorphism_inplace(apply_auto, &mut res_big, i);
158 }
159
160 match OP {
161 1 => module.vec_znx_big_add_small_inplace(&mut res_big, i, &lhs.data, i),
162 2 => module.vec_znx_big_sub_small_a_inplace(&mut res_big, i, &lhs.data, i),
163 3 => module.vec_znx_big_sub_small_b_inplace(&mut res_big, i, &lhs.data, i),
164 _ => {}
165 }
166 module.vec_znx_big_normalize(basek, &mut self.data, i, &res_big, i, scratch1);
167 });
168 }
169
170 pub(crate) fn keyswitch_from_fourier<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
171 &mut self,
172 module: &Module<FFT64>,
173 lhs: &FourierGLWECiphertext<DataLhs, FFT64>,
174 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
175 scratch: &mut Scratch,
176 ) {
177 let basek: usize = self.basek();
178
179 #[cfg(debug_assertions)]
180 {
181 assert_eq!(lhs.rank(), rhs.rank_in());
182 assert_eq!(self.rank(), rhs.rank_out());
183 assert_eq!(self.basek(), basek);
184 assert_eq!(lhs.basek(), basek);
185 assert_eq!(rhs.n(), module.n());
186 assert_eq!(self.n(), module.n());
187 assert_eq!(lhs.n(), module.n());
188 assert!(
189 scratch.available()
190 >= GLWECiphertext::keyswitch_from_fourier_scratch_space(
191 module,
192 self.basek(),
193 self.k(),
194 lhs.k(),
195 rhs.k(),
196 rhs.digits(),
197 rhs.rank_in(),
198 rhs.rank_out(),
199 )
200 );
201 }
202
203 let cols_in: usize = rhs.rank_in();
204 let cols_out: usize = rhs.rank_out() + 1;
205
206 // Buffer of the result of VMP in DFT
207 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols_out, rhs.size()); // Todo optimise
208
209 {
210 let digits = rhs.digits();
211
212 (0..digits).for_each(|di| {
213 // (lhs.size() + di) / digits = (a - (digit - di - 1) + digit - 1) / digits
214 let (mut ai_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols_in, (lhs.size() + di) / digits);
215
216 (0..cols_in).for_each(|col_i| {
217 module.vec_znx_dft_copy(
218 digits,
219 digits - 1 - di,
220 &mut ai_dft,
221 col_i,
222 &lhs.data,
223 col_i + 1,
224 );
225 });
226
227 if di == 0 {
228 module.vmp_apply(&mut res_dft, &ai_dft, &rhs.0.data, scratch2);
229 } else {
230 module.vmp_apply_add(&mut res_dft, &ai_dft, &rhs.0.data, di, scratch2);
231 }
232 });
233 }
234
235 module.vec_znx_dft_add_inplace(&mut res_dft, 0, &lhs.data, 0);
236
237 // Switches result of VMP outside of DFT
238 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume::<&mut [u8]>(res_dft);
239
240 (0..cols_out).for_each(|i| {
241 module.vec_znx_big_normalize(basek, &mut self.data, i, &res_big, i, scratch1);
242 });
243 }
244}