poulpy/poulpy-core/src/layouts/prepared/ggsw_ct.rs
poulpy-core/src/layouts/prepared/ggsw_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, GGSW, GGSWInfos, GGSWToMut, GGSWToRef, GLWEInfos, LWEInfos, Rank, TorusPrecision,
9};
10
11#[derive(PartialEq, Eq)]
12pub struct GGSWPrepared<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 GGSWPrepared<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 GGSWPrepared<D, B> {
38 fn rank(&self) -> Rank {
39 Rank(self.data.cols_out() as u32 - 1)
40 }
41}
42
43impl<D: Data, B: Backend> GGSWInfos for GGSWPrepared<D, B> {
44 fn dsize(&self) -> Dsize {
45 self.dsize
46 }
47
48 fn dnum(&self) -> Dnum {
49 Dnum(self.data.rows() as u32)
50 }
51}
52
53pub struct GGSWCiphertextPreparedBuilder<D: Data, B: Backend> {
54 data: Option<VmpPMat<D, B>>,
55 base2k: Option<Base2K>,
56 k: Option<TorusPrecision>,
57 dsize: Option<Dsize>,
58}
59
60impl<D: Data, B: Backend> GGSWPrepared<D, B> {
61 #[inline]
62 pub fn builder() -> GGSWCiphertextPreparedBuilder<D, B> {
63 GGSWCiphertextPreparedBuilder {
64 data: None,
65 base2k: None,
66 k: None,
67 dsize: None,
68 }
69 }
70}
71
72impl<B: Backend> GGSWCiphertextPreparedBuilder<Vec<u8>, B> {
73 #[inline]
74 pub fn layout<A>(mut self, infos: &A) -> Self
75 where
76 A: GGSWInfos,
77 B: VmpPMatAllocBytesImpl<B>,
78 {
79 debug_assert!(
80 infos.size() as u32 > infos.dsize().0,
81 "invalid ggsw: ceil(k/base2k): {} <= dsize: {}",
82 infos.size(),
83 infos.dsize()
84 );
85
86 assert!(
87 infos.dnum().0 * infos.dsize().0 <= infos.size() as u32,
88 "invalid ggsw: dnum: {} * dsize:{} > ceil(k/base2k): {}",
89 infos.dnum(),
90 infos.dsize(),
91 infos.size(),
92 );
93
94 self.data = Some(VmpPMat::alloc(
95 infos.n().into(),
96 infos.dnum().into(),
97 (infos.rank() + 1).into(),
98 (infos.rank() + 1).into(),
99 infos.size(),
100 ));
101 self.base2k = Some(infos.base2k());
102 self.k = Some(infos.k());
103 self.dsize = Some(infos.dsize());
104 self
105 }
106}
107
108impl<D: Data, B: Backend> GGSWCiphertextPreparedBuilder<D, B> {
109 #[inline]
110 pub fn data(mut self, data: VmpPMat<D, B>) -> Self {
111 self.data = Some(data);
112 self
113 }
114 #[inline]
115 pub fn base2k(mut self, base2k: Base2K) -> Self {
116 self.base2k = Some(base2k);
117 self
118 }
119 #[inline]
120 pub fn k(mut self, k: TorusPrecision) -> Self {
121 self.k = Some(k);
122 self
123 }
124
125 #[inline]
126 pub fn dsize(mut self, dsize: Dsize) -> Self {
127 self.dsize = Some(dsize);
128 self
129 }
130
131 pub fn build(self) -> Result<GGSWPrepared<D, B>, BuildError> {
132 let data: VmpPMat<D, B> = self.data.ok_or(BuildError::MissingData)?;
133 let base2k: Base2K = self.base2k.ok_or(BuildError::MissingBase2K)?;
134 let k: TorusPrecision = self.k.ok_or(BuildError::MissingK)?;
135 let dsize: Dsize = self.dsize.ok_or(BuildError::MissingDigits)?;
136
137 if base2k == 0_u32 {
138 return Err(BuildError::ZeroBase2K);
139 }
140
141 if dsize == 0_u32 {
142 return Err(BuildError::ZeroBase2K);
143 }
144
145 if k == 0_u32 {
146 return Err(BuildError::ZeroTorusPrecision);
147 }
148
149 if data.n() == 0 {
150 return Err(BuildError::ZeroDegree);
151 }
152
153 if data.cols() == 0 {
154 return Err(BuildError::ZeroCols);
155 }
156
157 if data.size() == 0 {
158 return Err(BuildError::ZeroLimbs);
159 }
160
161 Ok(GGSWPrepared {
162 data,
163 base2k,
164 k,
165 dsize,
166 })
167 }
168}
169
170impl<B: Backend> GGSWPrepared<Vec<u8>, B> {
171 pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
172 where
173 A: GGSWInfos,
174 Module<B>: VmpPMatAlloc<B>,
175 {
176 Self::alloc_with(
177 module,
178 infos.base2k(),
179 infos.k(),
180 infos.dnum(),
181 infos.dsize(),
182 infos.rank(),
183 )
184 }
185
186 pub fn alloc_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, rank: Rank) -> Self
187 where
188 Module<B>: VmpPMatAlloc<B>,
189 {
190 let size: usize = k.0.div_ceil(base2k.0) as usize;
191 debug_assert!(
192 size as u32 > dsize.0,
193 "invalid ggsw: ceil(k/base2k): {size} <= dsize: {}",
194 dsize.0
195 );
196
197 assert!(
198 dnum.0 * dsize.0 <= size as u32,
199 "invalid ggsw: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
200 dnum.0,
201 dsize.0,
202 );
203
204 Self {
205 data: module.vmp_pmat_alloc(
206 dnum.into(),
207 (rank + 1).into(),
208 (rank + 1).into(),
209 k.0.div_ceil(base2k.0) as usize,
210 ),
211 k,
212 base2k,
213 dsize,
214 }
215 }
216
217 pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
218 where
219 A: GGSWInfos,
220 Module<B>: VmpPMatAllocBytes,
221 {
222 Self::alloc_bytes_with(
223 module,
224 infos.base2k(),
225 infos.k(),
226 infos.dnum(),
227 infos.dsize(),
228 infos.rank(),
229 )
230 }
231
232 pub fn alloc_bytes_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, rank: Rank) -> usize
233 where
234 Module<B>: VmpPMatAllocBytes,
235 {
236 let size: usize = k.0.div_ceil(base2k.0) as usize;
237 debug_assert!(
238 size as u32 > dsize.0,
239 "invalid ggsw: ceil(k/base2k): {size} <= dsize: {}",
240 dsize.0
241 );
242
243 assert!(
244 dnum.0 * dsize.0 <= size as u32,
245 "invalid ggsw: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
246 dnum.0,
247 dsize.0,
248 );
249
250 module.vmp_pmat_alloc_bytes(dnum.into(), (rank + 1).into(), (rank + 1).into(), size)
251 }
252}
253
254impl<D: DataRef, B: Backend> GGSWPrepared<D, B> {
255 pub fn data(&self) -> &VmpPMat<D, B> {
256 &self.data
257 }
258}
259
260pub trait GGSWPrepareTmpBytes {
261 fn ggsw_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
262 where
263 A: GGSWInfos;
264}
265
266impl<B: Backend> GGSWPrepareTmpBytes for Module<B>
267where
268 Module<B>: VmpPrepareTmpBytes,
269{
270 fn ggsw_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
271 where
272 A: GGSWInfos,
273 {
274 self.vmp_prepare_tmp_bytes(
275 infos.dnum().into(),
276 (infos.rank() + 1).into(),
277 (infos.rank() + 1).into(),
278 infos.size(),
279 )
280 }
281}
282
283impl<B: Backend> GGSWPrepared<Vec<u8>, B>
284where
285 Module<B>: GGSWPrepareTmpBytes,
286{
287 pub fn prepare_tmp_bytes<A>(&self, module: Module<B>, infos: &A) -> usize
288 where
289 A: GGSWInfos,
290 {
291 module.ggsw_prepare_tmp_bytes(self)
292 }
293}
294
295pub trait GGSWPrepare<B: Backend> {
296 fn ggsw_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
297 where
298 R: GGSWPreparedToMut<B>,
299 O: GGSWToRef;
300}
301
302impl<B: Backend> GGSWPrepare<B> for Module<B>
303where
304 Module<B>: VmpPrepare<B>,
305{
306 fn ggsw_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
307 where
308 R: GGSWPreparedToMut<B>,
309 O: GGSWToRef,
310 {
311 let mut res: GGSWPrepared<&mut [u8], B> = res.to_mut();
312 let other: GGSW<&[u8]> = other.to_ref();
313 assert_eq!(res.k, other.k);
314 assert_eq!(res.base2k, other.base2k);
315 assert_eq!(res.dsize, other.dsize);
316 self.vmp_prepare(&mut res.data, &other.data, scratch);
317 }
318}
319
320impl<D: DataMut, B: Backend> GGSWPrepared<D, B>
321where
322 Module<B>: GGSWPrepare<B>,
323{
324 pub fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &mut Scratch<B>)
325 where
326 O: GGSWToRef,
327 {
328 module.ggsw_prepare(self, other, scratch);
329 }
330}
331
332pub trait GGSWPrepareAlloc<B: Backend> {
333 fn ggsw_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
334 where
335 O: GGSWToRef;
336}
337
338impl<B: Backend> GGSWPrepareAlloc<B> for Module<B>
339where
340 Module<B>: GGSWPrepare<B>,
341{
342 fn ggsw_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
343 where
344 O: GGSWToRef,
345 {
346 let mut ct_prepared: GGSWPrepared<Vec<u8>, B> = GGSWPrepared::alloc(self, other);
347 self.ggsw_prepare(&mut ct_prepared, other, scratch);
348 ct_prepared
349 }
350}
351
352impl<D: DataRef> GGSW<D> {
353 fn prepare_alloc<B: Backend>(&self, module: &Module<B>, scratch: &mut Scratch<B>)
354 where
355 Module<B>: GGSWPrepareAlloc<B>,
356 {
357 module.ggsw_prepare_alloc(self, scratch);
358 }
359}
360
361pub trait GGSWPreparedToMut<B: Backend> {
362 fn to_mut(&mut self) -> GGSWPrepared<&mut [u8], B>;
363}
364
365impl<D: DataMut, B: Backend> GGSWPreparedToMut<B> for GGSWPrepared<D, B> {
366 fn to_mut(&mut self) -> GGSWPrepared<&mut [u8], B> {
367 GGSWPrepared {
368 base2k: self.base2k,
369 k: self.k,
370 dsize: self.dsize,
371 data: self.data.to_mut(),
372 }
373 }
374}
375
376pub trait GGSWCiphertextPreparedToRef<B: Backend> {
377 fn to_ref(&self) -> GGSWPrepared<&[u8], B>;
378}
379
380impl<D: DataRef, B: Backend> GGSWCiphertextPreparedToRef<B> for GGSWPrepared<D, B> {
381 fn to_ref(&self) -> GGSWPrepared<&[u8], B> {
382 GGSWPrepared {
383 base2k: self.base2k,
384 k: self.k,
385 dsize: self.dsize,
386 data: self.data.to_mut(),
387 }
388 }
389}