poulpy/poulpy-core/src/conversion/gglwe_to_ggsw.rs
poulpy-core/src/conversion/gglwe_to_ggsw.rsRaw
1use poulpy_hal::{
2 api::{
3 ModuleN, ScratchAvailable, ScratchTakeBasic, VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxDftAddInplace, VecZnxDftApply,
4 VecZnxDftBytesOf, VecZnxDftCopy, VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
5 VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
6 },
7 layouts::{Backend, DataMut, Module, Scratch, VmpPMat, ZnxInfos},
8};
9
10use crate::{
11 GLWECopy, ScratchTakeCore,
12 layouts::{
13 GGLWE, GGLWEInfos, GGLWEToRef, GGSW, GGSWInfos, GGSWToMut, GLWEInfos, LWEInfos,
14 prepared::{TensorKeyPrepared, TensorKeyPreparedToRef},
15 },
16};
17
18impl GGLWE<Vec<u8>> {
19 pub fn from_gglw_tmp_bytes<R, A, M, BE: Backend>(module: &M, res_infos: &R, tsk_infos: &A) -> usize
20 where
21 M: GGSWFromGGLWE<BE>,
22 R: GGSWInfos,
23 A: GGLWEInfos,
24 {
25 module.ggsw_from_gglwe_tmp_bytes(res_infos, tsk_infos)
26 }
27}
28
29impl<D: DataMut> GGSW<D> {
30 pub fn from_gglwe<G, M, T, BE: Backend>(&mut self, module: &M, gglwe: &G, tsk: &T, scratch: &mut Scratch<BE>)
31 where
32 M: GGSWFromGGLWE<BE>,
33 G: GGLWEToRef,
34 T: TensorKeyPreparedToRef<BE>,
35 Scratch<BE>: ScratchTakeCore<BE>,
36 {
37 module.ggsw_from_gglwe(self, gglwe, tsk, scratch);
38 }
39}
40
41impl<BE: Backend> GGSWFromGGLWE<BE> for Module<BE>
42where
43 Self: GGSWExpandRows<BE> + GLWECopy,
44{
45 fn ggsw_from_gglwe_tmp_bytes<R, A>(&self, res_infos: &R, tsk_infos: &A) -> usize
46 where
47 R: GGSWInfos,
48 A: GGLWEInfos,
49 {
50 self.ggsw_expand_rows_tmp_bytes(res_infos, tsk_infos)
51 }
52
53 fn ggsw_from_gglwe<R, A, T>(&self, res: &mut R, a: &A, tsk: &T, scratch: &mut Scratch<BE>)
54 where
55 R: GGSWToMut,
56 A: GGLWEToRef,
57 T: TensorKeyPreparedToRef<BE>,
58 Scratch<BE>: ScratchTakeCore<BE>,
59 {
60 let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
61 let a: &GGLWE<&[u8]> = &a.to_ref();
62 let tsk: &TensorKeyPrepared<&[u8], BE> = &tsk.to_ref();
63
64 assert_eq!(res.rank(), a.rank_out());
65 assert_eq!(res.dnum(), a.dnum());
66 assert_eq!(res.n(), self.n() as u32);
67 assert_eq!(a.n(), self.n() as u32);
68 assert_eq!(tsk.n(), self.n() as u32);
69
70 for row in 0..res.dnum().into() {
71 self.glwe_copy(&mut res.at_mut(row, 0), &a.at(row, 0));
72 }
73
74 self.ggsw_expand_row(res, tsk, scratch);
75 }
76}
77
78pub trait GGSWFromGGLWE<BE: Backend> {
79 fn ggsw_from_gglwe_tmp_bytes<R, A>(&self, res_infos: &R, tsk_infos: &A) -> usize
80 where
81 R: GGSWInfos,
82 A: GGLWEInfos;
83
84 fn ggsw_from_gglwe<R, A, T>(&self, res: &mut R, a: &A, tsk: &T, scratch: &mut Scratch<BE>)
85 where
86 R: GGSWToMut,
87 A: GGLWEToRef,
88 T: TensorKeyPreparedToRef<BE>,
89 Scratch<BE>: ScratchTakeCore<BE>;
90}
91
92impl<BE: Backend> GGSWExpandRows<BE> for Module<BE> where
93 Self: Sized
94 + ModuleN
95 + VecZnxDftBytesOf
96 + VmpApplyDftToDftTmpBytes
97 + VecZnxBigBytesOf
98 + VecZnxNormalizeTmpBytes
99 + VecZnxDftBytesOf
100 + VmpApplyDftToDftTmpBytes
101 + VecZnxBigBytesOf
102 + VecZnxNormalizeTmpBytes
103 + VecZnxDftApply<BE>
104 + VecZnxDftCopy<BE>
105 + VmpApplyDftToDft<BE>
106 + VmpApplyDftToDftAdd<BE>
107 + VecZnxDftAddInplace<BE>
108 + VecZnxBigNormalize<BE>
109 + VecZnxIdftApplyTmpA<BE>
110 + VecZnxNormalize<BE>
111{
112}
113
114pub trait GGSWExpandRows<BE: Backend>
115where
116 Self: Sized
117 + ModuleN
118 + VecZnxDftBytesOf
119 + VmpApplyDftToDftTmpBytes
120 + VecZnxBigBytesOf
121 + VecZnxNormalizeTmpBytes
122 + VecZnxDftBytesOf
123 + VmpApplyDftToDftTmpBytes
124 + VecZnxBigBytesOf
125 + VecZnxNormalizeTmpBytes
126 + VecZnxDftApply<BE>
127 + VecZnxDftCopy<BE>
128 + VmpApplyDftToDft<BE>
129 + VmpApplyDftToDftAdd<BE>
130 + VecZnxDftAddInplace<BE>
131 + VecZnxBigNormalize<BE>
132 + VecZnxIdftApplyTmpA<BE>
133 + VecZnxNormalize<BE>,
134{
135 fn ggsw_expand_rows_tmp_bytes<R, A>(&self, res_infos: &R, tsk_infos: &A) -> usize
136 where
137 R: GGSWInfos,
138 A: GGLWEInfos,
139 {
140 let tsk_size: usize = tsk_infos.k().div_ceil(tsk_infos.base2k()) as usize;
141 let size_in: usize = res_infos
142 .k()
143 .div_ceil(tsk_infos.base2k())
144 .div_ceil(tsk_infos.dsize().into()) as usize;
145
146 let tmp_dft_i: usize = self.bytes_of_vec_znx_dft((tsk_infos.rank_out() + 1).into(), tsk_size);
147 let tmp_a: usize = self.bytes_of_vec_znx_dft(1, size_in);
148 let vmp: usize = self.vmp_apply_dft_to_dft_tmp_bytes(
149 tsk_size,
150 size_in,
151 size_in,
152 (tsk_infos.rank_in()).into(), // Verify if rank+1
153 (tsk_infos.rank_out()).into(), // Verify if rank+1
154 tsk_size,
155 );
156 let tmp_idft: usize = self.bytes_of_vec_znx_big(1, tsk_size);
157 let norm: usize = self.vec_znx_normalize_tmp_bytes();
158
159 tmp_dft_i + ((tmp_a + vmp) | (tmp_idft + norm))
160 }
161
162 fn ggsw_expand_row<R, T>(&self, res: &mut R, tsk: &T, scratch: &mut Scratch<BE>)
163 where
164 R: GGSWToMut,
165 T: TensorKeyPreparedToRef<BE>,
166 Scratch<BE>: ScratchTakeCore<BE>,
167 {
168 let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
169 let tsk: &TensorKeyPrepared<&[u8], BE> = &tsk.to_ref();
170
171 let basek_in: usize = res.base2k().into();
172 let basek_tsk: usize = tsk.base2k().into();
173
174 assert!(scratch.available() >= self.ggsw_expand_rows_tmp_bytes(res, tsk));
175
176 let rank: usize = res.rank().into();
177 let cols: usize = rank + 1;
178
179 let a_size: usize = (res.size() * basek_in).div_ceil(basek_tsk);
180
181 // Keyswitch the j-th row of the col 0
182 for row_i in 0..res.dnum().into() {
183 let a = &res.at(row_i, 0).data;
184
185 // Pre-compute DFT of (a0, a1, a2)
186 let (mut ci_dft, scratch_1) = scratch.take_vec_znx_dft(self, cols, a_size);
187
188 if basek_in == basek_tsk {
189 for i in 0..cols {
190 self.vec_znx_dft_apply(1, 0, &mut ci_dft, i, a, i);
191 }
192 } else {
193 let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(self, 1, a_size);
194 for i in 0..cols {
195 self.vec_znx_normalize(basek_tsk, &mut a_conv, 0, basek_in, a, i, scratch_2);
196 self.vec_znx_dft_apply(1, 0, &mut ci_dft, i, &a_conv, 0);
197 }
198 }
199
200 for col_j in 1..cols {
201 // Example for rank 3:
202 //
203 // Note: M is a vector (m, Bm, B^2m, B^3m, ...), so each column is
204 // actually composed of that many dnum and we focus on a specific row here
205 // implicitely given ci_dft.
206 //
207 // # Input
208 //
209 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0 , a1 , a2 )
210 // col 1: (0, 0, 0, 0)
211 // col 2: (0, 0, 0, 0)
212 // col 3: (0, 0, 0, 0)
213 //
214 // # Output
215 //
216 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0 , a1 , a2 )
217 // col 1: (-(b0s0 + b1s1 + b2s2) , b0 + M[i], b1 , b2 )
218 // col 2: (-(c0s0 + c1s1 + c2s2) , c0 , c1 + M[i], c2 )
219 // col 3: (-(d0s0 + d1s1 + d2s2) , d0 , d1 , d2 + M[i])
220
221 let dsize: usize = tsk.dsize().into();
222
223 let (mut tmp_dft_i, scratch_2) = scratch_1.take_vec_znx_dft(self, cols, tsk.size());
224 let (mut tmp_a, scratch_3) = scratch_2.take_vec_znx_dft(self, 1, ci_dft.size().div_ceil(dsize));
225
226 {
227 // Performs a key-switch for each combination of s[i]*s[j], i.e. for a0, a1, a2
228 //
229 // # Example for col=1
230 //
231 // a0 * (-(f0s0 + f1s1 + f1s2) + s0^2, f0, f1, f2) = (-(a0f0s0 + a0f1s1 + a0f1s2) + a0s0^2, a0f0, a0f1, a0f2)
232 // +
233 // a1 * (-(g0s0 + g1s1 + g1s2) + s0s1, g0, g1, g2) = (-(a1g0s0 + a1g1s1 + a1g1s2) + a1s0s1, a1g0, a1g1, a1g2)
234 // +
235 // a2 * (-(h0s0 + h1s1 + h1s2) + s0s2, h0, h1, h2) = (-(a2h0s0 + a2h1s1 + a2h1s2) + a2s0s2, a2h0, a2h1, a2h2)
236 // =
237 // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0, x1, x2)
238 for col_i in 1..cols {
239 let pmat: &VmpPMat<&[u8], BE> = &tsk.at(col_i - 1, col_j - 1).key.data; // Selects Enc(s[i]s[j])
240
241 // Extracts a[i] and multipies with Enc(s[i]s[j])
242 for di in 0..dsize {
243 tmp_a.set_size((ci_dft.size() + di) / dsize);
244
245 // Small optimization for dsize > 2
246 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
247 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(dsize-1) * B}.
248 // As such we can ignore the last dsize-2 limbs safely of the sum of vmp products.
249 // It is possible to further ignore the last dsize-1 limbs, but this introduce
250 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
251 // noise is kept with respect to the ideal functionality.
252 tmp_dft_i.set_size(tsk.size() - ((dsize - di) as isize - 2).max(0) as usize);
253
254 self.vec_znx_dft_copy(dsize, dsize - 1 - di, &mut tmp_a, 0, &ci_dft, col_i);
255 if di == 0 && col_i == 1 {
256 self.vmp_apply_dft_to_dft(&mut tmp_dft_i, &tmp_a, pmat, scratch_3);
257 } else {
258 self.vmp_apply_dft_to_dft_add(&mut tmp_dft_i, &tmp_a, pmat, di, scratch_3);
259 }
260 }
261 }
262 }
263
264 // Adds -(sum a[i] * s[i]) + m) on the i-th column of tmp_idft_i
265 //
266 // (-(x0s0 + x1s1 + x2s2) + a0s0s0 + a1s0s1 + a2s0s2, x0, x1, x2)
267 // +
268 // (0, -(a0s0 + a1s1 + a2s2) + M[i], 0, 0)
269 // =
270 // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0 -(a0s0 + a1s1 + a2s2) + M[i], x1, x2)
271 // =
272 // (-(x0s0 + x1s1 + x2s2), x0 + M[i], x1, x2)
273 self.vec_znx_dft_add_inplace(&mut tmp_dft_i, col_j, &ci_dft, 0);
274 let (mut tmp_idft, scratch_3) = scratch_2.take_vec_znx_big(self, 1, tsk.size());
275 for i in 0..cols {
276 self.vec_znx_idft_apply_tmpa(&mut tmp_idft, 0, &mut tmp_dft_i, i);
277 self.vec_znx_big_normalize(
278 basek_in,
279 &mut res.at_mut(row_i, col_j).data,
280 i,
281 basek_tsk,
282 &tmp_idft,
283 0,
284 scratch_3,
285 );
286 }
287 }
288 }
289 }
290}