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