poulpy_core/conversion/
gglwe_to_ggsw.rs

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::{GLWETensorKeyPrepared, GLWETensorKeyPreparedToRef},
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: GLWETensorKeyPreparedToRef<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: GLWETensorKeyPreparedToRef<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: &GLWETensorKeyPrepared<&[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: GLWETensorKeyPreparedToRef<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        + VecZnxDftApply<BE>
123        + VecZnxDftCopy<BE>
124        + VmpApplyDftToDft<BE>
125        + VmpApplyDftToDftAdd<BE>
126        + VecZnxDftAddInplace<BE>
127        + VecZnxBigNormalize<BE>
128        + VecZnxIdftApplyTmpA<BE>
129        + VecZnxNormalize<BE>,
130{
131    fn ggsw_expand_rows_tmp_bytes<R, A>(&self, res_infos: &R, tsk_infos: &A) -> usize
132    where
133        R: GGSWInfos,
134        A: GGLWEInfos,
135    {
136        let tsk_size: usize = tsk_infos.k().div_ceil(tsk_infos.base2k()) as usize;
137        let size_in: usize = res_infos
138            .k()
139            .div_ceil(tsk_infos.base2k())
140            .div_ceil(tsk_infos.dsize().into()) as usize;
141
142        let tmp_dft_i: usize = self.bytes_of_vec_znx_dft((tsk_infos.rank_out() + 1).into(), tsk_size);
143        let tmp_a: usize = self.bytes_of_vec_znx_dft(1, size_in);
144        let vmp: usize = self.vmp_apply_dft_to_dft_tmp_bytes(
145            tsk_size,
146            size_in,
147            size_in,
148            (tsk_infos.rank_in()).into(),  // Verify if rank+1
149            (tsk_infos.rank_out()).into(), // Verify if rank+1
150            tsk_size,
151        );
152        let tmp_idft: usize = self.bytes_of_vec_znx_big(1, tsk_size);
153        let norm: usize = self.vec_znx_normalize_tmp_bytes();
154
155        tmp_dft_i + ((tmp_a + vmp) | (tmp_idft + norm))
156    }
157
158    fn ggsw_expand_row<R, T>(&self, res: &mut R, tsk: &T, scratch: &mut Scratch<BE>)
159    where
160        R: GGSWToMut,
161        T: GLWETensorKeyPreparedToRef<BE>,
162        Scratch<BE>: ScratchTakeCore<BE>,
163    {
164        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
165        let tsk: &GLWETensorKeyPrepared<&[u8], BE> = &tsk.to_ref();
166
167        let basek_in: usize = res.base2k().into();
168        let basek_tsk: usize = tsk.base2k().into();
169
170        assert!(scratch.available() >= self.ggsw_expand_rows_tmp_bytes(res, tsk));
171
172        let rank: usize = res.rank().into();
173        let cols: usize = rank + 1;
174
175        let a_size: usize = (res.size() * basek_in).div_ceil(basek_tsk);
176
177        // Keyswitch the j-th row of the col 0
178        for row_i in 0..res.dnum().into() {
179            let a = &res.at(row_i, 0).data;
180
181            // Pre-compute DFT of (a0, a1, a2)
182            let (mut ci_dft, scratch_1) = scratch.take_vec_znx_dft(self, cols, a_size);
183
184            if basek_in == basek_tsk {
185                for i in 0..cols {
186                    self.vec_znx_dft_apply(1, 0, &mut ci_dft, i, a, i);
187                }
188            } else {
189                let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(self.n(), 1, a_size);
190                for i in 0..cols {
191                    self.vec_znx_normalize(basek_tsk, &mut a_conv, 0, basek_in, a, i, scratch_2);
192                    self.vec_znx_dft_apply(1, 0, &mut ci_dft, i, &a_conv, 0);
193                }
194            }
195
196            for col_j in 1..cols {
197                // Example for rank 3:
198                //
199                // Note: M is a vector (m, Bm, B^2m, B^3m, ...), so each column is
200                // actually composed of that many dnum and we focus on a specific row here
201                // implicitely given ci_dft.
202                //
203                // # Input
204                //
205                // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0    , a1    , a2    )
206                // col 1: (0, 0, 0, 0)
207                // col 2: (0, 0, 0, 0)
208                // col 3: (0, 0, 0, 0)
209                //
210                // # Output
211                //
212                // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0       , a1       , a2       )
213                // col 1: (-(b0s0 + b1s1 + b2s2)       , b0 + M[i], b1       , b2       )
214                // col 2: (-(c0s0 + c1s1 + c2s2)       , c0       , c1 + M[i], c2       )
215                // col 3: (-(d0s0 + d1s1 + d2s2)       , d0       , d1       , d2 + M[i])
216
217                let dsize: usize = tsk.dsize().into();
218
219                let (mut tmp_dft_i, scratch_2) = scratch_1.take_vec_znx_dft(self, cols, tsk.size());
220                let (mut tmp_a, scratch_3) = scratch_2.take_vec_znx_dft(self, 1, ci_dft.size().div_ceil(dsize));
221
222                {
223                    // Performs a key-switch for each combination of s[i]*s[j], i.e. for a0, a1, a2
224                    //
225                    // # Example for col=1
226                    //
227                    // a0 * (-(f0s0 + f1s1 + f1s2) + s0^2, f0, f1, f2) = (-(a0f0s0 + a0f1s1 + a0f1s2) + a0s0^2, a0f0, a0f1, a0f2)
228                    // +
229                    // a1 * (-(g0s0 + g1s1 + g1s2) + s0s1, g0, g1, g2) = (-(a1g0s0 + a1g1s1 + a1g1s2) + a1s0s1, a1g0, a1g1, a1g2)
230                    // +
231                    // a2 * (-(h0s0 + h1s1 + h1s2) + s0s2, h0, h1, h2) = (-(a2h0s0 + a2h1s1 + a2h1s2) + a2s0s2, a2h0, a2h1, a2h2)
232                    // =
233                    // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0, x1, x2)
234                    for col_i in 1..cols {
235                        let pmat: &VmpPMat<&[u8], BE> = &tsk.at(col_i - 1, col_j - 1).data; // Selects Enc(s[i]s[j])
236
237                        // Extracts a[i] and multipies with Enc(s[i]s[j])
238                        for di in 0..dsize {
239                            tmp_a.set_size((ci_dft.size() + di) / dsize);
240
241                            // Small optimization for dsize > 2
242                            // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
243                            // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(dsize-1) * B}.
244                            // As such we can ignore the last dsize-2 limbs safely of the sum of vmp products.
245                            // It is possible to further ignore the last dsize-1 limbs, but this introduce
246                            // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
247                            // noise is kept with respect to the ideal functionality.
248                            tmp_dft_i.set_size(tsk.size() - ((dsize - di) as isize - 2).max(0) as usize);
249
250                            self.vec_znx_dft_copy(dsize, dsize - 1 - di, &mut tmp_a, 0, &ci_dft, col_i);
251                            if di == 0 && col_i == 1 {
252                                self.vmp_apply_dft_to_dft(&mut tmp_dft_i, &tmp_a, pmat, scratch_3);
253                            } else {
254                                self.vmp_apply_dft_to_dft_add(&mut tmp_dft_i, &tmp_a, pmat, di, scratch_3);
255                            }
256                        }
257                    }
258                }
259
260                // Adds -(sum a[i] * s[i]) + m)  on the i-th column of tmp_idft_i
261                //
262                // (-(x0s0 + x1s1 + x2s2) + a0s0s0 + a1s0s1 + a2s0s2, x0, x1, x2)
263                // +
264                // (0, -(a0s0 + a1s1 + a2s2) + M[i], 0, 0)
265                // =
266                // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0 -(a0s0 + a1s1 + a2s2) + M[i], x1, x2)
267                // =
268                // (-(x0s0 + x1s1 + x2s2), x0 + M[i], x1, x2)
269                self.vec_znx_dft_add_inplace(&mut tmp_dft_i, col_j, &ci_dft, 0);
270                let (mut tmp_idft, scratch_3) = scratch_2.take_vec_znx_big(self, 1, tsk.size());
271                for i in 0..cols {
272                    self.vec_znx_idft_apply_tmpa(&mut tmp_idft, 0, &mut tmp_dft_i, i);
273                    self.vec_znx_big_normalize(
274                        basek_in,
275                        &mut res.at_mut(row_i, col_j).data,
276                        i,
277                        basek_tsk,
278                        &tmp_idft,
279                        0,
280                        scratch_3,
281                    );
282                }
283            }
284        }
285    }
286}
poulpy_core/conversion/gglwe_to_ggsw.rs

poulpy_core/conversion/
gglwe_to_ggsw.rs
