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
163impl<B: Backend> GGLWEPrepared<Vec<u8>, B> {
164 pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
165 where
166 A: GGLWEInfos,
167 Module<B>: VmpPMatAlloc<B>,
168 {
169 debug_assert_eq!(module.n(), infos.n().0 as usize, "module.n() != infos.n()");
170 Self::alloc_with(
171 module,
172 infos.base2k(),
173 infos.k(),
174 infos.rank_in(),
175 infos.rank_out(),
176 infos.dnum(),
177 infos.dsize(),
178 )
179 }
180
181 pub fn alloc_with(
182 module: &Module<B>,
183 base2k: Base2K,
184 k: TorusPrecision,
185 rank_in: Rank,
186 rank_out: Rank,
187 dnum: Dnum,
188 dsize: Dsize,
189 ) -> Self
190 where
191 Module<B>: VmpPMatAlloc<B>,
192 {
193 let size: usize = k.0.div_ceil(base2k.0) as usize;
194 debug_assert!(
195 size as u32 > dsize.0,
196 "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
197 dsize.0
198 );
199
200 assert!(
201 dnum.0 * dsize.0 <= size as u32,
202 "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
203 dnum.0,
204 dsize.0,
205 );
206
207 Self {
208 data: module.vmp_pmat_alloc(dnum.into(), rank_in.into(), (rank_out + 1).into(), size),
209 k,
210 base2k,
211 dsize,
212 }
213 }
214
215 pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
216 where
217 A: GGLWEInfos,
218 Module<B>: VmpPMatAllocBytes,
219 {
220 debug_assert_eq!(module.n(), infos.n().0 as usize, "module.n() != infos.n()");
221 Self::alloc_bytes_with(
222 module,
223 infos.base2k(),
224 infos.k(),
225 infos.rank_in(),
226 infos.rank_out(),
227 infos.dnum(),
228 infos.dsize(),
229 )
230 }
231
232 pub fn alloc_bytes_with(
233 module: &Module<B>,
234 base2k: Base2K,
235 k: TorusPrecision,
236 rank_in: Rank,
237 rank_out: Rank,
238 dnum: Dnum,
239 dsize: Dsize,
240 ) -> usize
241 where
242 Module<B>: VmpPMatAllocBytes,
243 {
244 let size: usize = k.0.div_ceil(base2k.0) as usize;
245 debug_assert!(
246 size as u32 > dsize.0,
247 "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
248 dsize.0
249 );
250
251 assert!(
252 dnum.0 * dsize.0 <= size as u32,
253 "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
254 dnum.0,
255 dsize.0,
256 );
257
258 module.vmp_pmat_alloc_bytes(dnum.into(), rank_in.into(), (rank_out + 1).into(), size)
259 }
260}
261
262pub trait GGLWEPrepareTmpBytes {
263 fn gglwe_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
264 where
265 A: GGLWEInfos;
266}
267
268impl<B: Backend> GGLWEPrepareTmpBytes for Module<B>
269where
270 Module<B>: VmpPrepareTmpBytes,
271{
272 fn gglwe_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
273 where
274 A: GGLWEInfos,
275 {
276 self.vmp_prepare_tmp_bytes(
277 infos.dnum().into(),
278 infos.rank_in().into(),
279 (infos.rank() + 1).into(),
280 infos.size(),
281 )
282 }
283}
284
285impl<B: Backend> GGLWEPrepared<Vec<u8>, B> {
286 pub fn prepare_tmp_bytes(&self, module: &Module<B>)
287 where
288 Module<B>: GGLWEPrepareTmpBytes,
289 {
290 module.gglwe_prepare_tmp_bytes(self)
291 }
292}
293
294pub trait GGLWEPrepare<B: Backend> {
295 fn gglwe_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
296 where
297 R: GGLWEPreparedToMut<B>,
298 O: GGLWEToRef;
299}
300
301impl<B: Backend> GGLWEPrepare<B> for Module<B>
302where
303 Module<B>: VmpPrepare<B>,
304{
305 fn gglwe_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
306 where
307 R: GGLWEPreparedToMut<B>,
308 O: GGLWEToRef,
309 {
310 let mut res: GGLWEPrepared<&mut [u8], B> = self.to_mut();
311 let other: GGLWE<&[u8]> = other.to_ref();
312
313 assert_eq!(res.n(), self.n() as u32);
314 assert_eq!(other.n(), self.n() as u32);
315 assert_eq!(res.base2k, other.base2k);
316 assert_eq!(res.k, other.k);
317 assert_eq!(res.dsize, other.dsize);
318
319 self.vmp_prepare(&mut res.data, &other.data, scratch);
320 }
321}
322
323impl<D: DataMut, B: Backend> GGLWEPrepared<D, B>
324where
325 Module<B>: GGLWEPrepare<B>,
326{
327 fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &mut Scratch<B>)
328 where
329 O: GGLWEToRef,
330 {
331 module.gglwe_prepare(self, other, scratch);
332 }
333}
334
335pub trait GGLWEPrepareAlloc<B: Backend> {
336 fn gglwe_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>) -> GGLWEPrepared<Vec<u8>, B>;
337}
338
339impl<B: Backend> GGLWEPrepareAlloc<B> for Module<B>
340where
341 Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
342{
343 fn gglwe_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>) -> GGLWEPrepared<Vec<u8>, B> {
344 let mut ct_prepared: GGLWEPrepared<Vec<u8>, B> = GGLWEPrepared::alloc(self, &other.to_ref());
345 ct_prepared.prepare(self, &other.to_ref(), scratch);
346 ct_prepared
347 }
348}
349
350impl<D: DataRef> GGLWE<D> {
351 fn prepare_alloc<B: Backend>(&self, module: &Module<B>, scratch: &Scratch<B>) -> GGLWEPrepared<Vec<u8>, B>
352 where
353 Module<B>: GGLWEPrepareAlloc<B>,
354 {
355 module.gglwe_prepare_alloc(self, scratch)
356 }
357}
358
359pub trait GGLWEPreparedToMut<B: Backend> {
360 fn to_mut(&mut self) -> GGLWEPrepared<&mut [u8], B>;
361}
362
363impl<D: DataMut, B: Backend> GGLWEPreparedToMut<B> for GGLWEPrepared<D, B> {
364 fn to_mut(&mut self) -> GGLWEPrepared<&mut [u8], B> {
365 GGLWEPrepared {
366 k: self.k,
367 base2k: self.base2k,
368 dsize: self.dsize,
369 data: self.data.to_mut(),
370 }
371 }
372}
373
374pub trait GGLWEPreparedToRef<B: Backend> {
375 fn to_ref(&self) -> GGLWEPrepared<&[u8], B>;
376}
377
378impl<D: DataRef, B: Backend> GGLWEPreparedToRef<B> for GGLWEPrepared<D, B> {
379 fn to_ref(&self) -> GGLWEPrepared<&[u8], B> {
380 GGLWEPrepared {
381 k: self.k,
382 base2k: self.base2k,
383 dsize: self.dsize,
384 data: self.data.to_ref(),
385 }
386 }
387}