poulpy/poulpy-core/src/layouts/prepared/gglwe_ct.rs
poulpy-core/src/layouts/prepared/gglwe_ct.rsRaw
1use poulpy_hal::{
2 api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare, VmpPrepareTmpBytes},
3 layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat, VmpPMatToMut, VmpPMatToRef, ZnxInfos},
4 oep::VmpPMatAllocBytesImpl,
5};
6
7use crate::layouts::{
8 Base2K, BuildError, Degree, Dnum, Dsize, GGLWE, GGLWEInfos, GGLWEToRef, GLWEInfos, LWEInfos, Rank, TorusPrecision,
9};
10
11#[derive(PartialEq, Eq)]
12pub struct GGLWEPrepared<D: Data, B: Backend> {
13 pub(crate) data: VmpPMat<D, B>,
14 pub(crate) k: TorusPrecision,
15 pub(crate) base2k: Base2K,
16 pub(crate) dsize: Dsize,
17}
18
19impl<D: Data, B: Backend> LWEInfos for GGLWEPrepared<D, B> {
20 fn n(&self) -> Degree {
21 Degree(self.data.n() as u32)
22 }
23
24 fn base2k(&self) -> Base2K {
25 self.base2k
26 }
27
28 fn k(&self) -> TorusPrecision {
29 self.k
30 }
31
32 fn size(&self) -> usize {
33 self.data.size()
34 }
35}
36
37impl<D: Data, B: Backend> GLWEInfos for GGLWEPrepared<D, B> {
38 fn rank(&self) -> Rank {
39 self.rank_out()
40 }
41}
42
43impl<D: Data, B: Backend> GGLWEInfos for GGLWEPrepared<D, B> {
44 fn rank_in(&self) -> Rank {
45 Rank(self.data.cols_in() as u32)
46 }
47
48 fn rank_out(&self) -> Rank {
49 Rank(self.data.cols_out() as u32 - 1)
50 }
51
52 fn dsize(&self) -> Dsize {
53 self.dsize
54 }
55
56 fn dnum(&self) -> Dnum {
57 Dnum(self.data.rows() as u32)
58 }
59}
60
61pub struct GGLWEPreparedBuilder<D: Data, B: Backend> {
62 data: Option<VmpPMat<D, B>>,
63 base2k: Option<Base2K>,
64 k: Option<TorusPrecision>,
65 dsize: Option<Dsize>,
66}
67
68impl<D: Data, B: Backend> GGLWEPrepared<D, B> {
69 #[inline]
70 pub fn builder() -> GGLWEPreparedBuilder<D, B> {
71 GGLWEPreparedBuilder {
72 data: None,
73 base2k: None,
74 k: None,
75 dsize: None,
76 }
77 }
78}
79
80impl<B: Backend> GGLWEPreparedBuilder<Vec<u8>, B> {
81 #[inline]
82 pub fn layout<A>(mut self, infos: &A) -> Self
83 where
84 A: GGLWEInfos,
85 B: VmpPMatAllocBytesImpl<B>,
86 {
87 self.data = Some(VmpPMat::alloc(
88 infos.n().into(),
89 infos.dnum().into(),
90 infos.rank_in().into(),
91 (infos.rank_out() + 1).into(),
92 infos.size(),
93 ));
94 self.base2k = Some(infos.base2k());
95 self.k = Some(infos.k());
96 self.dsize = Some(infos.dsize());
97 self
98 }
99}
100
101impl<D: Data, B: Backend> GGLWEPreparedBuilder<D, B> {
102 #[inline]
103 pub fn data(mut self, data: VmpPMat<D, B>) -> Self {
104 self.data = Some(data);
105 self
106 }
107 #[inline]
108 pub fn base2k(mut self, base2k: Base2K) -> Self {
109 self.base2k = Some(base2k);
110 self
111 }
112 #[inline]
113 pub fn k(mut self, k: TorusPrecision) -> Self {
114 self.k = Some(k);
115 self
116 }
117
118 #[inline]
119 pub fn dsize(mut self, dsize: Dsize) -> Self {
120 self.dsize = Some(dsize);
121 self
122 }
123
124 pub fn build(self) -> Result<GGLWEPrepared<D, B>, BuildError> {
125 let data: VmpPMat<D, B> = self.data.ok_or(BuildError::MissingData)?;
126 let base2k: Base2K = self.base2k.ok_or(BuildError::MissingBase2K)?;
127 let k: TorusPrecision = self.k.ok_or(BuildError::MissingK)?;
128 let dsize: Dsize = self.dsize.ok_or(BuildError::MissingDigits)?;
129
130 if base2k == 0_u32 {
131 return Err(BuildError::ZeroBase2K);
132 }
133
134 if dsize == 0_u32 {
135 return Err(BuildError::ZeroBase2K);
136 }
137
138 if k == 0_u32 {
139 return Err(BuildError::ZeroTorusPrecision);
140 }
141
142 if data.n() == 0 {
143 return Err(BuildError::ZeroDegree);
144 }
145
146 if data.cols() == 0 {
147 return Err(BuildError::ZeroCols);
148 }
149
150 if data.size() == 0 {
151 return Err(BuildError::ZeroLimbs);
152 }
153
154 Ok(GGLWEPrepared {
155 data,
156 base2k,
157 k,
158 dsize,
159 })
160 }
161}
162
163pub trait GGLWEPreparedAlloc<B: Backend> {
164 fn gglwe_prepared_alloc(
165 &self,
166 base2k: Base2K,
167 k: TorusPrecision,
168 rank_in: Rank,
169 rank_out: Rank,
170 dnum: Dnum,
171 dsize: Dsize,
172 ) -> GGLWEPrepared<Vec<u8>, B>;
173}
174
175impl<B: Backend> GGLWEPreparedAlloc<B> for Module<B>
176where
177 Module<B>: VmpPMatAlloc<B>,
178{
179 fn gglwe_prepared_alloc(
180 &self,
181 base2k: Base2K,
182 k: TorusPrecision,
183 rank_in: Rank,
184 rank_out: Rank,
185 dnum: Dnum,
186 dsize: Dsize,
187 ) -> GGLWEPrepared<Vec<u8>, B> {
188 let size: usize = k.0.div_ceil(base2k.0) as usize;
189 debug_assert!(
190 size as u32 > dsize.0,
191 "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
192 dsize.0
193 );
194
195 assert!(
196 dnum.0 * dsize.0 <= size as u32,
197 "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
198 dnum.0,
199 dsize.0,
200 );
201
202 GGLWEPrepared {
203 data: self.vmp_pmat_alloc(dnum.into(), rank_in.into(), (rank_out + 1).into(), size),
204 k,
205 base2k,
206 dsize,
207 }
208 }
209}
210
211pub trait GGLWEPreparedAllocFromInfos<B: Backend> {
212 fn gglwe_prepared_alloc_from_infos<A>(&self, infos: &A) -> GGLWEPrepared<Vec<u8>, B>
213 where
214 A: GGLWEInfos;
215}
216
217impl<B: Backend> GGLWEPreparedAllocFromInfos<B> for Module<B>
218where
219 Module<B>: GGLWEPreparedAlloc<B>,
220{
221 fn gglwe_prepared_alloc_from_infos<A>(&self, infos: &A) -> GGLWEPrepared<Vec<u8>, B>
222 where
223 A: GGLWEInfos,
224 {
225 assert_eq!(self.n() as u32, infos.n(), "module.n() != infos.n()");
226 self.gglwe_prepared_alloc(
227 infos.base2k(),
228 infos.k(),
229 infos.rank_in(),
230 infos.rank_out(),
231 infos.dnum(),
232 infos.dsize(),
233 )
234 }
235}
236
237pub trait GGLWEPreparedAllocBytes<B: Backend> {
238 fn gglwe_prepared_alloc_bytes(
239 &self,
240 base2k: Base2K,
241 k: TorusPrecision,
242 rank_in: Rank,
243 rank_out: Rank,
244 dnum: Dnum,
245 dsize: Dsize,
246 ) -> usize;
247}
248
249impl<B: Backend> GGLWEPreparedAllocBytes<B> for Module<B>
250where
251 Module<B>: VmpPMatAllocBytes,
252{
253 fn gglwe_prepared_alloc_bytes(
254 &self,
255 base2k: Base2K,
256 k: TorusPrecision,
257 rank_in: Rank,
258 rank_out: Rank,
259 dnum: Dnum,
260 dsize: Dsize,
261 ) -> usize {
262 let size: usize = k.0.div_ceil(base2k.0) as usize;
263 debug_assert!(
264 size as u32 > dsize.0,
265 "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
266 dsize.0
267 );
268
269 assert!(
270 dnum.0 * dsize.0 <= size as u32,
271 "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
272 dnum.0,
273 dsize.0,
274 );
275
276 self.vmp_pmat_alloc_bytes(dnum.into(), rank_in.into(), (rank_out + 1).into(), size)
277 }
278}
279
280pub trait GGLWEPreparedAllocBytesFromInfos<B: Backend> {
281 fn gglwe_prepared_alloc_bytes_from_infos<A>(&self, infos: &A) -> usize
282 where
283 A: GGLWEInfos;
284}
285
286impl<B: Backend> GGLWEPreparedAllocBytesFromInfos<B> for Module<B>
287where
288 Module<B>: GGLWEPreparedAllocBytes<B>,
289{
290 fn gglwe_prepared_alloc_bytes_from_infos<A>(&self, infos: &A) -> usize
291 where
292 A: GGLWEInfos,
293 {
294 assert_eq!(self.n() as u32, infos.n(), "module.n() != infos.n()");
295 self.gglwe_prepared_alloc_bytes(
296 infos.base2k(),
297 infos.k(),
298 infos.rank_in(),
299 infos.rank_out(),
300 infos.dnum(),
301 infos.dsize(),
302 )
303 }
304}
305
306impl<B: Backend> GGLWEPrepared<Vec<u8>, B> {
307 pub fn alloc_from_infos<A>(module: &Module<B>, infos: &A) -> Self
308 where
309 A: GGLWEInfos,
310 Module<B>: GGLWEPreparedAllocFromInfos<B>,
311 {
312 module.gglwe_prepared_alloc_from_infos(infos)
313 }
314
315 pub fn alloc(
316 module: &Module<B>,
317 base2k: Base2K,
318 k: TorusPrecision,
319 rank_in: Rank,
320 rank_out: Rank,
321 dnum: Dnum,
322 dsize: Dsize,
323 ) -> Self
324 where
325 Module<B>: GGLWEPreparedAlloc<B>,
326 {
327 module.gglwe_prepared_alloc(base2k, k, rank_in, rank_out, dnum, dsize)
328 }
329
330 pub fn alloc_bytes_from_infos<A>(module: &Module<B>, infos: &A) -> usize
331 where
332 A: GGLWEInfos,
333 Module<B>: GGLWEPreparedAllocBytesFromInfos<B>,
334 {
335 module.gglwe_prepared_alloc_bytes_from_infos(infos)
336 }
337
338 pub fn alloc_bytes(
339 module: &Module<B>,
340 base2k: Base2K,
341 k: TorusPrecision,
342 rank_in: Rank,
343 rank_out: Rank,
344 dnum: Dnum,
345 dsize: Dsize,
346 ) -> usize
347 where
348 Module<B>: GGLWEPreparedAllocBytes<B>,
349 {
350 module.gglwe_prepared_alloc_bytes(base2k, k, rank_in, rank_out, dnum, dsize)
351 }
352}
353
354pub trait GGLWEPrepareTmpBytes {
355 fn gglwe_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
356 where
357 A: GGLWEInfos;
358}
359
360impl<B: Backend> GGLWEPrepareTmpBytes for Module<B>
361where
362 Module<B>: VmpPrepareTmpBytes,
363{
364 fn gglwe_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
365 where
366 A: GGLWEInfos,
367 {
368 self.vmp_prepare_tmp_bytes(
369 infos.dnum().into(),
370 infos.rank_in().into(),
371 (infos.rank() + 1).into(),
372 infos.size(),
373 )
374 }
375}
376
377impl<B: Backend> GGLWEPrepared<Vec<u8>, B> {
378 pub fn prepare_tmp_bytes(&self, module: &Module<B>) -> usize
379 where
380 Module<B>: GGLWEPrepareTmpBytes,
381 {
382 module.gglwe_prepare_tmp_bytes(self)
383 }
384}
385
386pub trait GGLWEPrepare<B: Backend> {
387 fn gglwe_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
388 where
389 R: GGLWEPreparedToMut<B>,
390 O: GGLWEToRef;
391}
392
393impl<B: Backend> GGLWEPrepare<B> for Module<B>
394where
395 Module<B>: VmpPrepare<B>,
396{
397 fn gglwe_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
398 where
399 R: GGLWEPreparedToMut<B>,
400 O: GGLWEToRef,
401 {
402 let mut res: GGLWEPrepared<&mut [u8], B> = res.to_mut();
403 let other: GGLWE<&[u8]> = other.to_ref();
404
405 assert_eq!(res.n(), self.n() as u32);
406 assert_eq!(other.n(), self.n() as u32);
407 assert_eq!(res.base2k, other.base2k);
408 assert_eq!(res.k, other.k);
409 assert_eq!(res.dsize, other.dsize);
410
411 self.vmp_prepare(&mut res.data, &other.data, scratch);
412 }
413}
414
415impl<D: DataMut, B: Backend> GGLWEPrepared<D, B>
416where
417 Module<B>: GGLWEPrepare<B>,
418{
419 fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &mut Scratch<B>)
420 where
421 O: GGLWEToRef,
422 {
423 module.gglwe_prepare(self, other, scratch);
424 }
425}
426
427pub trait GGLWEPreparedToMut<B: Backend> {
428 fn to_mut(&mut self) -> GGLWEPrepared<&mut [u8], B>;
429}
430
431impl<D: DataMut, B: Backend> GGLWEPreparedToMut<B> for GGLWEPrepared<D, B> {
432 fn to_mut(&mut self) -> GGLWEPrepared<&mut [u8], B> {
433 GGLWEPrepared {
434 k: self.k,
435 base2k: self.base2k,
436 dsize: self.dsize,
437 data: self.data.to_mut(),
438 }
439 }
440}
441
442pub trait GGLWEPreparedToRef<B: Backend> {
443 fn to_ref(&self) -> GGLWEPrepared<&[u8], B>;
444}
445
446impl<D: DataRef, B: Backend> GGLWEPreparedToRef<B> for GGLWEPrepared<D, B> {
447 fn to_ref(&self) -> GGLWEPrepared<&[u8], B> {
448 GGLWEPrepared {
449 k: self.k,
450 base2k: self.base2k,
451 dsize: self.dsize,
452 data: self.data.to_ref(),
453 }
454 }
455}