poulpy/poulpy-core/src/automorphism/gglwe_atk.rs
poulpy-core/src/automorphism/gglwe_atk.rsRaw
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}