poulpy/core/src/ggsw/keyswitch.rs
core/src/ggsw/keyswitch.rsRaw
1use backend::hal::{
2 api::{
3 ScratchAvailable, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxBigAllocBytes, VecZnxCopy, VecZnxDftAddInplace,
4 VecZnxDftCopy, VecZnxDftToVecZnxBigTmpA, VecZnxNormalizeTmpBytes, ZnxInfos,
5 },
6 layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnx, VmpPMat},
7};
8
9use crate::{
10 GGLWECiphertext, GGSWCiphertext, GLWECiphertext, GLWEKeyswitchFamily, GLWEOps, GLWESwitchingKeyExec, GLWETensorKeyExec, Infos,
11};
12
13pub trait GGSWKeySwitchFamily<B> =
14 GLWEKeyswitchFamily<B> + VecZnxBigAllocBytes + VecZnxDftCopy<B> + VecZnxDftAddInplace<B> + VecZnxDftToVecZnxBigTmpA<B>;
15
16impl GGSWCiphertext<Vec<u8>> {
17 pub(crate) fn expand_row_scratch_space<B: Backend>(
18 module: &Module<B>,
19 n: usize,
20 basek: usize,
21 self_k: usize,
22 k_tsk: usize,
23 digits: usize,
24 rank: usize,
25 ) -> usize
26 where
27 Module<B>: GGSWKeySwitchFamily<B> + VecZnxNormalizeTmpBytes,
28 {
29 let tsk_size: usize = k_tsk.div_ceil(basek);
30 let self_size_out: usize = self_k.div_ceil(basek);
31 let self_size_in: usize = self_size_out.div_ceil(digits);
32
33 let tmp_dft_i: usize = module.vec_znx_dft_alloc_bytes(n, rank + 1, tsk_size);
34 let tmp_a: usize = module.vec_znx_dft_alloc_bytes(n, 1, self_size_in);
35 let vmp: usize = module.vmp_apply_tmp_bytes(
36 n,
37 self_size_out,
38 self_size_in,
39 self_size_in,
40 rank,
41 rank,
42 tsk_size,
43 );
44 let tmp_idft: usize = module.vec_znx_big_alloc_bytes(n, 1, tsk_size);
45 let norm: usize = module.vec_znx_normalize_tmp_bytes(n);
46 tmp_dft_i + ((tmp_a + vmp) | (tmp_idft + norm))
47 }
48
49 pub fn keyswitch_scratch_space<B: Backend>(
50 module: &Module<B>,
51 n: usize,
52 basek: usize,
53 k_out: usize,
54 k_in: usize,
55 k_ksk: usize,
56 digits_ksk: usize,
57 k_tsk: usize,
58 digits_tsk: usize,
59 rank: usize,
60 ) -> usize
61 where
62 Module<B>: GLWEKeyswitchFamily<B> + GGSWKeySwitchFamily<B> + VecZnxNormalizeTmpBytes,
63 {
64 let out_size: usize = k_out.div_ceil(basek);
65 let res_znx: usize = VecZnx::alloc_bytes(n, rank + 1, out_size);
66 let ci_dft: usize = module.vec_znx_dft_alloc_bytes(n, rank + 1, out_size);
67 let ks: usize = GLWECiphertext::keyswitch_scratch_space(module, n, basek, k_out, k_in, k_ksk, digits_ksk, rank, rank);
68 let expand_rows: usize = GGSWCiphertext::expand_row_scratch_space(module, n, basek, k_out, k_tsk, digits_tsk, rank);
69 let res_dft: usize = module.vec_znx_dft_alloc_bytes(n, rank + 1, out_size);
70 res_znx + ci_dft + (ks | expand_rows | res_dft)
71 }
72
73 pub fn keyswitch_inplace_scratch_space<B: Backend>(
74 module: &Module<B>,
75 n: usize,
76 basek: usize,
77 k_out: usize,
78 k_ksk: usize,
79 digits_ksk: usize,
80 k_tsk: usize,
81 digits_tsk: usize,
82 rank: usize,
83 ) -> usize
84 where
85 Module<B>: GLWEKeyswitchFamily<B> + GGSWKeySwitchFamily<B> + VecZnxNormalizeTmpBytes,
86 {
87 GGSWCiphertext::keyswitch_scratch_space(
88 module, n, basek, k_out, k_out, k_ksk, digits_ksk, k_tsk, digits_tsk, rank,
89 )
90 }
91}
92
93impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
94 pub fn from_gglwe<DataA, DataTsk, B: Backend>(
95 &mut self,
96 module: &Module<B>,
97 a: &GGLWECiphertext<DataA>,
98 tsk: &GLWETensorKeyExec<DataTsk, B>,
99 scratch: &mut Scratch<B>,
100 ) where
101 DataA: DataRef,
102 DataTsk: DataRef,
103 Module<B>: GGSWKeySwitchFamily<B> + VecZnxNormalizeTmpBytes + VecZnxCopy,
104 Scratch<B>: TakeVecZnxDft<B> + TakeVecZnxBig<B> + ScratchAvailable + TakeVecZnx,
105 {
106 #[cfg(debug_assertions)]
107 {
108 assert_eq!(self.rank(), a.rank());
109 assert_eq!(self.rows(), a.rows());
110 assert_eq!(self.n(), module.n());
111 assert_eq!(a.n(), module.n());
112 assert_eq!(tsk.n(), module.n());
113 }
114 (0..self.rows()).for_each(|row_i| {
115 self.at_mut(row_i, 0).copy(module, &a.at(row_i, 0));
116 });
117 self.expand_row(module, tsk, scratch);
118 }
119
120 pub fn keyswitch<DataLhs: DataRef, DataKsk: DataRef, DataTsk: DataRef, B: Backend>(
121 &mut self,
122 module: &Module<B>,
123 lhs: &GGSWCiphertext<DataLhs>,
124 ksk: &GLWESwitchingKeyExec<DataKsk, B>,
125 tsk: &GLWETensorKeyExec<DataTsk, B>,
126 scratch: &mut Scratch<B>,
127 ) where
128 Module<B>: GLWEKeyswitchFamily<B> + GGSWKeySwitchFamily<B> + VecZnxNormalizeTmpBytes,
129 Scratch<B>: TakeVecZnxDft<B> + TakeVecZnxBig<B> + ScratchAvailable,
130 {
131 self.keyswitch_internal(module, lhs, ksk, scratch);
132 self.expand_row(module, tsk, scratch);
133 }
134
135 pub fn keyswitch_inplace<DataKsk: DataRef, DataTsk: DataRef, B: Backend>(
136 &mut self,
137 module: &Module<B>,
138 ksk: &GLWESwitchingKeyExec<DataKsk, B>,
139 tsk: &GLWETensorKeyExec<DataTsk, B>,
140 scratch: &mut Scratch<B>,
141 ) where
142 Module<B>: GLWEKeyswitchFamily<B> + GGSWKeySwitchFamily<B> + VecZnxNormalizeTmpBytes,
143 Scratch<B>: TakeVecZnxDft<B> + TakeVecZnxBig<B> + ScratchAvailable,
144 {
145 unsafe {
146 let self_ptr: *mut GGSWCiphertext<DataSelf> = self as *mut GGSWCiphertext<DataSelf>;
147 self.keyswitch(module, &*self_ptr, ksk, tsk, scratch);
148 }
149 }
150
151 pub fn expand_row<DataTsk: DataRef, B: Backend>(
152 &mut self,
153 module: &Module<B>,
154 tsk: &GLWETensorKeyExec<DataTsk, B>,
155 scratch: &mut Scratch<B>,
156 ) where
157 Module<B>: GGSWKeySwitchFamily<B> + VecZnxNormalizeTmpBytes,
158 Scratch<B>: TakeVecZnxDft<B> + TakeVecZnxBig<B> + ScratchAvailable,
159 {
160 assert!(
161 scratch.available()
162 >= GGSWCiphertext::expand_row_scratch_space(
163 module,
164 self.n(),
165 self.basek(),
166 self.k(),
167 tsk.k(),
168 tsk.digits(),
169 tsk.rank()
170 )
171 );
172
173 let n: usize = self.n();
174 let rank: usize = self.rank();
175 let cols: usize = rank + 1;
176
177 // Keyswitch the j-th row of the col 0
178 (0..self.rows()).for_each(|row_i| {
179 // Pre-compute DFT of (a0, a1, a2)
180 let (mut ci_dft, scratch1) = scratch.take_vec_znx_dft(n, cols, self.size());
181 (0..cols).for_each(|i| {
182 module.vec_znx_dft_from_vec_znx(1, 0, &mut ci_dft, i, &self.at(row_i, 0).data, i);
183 });
184
185 (1..cols).for_each(|col_j| {
186 // Example for rank 3:
187 //
188 // Note: M is a vector (m, Bm, B^2m, B^3m, ...), so each column is
189 // actually composed of that many rows and we focus on a specific row here
190 // implicitely given ci_dft.
191 //
192 // # Input
193 //
194 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0 , a1 , a2 )
195 // col 1: (0, 0, 0, 0)
196 // col 2: (0, 0, 0, 0)
197 // col 3: (0, 0, 0, 0)
198 //
199 // # Output
200 //
201 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0 , a1 , a2 )
202 // col 1: (-(b0s0 + b1s1 + b2s2) , b0 + M[i], b1 , b2 )
203 // col 2: (-(c0s0 + c1s1 + c2s2) , c0 , c1 + M[i], c2 )
204 // col 3: (-(d0s0 + d1s1 + d2s2) , d0 , d1 , d2 + M[i])
205
206 let digits: usize = tsk.digits();
207
208 let (mut tmp_dft_i, scratch2) = scratch1.take_vec_znx_dft(n, cols, tsk.size());
209 let (mut tmp_a, scratch3) = scratch2.take_vec_znx_dft(n, 1, ci_dft.size().div_ceil(digits));
210
211 {
212 // Performs a key-switch for each combination of s[i]*s[j], i.e. for a0, a1, a2
213 //
214 // # Example for col=1
215 //
216 // a0 * (-(f0s0 + f1s1 + f1s2) + s0^2, f0, f1, f2) = (-(a0f0s0 + a0f1s1 + a0f1s2) + a0s0^2, a0f0, a0f1, a0f2)
217 // +
218 // a1 * (-(g0s0 + g1s1 + g1s2) + s0s1, g0, g1, g2) = (-(a1g0s0 + a1g1s1 + a1g1s2) + a1s0s1, a1g0, a1g1, a1g2)
219 // +
220 // a2 * (-(h0s0 + h1s1 + h1s2) + s0s2, h0, h1, h2) = (-(a2h0s0 + a2h1s1 + a2h1s2) + a2s0s2, a2h0, a2h1, a2h2)
221 // =
222 // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0, x1, x2)
223 (1..cols).for_each(|col_i| {
224 let pmat: &VmpPMat<DataTsk, B> = &tsk.at(col_i - 1, col_j - 1).key.data; // Selects Enc(s[i]s[j])
225
226 // Extracts a[i] and multipies with Enc(s[i]s[j])
227 (0..digits).for_each(|di| {
228 tmp_a.set_size((ci_dft.size() + di) / digits);
229
230 // Small optimization for digits > 2
231 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
232 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(digits-1) * B}.
233 // As such we can ignore the last digits-2 limbs safely of the sum of vmp products.
234 // It is possible to further ignore the last digits-1 limbs, but this introduce
235 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
236 // noise is kept with respect to the ideal functionality.
237 tmp_dft_i.set_size(tsk.size() - ((digits - di) as isize - 2).max(0) as usize);
238
239 module.vec_znx_dft_copy(digits, digits - 1 - di, &mut tmp_a, 0, &ci_dft, col_i);
240 if di == 0 && col_i == 1 {
241 module.vmp_apply(&mut tmp_dft_i, &tmp_a, pmat, scratch3);
242 } else {
243 module.vmp_apply_add(&mut tmp_dft_i, &tmp_a, pmat, di, scratch3);
244 }
245 });
246 });
247 }
248
249 // Adds -(sum a[i] * s[i]) + m) on the i-th column of tmp_idft_i
250 //
251 // (-(x0s0 + x1s1 + x2s2) + a0s0s0 + a1s0s1 + a2s0s2, x0, x1, x2)
252 // +
253 // (0, -(a0s0 + a1s1 + a2s2) + M[i], 0, 0)
254 // =
255 // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0 -(a0s0 + a1s1 + a2s2) + M[i], x1, x2)
256 // =
257 // (-(x0s0 + x1s1 + x2s2), x0 + M[i], x1, x2)
258 module.vec_znx_dft_add_inplace(&mut tmp_dft_i, col_j, &ci_dft, 0);
259 let (mut tmp_idft, scratch3) = scratch2.take_vec_znx_big(n, 1, tsk.size());
260 (0..cols).for_each(|i| {
261 module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut tmp_idft, 0, &mut tmp_dft_i, i);
262 module.vec_znx_big_normalize(
263 self.basek(),
264 &mut self.at_mut(row_i, col_j).data,
265 i,
266 &tmp_idft,
267 0,
268 scratch3,
269 );
270 });
271 })
272 })
273 }
274
275 fn keyswitch_internal<DataLhs: DataRef, DataKsk: DataRef, B: Backend>(
276 &mut self,
277 module: &Module<B>,
278 lhs: &GGSWCiphertext<DataLhs>,
279 ksk: &GLWESwitchingKeyExec<DataKsk, B>,
280 scratch: &mut Scratch<B>,
281 ) where
282 Module<B>: GLWEKeyswitchFamily<B> + GGSWKeySwitchFamily<B> + VecZnxNormalizeTmpBytes,
283 Scratch<B>: TakeVecZnxDft<B> + TakeVecZnxBig<B> + ScratchAvailable,
284 {
285 // Keyswitch the j-th row of the col 0
286 (0..lhs.rows()).for_each(|row_i| {
287 // Key-switch column 0, i.e.
288 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0s0' + a1s1' + a2s2') + M[i], a0, a1, a2)
289 self.at_mut(row_i, 0)
290 .keyswitch(module, &lhs.at(row_i, 0), ksk, scratch);
291 })
292 }
293}