poulpy/core/src/glwe/keyswitch.rs
core/src/glwe/keyswitch.rsRaw
1use backend::hal::{
2 api::{
3 DataViewMut, ScratchAvailable, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
4 VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VmpApply, VmpApplyAdd, VmpApplyTmpBytes, ZnxInfos,
5 },
6 layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnx, VecZnxBig, VecZnxDft, VmpPMat},
7};
8
9use crate::{GLWECiphertext, GLWESwitchingKeyExec, Infos};
10
11pub trait GLWEKeyswitchFamily<B: Backend> = VecZnxDftAllocBytes
12 + VmpApplyTmpBytes
13 + VecZnxBigNormalizeTmpBytes
14 + VmpApplyTmpBytes
15 + VmpApply<B>
16 + VmpApplyAdd<B>
17 + VecZnxDftFromVecZnx<B>
18 + VecZnxDftToVecZnxBigConsume<B>
19 + VecZnxBigAddSmallInplace<B>
20 + VecZnxBigNormalize<B>;
21
22impl GLWECiphertext<Vec<u8>> {
23 pub fn keyswitch_scratch_space<B: Backend>(
24 module: &Module<B>,
25 n: usize,
26 basek: usize,
27 k_out: usize,
28 k_in: usize,
29 k_ksk: usize,
30 digits: usize,
31 rank_in: usize,
32 rank_out: usize,
33 ) -> usize
34 where
35 Module<B>: GLWEKeyswitchFamily<B>,
36 {
37 let in_size: usize = k_in.div_ceil(basek).div_ceil(digits);
38 let out_size: usize = k_out.div_ceil(basek);
39 let ksk_size: usize = k_ksk.div_ceil(basek);
40 let res_dft: usize = module.vec_znx_dft_alloc_bytes(n, rank_out + 1, ksk_size); // TODO OPTIMIZE
41 let ai_dft: usize = module.vec_znx_dft_alloc_bytes(n, rank_in, in_size);
42 let vmp: usize = module.vmp_apply_tmp_bytes(
43 n,
44 out_size,
45 in_size,
46 in_size,
47 rank_in,
48 rank_out + 1,
49 ksk_size,
50 ) + module.vec_znx_dft_alloc_bytes(n, rank_in, in_size);
51 let normalize: usize = module.vec_znx_big_normalize_tmp_bytes(n);
52 return res_dft + ((ai_dft + vmp) | normalize);
53 }
54
55 pub fn keyswitch_from_fourier_scratch_space<B: Backend>(
56 module: &Module<B>,
57 n: usize,
58 basek: usize,
59 k_out: usize,
60 k_in: usize,
61 k_ksk: usize,
62 digits: usize,
63 rank_in: usize,
64 rank_out: usize,
65 ) -> usize
66 where
67 Module<B>: GLWEKeyswitchFamily<B>,
68 {
69 Self::keyswitch_scratch_space(
70 module, n, basek, k_out, k_in, k_ksk, digits, rank_in, rank_out,
71 )
72 }
73
74 pub fn keyswitch_inplace_scratch_space<B: Backend>(
75 module: &Module<B>,
76 n: usize,
77 basek: usize,
78 k_out: usize,
79 k_ksk: usize,
80 digits: usize,
81 rank: usize,
82 ) -> usize
83 where
84 Module<B>: GLWEKeyswitchFamily<B>,
85 {
86 Self::keyswitch_scratch_space(module, n, basek, k_out, k_out, k_ksk, digits, rank, rank)
87 }
88}
89
90impl<DataSelf: DataRef> GLWECiphertext<DataSelf> {
91 pub(crate) fn assert_keyswitch<B: Backend, DataLhs, DataRhs>(
92 &self,
93 module: &Module<B>,
94 lhs: &GLWECiphertext<DataLhs>,
95 rhs: &GLWESwitchingKeyExec<DataRhs, B>,
96 scratch: &Scratch<B>,
97 ) where
98 DataLhs: DataRef,
99 DataRhs: DataRef,
100 Module<B>: GLWEKeyswitchFamily<B>,
101 Scratch<B>: ScratchAvailable,
102 {
103 let basek: usize = self.basek();
104 assert_eq!(
105 lhs.rank(),
106 rhs.rank_in(),
107 "lhs.rank(): {} != rhs.rank_in(): {}",
108 lhs.rank(),
109 rhs.rank_in()
110 );
111 assert_eq!(
112 self.rank(),
113 rhs.rank_out(),
114 "self.rank(): {} != rhs.rank_out(): {}",
115 self.rank(),
116 rhs.rank_out()
117 );
118 assert_eq!(self.basek(), basek);
119 assert_eq!(lhs.basek(), basek);
120 assert_eq!(rhs.n(), self.n());
121 assert_eq!(lhs.n(), self.n());
122 assert!(
123 scratch.available()
124 >= GLWECiphertext::keyswitch_scratch_space(
125 module,
126 self.n(),
127 self.basek(),
128 self.k(),
129 lhs.k(),
130 rhs.k(),
131 rhs.digits(),
132 rhs.rank_in(),
133 rhs.rank_out(),
134 ),
135 "scratch.available()={} < GLWECiphertext::keyswitch_scratch_space(
136 module,
137 self.basek(),
138 self.k(),
139 lhs.k(),
140 rhs.k(),
141 rhs.digits(),
142 rhs.rank_in(),
143 rhs.rank_out(),
144 )={}",
145 scratch.available(),
146 GLWECiphertext::keyswitch_scratch_space(
147 module,
148 self.n(),
149 self.basek(),
150 self.k(),
151 lhs.k(),
152 rhs.k(),
153 rhs.digits(),
154 rhs.rank_in(),
155 rhs.rank_out(),
156 )
157 );
158 }
159}
160
161impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
162 pub fn keyswitch<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
163 &mut self,
164 module: &Module<B>,
165 lhs: &GLWECiphertext<DataLhs>,
166 rhs: &GLWESwitchingKeyExec<DataRhs, B>,
167 scratch: &mut Scratch<B>,
168 ) where
169 Module<B>: GLWEKeyswitchFamily<B>,
170 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
171 {
172 #[cfg(debug_assertions)]
173 {
174 self.assert_keyswitch(module, lhs, rhs, scratch);
175 }
176 let (res_dft, scratch1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // Todo optimise
177 let res_big: VecZnxBig<_, B> = keyswitch(module, res_dft, lhs, rhs, scratch1);
178 (0..self.cols()).for_each(|i| {
179 module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch1);
180 })
181 }
182
183 pub fn keyswitch_inplace<DataRhs: DataRef, B: Backend>(
184 &mut self,
185 module: &Module<B>,
186 rhs: &GLWESwitchingKeyExec<DataRhs, B>,
187 scratch: &mut Scratch<B>,
188 ) where
189 Module<B>: GLWEKeyswitchFamily<B>,
190 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
191 {
192 unsafe {
193 let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
194 self.keyswitch(&module, &*self_ptr, rhs, scratch);
195 }
196 }
197}
198
199pub(crate) fn keyswitch<B: Backend, DataRes, DataIn, DataKey>(
200 module: &Module<B>,
201 res_dft: VecZnxDft<DataRes, B>,
202 lhs: &GLWECiphertext<DataIn>,
203 rhs: &GLWESwitchingKeyExec<DataKey, B>,
204 scratch: &mut Scratch<B>,
205) -> VecZnxBig<DataRes, B>
206where
207 DataRes: DataMut,
208 DataIn: DataRef,
209 DataKey: DataRef,
210 Module<B>: GLWEKeyswitchFamily<B>,
211 Scratch<B>: TakeVecZnxDft<B>,
212{
213 if rhs.digits() == 1 {
214 return keyswitch_vmp_one_digit(module, res_dft, &lhs.data, &rhs.key.data, scratch);
215 }
216
217 keyswitch_vmp_multiple_digits(
218 module,
219 res_dft,
220 &lhs.data,
221 &rhs.key.data,
222 rhs.digits(),
223 scratch,
224 )
225}
226
227fn keyswitch_vmp_one_digit<B: Backend, DataRes, DataIn, DataVmp>(
228 module: &Module<B>,
229 mut res_dft: VecZnxDft<DataRes, B>,
230 a: &VecZnx<DataIn>,
231 mat: &VmpPMat<DataVmp, B>,
232 scratch: &mut Scratch<B>,
233) -> VecZnxBig<DataRes, B>
234where
235 DataRes: DataMut,
236 DataIn: DataRef,
237 DataVmp: DataRef,
238 Module<B>:
239 VecZnxDftAllocBytes + VecZnxDftFromVecZnx<B> + VmpApply<B> + VecZnxDftToVecZnxBigConsume<B> + VecZnxBigAddSmallInplace<B>,
240 Scratch<B>: TakeVecZnxDft<B>,
241{
242 let cols: usize = a.cols();
243 let (mut ai_dft, scratch1) = scratch.take_vec_znx_dft(a.n(), cols - 1, a.size());
244 (0..cols - 1).for_each(|col_i| {
245 module.vec_znx_dft_from_vec_znx(1, 0, &mut ai_dft, col_i, a, col_i + 1);
246 });
247 module.vmp_apply(&mut res_dft, &ai_dft, mat, scratch1);
248 let mut res_big: VecZnxBig<DataRes, B> = module.vec_znx_dft_to_vec_znx_big_consume(res_dft);
249 module.vec_znx_big_add_small_inplace(&mut res_big, 0, a, 0);
250 res_big
251}
252
253fn keyswitch_vmp_multiple_digits<B: Backend, DataRes, DataIn, DataVmp>(
254 module: &Module<B>,
255 mut res_dft: VecZnxDft<DataRes, B>,
256 a: &VecZnx<DataIn>,
257 mat: &VmpPMat<DataVmp, B>,
258 digits: usize,
259 scratch: &mut Scratch<B>,
260) -> VecZnxBig<DataRes, B>
261where
262 DataRes: DataMut,
263 DataIn: DataRef,
264 DataVmp: DataRef,
265 Module<B>: VecZnxDftAllocBytes
266 + VecZnxDftFromVecZnx<B>
267 + VmpApply<B>
268 + VmpApplyAdd<B>
269 + VecZnxDftToVecZnxBigConsume<B>
270 + VecZnxBigAddSmallInplace<B>,
271 Scratch<B>: TakeVecZnxDft<B>,
272{
273 let cols: usize = a.cols();
274 let size: usize = a.size();
275 let (mut ai_dft, scratch1) = scratch.take_vec_znx_dft(a.n(), cols - 1, size.div_ceil(digits));
276
277 ai_dft.data_mut().fill(0);
278
279 (0..digits).for_each(|di| {
280 ai_dft.set_size((size + di) / digits);
281
282 // Small optimization for digits > 2
283 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
284 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(digits-1) * B}.
285 // As such we can ignore the last digits-2 limbs safely of the sum of vmp products.
286 // It is possible to further ignore the last digits-1 limbs, but this introduce
287 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
288 // noise is kept with respect to the ideal functionality.
289 res_dft.set_size(mat.size() - ((digits - di) as isize - 2).max(0) as usize);
290
291 (0..cols - 1).for_each(|col_i| {
292 module.vec_znx_dft_from_vec_znx(digits, digits - di - 1, &mut ai_dft, col_i, a, col_i + 1);
293 });
294
295 if di == 0 {
296 module.vmp_apply(&mut res_dft, &ai_dft, mat, scratch1);
297 } else {
298 module.vmp_apply_add(&mut res_dft, &ai_dft, mat, di, scratch1);
299 }
300 });
301
302 res_dft.set_size(res_dft.max_size());
303 let mut res_big: VecZnxBig<DataRes, B> = module.vec_znx_dft_to_vec_znx_big_consume(res_dft);
304 module.vec_znx_big_add_small_inplace(&mut res_big, 0, a, 0);
305 res_big
306}