poulpy/poulpy-core/src/scratch.rs
poulpy-core/src/scratch.rsRaw
1use poulpy_hal::{
2 api::{ModuleN, ScratchAvailable, ScratchTakeBasic, SvpPPolBytesOf, VecZnxDftBytesOf, VmpPMatBytesOf},
3 layouts::{Backend, Scratch},
4};
5
6use crate::{
7 dist::Distribution,
8 layouts::{
9 AutomorphismKey, Degree, GGLWE, GGLWEInfos, GGLWELayout, GGSW, GGSWInfos, GLWE, GLWEInfos, GLWEPlaintext, GLWEPublicKey,
10 GLWESecret, GLWESwitchingKey, Rank, TensorKey,
11 prepared::{
12 AutomorphismKeyPrepared, GGLWEPrepared, GGSWPrepared, GLWEPublicKeyPrepared, GLWESecretPrepared,
13 GLWESwitchingKeyPrepared, TensorKeyPrepared,
14 },
15 },
16};
17
18pub trait ScratchTakeCore<B: Backend>
19where
20 Self: ScratchTakeBasic + ScratchAvailable,
21{
22 fn take_glwe_ct<A, M>(&mut self, module: &M, infos: &A) -> (GLWE<&mut [u8]>, &mut Self)
23 where
24 A: GLWEInfos,
25 M: ModuleN,
26 {
27 assert_eq!(module.n() as u32, infos.n());
28 let (data, scratch) = self.take_vec_znx(module, (infos.rank() + 1).into(), infos.size());
29 (
30 GLWE {
31 k: infos.k(),
32 base2k: infos.base2k(),
33 data,
34 },
35 scratch,
36 )
37 }
38
39 fn take_glwe_ct_slice<A, M>(&mut self, module: &M, size: usize, infos: &A) -> (Vec<GLWE<&mut [u8]>>, &mut Self)
40 where
41 A: GLWEInfos,
42 M: ModuleN,
43 {
44 let mut scratch: &mut Self = self;
45 let mut cts: Vec<GLWE<&mut [u8]>> = Vec::with_capacity(size);
46 for _ in 0..size {
47 let (ct, new_scratch) = scratch.take_glwe_ct(module, infos);
48 scratch = new_scratch;
49 cts.push(ct);
50 }
51 (cts, scratch)
52 }
53
54 fn take_glwe_pt<A, M>(&mut self, module: &M, infos: &A) -> (GLWEPlaintext<&mut [u8]>, &mut Self)
55 where
56 A: GLWEInfos,
57 M: ModuleN,
58 {
59 assert_eq!(module.n() as u32, infos.n());
60 let (data, scratch) = self.take_vec_znx(module, 1, infos.size());
61 (
62 GLWEPlaintext {
63 k: infos.k(),
64 base2k: infos.base2k(),
65 data,
66 },
67 scratch,
68 )
69 }
70
71 fn take_gglwe<A, M>(&mut self, module: &M, infos: &A) -> (GGLWE<&mut [u8]>, &mut Self)
72 where
73 A: GGLWEInfos,
74 M: ModuleN,
75 {
76 assert_eq!(module.n() as u32, infos.n());
77 let (data, scratch) = self.take_mat_znx(
78 module,
79 infos.dnum().0.div_ceil(infos.dsize().0) as usize,
80 infos.rank_in().into(),
81 (infos.rank_out() + 1).into(),
82 infos.size(),
83 );
84 (
85 GGLWE {
86 k: infos.k(),
87 base2k: infos.base2k(),
88 dsize: infos.dsize(),
89 data,
90 },
91 scratch,
92 )
93 }
94
95 fn take_gglwe_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GGLWEPrepared<&mut [u8], B>, &mut Self)
96 where
97 A: GGLWEInfos,
98 M: ModuleN + VmpPMatBytesOf,
99 {
100 assert_eq!(module.n() as u32, infos.n());
101 let (data, scratch) = self.take_vmp_pmat(
102 module,
103 infos.dnum().into(),
104 infos.rank_in().into(),
105 (infos.rank_out() + 1).into(),
106 infos.size(),
107 );
108 (
109 GGLWEPrepared {
110 k: infos.k(),
111 base2k: infos.base2k(),
112 dsize: infos.dsize(),
113 data,
114 },
115 scratch,
116 )
117 }
118
119 fn take_ggsw<A, M>(&mut self, module: &M, infos: &A) -> (GGSW<&mut [u8]>, &mut Self)
120 where
121 A: GGSWInfos,
122 M: ModuleN,
123 {
124 assert_eq!(module.n() as u32, infos.n());
125 let (data, scratch) = self.take_mat_znx(
126 module,
127 infos.dnum().into(),
128 (infos.rank() + 1).into(),
129 (infos.rank() + 1).into(),
130 infos.size(),
131 );
132 (
133 GGSW {
134 k: infos.k(),
135 base2k: infos.base2k(),
136 dsize: infos.dsize(),
137 data,
138 },
139 scratch,
140 )
141 }
142
143 fn take_ggsw_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GGSWPrepared<&mut [u8], B>, &mut Self)
144 where
145 A: GGSWInfos,
146 M: ModuleN + VmpPMatBytesOf,
147 {
148 assert_eq!(module.n() as u32, infos.n());
149 let (data, scratch) = self.take_vmp_pmat(
150 module,
151 infos.dnum().into(),
152 (infos.rank() + 1).into(),
153 (infos.rank() + 1).into(),
154 infos.size(),
155 );
156 (
157 GGSWPrepared {
158 k: infos.k(),
159 base2k: infos.base2k(),
160 dsize: infos.dsize(),
161 data,
162 },
163 scratch,
164 )
165 }
166
167 fn take_ggsw_prepared_slice<A, M>(
168 &mut self,
169 module: &M,
170 size: usize,
171 infos: &A,
172 ) -> (Vec<GGSWPrepared<&mut [u8], B>>, &mut Self)
173 where
174 A: GGSWInfos,
175 M: ModuleN + VmpPMatBytesOf,
176 {
177 let mut scratch: &mut Self = self;
178 let mut cts: Vec<GGSWPrepared<&mut [u8], B>> = Vec::with_capacity(size);
179 for _ in 0..size {
180 let (ct, new_scratch) = scratch.take_ggsw_prepared(module, infos);
181 scratch = new_scratch;
182 cts.push(ct)
183 }
184 (cts, scratch)
185 }
186
187 fn take_glwe_pk<A, M>(&mut self, module: &M, infos: &A) -> (GLWEPublicKey<&mut [u8]>, &mut Self)
188 where
189 A: GLWEInfos,
190 M: ModuleN,
191 {
192 assert_eq!(module.n() as u32, infos.n());
193 let (data, scratch) = self.take_vec_znx(module, (infos.rank() + 1).into(), infos.size());
194 (
195 GLWEPublicKey {
196 k: infos.k(),
197 dist: Distribution::NONE,
198 base2k: infos.base2k(),
199 data,
200 },
201 scratch,
202 )
203 }
204
205 fn take_glwe_pk_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GLWEPublicKeyPrepared<&mut [u8], B>, &mut Self)
206 where
207 A: GLWEInfos,
208 M: ModuleN + VecZnxDftBytesOf,
209 {
210 assert_eq!(module.n() as u32, infos.n());
211 let (data, scratch) = self.take_vec_znx_dft(module, (infos.rank() + 1).into(), infos.size());
212 (
213 GLWEPublicKeyPrepared {
214 k: infos.k(),
215 dist: Distribution::NONE,
216 base2k: infos.base2k(),
217 data,
218 },
219 scratch,
220 )
221 }
222
223 fn take_glwe_secret<M>(&mut self, module: &M, rank: Rank) -> (GLWESecret<&mut [u8]>, &mut Self)
224 where
225 M: ModuleN,
226 {
227 let (data, scratch) = self.take_scalar_znx(module, rank.into());
228 (
229 GLWESecret {
230 data,
231 dist: Distribution::NONE,
232 },
233 scratch,
234 )
235 }
236
237 fn take_glwe_secret_prepared<M>(&mut self, module: &M, rank: Rank) -> (GLWESecretPrepared<&mut [u8], B>, &mut Self)
238 where
239 M: ModuleN + SvpPPolBytesOf,
240 {
241 let (data, scratch) = self.take_svp_ppol(module, rank.into());
242 (
243 GLWESecretPrepared {
244 data,
245 dist: Distribution::NONE,
246 },
247 scratch,
248 )
249 }
250
251 fn take_glwe_switching_key<A, M>(&mut self, module: &M, infos: &A) -> (GLWESwitchingKey<&mut [u8]>, &mut Self)
252 where
253 A: GGLWEInfos,
254 M: ModuleN,
255 {
256 assert_eq!(module.n() as u32, infos.n());
257 let (data, scratch) = self.take_gglwe(module, infos);
258 (
259 GLWESwitchingKey {
260 key: data,
261 input_degree: Degree(0),
262 output_degree: Degree(0),
263 },
264 scratch,
265 )
266 }
267
268 fn take_gglwe_switching_key_prepared<A, M>(
269 &mut self,
270 module: &M,
271 infos: &A,
272 ) -> (GLWESwitchingKeyPrepared<&mut [u8], B>, &mut Self)
273 where
274 A: GGLWEInfos,
275 M: ModuleN + VmpPMatBytesOf,
276 {
277 assert_eq!(module.n() as u32, infos.n());
278 let (data, scratch) = self.take_gglwe_prepared(module, infos);
279 (
280 GLWESwitchingKeyPrepared {
281 key: data,
282 input_degree: Degree(0),
283 output_degree: Degree(0),
284 },
285 scratch,
286 )
287 }
288
289 fn take_gglwe_automorphism_key<A, M>(&mut self, module: &M, infos: &A) -> (AutomorphismKey<&mut [u8]>, &mut Self)
290 where
291 A: GGLWEInfos,
292 M: ModuleN,
293 {
294 assert_eq!(module.n() as u32, infos.n());
295 let (data, scratch) = self.take_gglwe(module, infos);
296 (AutomorphismKey { key: data, p: 0 }, scratch)
297 }
298
299 fn take_gglwe_automorphism_key_prepared<A, M>(
300 &mut self,
301 module: &M,
302 infos: &A,
303 ) -> (AutomorphismKeyPrepared<&mut [u8], B>, &mut Self)
304 where
305 A: GGLWEInfos,
306 M: ModuleN + VmpPMatBytesOf,
307 {
308 assert_eq!(module.n() as u32, infos.n());
309 let (data, scratch) = self.take_gglwe_prepared(module, infos);
310 (AutomorphismKeyPrepared { key: data, p: 0 }, scratch)
311 }
312
313 fn take_tensor_key<A, M>(&mut self, module: &M, infos: &A) -> (TensorKey<&mut [u8]>, &mut Self)
314 where
315 A: GGLWEInfos,
316 M: ModuleN,
317 {
318 assert_eq!(module.n() as u32, infos.n());
319 assert_eq!(
320 infos.rank_in(),
321 infos.rank_out(),
322 "rank_in != rank_out is not supported for GGLWETensorKey"
323 );
324 let mut keys: Vec<GGLWE<&mut [u8]>> = Vec::new();
325 let pairs: usize = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1) as usize;
326
327 let mut scratch: &mut Self = self;
328
329 let mut ksk_infos: GGLWELayout = infos.gglwe_layout();
330 ksk_infos.rank_in = Rank(1);
331
332 if pairs != 0 {
333 let (gglwe, s) = scratch.take_gglwe(module, &ksk_infos);
334 scratch = s;
335 keys.push(gglwe);
336 }
337 for _ in 1..pairs {
338 let (gglwe, s) = scratch.take_gglwe(module, &ksk_infos);
339 scratch = s;
340 keys.push(gglwe);
341 }
342 (TensorKey { keys }, scratch)
343 }
344
345 fn take_gglwe_tensor_key_prepared<A, M>(&mut self, module: &M, infos: &A) -> (TensorKeyPrepared<&mut [u8], B>, &mut Self)
346 where
347 A: GGLWEInfos,
348 M: ModuleN + VmpPMatBytesOf,
349 {
350 assert_eq!(module.n() as u32, infos.n());
351 assert_eq!(
352 infos.rank_in(),
353 infos.rank_out(),
354 "rank_in != rank_out is not supported for GGLWETensorKeyPrepared"
355 );
356
357 let mut keys: Vec<GGLWEPrepared<&mut [u8], B>> = Vec::new();
358 let pairs: usize = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1) as usize;
359
360 let mut scratch: &mut Self = self;
361
362 let mut ksk_infos: GGLWELayout = infos.gglwe_layout();
363 ksk_infos.rank_in = Rank(1);
364
365 if pairs != 0 {
366 let (gglwe, s) = scratch.take_gglwe_prepared(module, &ksk_infos);
367 scratch = s;
368 keys.push(gglwe);
369 }
370 for _ in 1..pairs {
371 let (gglwe, s) = scratch.take_gglwe_prepared(module, &ksk_infos);
372 scratch = s;
373 keys.push(gglwe);
374 }
375 (TensorKeyPrepared { keys }, scratch)
376 }
377}
378
379impl<B: Backend> ScratchTakeCore<B> for Scratch<B> where Self: ScratchTakeBasic + ScratchAvailable {}