poulpy/poulpy-core/src/automorphism/glwe_ct.rs
poulpy-core/src/automorphism/glwe_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchAvailable, TakeVecZnxDft, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace,
4 VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallAInplace, VecZnxBigSubSmallBInplace,
5 VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VmpApply, VmpApplyAdd, VmpApplyTmpBytes,
6 },
7 layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnxBig},
8};
9
10use crate::layouts::{GLWECiphertext, Infos, prepared::GGLWEAutomorphismKeyPrepared};
11
12impl GLWECiphertext<Vec<u8>> {
13 #[allow(clippy::too_many_arguments)]
14 pub fn automorphism_scratch_space<B: Backend>(
15 module: &Module<B>,
16 n: usize,
17 basek: usize,
18 k_out: usize,
19 k_in: usize,
20 k_ksk: usize,
21 digits: usize,
22 rank: usize,
23 ) -> usize
24 where
25 Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigNormalizeTmpBytes,
26 {
27 Self::keyswitch_scratch_space(module, n, basek, k_out, k_in, k_ksk, digits, rank, rank)
28 }
29
30 pub fn automorphism_inplace_scratch_space<B: Backend>(
31 module: &Module<B>,
32 n: usize,
33 basek: usize,
34 k_out: usize,
35 k_ksk: usize,
36 digits: usize,
37 rank: usize,
38 ) -> usize
39 where
40 Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigNormalizeTmpBytes,
41 {
42 Self::keyswitch_inplace_scratch_space(module, n, basek, k_out, k_ksk, digits, rank)
43 }
44}
45
46impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
47 pub fn automorphism<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
48 &mut self,
49 module: &Module<B>,
50 lhs: &GLWECiphertext<DataLhs>,
51 rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
52 scratch: &mut Scratch<B>,
53 ) where
54 Module<B>: VecZnxDftAllocBytes
55 + VmpApplyTmpBytes
56 + VecZnxBigNormalizeTmpBytes
57 + VmpApply<B>
58 + VmpApplyAdd<B>
59 + VecZnxDftFromVecZnx<B>
60 + VecZnxDftToVecZnxBigConsume<B>
61 + VecZnxBigAddSmallInplace<B>
62 + VecZnxBigNormalize<B>
63 + VecZnxAutomorphismInplace,
64 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
65 {
66 self.keyswitch(module, lhs, &rhs.key, scratch);
67 (0..self.rank() + 1).for_each(|i| {
68 module.vec_znx_automorphism_inplace(rhs.p(), &mut self.data, i);
69 })
70 }
71
72 pub fn automorphism_inplace<DataRhs: DataRef, B: Backend>(
73 &mut self,
74 module: &Module<B>,
75 rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
76 scratch: &mut Scratch<B>,
77 ) where
78 Module<B>: VecZnxDftAllocBytes
79 + VmpApplyTmpBytes
80 + VecZnxBigNormalizeTmpBytes
81 + VmpApply<B>
82 + VmpApplyAdd<B>
83 + VecZnxDftFromVecZnx<B>
84 + VecZnxDftToVecZnxBigConsume<B>
85 + VecZnxBigAddSmallInplace<B>
86 + VecZnxBigNormalize<B>
87 + VecZnxAutomorphismInplace,
88 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
89 {
90 self.keyswitch_inplace(module, &rhs.key, scratch);
91 (0..self.rank() + 1).for_each(|i| {
92 module.vec_znx_automorphism_inplace(rhs.p(), &mut self.data, i);
93 })
94 }
95
96 pub fn automorphism_add<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
97 &mut self,
98 module: &Module<B>,
99 lhs: &GLWECiphertext<DataLhs>,
100 rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
101 scratch: &mut Scratch<B>,
102 ) where
103 Module<B>: VecZnxDftAllocBytes
104 + VmpApplyTmpBytes
105 + VecZnxBigNormalizeTmpBytes
106 + VmpApply<B>
107 + VmpApplyAdd<B>
108 + VecZnxDftFromVecZnx<B>
109 + VecZnxDftToVecZnxBigConsume<B>
110 + VecZnxBigAddSmallInplace<B>
111 + VecZnxBigNormalize<B>
112 + VecZnxBigAutomorphismInplace<B>,
113 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
114 {
115 #[cfg(debug_assertions)]
116 {
117 self.assert_keyswitch(module, lhs, &rhs.key, scratch);
118 }
119 let (res_dft, scratch1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
120 let mut res_big: VecZnxBig<_, B> = lhs.keyswitch_internal(module, res_dft, &rhs.key, scratch1);
121 (0..self.cols()).for_each(|i| {
122 module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i);
123 module.vec_znx_big_add_small_inplace(&mut res_big, i, &lhs.data, i);
124 module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch1);
125 })
126 }
127
128 pub fn automorphism_add_inplace<DataRhs: DataRef, B: Backend>(
129 &mut self,
130 module: &Module<B>,
131 rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
132 scratch: &mut Scratch<B>,
133 ) where
134 Module<B>: VecZnxDftAllocBytes
135 + VmpApplyTmpBytes
136 + VecZnxBigNormalizeTmpBytes
137 + VmpApply<B>
138 + VmpApplyAdd<B>
139 + VecZnxDftFromVecZnx<B>
140 + VecZnxDftToVecZnxBigConsume<B>
141 + VecZnxBigAddSmallInplace<B>
142 + VecZnxBigNormalize<B>
143 + VecZnxBigAutomorphismInplace<B>,
144 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
145 {
146 unsafe {
147 let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
148 self.automorphism_add(module, &*self_ptr, rhs, scratch);
149 }
150 }
151
152 pub fn automorphism_sub_ab<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
153 &mut self,
154 module: &Module<B>,
155 lhs: &GLWECiphertext<DataLhs>,
156 rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
157 scratch: &mut Scratch<B>,
158 ) where
159 Module<B>: VecZnxDftAllocBytes
160 + VmpApplyTmpBytes
161 + VecZnxBigNormalizeTmpBytes
162 + VmpApply<B>
163 + VmpApplyAdd<B>
164 + VecZnxDftFromVecZnx<B>
165 + VecZnxDftToVecZnxBigConsume<B>
166 + VecZnxBigAddSmallInplace<B>
167 + VecZnxBigNormalize<B>
168 + VecZnxBigAutomorphismInplace<B>
169 + VecZnxBigSubSmallAInplace<B>,
170 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
171 {
172 #[cfg(debug_assertions)]
173 {
174 self.assert_keyswitch(module, lhs, &rhs.key, scratch);
175 }
176 let (res_dft, scratch1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
177 let mut res_big: VecZnxBig<_, B> = lhs.keyswitch_internal(module, res_dft, &rhs.key, scratch1);
178 (0..self.cols()).for_each(|i| {
179 module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i);
180 module.vec_znx_big_sub_small_a_inplace(&mut res_big, i, &lhs.data, i);
181 module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch1);
182 })
183 }
184
185 pub fn automorphism_sub_ab_inplace<DataRhs: DataRef, B: Backend>(
186 &mut self,
187 module: &Module<B>,
188 rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
189 scratch: &mut Scratch<B>,
190 ) where
191 Module<B>: VecZnxDftAllocBytes
192 + VmpApplyTmpBytes
193 + VecZnxBigNormalizeTmpBytes
194 + VmpApply<B>
195 + VmpApplyAdd<B>
196 + VecZnxDftFromVecZnx<B>
197 + VecZnxDftToVecZnxBigConsume<B>
198 + VecZnxBigAddSmallInplace<B>
199 + VecZnxBigNormalize<B>
200 + VecZnxBigAutomorphismInplace<B>
201 + VecZnxBigSubSmallAInplace<B>,
202 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
203 {
204 unsafe {
205 let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
206 self.automorphism_sub_ab(module, &*self_ptr, rhs, scratch);
207 }
208 }
209
210 pub fn automorphism_sub_ba<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
211 &mut self,
212 module: &Module<B>,
213 lhs: &GLWECiphertext<DataLhs>,
214 rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
215 scratch: &mut Scratch<B>,
216 ) where
217 Module<B>: VecZnxDftAllocBytes
218 + VmpApplyTmpBytes
219 + VecZnxBigNormalizeTmpBytes
220 + VmpApply<B>
221 + VmpApplyAdd<B>
222 + VecZnxDftFromVecZnx<B>
223 + VecZnxDftToVecZnxBigConsume<B>
224 + VecZnxBigAddSmallInplace<B>
225 + VecZnxBigNormalize<B>
226 + VecZnxBigAutomorphismInplace<B>
227 + VecZnxBigSubSmallBInplace<B>,
228 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
229 {
230 #[cfg(debug_assertions)]
231 {
232 self.assert_keyswitch(module, lhs, &rhs.key, scratch);
233 }
234 let (res_dft, scratch1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
235 let mut res_big: VecZnxBig<_, B> = lhs.keyswitch_internal(module, res_dft, &rhs.key, scratch1);
236 (0..self.cols()).for_each(|i| {
237 module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i);
238 module.vec_znx_big_sub_small_b_inplace(&mut res_big, i, &lhs.data, i);
239 module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch1);
240 })
241 }
242
243 pub fn automorphism_sub_ba_inplace<DataRhs: DataRef, B: Backend>(
244 &mut self,
245 module: &Module<B>,
246 rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
247 scratch: &mut Scratch<B>,
248 ) where
249 Module<B>: VecZnxDftAllocBytes
250 + VmpApplyTmpBytes
251 + VecZnxBigNormalizeTmpBytes
252 + VmpApply<B>
253 + VmpApplyAdd<B>
254 + VecZnxDftFromVecZnx<B>
255 + VecZnxDftToVecZnxBigConsume<B>
256 + VecZnxBigAddSmallInplace<B>
257 + VecZnxBigNormalize<B>
258 + VecZnxBigAutomorphismInplace<B>
259 + VecZnxBigSubSmallBInplace<B>,
260 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
261 {
262 unsafe {
263 let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
264 self.automorphism_sub_ba(module, &*self_ptr, rhs, scratch);
265 }
266 }
267}