Skip to main content

poulpy_core/layouts/
gglwe.rs

1use poulpy_hal::{
2    layouts::{
3        Backend, Data, FillUniform, HostDataMut, HostDataRef, MatZnx, MatZnxAtBackendMut, MatZnxAtBackendRef, MatZnxToBackendMut,
4        MatZnxToBackendRef, Module, ReaderFrom, TransferFrom, WriterTo,
5    },
6    source::Source,
7};
8
9use crate::api::ModuleTransfer;
10use crate::layouts::{
11    Base2K, Degree, Dnum, Dsize, GGSWAtViewRef, GLWE, GLWEInfos, GLWEViewMut, GLWEViewRef, LWEInfos, Rank, TorusPrecision,
12};
13use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
14
15use std::{
16    fmt,
17    ops::{Deref, DerefMut},
18};
19
20pub trait GGLWEInfos
21where
22    Self: GLWEInfos,
23{
24    fn dnum(&self) -> Dnum;
25    fn dsize(&self) -> Dsize;
26    fn rank_in(&self) -> Rank;
27    fn rank_out(&self) -> Rank;
28    fn gglwe_layout(&self) -> GGLWELayout {
29        GGLWELayout {
30            n: self.n(),
31            base2k: self.base2k(),
32            k: self.max_k(),
33            rank_in: self.rank_in(),
34            rank_out: self.rank_out(),
35            dsize: self.dsize(),
36            dnum: self.dnum(),
37        }
38    }
39}
40
41pub trait SetGGLWEInfos {
42    fn set_dsize(&mut self, dsize: usize);
43}
44
45#[derive(PartialEq, Eq, Copy, Clone, Debug)]
46pub struct GGLWELayout {
47    pub n: Degree,
48    pub base2k: Base2K,
49    pub k: TorusPrecision,
50    pub rank_in: Rank,
51    pub rank_out: Rank,
52    pub dnum: Dnum,
53    pub dsize: Dsize,
54}
55
56impl LWEInfos for GGLWELayout {
57    fn base2k(&self) -> Base2K {
58        self.base2k
59    }
60
61    fn n(&self) -> Degree {
62        self.n
63    }
64
65    fn size(&self) -> usize {
66        self.k.as_usize().div_ceil(self.base2k.as_usize())
67    }
68}
69
70impl GLWEInfos for GGLWELayout {
71    fn rank(&self) -> Rank {
72        self.rank_out
73    }
74}
75
76impl GGLWEInfos for GGLWELayout {
77    fn rank_in(&self) -> Rank {
78        self.rank_in
79    }
80
81    fn dsize(&self) -> Dsize {
82        self.dsize
83    }
84
85    fn rank_out(&self) -> Rank {
86        self.rank_out
87    }
88
89    fn dnum(&self) -> Dnum {
90        self.dnum
91    }
92}
93
94#[derive(PartialEq, Eq, Clone)]
95pub struct GGLWE<D: Data> {
96    pub(crate) data: MatZnx<D>,
97    pub(crate) base2k: Base2K,
98    pub(crate) dsize: Dsize,
99}
100
101pub struct GGLWEBackendRef<'a, BE: Backend + 'a> {
102    inner: GGLWE<BE::BufRef<'a>>,
103}
104
105impl<'a, BE: Backend + 'a> GGLWEBackendRef<'a, BE> {
106    pub fn from_inner(inner: GGLWE<BE::BufRef<'a>>) -> Self {
107        Self { inner }
108    }
109
110    pub fn into_inner(self) -> GGLWE<BE::BufRef<'a>> {
111        self.inner
112    }
113
114    pub fn at_view(&self, row: usize, col: usize) -> GLWEViewRef<'_, BE> {
115        GLWEViewRef::from_inner(gglwe_at_backend_ref_from_ref::<BE>(&self.inner, row, col))
116    }
117}
118
119impl<'a, BE: Backend + 'a> Deref for GGLWEBackendRef<'a, BE> {
120    type Target = GGLWE<BE::BufRef<'a>>;
121
122    fn deref(&self) -> &Self::Target {
123        &self.inner
124    }
125}
126
127pub struct GGLWEBackendMut<'a, BE: Backend + 'a> {
128    inner: GGLWE<BE::BufMut<'a>>,
129}
130
131impl<'a, BE: Backend + 'a> GGLWEBackendMut<'a, BE> {
132    pub fn from_inner(inner: GGLWE<BE::BufMut<'a>>) -> Self {
133        Self { inner }
134    }
135
136    pub fn into_inner(self) -> GGLWE<BE::BufMut<'a>> {
137        self.inner
138    }
139
140    pub fn at_view(&self, row: usize, col: usize) -> GLWEViewRef<'_, BE> {
141        GLWEViewRef::from_inner(gglwe_at_backend_ref_from_mut::<BE>(&self.inner, row, col))
142    }
143
144    pub fn at_view_mut(&mut self, row: usize, col: usize) -> GLWEViewMut<'_, BE> {
145        GLWEViewMut::from_inner(gglwe_at_backend_mut_from_mut::<BE>(&mut self.inner, row, col))
146    }
147}
148
149impl<'a, BE: Backend + 'a> Deref for GGLWEBackendMut<'a, BE> {
150    type Target = GGLWE<BE::BufMut<'a>>;
151
152    fn deref(&self) -> &Self::Target {
153        &self.inner
154    }
155}
156
157impl<'a, BE: Backend + 'a> DerefMut for GGLWEBackendMut<'a, BE> {
158    fn deref_mut(&mut self) -> &mut Self::Target {
159        &mut self.inner
160    }
161}
162
163impl_gglwe_infos_for_inner!(GGLWEBackendRef<'a, BE>, ['a, BE: Backend + 'a]; inner);
164impl_gglwe_infos_for_inner!(GGLWEBackendMut<'a, BE>, ['a, BE: Backend + 'a]; inner);
165
166impl<'a, BE: Backend + 'a> GGLWEToBackendRef<BE> for GGLWEBackendRef<'a, BE> {
167    fn to_backend_ref(&self) -> GGLWEBackendRef<'_, BE> {
168        GGLWEBackendRef::from_inner(GGLWE {
169            base2k: self.inner.base2k,
170            dsize: self.inner.dsize,
171            data: poulpy_hal::layouts::mat_znx_backend_ref_from_ref::<BE>(&self.inner.data),
172        })
173    }
174}
175
176impl<'a, BE: Backend + 'a> GGSWAtViewRef<BE> for GGLWEBackendRef<'a, BE> {
177    fn at_view(&self, row: usize, col: usize) -> GLWEViewRef<'_, BE> {
178        GGLWEBackendRef::at_view(self, row, col)
179    }
180}
181
182impl<'a, BE: Backend + 'a> GGLWEToBackendRef<BE> for GGLWEBackendMut<'a, BE> {
183    fn to_backend_ref(&self) -> GGLWEBackendRef<'_, BE> {
184        GGLWEBackendRef::from_inner(GGLWE {
185            base2k: self.inner.base2k,
186            dsize: self.inner.dsize,
187            data: poulpy_hal::layouts::mat_znx_backend_ref_from_mut::<BE>(&self.inner.data),
188        })
189    }
190}
191
192impl<'a, BE: Backend + 'a> GGLWEToBackendMut<BE> for GGLWEBackendMut<'a, BE> {
193    fn to_backend_mut(&mut self) -> GGLWEBackendMut<'_, BE> {
194        GGLWEBackendMut::from_inner(GGLWE {
195            base2k: self.inner.base2k,
196            dsize: self.inner.dsize,
197            data: poulpy_hal::layouts::mat_znx_backend_mut_from_mut::<BE>(&mut self.inner.data),
198        })
199    }
200}
201
202impl<D: Data> LWEInfos for GGLWE<D> {
203    fn base2k(&self) -> Base2K {
204        self.base2k
205    }
206
207    fn n(&self) -> Degree {
208        Degree(self.data.n() as u32)
209    }
210
211    fn size(&self) -> usize {
212        self.data.size()
213    }
214}
215
216impl<D: Data> LWEInfos for &GGLWE<D> {
217    fn base2k(&self) -> Base2K {
218        self.base2k
219    }
220
221    fn n(&self) -> Degree {
222        Degree(self.data.n() as u32)
223    }
224
225    fn size(&self) -> usize {
226        self.data.size()
227    }
228}
229
230impl<D: Data> GLWEInfos for GGLWE<D> {
231    fn rank(&self) -> Rank {
232        self.rank_out()
233    }
234}
235
236impl<D: Data> GLWEInfos for &GGLWE<D> {
237    fn rank(&self) -> Rank {
238        self.rank_out()
239    }
240}
241
242impl<D: Data> GGLWEInfos for GGLWE<D> {
243    fn rank_in(&self) -> Rank {
244        Rank(self.data.cols_in() as u32)
245    }
246
247    fn rank_out(&self) -> Rank {
248        Rank(self.data.cols_out() as u32 - 1)
249    }
250
251    fn dsize(&self) -> Dsize {
252        self.dsize
253    }
254
255    fn dnum(&self) -> Dnum {
256        Dnum(self.data.rows() as u32)
257    }
258}
259
260impl<D: Data> GGLWEInfos for &GGLWE<D> {
261    fn rank_in(&self) -> Rank {
262        Rank(self.data.cols_in() as u32)
263    }
264
265    fn rank_out(&self) -> Rank {
266        Rank(self.data.cols_out() as u32 - 1)
267    }
268
269    fn dsize(&self) -> Dsize {
270        self.dsize
271    }
272
273    fn dnum(&self) -> Dnum {
274        Dnum(self.data.rows() as u32)
275    }
276}
277
278impl<D: Data> LWEInfos for &mut GGLWE<D> {
279    fn base2k(&self) -> Base2K {
280        self.base2k
281    }
282
283    fn n(&self) -> Degree {
284        Degree(self.data.n() as u32)
285    }
286
287    fn size(&self) -> usize {
288        self.data.size()
289    }
290}
291
292impl<D: Data> GLWEInfos for &mut GGLWE<D> {
293    fn rank(&self) -> Rank {
294        self.rank_out()
295    }
296}
297
298impl<D: Data> GGLWEInfos for &mut GGLWE<D> {
299    fn rank_in(&self) -> Rank {
300        Rank(self.data.cols_in() as u32)
301    }
302
303    fn rank_out(&self) -> Rank {
304        Rank(self.data.cols_out() as u32 - 1)
305    }
306
307    fn dsize(&self) -> Dsize {
308        self.dsize
309    }
310
311    fn dnum(&self) -> Dnum {
312        Dnum(self.data.rows() as u32)
313    }
314}
315
316impl<D: HostDataRef> GGLWE<D> {
317    pub fn data(&self) -> &MatZnx<D> {
318        &self.data
319    }
320}
321
322pub(crate) trait GGLWEAtBackendRef<BE: Backend> {
323    fn at_backend(&self, row: usize, col: usize) -> GLWE<BE::BufRef<'_>>;
324}
325
326impl<BE: Backend> GGLWEAtBackendRef<BE> for GGLWE<BE::OwnedBuf> {
327    fn at_backend(&self, row: usize, col: usize) -> GLWE<BE::BufRef<'_>> {
328        let data = <MatZnx<BE::OwnedBuf> as MatZnxAtBackendRef<BE>>::at_backend(&self.data, row, col);
329        GLWE {
330            base2k: self.base2k,
331            data,
332        }
333    }
334}
335
336pub(crate) fn gglwe_at_backend_ref_from_ref<'a, 'b, BE: Backend>(
337    gglwe: &'a GGLWE<BE::BufRef<'b>>,
338    row: usize,
339    col: usize,
340) -> GLWE<BE::BufRef<'a>> {
341    let data = poulpy_hal::layouts::mat_znx_at_backend_ref_from_ref::<BE>(&gglwe.data, row, col);
342    GLWE {
343        base2k: gglwe.base2k,
344        data,
345    }
346}
347
348pub trait GGLWEAtViewRef<BE: Backend> {
349    fn at_view(&self, row: usize, col: usize) -> GLWEViewRef<'_, BE>;
350}
351
352impl<BE: Backend> GGLWEAtViewRef<BE> for GGLWE<BE::OwnedBuf> {
353    fn at_view(&self, row: usize, col: usize) -> GLWEViewRef<'_, BE> {
354        GLWEViewRef::from_inner(<GGLWE<BE::OwnedBuf> as GGLWEAtBackendRef<BE>>::at_backend(self, row, col))
355    }
356}
357
358impl<'b, BE: Backend + 'b> GGLWEAtViewRef<BE> for &GGLWE<BE::BufRef<'b>> {
359    fn at_view(&self, row: usize, col: usize) -> GLWEViewRef<'_, BE> {
360        GLWEViewRef::from_inner(gglwe_at_backend_ref_from_ref::<BE>(self, row, col))
361    }
362}
363
364pub(crate) fn gglwe_at_backend_ref_from_mut<'a, 'b, BE: Backend>(
365    gglwe: &'a GGLWE<BE::BufMut<'b>>,
366    row: usize,
367    col: usize,
368) -> GLWE<BE::BufRef<'a>> {
369    let data = poulpy_hal::layouts::mat_znx_at_backend_ref_from_mut::<BE>(&gglwe.data, row, col);
370    GLWE {
371        base2k: gglwe.base2k,
372        data,
373    }
374}
375
376impl<D: HostDataMut> GGLWE<D> {
377    pub fn data_mut(&mut self) -> &mut MatZnx<D> {
378        &mut self.data
379    }
380}
381
382pub(crate) trait GGLWEAtBackendMut<BE: Backend> {
383    fn at_backend_mut(&mut self, row: usize, col: usize) -> GLWE<BE::BufMut<'_>>;
384}
385
386impl<BE: Backend> GGLWEAtBackendMut<BE> for GGLWE<BE::OwnedBuf> {
387    fn at_backend_mut(&mut self, row: usize, col: usize) -> GLWE<BE::BufMut<'_>> {
388        let base2k = self.base2k;
389        let data = <MatZnx<BE::OwnedBuf> as MatZnxAtBackendMut<BE>>::at_backend_mut(&mut self.data, row, col);
390        GLWE { base2k, data }
391    }
392}
393
394pub(crate) fn gglwe_at_backend_mut_from_mut<'a, 'b, BE: Backend>(
395    gglwe: &'a mut GGLWE<BE::BufMut<'b>>,
396    row: usize,
397    col: usize,
398) -> GLWE<BE::BufMut<'a>> {
399    let base2k = gglwe.base2k;
400    let data = poulpy_hal::layouts::mat_znx_at_backend_mut_from_mut::<BE>(&mut gglwe.data, row, col);
401    GLWE { base2k, data }
402}
403
404pub trait GGLWEAtViewMut<BE: Backend> {
405    fn at_view_mut(&mut self, row: usize, col: usize) -> GLWEViewMut<'_, BE>;
406}
407
408impl<BE: Backend> GGLWEAtViewMut<BE> for GGLWE<BE::OwnedBuf> {
409    fn at_view_mut(&mut self, row: usize, col: usize) -> GLWEViewMut<'_, BE> {
410        GLWEViewMut::from_inner(<GGLWE<BE::OwnedBuf> as GGLWEAtBackendMut<BE>>::at_backend_mut(self, row, col))
411    }
412}
413
414impl<'b, BE: Backend + 'b> GGLWEAtViewMut<BE> for &mut GGLWE<BE::BufMut<'b>> {
415    fn at_view_mut(&mut self, row: usize, col: usize) -> GLWEViewMut<'_, BE> {
416        GLWEViewMut::from_inner(gglwe_at_backend_mut_from_mut::<BE>(*self, row, col))
417    }
418}
419
420impl<D: HostDataRef> fmt::Debug for GGLWE<D> {
421    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
422        write!(f, "{self}")
423    }
424}
425
426impl<D: HostDataMut> FillUniform for GGLWE<D> {
427    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
428        self.data.fill_uniform(log_bound, source);
429    }
430}
431
432impl<D: HostDataRef> fmt::Display for GGLWE<D> {
433    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
434        write!(
435            f,
436            "(GGLWE: k={} base2k={} dsize={}) {}",
437            self.max_k().0,
438            self.base2k().0,
439            self.dsize().0,
440            self.data
441        )
442    }
443}
444
445impl<D: HostDataRef> GGLWE<D> {
446    pub fn at(&self, row: usize, col: usize) -> GLWE<&[u8]> {
447        let data = self.data.at(row, col);
448        GLWE {
449            base2k: self.base2k,
450            data,
451        }
452    }
453}
454
455impl<D: HostDataMut> GGLWE<D> {
456    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWE<&mut [u8]> {
457        let base2k = self.base2k;
458        let data = self.data.at_mut(row, col);
459        GLWE { base2k, data }
460    }
461}
462
463impl<D: HostDataRef> GGLWE<D> {
464    /// Copies this ciphertext's backing bytes into an owned buffer of
465    /// backend `To`, routing via host bytes.
466    pub fn to_backend<BE, To>(&self, dst: &Module<To>) -> GGLWE<To::OwnedBuf>
467    where
468        BE: Backend<OwnedBuf = D>,
469        To: Backend,
470        To: TransferFrom<BE>,
471    {
472        dst.upload_gglwe(self)
473    }
474}
475
476impl<D: Data> GGLWE<D> {
477    /// Zero-cost rename when both backends share the same `OwnedBuf`.
478    pub fn reinterpret<To>(self) -> GGLWE<To::OwnedBuf>
479    where
480        To: Backend<OwnedBuf = D>,
481    {
482        let (n, rows, cols_in, cols_out, size) = (
483            self.data.n(),
484            self.data.rows(),
485            self.data.cols_in(),
486            self.data.cols_out(),
487            self.data.size(),
488        );
489        GGLWE {
490            data: MatZnx::from_data(self.data.into_data(), n, rows, cols_in, cols_out, size),
491            base2k: self.base2k,
492            dsize: self.dsize,
493        }
494    }
495}
496
497#[expect(
498    dead_code,
499    reason = "host-owned constructors are kept for serialization and host-only staging"
500)]
501impl GGLWE<Vec<u8>> {
502    pub(crate) fn alloc_from_infos<A>(infos: &A) -> Self
503    where
504        A: GGLWEInfos,
505    {
506        Self::alloc(
507            infos.n(),
508            infos.base2k(),
509            infos.max_k(),
510            infos.rank_in(),
511            infos.rank_out(),
512            infos.dnum(),
513            infos.dsize(),
514        )
515    }
516
517    pub(crate) fn alloc(
518        n: Degree,
519        base2k: Base2K,
520        k: TorusPrecision,
521        rank_in: Rank,
522        rank_out: Rank,
523        dnum: Dnum,
524        dsize: Dsize,
525    ) -> Self {
526        let size: usize = k.0.div_ceil(base2k.0) as usize;
527        assert!(
528            size as u32 > dsize.0,
529            "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
530            dsize.0
531        );
532
533        assert!(
534            dnum.0 * dsize.0 <= size as u32,
535            "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
536            dnum.0,
537            dsize.0,
538        );
539
540        GGLWE {
541            data: MatZnx::from_data(
542                poulpy_hal::layouts::HostBytesBackend::alloc_bytes(MatZnx::<Vec<u8>>::bytes_of(
543                    n.into(),
544                    dnum.into(),
545                    rank_in.into(),
546                    (rank_out + 1).into(),
547                    size,
548                )),
549                n.into(),
550                dnum.into(),
551                rank_in.into(),
552                (rank_out + 1).into(),
553                size,
554            ),
555            base2k,
556            dsize,
557        }
558    }
559
560    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
561    where
562        A: GGLWEInfos,
563    {
564        Self::bytes_of(
565            infos.n(),
566            infos.base2k(),
567            infos.max_k(),
568            infos.rank_in(),
569            infos.rank_out(),
570            infos.dnum(),
571            infos.dsize(),
572        )
573    }
574
575    pub fn bytes_of(
576        n: Degree,
577        base2k: Base2K,
578        k: TorusPrecision,
579        rank_in: Rank,
580        rank_out: Rank,
581        dnum: Dnum,
582        dsize: Dsize,
583    ) -> usize {
584        let size: usize = k.0.div_ceil(base2k.0) as usize;
585        assert!(
586            size as u32 > dsize.0,
587            "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
588            dsize.0
589        );
590
591        assert!(
592            dnum.0 * dsize.0 <= size as u32,
593            "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
594            dnum.0,
595            dsize.0,
596        );
597
598        MatZnx::bytes_of(
599            n.into(),
600            dnum.into(),
601            rank_in.into(),
602            (rank_out + 1).into(),
603            k.0.div_ceil(base2k.0) as usize,
604        )
605    }
606}
607
608pub trait GGLWEToBackendMut<BE: Backend>: GGLWEToBackendRef<BE> {
609    fn to_backend_mut(&mut self) -> GGLWEBackendMut<'_, BE>;
610}
611
612impl<BE: Backend, D: Data> GGLWEToBackendMut<BE> for GGLWE<D>
613where
614    MatZnx<D>: MatZnxToBackendRef<BE> + MatZnxToBackendMut<BE>,
615{
616    fn to_backend_mut(&mut self) -> GGLWEBackendMut<'_, BE> {
617        GGLWEBackendMut::from_inner(GGLWE {
618            base2k: self.base2k(),
619            dsize: self.dsize(),
620            data: self.data.to_backend_mut(),
621        })
622    }
623}
624
625impl<'b, BE: Backend + 'b> GGLWEToBackendRef<BE> for &mut GGLWE<BE::BufMut<'b>> {
626    fn to_backend_ref(&self) -> GGLWEBackendRef<'_, BE> {
627        GGLWEBackendRef::from_inner(GGLWE {
628            base2k: self.base2k(),
629            dsize: self.dsize(),
630            data: poulpy_hal::layouts::mat_znx_backend_ref_from_mut::<BE>(&self.data),
631        })
632    }
633}
634
635impl<'b, BE: Backend + 'b> GGLWEToBackendMut<BE> for &mut GGLWE<BE::BufMut<'b>> {
636    fn to_backend_mut(&mut self) -> GGLWEBackendMut<'_, BE> {
637        GGLWEBackendMut::from_inner(GGLWE {
638            base2k: self.base2k(),
639            dsize: self.dsize(),
640            data: poulpy_hal::layouts::mat_znx_backend_mut_from_mut::<BE>(&mut self.data),
641        })
642    }
643}
644
645pub trait GGLWEToBackendRef<BE: Backend> {
646    fn to_backend_ref(&self) -> GGLWEBackendRef<'_, BE>;
647}
648
649impl<BE: Backend, D: Data> GGLWEToBackendRef<BE> for GGLWE<D>
650where
651    MatZnx<D>: MatZnxToBackendRef<BE>,
652{
653    fn to_backend_ref(&self) -> GGLWEBackendRef<'_, BE> {
654        GGLWEBackendRef::from_inner(GGLWE {
655            base2k: self.base2k(),
656            dsize: self.dsize(),
657            data: self.data.to_backend_ref(),
658        })
659    }
660}
661
662impl<D: HostDataMut> ReaderFrom for GGLWE<D> {
663    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
664        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
665        self.dsize = Dsize(reader.read_u32::<LittleEndian>()?);
666        self.data.read_from(reader)
667    }
668}
669
670impl<D: HostDataRef> WriterTo for GGLWE<D> {
671    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
672        writer.write_u32::<LittleEndian>(self.base2k.0)?;
673        writer.write_u32::<LittleEndian>(self.dsize.0)?;
674        self.data.write_to(writer)
675    }
676}