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