poulpy_core/automorphism/
gglwe_atk.rs

1use poulpy_hal::{
2    api::{VecZnxAutomorphism, VecZnxAutomorphismInplace},
3    layouts::{Backend, CyclotomicOrder, DataMut, GaloisElement, Module, Scratch},
4};
5
6use crate::{
7    GLWEKeyswitch, ScratchTakeCore,
8    layouts::{
9        GGLWE, GGLWEInfos, GGLWEPreparedToRef, GGLWEToMut, GGLWEToRef, GLWE, GLWEAutomorphismKey, GetGaloisElement,
10        SetGaloisElement,
11    },
12};
13
14impl GLWEAutomorphismKey<Vec<u8>> {
15    pub fn automorphism_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
16    where
17        R: GGLWEInfos,
18        A: GGLWEInfos,
19        K: GGLWEInfos,
20        M: GLWEAutomorphismKeyAutomorphism<BE>,
21    {
22        module.glwe_automorphism_key_automorphism_tmp_bytes(res_infos, a_infos, key_infos)
23    }
24}
25
26impl<DataSelf: DataMut> GLWEAutomorphismKey<DataSelf> {
27    pub fn automorphism<A, K, M, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
28    where
29        A: GGLWEToRef + GetGaloisElement + GGLWEInfos,
30        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
31        Scratch<BE>: ScratchTakeCore<BE>,
32        M: GLWEAutomorphismKeyAutomorphism<BE>,
33    {
34        module.glwe_automorphism_key_automorphism(self, a, key, scratch);
35    }
36
37    pub fn automorphism_inplace<K, M, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
38    where
39        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
40        Scratch<BE>: ScratchTakeCore<BE>,
41        M: GLWEAutomorphismKeyAutomorphism<BE>,
42    {
43        module.glwe_automorphism_key_automorphism_inplace(self, key, scratch);
44    }
45}
46
47impl<BE: Backend> GLWEAutomorphismKeyAutomorphism<BE> for Module<BE>
48where
49    Self: GaloisElement + GLWEKeyswitch<BE> + VecZnxAutomorphism + VecZnxAutomorphismInplace<BE> + CyclotomicOrder,
50    Scratch<BE>: ScratchTakeCore<BE>,
51{
52    fn glwe_automorphism_key_automorphism_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
53    where
54        R: GGLWEInfos,
55        A: GGLWEInfos,
56        K: GGLWEInfos,
57    {
58        if res_infos.glwe_layout() == a_infos.glwe_layout() {
59            self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
60        } else {
61            self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos) + GLWE::bytes_of_from_infos(a_infos)
62        }
63    }
64
65    fn glwe_automorphism_key_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
66    where
67        R: GGLWEToMut + SetGaloisElement + GGLWEInfos,
68        A: GGLWEToRef + GetGaloisElement + GGLWEInfos,
69        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
70    {
71        assert!(
72            res.dnum().as_u32() <= a.dnum().as_u32(),
73            "res dnum: {} > a dnum: {}",
74            res.dnum(),
75            a.dnum()
76        );
77
78        assert_eq!(
79            res.dsize(),
80            a.dsize(),
81            "res dnum: {} != a dnum: {}",
82            res.dsize(),
83            a.dsize()
84        );
85
86        assert_eq!(res.base2k(), a.base2k());
87
88        let cols_out: usize = (key.rank_out() + 1).into();
89        let cols_in: usize = key.rank_in().into();
90
91        let p: i64 = a.p();
92        let p_inv: i64 = self.galois_element_inv(p);
93
94        let same_layout: bool = res.glwe_layout() == a.glwe_layout();
95
96        {
97            let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
98            let a: &GGLWE<&[u8]> = &a.to_ref();
99
100            for row in 0..res.dnum().as_usize() {
101                for col in 0..cols_in {
102                    let mut res_tmp: GLWE<&mut [u8]> = res.at_mut(row, col);
103                    let a_ct: GLWE<&[u8]> = a.at(row, col);
104
105                    if same_layout {
106                        // Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
107                        for i in 0..cols_out {
108                            self.vec_znx_automorphism(p, res_tmp.data_mut(), i, &a_ct.data, i);
109                        }
110
111                        // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
112                        self.glwe_keyswitch_inplace(&mut res_tmp, key, scratch);
113                    } else {
114                        let (mut tmp_glwe, scratch_1) = scratch.take_glwe(a);
115
116                        // Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
117                        for i in 0..cols_out {
118                            self.vec_znx_automorphism(p, tmp_glwe.data_mut(), i, &a_ct.data, i);
119                        }
120
121                        // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
122                        self.glwe_keyswitch(&mut res_tmp, &tmp_glwe, key, scratch_1);
123                    }
124
125                    // Applies back the automorphism X^{-k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) to (-pi^{-1}_{k'+k}(s)a + s, a)
126                    for i in 0..cols_out {
127                        self.vec_znx_automorphism_inplace(p_inv, res_tmp.data_mut(), i, scratch);
128                    }
129                }
130            }
131        }
132
133        res.set_p((p * key.p()) % self.cyclotomic_order());
134    }
135
136    fn glwe_automorphism_key_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
137    where
138        R: GGLWEToMut + SetGaloisElement + GetGaloisElement + GGLWEInfos,
139        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
140        Scratch<BE>: ScratchTakeCore<BE>,
141    {
142        assert_eq!(
143            res.rank(),
144            key.rank(),
145            "key rank: {} != key rank: {}",
146            res.rank(),
147            key.rank()
148        );
149
150        let cols_out: usize = (key.rank_out() + 1).into();
151        let cols_in: usize = key.rank_in().into();
152        let p: i64 = res.p();
153        let p_inv: i64 = self.galois_element_inv(p);
154
155        {
156            let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
157            for row in 0..res.dnum().as_usize() {
158                for col in 0..cols_in {
159                    let mut res_tmp: GLWE<&mut [u8]> = res.at_mut(row, col);
160
161                    // Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
162                    for i in 0..cols_out {
163                        self.vec_znx_automorphism_inplace(p, res_tmp.data_mut(), i, scratch);
164                    }
165
166                    // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
167                    self.glwe_keyswitch_inplace(&mut res_tmp, key, scratch);
168
169                    // Applies back the automorphism X^{-k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) to (-pi^{-1}_{k'+k}(s)a + s, a)
170                    for i in 0..cols_out {
171                        self.vec_znx_automorphism_inplace(p_inv, res_tmp.data_mut(), i, scratch);
172                    }
173                }
174            }
175        }
176
177        res.set_p((res.p() * key.p()) % self.cyclotomic_order());
178    }
179}
180
181pub trait GLWEAutomorphismKeyAutomorphism<BE: Backend> {
182    fn glwe_automorphism_key_automorphism_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
183    where
184        R: GGLWEInfos,
185        A: GGLWEInfos,
186        K: GGLWEInfos;
187
188    fn glwe_automorphism_key_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
189    where
190        R: GGLWEToMut + SetGaloisElement + GGLWEInfos,
191        A: GGLWEToRef + GetGaloisElement + GGLWEInfos,
192        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos;
193
194    fn glwe_automorphism_key_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
195    where
196        R: GGLWEToMut + SetGaloisElement + GetGaloisElement + GGLWEInfos,
197        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos;
198}
poulpy_core/automorphism/gglwe_atk.rs

poulpy_core/automorphism/
gglwe_atk.rs
