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, VecZnxDftApply, VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalize,
5 },
6 layouts::{Backend, DataMut, Module, Scratch, VecZnxBig},
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
69 for row in 0..res.dnum().into() {
70 self.glwe_copy(&mut res.at_mut(row, 0), &a.at(row, 0));
71 }
72
73 self.ggsw_expand_row(res, tsk, scratch);
74 }
75}
76
77pub trait GGSWFromGGLWE<BE: Backend> {
78 fn ggsw_from_gglwe_tmp_bytes<R, A>(&self, res_infos: &R, tsk_infos: &A) -> usize
79 where
80 R: GGSWInfos,
81 A: GGLWEInfos;
82
83 fn ggsw_from_gglwe<R, A, T>(&self, res: &mut R, a: &A, tsk: &T, scratch: &mut Scratch<BE>)
84 where
85 R: GGSWToMut,
86 A: GGLWEToRef,
87 T: GGLWEToGGSWKeyPreparedToRef<BE>,
88 Scratch<BE>: ScratchTakeCore<BE>;
89}
90
91pub trait GGSWExpandRows<BE: Backend> {
92 fn ggsw_expand_rows_tmp_bytes<R, A>(&self, res_infos: &R, tsk_infos: &A) -> usize
93 where
94 R: GGSWInfos,
95 A: GGLWEInfos;
96
97 fn ggsw_expand_row<R, T>(&self, res: &mut R, tsk: &T, scratch: &mut Scratch<BE>)
98 where
99 R: GGSWToMut,
100 T: GGLWEToGGSWKeyPreparedToRef<BE>,
101 Scratch<BE>: ScratchTakeCore<BE>;
102}
103
104impl<BE: Backend> GGSWExpandRows<BE> for Module<BE>
105where
106 Self: GGLWEProduct<BE>
107 + VecZnxBigNormalize<BE>
108 + VecZnxBigNormalizeTmpBytes
109 + VecZnxBigBytesOf
110 + VecZnxDftBytesOf
111 + VecZnxDftApply<BE>
112 + VecZnxNormalize<BE>
113 + VecZnxBigAddSmallInplace<BE>
114 + VecZnxIdftApplyConsume<BE>,
115{
116 fn ggsw_expand_rows_tmp_bytes<R, A>(&self, res_infos: &R, tsk_infos: &A) -> usize
117 where
118 R: GGSWInfos,
119 A: GGLWEInfos,
120 {
121 let base2k_in: usize = res_infos.base2k().into();
122 let base2k_tsk: usize = tsk_infos.base2k().into();
123
124 let rank: usize = res_infos.rank().into();
125 let cols: usize = rank + 1;
126
127 let res_size = res_infos.size();
128 let a_size: usize = (res_infos.size() * base2k_in).div_ceil(base2k_tsk);
129
130 let a_dft = self.bytes_of_vec_znx_dft(cols - 1, a_size);
131 let res_dft = self.bytes_of_vec_znx_dft(cols, a_size);
132 let gglwe_prod: usize = self.gglwe_product_dft_tmp_bytes(res_size, a_size, tsk_infos);
133 let normalize = self.vec_znx_big_normalize_tmp_bytes();
134
135 (a_dft + res_dft + gglwe_prod).max(normalize)
136 }
137
138 fn ggsw_expand_row<R, T>(&self, res: &mut R, tsk: &T, scratch: &mut Scratch<BE>)
139 where
140 R: GGSWToMut,
141 T: GGLWEToGGSWKeyPreparedToRef<BE>,
142 Scratch<BE>: ScratchTakeCore<BE>,
143 {
144 let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
145 let tsk: &GGLWEToGGSWKeyPrepared<&[u8], BE> = &tsk.to_ref();
146
147 let base2k_in: usize = res.base2k().into();
148 let base2k_tsk: usize = tsk.base2k().into();
149
150 assert!(scratch.available() >= self.ggsw_expand_rows_tmp_bytes(res, tsk));
151
152 let rank: usize = res.rank().into();
153 let cols: usize = rank + 1;
154
155 let a_size: usize = (res.size() * base2k_in).div_ceil(base2k_tsk);
156
157 // Keyswitch the j-th row of the col 0
158 for row in 0..res.dnum().as_usize() {
159 let (mut a_dft, scratch_1) = scratch.take_vec_znx_dft(self, cols - 1, a_size);
160
161 {
162 let glwe_mi_1: &GLWE<&[u8]> = &res.at(row, 0);
163
164 if base2k_in == base2k_tsk {
165 for col_i in 0..cols - 1 {
166 self.vec_znx_dft_apply(1, 0, &mut a_dft, col_i, glwe_mi_1.data(), col_i + 1);
167 }
168 } else {
169 let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(self.n(), 1, a_size);
170 for i in 0..cols - 1 {
171 self.vec_znx_normalize(
172 base2k_tsk,
173 &mut a_conv,
174 0,
175 base2k_in,
176 glwe_mi_1.data(),
177 i + 1,
178 scratch_2,
179 );
180 self.vec_znx_dft_apply(1, 0, &mut a_dft, i, &a_conv, 0);
181 }
182 }
183 }
184
185 // Example for rank 3:
186 //
187 // Note: M is a vector (m, Bm, B^2m, B^3m, ...), so each column is
188 // actually composed of that many dnum and we focus on a specific row here
189 // implicitely given ci_dft.
190 //
191 // # Input
192 //
193 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0 , a1 , a2 )
194 // col 1: (0, 0, 0, 0)
195 // col 2: (0, 0, 0, 0)
196 // col 3: (0, 0, 0, 0)
197 //
198 // # Output
199 //
200 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0 , a1 , a2 )
201 // col 1: (-(b0s0 + b1s1 + b2s2) , b0 + M[i], b1 , b2 )
202 // col 2: (-(c0s0 + c1s1 + c2s2) , c0 , c1 + M[i], c2 )
203 // col 3: (-(d0s0 + d1s1 + d2s2) , d0 , d1 , d2 + M[i])
204 for col in 1..cols {
205 let (mut res_dft, scratch_2) = scratch_1.take_vec_znx_dft(self, cols, tsk.size()); // Todo optimise
206
207 // Performs a key-switch for each combination of s[i]*s[j], i.e. for a0, a1, a2
208 //
209 // # Example for col=1
210 //
211 // a0 * (-(f0s0 + f1s1 + f1s2) + s0^2, f0, f1, f2) = (-(a0f0s0 + a0f1s1 + a0f1s2) + a0s0^2, a0f0, a0f1, a0f2)
212 // +
213 // a1 * (-(g0s0 + g1s1 + g1s2) + s0s1, g0, g1, g2) = (-(a1g0s0 + a1g1s1 + a1g1s2) + a1s0s1, a1g0, a1g1, a1g2)
214 // +
215 // a2 * (-(h0s0 + h1s1 + h1s2) + s0s2, h0, h1, h2) = (-(a2h0s0 + a2h1s1 + a2h1s2) + a2s0s2, a2h0, a2h1, a2h2)
216 // =
217 // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0, x1, x2)
218 self.gglwe_product_dft(&mut res_dft, &a_dft, tsk.at(col - 1), scratch_2);
219
220 let mut res_big: VecZnxBig<&mut [u8], BE> = self.vec_znx_idft_apply_consume(res_dft);
221
222 // Adds -(sum a[i] * s[i]) + m) on the i-th column of tmp_idft_i
223 //
224 // (-(x0s0 + x1s1 + x2s2) + a0s0s0 + a1s0s1 + a2s0s2, x0, x1, x2)
225 // +
226 // (0, -(a0s0 + a1s1 + a2s2) + M[i], 0, 0)
227 // =
228 // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0 -(a0s0 + a1s1 + a2s2) + M[i], x1, x2)
229 // =
230 // (-(x0s0 + x1s1 + x2s2), x0 + M[i], x1, x2)
231 self.vec_znx_big_add_small_inplace(&mut res_big, col, res.at(row, 0).data(), 0);
232
233 for j in 0..cols {
234 self.vec_znx_big_normalize(
235 res.base2k().as_usize(),
236 res.at_mut(row, col).data_mut(),
237 j,
238 tsk.base2k().as_usize(),
239 &res_big,
240 j,
241 scratch_2,
242 );
243 }
244 }
245 }
246 }
247}