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, GLWEPrepared,
10 GLWEPublicKey, GLWESecret, GLWESwitchingKey, Rank, TensorKey,
11 prepared::{
12 GGLWEPrepared, GGSWPrepared, GLWEAutomorphismKeyPrepared, GLWEPublicKeyPrepared, GLWESecretPrepared,
13 GLWESwitchingKeyPrepared, TensorKeyPrepared,
14 },
15 },
16};
17
18pub trait ScratchTakeCore<B: Backend>
19where
20 Self: ScratchTakeBasic + ScratchAvailable,
21{
22 fn take_glwe<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_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(module, infos);
48 scratch = new_scratch;
49 cts.push(ct);
50 }
51 (cts, scratch)
52 }
53
54 fn take_glwe_plaintext<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_public_key<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_glwe(module, infos);
194 (
195 GLWEPublicKey {
196 dist: Distribution::NONE,
197 key: data,
198 },
199 scratch,
200 )
201 }
202
203 fn take_glwe_public_key_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GLWEPublicKeyPrepared<&mut [u8], B>, &mut Self)
204 where
205 A: GLWEInfos,
206 M: ModuleN + VecZnxDftBytesOf,
207 {
208 let (data, scratch) = self.take_glwe_prepared(module, infos);
209 (
210 GLWEPublicKeyPrepared {
211 dist: Distribution::NONE,
212 key: data,
213 },
214 scratch,
215 )
216 }
217
218 fn take_glwe_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GLWEPrepared<&mut [u8], B>, &mut Self)
219 where
220 A: GLWEInfos,
221 M: ModuleN + VecZnxDftBytesOf,
222 {
223 assert_eq!(module.n() as u32, infos.n());
224 let (data, scratch) = self.take_vec_znx_dft(module, (infos.rank() + 1).into(), infos.size());
225 (
226 GLWEPrepared {
227 k: infos.k(),
228 base2k: infos.base2k(),
229 data,
230 },
231 scratch,
232 )
233 }
234
235 fn take_glwe_secret<M>(&mut self, module: &M, rank: Rank) -> (GLWESecret<&mut [u8]>, &mut Self)
236 where
237 M: ModuleN,
238 {
239 let (data, scratch) = self.take_scalar_znx(module, rank.into());
240 (
241 GLWESecret {
242 data,
243 dist: Distribution::NONE,
244 },
245 scratch,
246 )
247 }
248
249 fn take_glwe_secret_prepared<M>(&mut self, module: &M, rank: Rank) -> (GLWESecretPrepared<&mut [u8], B>, &mut Self)
250 where
251 M: ModuleN + SvpPPolBytesOf,
252 {
253 let (data, scratch) = self.take_svp_ppol(module, rank.into());
254 (
255 GLWESecretPrepared {
256 data,
257 dist: Distribution::NONE,
258 },
259 scratch,
260 )
261 }
262
263 fn take_glwe_switching_key<A, M>(&mut self, module: &M, infos: &A) -> (GLWESwitchingKey<&mut [u8]>, &mut Self)
264 where
265 A: GGLWEInfos,
266 M: ModuleN,
267 {
268 assert_eq!(module.n() as u32, infos.n());
269 let (data, scratch) = self.take_gglwe(module, infos);
270 (
271 GLWESwitchingKey {
272 key: data,
273 input_degree: Degree(0),
274 output_degree: Degree(0),
275 },
276 scratch,
277 )
278 }
279
280 fn take_glwe_switching_key_prepared<A, M>(
281 &mut self,
282 module: &M,
283 infos: &A,
284 ) -> (GLWESwitchingKeyPrepared<&mut [u8], B>, &mut Self)
285 where
286 A: GGLWEInfos,
287 M: ModuleN + VmpPMatBytesOf,
288 {
289 assert_eq!(module.n() as u32, infos.n());
290 let (data, scratch) = self.take_gglwe_prepared(module, infos);
291 (
292 GLWESwitchingKeyPrepared {
293 key: data,
294 input_degree: Degree(0),
295 output_degree: Degree(0),
296 },
297 scratch,
298 )
299 }
300
301 fn take_gglwe_automorphism_key<A, M>(&mut self, module: &M, infos: &A) -> (AutomorphismKey<&mut [u8]>, &mut Self)
302 where
303 A: GGLWEInfos,
304 M: ModuleN,
305 {
306 assert_eq!(module.n() as u32, infos.n());
307 let (data, scratch) = self.take_gglwe(module, infos);
308 (AutomorphismKey { key: data, p: 0 }, scratch)
309 }
310
311 fn take_gglwe_automorphism_key_prepared<A, M>(
312 &mut self,
313 module: &M,
314 infos: &A,
315 ) -> (GLWEAutomorphismKeyPrepared<&mut [u8], B>, &mut Self)
316 where
317 A: GGLWEInfos,
318 M: ModuleN + VmpPMatBytesOf,
319 {
320 assert_eq!(module.n() as u32, infos.n());
321 let (data, scratch) = self.take_gglwe_prepared(module, infos);
322 (GLWEAutomorphismKeyPrepared { key: data, p: 0 }, scratch)
323 }
324
325 fn take_tensor_key<A, M>(&mut self, module: &M, infos: &A) -> (TensorKey<&mut [u8]>, &mut Self)
326 where
327 A: GGLWEInfos,
328 M: ModuleN,
329 {
330 assert_eq!(module.n() as u32, infos.n());
331 assert_eq!(
332 infos.rank_in(),
333 infos.rank_out(),
334 "rank_in != rank_out is not supported for GGLWETensorKey"
335 );
336 let mut keys: Vec<GGLWE<&mut [u8]>> = Vec::new();
337 let pairs: usize = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1) as usize;
338
339 let mut scratch: &mut Self = self;
340
341 let mut ksk_infos: GGLWELayout = infos.gglwe_layout();
342 ksk_infos.rank_in = Rank(1);
343
344 if pairs != 0 {
345 let (gglwe, s) = scratch.take_gglwe(module, &ksk_infos);
346 scratch = s;
347 keys.push(gglwe);
348 }
349 for _ in 1..pairs {
350 let (gglwe, s) = scratch.take_gglwe(module, &ksk_infos);
351 scratch = s;
352 keys.push(gglwe);
353 }
354 (TensorKey { keys }, scratch)
355 }
356
357 fn take_gglwe_tensor_key_prepared<A, M>(&mut self, module: &M, infos: &A) -> (TensorKeyPrepared<&mut [u8], B>, &mut Self)
358 where
359 A: GGLWEInfos,
360 M: ModuleN + VmpPMatBytesOf,
361 {
362 assert_eq!(module.n() as u32, infos.n());
363 assert_eq!(
364 infos.rank_in(),
365 infos.rank_out(),
366 "rank_in != rank_out is not supported for GGLWETensorKeyPrepared"
367 );
368
369 let mut keys: Vec<GGLWEPrepared<&mut [u8], B>> = Vec::new();
370 let pairs: usize = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1) as usize;
371
372 let mut scratch: &mut Self = self;
373
374 let mut ksk_infos: GGLWELayout = infos.gglwe_layout();
375 ksk_infos.rank_in = Rank(1);
376
377 if pairs != 0 {
378 let (gglwe, s) = scratch.take_gglwe_prepared(module, &ksk_infos);
379 scratch = s;
380 keys.push(gglwe);
381 }
382 for _ in 1..pairs {
383 let (gglwe, s) = scratch.take_gglwe_prepared(module, &ksk_infos);
384 scratch = s;
385 keys.push(gglwe);
386 }
387 (TensorKeyPrepared { keys }, scratch)
388 }
389}
390
391impl<B: Backend> ScratchTakeCore<B> for Scratch<B> where Self: ScratchTakeBasic + ScratchAvailable {}