poulpy_core/automorphism/
ggsw_ct.rs

1use poulpy_hal::{
2    api::ScratchAvailable,
3    layouts::{Backend, DataMut, Module, Scratch},
4};
5
6use crate::{
7    GGSWExpandRows, ScratchTakeCore,
8    automorphism::glwe_ct::GLWEAutomorphism,
9    layouts::{
10        GGLWEInfos, GGLWEPreparedToRef, GGSW, GGSWInfos, GGSWToMut, GGSWToRef, GetGaloisElement,
11        prepared::{GLWETensorKeyPrepared, GLWETensorKeyPreparedToRef},
12    },
13};
14
15impl GGSW<Vec<u8>> {
16    pub fn automorphism_tmp_bytes<R, A, K, T, M, BE: Backend>(
17        module: &M,
18        res_infos: &R,
19        a_infos: &A,
20        key_infos: &K,
21        tsk_infos: &T,
22    ) -> usize
23    where
24        R: GGSWInfos,
25        A: GGSWInfos,
26        K: GGLWEInfos,
27        T: GGLWEInfos,
28        M: GGSWAutomorphism<BE>,
29    {
30        module.ggsw_automorphism_tmp_bytes(res_infos, a_infos, key_infos, tsk_infos)
31    }
32}
33
34impl<D: DataMut> GGSW<D> {
35    pub fn automorphism<A, K, T, M, BE: Backend>(&mut self, module: &M, a: &A, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
36    where
37        A: GGSWToRef,
38        K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
39        T: GLWETensorKeyPreparedToRef<BE>,
40        Scratch<BE>: ScratchTakeCore<BE>,
41        M: GGSWAutomorphism<BE>,
42    {
43        module.ggsw_automorphism(self, a, key, tsk, scratch);
44    }
45
46    pub fn automorphism_inplace<K, T, M, BE: Backend>(&mut self, module: &M, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
47    where
48        K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
49        T: GLWETensorKeyPreparedToRef<BE>,
50        Scratch<BE>: ScratchTakeCore<BE>,
51        M: GGSWAutomorphism<BE>,
52    {
53        module.ggsw_automorphism_inplace(self, key, tsk, scratch);
54    }
55}
56
57impl<BE: Backend> GGSWAutomorphism<BE> for Module<BE> where Self: GLWEAutomorphism<BE> + GGSWExpandRows<BE> {}
58
59pub trait GGSWAutomorphism<BE: Backend>
60where
61    Self: GLWEAutomorphism<BE> + GGSWExpandRows<BE>,
62{
63    fn ggsw_automorphism_tmp_bytes<R, A, K, T>(&self, res_infos: &R, a_infos: &A, key_infos: &K, tsk_infos: &T) -> usize
64    where
65        R: GGSWInfos,
66        A: GGSWInfos,
67        K: GGLWEInfos,
68        T: GGLWEInfos,
69    {
70        let out_size: usize = res_infos.size();
71        let ci_dft: usize = self.bytes_of_vec_znx_dft((key_infos.rank_out() + 1).into(), out_size);
72        let ks_internal: usize = self.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos);
73        let expand: usize = self.ggsw_expand_rows_tmp_bytes(res_infos, tsk_infos);
74        ci_dft + (ks_internal.max(expand))
75    }
76
77    fn ggsw_automorphism<R, A, K, T>(&self, res: &mut R, a: &A, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
78    where
79        R: GGSWToMut,
80        A: GGSWToRef,
81        K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
82        T: GLWETensorKeyPreparedToRef<BE>,
83        Scratch<BE>: ScratchTakeCore<BE>,
84    {
85        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
86        let a: &GGSW<&[u8]> = &a.to_ref();
87        let tsk: &GLWETensorKeyPrepared<&[u8], BE> = &tsk.to_ref();
88
89        assert_eq!(res.dsize(), a.dsize());
90        assert!(res.dnum() <= a.dnum());
91        assert!(scratch.available() >= self.ggsw_automorphism_tmp_bytes(res, a, key, tsk));
92
93        // Keyswitch the j-th row of the col 0
94        for row in 0..res.dnum().as_usize() {
95            // Key-switch column 0, i.e.
96            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2)
97            self.glwe_automorphism(&mut res.at_mut(row, 0), &a.at(row, 0), key, scratch);
98        }
99
100        self.ggsw_expand_row(res, tsk, scratch);
101    }
102
103    fn ggsw_automorphism_inplace<R, K, T>(&self, res: &mut R, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
104    where
105        R: GGSWToMut,
106        K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
107        T: GLWETensorKeyPreparedToRef<BE>,
108        Scratch<BE>: ScratchTakeCore<BE>,
109    {
110        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
111        let tsk: &GLWETensorKeyPrepared<&[u8], BE> = &tsk.to_ref();
112
113        // Keyswitch the j-th row of the col 0
114        for row in 0..res.dnum().as_usize() {
115            // Key-switch column 0, i.e.
116            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2)
117            self.glwe_automorphism_inplace(&mut res.at_mut(row, 0), key, scratch);
118        }
119
120        self.ggsw_expand_row(res, tsk, scratch);
121    }
122}
123
124impl<DataSelf: DataMut> GGSW<DataSelf> {}
poulpy_core/automorphism/ggsw_ct.rs

poulpy_core/automorphism/
ggsw_ct.rs
