poulpy/poulpy-core/src/conversion/gglwe_to_ggsw.rs
poulpy-core/src/conversion/gglwe_to_ggsw.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchAvailable, ScratchTakeBasic, VecZnxBigAddSmallInplace, VecZnxBigBytesOf, VecZnxBigNormalize,
4 VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftApply, VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalize,
5 },
6 layouts::{Backend, DataMut, Module, Scratch, VecZnx, VecZnxBig, VecZnxDft, VecZnxDftToRef, VecZnxToRef},
7};
8
9use crate::{
10 GGLWEProduct, GLWECopy, ScratchTakeCore,
11 layouts::{
12 GGLWE, GGLWEInfos, GGLWEToGGSWKeyPrepared, GGLWEToGGSWKeyPreparedToRef, GGLWEToRef, GGSW, GGSWInfos, GGSWToMut, GLWE,
13 GLWEInfos, LWEInfos,
14 },
15};
16
17impl GGLWE<Vec<u8>> {
18 pub fn from_gglw_tmp_bytes<R, A, M, BE: Backend>(module: &M, res_infos: &R, tsk_infos: &A) -> usize
19 where
20 M: GGSWFromGGLWE<BE>,
21 R: GGSWInfos,
22 A: GGLWEInfos,
23 {
24 module.ggsw_from_gglwe_tmp_bytes(res_infos, tsk_infos)
25 }
26}
27
28impl<D: DataMut> GGSW<D> {
29 pub fn from_gglwe<G, M, T, BE: Backend>(&mut self, module: &M, gglwe: &G, tsk: &T, scratch: &mut Scratch<BE>)
30 where
31 M: GGSWFromGGLWE<BE>,
32 G: GGLWEToRef,
33 T: GGLWEToGGSWKeyPreparedToRef<BE>,
34 Scratch<BE>: ScratchTakeCore<BE>,
35 {
36 module.ggsw_from_gglwe(self, gglwe, tsk, scratch);
37 }
38}
39
40impl<BE: Backend> GGSWFromGGLWE<BE> for Module<BE>
41where
42 Self: GGSWExpandRows<BE> + GLWECopy,
43{
44 fn ggsw_from_gglwe_tmp_bytes<R, A>(&self, res_infos: &R, tsk_infos: &A) -> usize
45 where
46 R: GGSWInfos,
47 A: GGLWEInfos,
48 {
49 self.ggsw_expand_rows_tmp_bytes(res_infos, tsk_infos)
50 }
51
52 fn ggsw_from_gglwe<R, A, T>(&self, res: &mut R, a: &A, tsk: &T, scratch: &mut Scratch<BE>)
53 where
54 R: GGSWToMut,
55 A: GGLWEToRef,
56 T: GGLWEToGGSWKeyPreparedToRef<BE>,
57 Scratch<BE>: ScratchTakeCore<BE>,
58 {
59 let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
60 let a: &GGLWE<&[u8]> = &a.to_ref();
61 let tsk: &GGLWEToGGSWKeyPrepared<&[u8], BE> = &tsk.to_ref();
62
63 assert_eq!(res.rank(), a.rank_out());
64 assert_eq!(res.dnum(), a.dnum());
65 assert_eq!(res.n(), self.n() as u32);
66 assert_eq!(a.n(), self.n() as u32);
67 assert_eq!(tsk.n(), self.n() as u32);
68 assert_eq!(res.base2k(), a.base2k());
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: GGLWEToGGSWKeyPreparedToRef<BE>,
89 Scratch<BE>: ScratchTakeCore<BE>;
90}
91
92pub trait GGSWExpandRows<BE: Backend> {
93 fn ggsw_expand_rows_tmp_bytes<R, A>(&self, res_infos: &R, tsk_infos: &A) -> usize
94 where
95 R: GGSWInfos,
96 A: GGLWEInfos;
97
98 fn ggsw_expand_row<R, T>(&self, res: &mut R, tsk: &T, scratch: &mut Scratch<BE>)
99 where
100 R: GGSWToMut,
101 T: GGLWEToGGSWKeyPreparedToRef<BE>,
102 Scratch<BE>: ScratchTakeCore<BE>;
103}
104
105impl<BE: Backend> GGSWExpandRows<BE> for Module<BE>
106where
107 Self: GGLWEProduct<BE>
108 + VecZnxBigNormalize<BE>
109 + VecZnxBigNormalizeTmpBytes
110 + VecZnxBigBytesOf
111 + VecZnxDftBytesOf
112 + VecZnxDftApply<BE>
113 + VecZnxNormalize<BE>
114 + VecZnxBigAddSmallInplace<BE>
115 + VecZnxIdftApplyConsume<BE>
116 + VecZnxCopy,
117{
118 fn ggsw_expand_rows_tmp_bytes<R, A>(&self, res_infos: &R, tsk_infos: &A) -> usize
119 where
120 R: GGSWInfos,
121 A: GGLWEInfos,
122 {
123 let tsk_base2k: usize = tsk_infos.base2k().into();
124
125 let rank: usize = res_infos.rank().into();
126 let cols: usize = rank + 1;
127
128 let res_size: usize = res_infos.size();
129 let a_size: usize = res_infos.max_k().as_usize().div_ceil(tsk_base2k);
130
131 let a_0: usize = VecZnx::bytes_of(self.n(), 1, a_size);
132 let a_dft: usize = self.bytes_of_vec_znx_dft(cols - 1, a_size);
133 let res_dft: usize = self.bytes_of_vec_znx_dft(cols, a_size);
134 let gglwe_prod: usize = self.gglwe_product_dft_tmp_bytes(res_size, a_size, tsk_infos);
135 let normalize: usize = self.vec_znx_big_normalize_tmp_bytes();
136
137 (a_0 + a_dft + res_dft + gglwe_prod).max(normalize)
138 }
139
140 fn ggsw_expand_row<R, T>(&self, res: &mut R, tsk: &T, scratch: &mut Scratch<BE>)
141 where
142 R: GGSWToMut,
143 T: GGLWEToGGSWKeyPreparedToRef<BE>,
144 Scratch<BE>: ScratchTakeCore<BE>,
145 {
146 let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
147 let tsk: &GGLWEToGGSWKeyPrepared<&[u8], BE> = &tsk.to_ref();
148
149 let res_base2k: usize = res.base2k().into();
150 let tsk_base2k: usize = tsk.base2k().into();
151
152 assert!(scratch.available() >= self.ggsw_expand_rows_tmp_bytes(res, tsk));
153
154 let rank: usize = res.rank().into();
155 let cols: usize = rank + 1;
156
157 let res_conv_size: usize = res.max_k().as_usize().div_ceil(tsk_base2k);
158
159 let (mut a_dft, scratch_1) = scratch.take_vec_znx_dft(self, cols - 1, res_conv_size);
160 let (mut a_0, scratch_2) = scratch_1.take_vec_znx(self.n(), 1, res_conv_size);
161
162 // Keyswitch the j-th row of the col 0
163 for row in 0..res.dnum().as_usize() {
164 let glwe_mi_1: &GLWE<&[u8]> = &res.at(row, 0);
165
166 if res_base2k == tsk_base2k {
167 for col_i in 0..cols - 1 {
168 self.vec_znx_dft_apply(1, 0, &mut a_dft, col_i, glwe_mi_1.data(), col_i + 1);
169 }
170 self.vec_znx_copy(&mut a_0, 0, glwe_mi_1.data(), 0);
171 } else {
172 for i in 0..cols - 1 {
173 self.vec_znx_normalize(&mut a_0, tsk_base2k, 0, 0, glwe_mi_1.data(), res_base2k, i + 1, scratch_2);
174 self.vec_znx_dft_apply(1, 0, &mut a_dft, i, &a_0, 0);
175 }
176 self.vec_znx_normalize(&mut a_0, tsk_base2k, 0, 0, glwe_mi_1.data(), res_base2k, 0, scratch_2);
177 }
178
179 ggsw_expand_rows_internal(self, row, res, &a_0, &a_dft, tsk, scratch_2)
180 }
181 }
182}
183
184fn ggsw_expand_rows_internal<M, R, C, A, T, BE: Backend>(
185 module: &M,
186 row: usize,
187 res: &mut R,
188 a_0: &C,
189 a_dft: &A,
190 tsk: &T,
191 scratch: &mut Scratch<BE>,
192) where
193 R: GGSWToMut,
194 C: VecZnxToRef,
195 A: VecZnxDftToRef<BE>,
196 M: GGLWEProduct<BE> + VecZnxIdftApplyConsume<BE> + VecZnxBigAddSmallInplace<BE> + VecZnxBigNormalize<BE>,
197 T: GGLWEToGGSWKeyPreparedToRef<BE>,
198 Scratch<BE>: ScratchTakeCore<BE>,
199{
200 let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
201 let a_0: &VecZnx<&[u8]> = &a_0.to_ref();
202 let a_dft: &VecZnxDft<&[u8], BE> = &a_dft.to_ref();
203 let tsk: &GGLWEToGGSWKeyPrepared<&[u8], BE> = &tsk.to_ref();
204 let cols: usize = res.rank().as_usize() + 1;
205
206 // Example for rank 3:
207 //
208 // Note: M is a vector (m, Bm, B^2m, B^3m, ...), so each column is
209 // actually composed of that many dnum and we focus on a specific row here
210 // implicitely given ci_dft.
211 //
212 // # Input
213 //
214 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0 , a1 , a2 )
215 // col 1: (0, 0, 0, 0)
216 // col 2: (0, 0, 0, 0)
217 // col 3: (0, 0, 0, 0)
218 //
219 // # Output
220 //
221 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0 , a1 , a2 )
222 // col 1: (-(b0s0 + b1s1 + b2s2) , b0 + M[i], b1 , b2 )
223 // col 2: (-(c0s0 + c1s1 + c2s2) , c0 , c1 + M[i], c2 )
224 // col 3: (-(d0s0 + d1s1 + d2s2) , d0 , d1 , d2 + M[i])
225 for col in 1..cols {
226 let (mut res_dft, scratch_1) = scratch.take_vec_znx_dft(module, cols, tsk.size()); // Todo optimise
227
228 // Performs a key-switch for each combination of s[i]*s[j], i.e. for a0, a1, a2
229 //
230 // # Example for col=1
231 //
232 // a0 * (-(f0s0 + f1s1 + f1s2) + s0^2, f0, f1, f2) = (-(a0f0s0 + a0f1s1 + a0f1s2) + a0s0^2, a0f0, a0f1, a0f2)
233 // +
234 // a1 * (-(g0s0 + g1s1 + g1s2) + s0s1, g0, g1, g2) = (-(a1g0s0 + a1g1s1 + a1g1s2) + a1s0s1, a1g0, a1g1, a1g2)
235 // +
236 // a2 * (-(h0s0 + h1s1 + h1s2) + s0s2, h0, h1, h2) = (-(a2h0s0 + a2h1s1 + a2h1s2) + a2s0s2, a2h0, a2h1, a2h2)
237 // =
238 // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0, x1, x2)
239 module.gglwe_product_dft(&mut res_dft, a_dft, tsk.at(col - 1), scratch_1);
240
241 let mut res_big: VecZnxBig<&mut [u8], BE> = module.vec_znx_idft_apply_consume(res_dft);
242
243 // Adds -(sum a[i] * s[i]) + m) on the i-th column of tmp_idft_i
244 //
245 // (-(x0s0 + x1s1 + x2s2) + a0s0s0 + a1s0s1 + a2s0s2, x0, x1, x2)
246 // +
247 // (0, -(a0s0 + a1s1 + a2s2) + M[i], 0, 0)
248 // =
249 // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0 -(a0s0 + a1s1 + a2s2) + M[i], x1, x2)
250 // =
251 // (-(x0s0 + x1s1 + x2s2), x0 + M[i], x1, x2)
252 module.vec_znx_big_add_small_inplace(&mut res_big, col, a_0, 0);
253
254 let res_base2k: usize = res.base2k().as_usize();
255
256 for j in 0..cols {
257 module.vec_znx_big_normalize(
258 res.at_mut(row, col).data_mut(),
259 res_base2k,
260 0,
261 j,
262 &res_big,
263 tsk.base2k().as_usize(),
264 j,
265 scratch_1,
266 );
267 }
268 }
269}