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