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