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