poulpy/core/src/gglwe/layouts_exec.rs
core/src/gglwe/layouts_exec.rsRaw
1use backend::hal::{
2 api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPMatPrepare},
3 layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat},
4};
5
6use crate::{AutomorphismKey, GGLWECiphertext, GGLWEExecLayoutFamily, GLWESwitchingKey, GLWETensorKey, Infos};
7
8#[derive(PartialEq, Eq)]
9pub struct GGLWECiphertextExec<D: Data, B: Backend> {
10 pub(crate) data: VmpPMat<D, B>,
11 pub(crate) basek: usize,
12 pub(crate) k: usize,
13 pub(crate) digits: usize,
14}
15
16impl<B: Backend> GGLWECiphertextExec<Vec<u8>, B> {
17 pub fn alloc(
18 module: &Module<B>,
19 n: usize,
20 basek: usize,
21 k: usize,
22 rows: usize,
23 digits: usize,
24 rank_in: usize,
25 rank_out: usize,
26 ) -> Self
27 where
28 Module<B>: GGLWEExecLayoutFamily<B>,
29 {
30 let size: usize = k.div_ceil(basek);
31 debug_assert!(
32 size > digits,
33 "invalid gglwe: ceil(k/basek): {} <= digits: {}",
34 size,
35 digits
36 );
37
38 assert!(
39 rows * digits <= size,
40 "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
41 rows,
42 digits,
43 size
44 );
45
46 Self {
47 data: module.vmp_pmat_alloc(n, rows, rank_in, rank_out + 1, size),
48 basek: basek,
49 k,
50 digits,
51 }
52 }
53
54 pub fn bytes_of(
55 module: &Module<B>,
56 n: usize,
57 basek: usize,
58 k: usize,
59 rows: usize,
60 digits: usize,
61 rank_in: usize,
62 rank_out: usize,
63 ) -> usize
64 where
65 Module<B>: GGLWEExecLayoutFamily<B>,
66 {
67 let size: usize = k.div_ceil(basek);
68 debug_assert!(
69 size > digits,
70 "invalid gglwe: ceil(k/basek): {} <= digits: {}",
71 size,
72 digits
73 );
74
75 assert!(
76 rows * digits <= size,
77 "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
78 rows,
79 digits,
80 size
81 );
82
83 module.vmp_pmat_alloc_bytes(n, rows, rank_in, rank_out + 1, rows)
84 }
85}
86
87impl<D: Data, B: Backend> Infos for GGLWECiphertextExec<D, B> {
88 type Inner = VmpPMat<D, B>;
89
90 fn inner(&self) -> &Self::Inner {
91 &self.data
92 }
93
94 fn basek(&self) -> usize {
95 self.basek
96 }
97
98 fn k(&self) -> usize {
99 self.k
100 }
101}
102
103impl<D: Data, B: Backend> GGLWECiphertextExec<D, B> {
104 pub fn rank(&self) -> usize {
105 self.data.cols_out() - 1
106 }
107
108 pub fn digits(&self) -> usize {
109 self.digits
110 }
111
112 pub fn rank_in(&self) -> usize {
113 self.data.cols_in()
114 }
115
116 pub fn rank_out(&self) -> usize {
117 self.data.cols_out() - 1
118 }
119}
120
121impl<D: DataMut, B: Backend> GGLWECiphertextExec<D, B> {
122 pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GGLWECiphertext<DataOther>, scratch: &mut Scratch<B>)
123 where
124 DataOther: DataRef,
125 Module<B>: GGLWEExecLayoutFamily<B>,
126 {
127 module.vmp_prepare(&mut self.data, &other.data, scratch);
128 self.basek = other.basek;
129 self.k = other.k;
130 self.digits = other.digits;
131 }
132}
133
134#[derive(PartialEq, Eq)]
135pub struct GLWESwitchingKeyExec<D: Data, B: Backend> {
136 pub(crate) key: GGLWECiphertextExec<D, B>,
137 pub(crate) sk_in_n: usize, // Degree of sk_in
138 pub(crate) sk_out_n: usize, // Degree of sk_out
139}
140
141impl<B: Backend> GLWESwitchingKeyExec<Vec<u8>, B> {
142 pub fn alloc(
143 module: &Module<B>,
144 n: usize,
145 basek: usize,
146 k: usize,
147 rows: usize,
148 digits: usize,
149 rank_in: usize,
150 rank_out: usize,
151 ) -> Self
152 where
153 Module<B>: GGLWEExecLayoutFamily<B>,
154 {
155 GLWESwitchingKeyExec::<Vec<u8>, B> {
156 key: GGLWECiphertextExec::alloc(module, n, basek, k, rows, digits, rank_in, rank_out),
157 sk_in_n: 0,
158 sk_out_n: 0,
159 }
160 }
161
162 pub fn bytes_of(
163 module: &Module<B>,
164 n: usize,
165 basek: usize,
166 k: usize,
167 rows: usize,
168 digits: usize,
169 rank_in: usize,
170 rank_out: usize,
171 ) -> usize
172 where
173 Module<B>: GGLWEExecLayoutFamily<B>,
174 {
175 GGLWECiphertextExec::bytes_of(module, n, basek, k, rows, digits, rank_in, rank_out)
176 }
177
178 pub fn from<DataOther: DataRef>(module: &Module<B>, other: &GLWESwitchingKey<DataOther>, scratch: &mut Scratch<B>) -> Self
179 where
180 Module<B>: GGLWEExecLayoutFamily<B>,
181 {
182 let mut ksk_exec: GLWESwitchingKeyExec<Vec<u8>, B> = Self::alloc(
183 module,
184 other.n(),
185 other.basek(),
186 other.k(),
187 other.rows(),
188 other.digits(),
189 other.rank_in(),
190 other.rank_out(),
191 );
192 ksk_exec.prepare(module, other, scratch);
193 ksk_exec
194 }
195}
196
197impl<D: Data, B: Backend> Infos for GLWESwitchingKeyExec<D, B> {
198 type Inner = VmpPMat<D, B>;
199
200 fn inner(&self) -> &Self::Inner {
201 self.key.inner()
202 }
203
204 fn basek(&self) -> usize {
205 self.key.basek()
206 }
207
208 fn k(&self) -> usize {
209 self.key.k()
210 }
211}
212
213impl<D: Data, B: Backend> GLWESwitchingKeyExec<D, B> {
214 pub fn rank(&self) -> usize {
215 self.key.data.cols_out() - 1
216 }
217
218 pub fn rank_in(&self) -> usize {
219 self.key.data.cols_in()
220 }
221
222 pub fn rank_out(&self) -> usize {
223 self.key.data.cols_out() - 1
224 }
225
226 pub fn digits(&self) -> usize {
227 self.key.digits()
228 }
229
230 pub fn sk_degree_in(&self) -> usize {
231 self.sk_in_n
232 }
233
234 pub fn sk_degree_out(&self) -> usize {
235 self.sk_out_n
236 }
237}
238
239impl<D: DataMut, B: Backend> GLWESwitchingKeyExec<D, B> {
240 pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GLWESwitchingKey<DataOther>, scratch: &mut Scratch<B>)
241 where
242 DataOther: DataRef,
243 Module<B>: GGLWEExecLayoutFamily<B>,
244 {
245 self.key.prepare(module, &other.key, scratch);
246 self.sk_in_n = other.sk_in_n;
247 self.sk_out_n = other.sk_out_n;
248 }
249}
250
251#[derive(PartialEq, Eq)]
252pub struct AutomorphismKeyExec<D: Data, B: Backend> {
253 pub(crate) key: GLWESwitchingKeyExec<D, B>,
254 pub(crate) p: i64,
255}
256
257impl<B: Backend> AutomorphismKeyExec<Vec<u8>, B> {
258 pub fn alloc(module: &Module<B>, n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self
259 where
260 Module<B>: GGLWEExecLayoutFamily<B>,
261 {
262 AutomorphismKeyExec::<Vec<u8>, B> {
263 key: GLWESwitchingKeyExec::alloc(module, n, basek, k, rows, digits, rank, rank),
264 p: 0,
265 }
266 }
267
268 pub fn bytes_of(module: &Module<B>, n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize
269 where
270 Module<B>: GGLWEExecLayoutFamily<B>,
271 {
272 GLWESwitchingKeyExec::bytes_of(module, n, basek, k, rows, digits, rank, rank)
273 }
274
275 pub fn from<DataOther: DataRef>(module: &Module<B>, other: &AutomorphismKey<DataOther>, scratch: &mut Scratch<B>) -> Self
276 where
277 Module<B>: GGLWEExecLayoutFamily<B>,
278 {
279 let mut atk_exec: AutomorphismKeyExec<Vec<u8>, B> = Self::alloc(
280 module,
281 other.n(),
282 other.basek(),
283 other.k(),
284 other.rows(),
285 other.digits(),
286 other.rank(),
287 );
288 atk_exec.prepare(module, other, scratch);
289 atk_exec
290 }
291}
292
293impl<D: DataMut, B: Backend> AutomorphismKeyExec<D, B> {
294 pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &AutomorphismKey<DataOther>, scratch: &mut Scratch<B>)
295 where
296 DataOther: DataRef,
297 Module<B>: GGLWEExecLayoutFamily<B>,
298 {
299 self.key.prepare(module, &other.key, scratch);
300 self.p = other.p;
301 }
302}
303
304impl<D: Data, B: Backend> Infos for AutomorphismKeyExec<D, B> {
305 type Inner = VmpPMat<D, B>;
306
307 fn inner(&self) -> &Self::Inner {
308 &self.key.inner()
309 }
310
311 fn basek(&self) -> usize {
312 self.key.basek()
313 }
314
315 fn k(&self) -> usize {
316 self.key.k()
317 }
318}
319
320impl<D: Data, B: Backend> AutomorphismKeyExec<D, B> {
321 pub fn p(&self) -> i64 {
322 self.p
323 }
324
325 pub fn digits(&self) -> usize {
326 self.key.digits()
327 }
328
329 pub fn rank(&self) -> usize {
330 self.key.rank()
331 }
332
333 pub fn rank_in(&self) -> usize {
334 self.key.rank_in()
335 }
336
337 pub fn rank_out(&self) -> usize {
338 self.key.rank_out()
339 }
340}
341
342#[derive(PartialEq, Eq)]
343pub struct GLWETensorKeyExec<D: Data, B: Backend> {
344 pub(crate) keys: Vec<GLWESwitchingKeyExec<D, B>>,
345}
346
347impl<B: Backend> GLWETensorKeyExec<Vec<u8>, B> {
348 pub fn alloc(module: &Module<B>, n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self
349 where
350 Module<B>: GGLWEExecLayoutFamily<B>,
351 {
352 let mut keys: Vec<GLWESwitchingKeyExec<Vec<u8>, B>> = Vec::new();
353 let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
354 (0..pairs).for_each(|_| {
355 keys.push(GLWESwitchingKeyExec::alloc(
356 module, n, basek, k, rows, digits, 1, rank,
357 ));
358 });
359 Self { keys }
360 }
361
362 pub fn bytes_of(module: &Module<B>, n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize
363 where
364 Module<B>: GGLWEExecLayoutFamily<B>,
365 {
366 let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
367 pairs * GLWESwitchingKeyExec::bytes_of(module, n, basek, k, rows, digits, 1, rank)
368 }
369
370 pub fn from<D: DataRef>(
371 module: &Module<B>,
372 other: &GLWETensorKey<D>,
373 scratch: &mut Scratch<B>,
374 ) -> GLWETensorKeyExec<Vec<u8>, B>
375 where
376 Module<B>: GGLWEExecLayoutFamily<B>,
377 {
378 let mut tsk_exec: GLWETensorKeyExec<Vec<u8>, B> = Self::alloc(
379 module,
380 other.n(),
381 other.basek(),
382 other.k(),
383 other.rows(),
384 other.digits(),
385 other.rank(),
386 );
387 tsk_exec.prepare(module, other, scratch);
388 tsk_exec
389 }
390}
391
392impl<D: Data, B: Backend> Infos for GLWETensorKeyExec<D, B> {
393 type Inner = VmpPMat<D, B>;
394
395 fn inner(&self) -> &Self::Inner {
396 &self.keys[0].inner()
397 }
398
399 fn basek(&self) -> usize {
400 self.keys[0].basek()
401 }
402
403 fn k(&self) -> usize {
404 self.keys[0].k()
405 }
406}
407
408impl<D: Data, B: Backend> GLWETensorKeyExec<D, B> {
409 pub fn rank(&self) -> usize {
410 self.keys[0].rank()
411 }
412
413 pub fn rank_in(&self) -> usize {
414 self.keys[0].rank_in()
415 }
416
417 pub fn rank_out(&self) -> usize {
418 self.keys[0].rank_out()
419 }
420
421 pub fn digits(&self) -> usize {
422 self.keys[0].digits()
423 }
424}
425
426impl<D: DataMut, B: Backend> GLWETensorKeyExec<D, B> {
427 // Returns a mutable reference to GLWESwitchingKey_{s}(s[i] * s[j])
428 pub fn at_mut(&mut self, mut i: usize, mut j: usize) -> &mut GLWESwitchingKeyExec<D, B> {
429 if i > j {
430 std::mem::swap(&mut i, &mut j);
431 };
432 let rank: usize = self.rank();
433 &mut self.keys[i * rank + j - (i * (i + 1) / 2)]
434 }
435}
436
437impl<D: DataRef, B: Backend> GLWETensorKeyExec<D, B> {
438 // Returns a reference to GLWESwitchingKey_{s}(s[i] * s[j])
439 pub fn at(&self, mut i: usize, mut j: usize) -> &GLWESwitchingKeyExec<D, B> {
440 if i > j {
441 std::mem::swap(&mut i, &mut j);
442 };
443 let rank: usize = self.rank();
444 &self.keys[i * rank + j - (i * (i + 1) / 2)]
445 }
446}
447
448impl<D: DataMut, B: Backend> GLWETensorKeyExec<D, B> {
449 pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GLWETensorKey<DataOther>, scratch: &mut Scratch<B>)
450 where
451 DataOther: DataRef,
452 Module<B>: GGLWEExecLayoutFamily<B>,
453 {
454 #[cfg(debug_assertions)]
455 {
456 assert_eq!(self.keys.len(), other.keys.len());
457 }
458 self.keys
459 .iter_mut()
460 .zip(other.keys.iter())
461 .for_each(|(a, b)| {
462 a.prepare(module, b, scratch);
463 });
464 }
465}