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