poulpy/core/src/glwe/ops.rs
core/src/glwe/ops.rsRaw
1use backend::hal::{
2 api::{
3 VecZnxAdd, VecZnxAddInplace, VecZnxCopy, VecZnxMulXpMinusOne, VecZnxMulXpMinusOneInplace, VecZnxNegateInplace,
4 VecZnxNormalize, VecZnxNormalizeInplace, VecZnxRotate, VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub,
5 VecZnxSubABInplace, VecZnxSubBAInplace, ZnxZero,
6 },
7 layouts::{Backend, Module, Scratch, VecZnx},
8};
9
10use crate::{GLWECiphertext, GLWECiphertextToMut, GLWECiphertextToRef, Infos, SetMetaData};
11
12pub trait GLWEOps: GLWECiphertextToMut + SetMetaData + Sized {
13 fn add<A, B, BACKEND: Backend>(&mut self, module: &Module<BACKEND>, a: &A, b: &B)
14 where
15 A: GLWECiphertextToRef,
16 B: GLWECiphertextToRef,
17 Module<BACKEND>: VecZnxAdd + VecZnxCopy,
18 {
19 #[cfg(debug_assertions)]
20 {
21 assert_eq!(a.n(), self.n());
22 assert_eq!(b.n(), self.n());
23 assert_eq!(a.basek(), b.basek());
24 assert!(self.rank() >= a.rank().max(b.rank()));
25 }
26
27 let min_col: usize = a.rank().min(b.rank()) + 1;
28 let max_col: usize = a.rank().max(b.rank() + 1);
29 let self_col: usize = self.rank() + 1;
30
31 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
32 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
33 let b_ref: &GLWECiphertext<&[u8]> = &b.to_ref();
34
35 (0..min_col).for_each(|i| {
36 module.vec_znx_add(&mut self_mut.data, i, &a_ref.data, i, &b_ref.data, i);
37 });
38
39 if a.rank() > b.rank() {
40 (min_col..max_col).for_each(|i| {
41 module.vec_znx_copy(&mut self_mut.data, i, &a_ref.data, i);
42 });
43 } else {
44 (min_col..max_col).for_each(|i| {
45 module.vec_znx_copy(&mut self_mut.data, i, &b_ref.data, i);
46 });
47 }
48
49 let size: usize = self_mut.size();
50 (max_col..self_col).for_each(|i| {
51 (0..size).for_each(|j| {
52 self_mut.data.zero_at(i, j);
53 });
54 });
55
56 self.set_basek(a.basek());
57 self.set_k(set_k_binary(self, a, b));
58 }
59
60 fn add_inplace<A, BACKEND: Backend>(&mut self, module: &Module<BACKEND>, a: &A)
61 where
62 A: GLWECiphertextToRef + Infos,
63 Module<BACKEND>: VecZnxAddInplace,
64 {
65 #[cfg(debug_assertions)]
66 {
67 assert_eq!(a.n(), self.n());
68 assert_eq!(self.basek(), a.basek());
69 assert!(self.rank() >= a.rank())
70 }
71
72 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
73 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
74
75 (0..a.rank() + 1).for_each(|i| {
76 module.vec_znx_add_inplace(&mut self_mut.data, i, &a_ref.data, i);
77 });
78
79 self.set_k(set_k_unary(self, a))
80 }
81
82 fn sub<A, B, BACKEND: Backend>(&mut self, module: &Module<BACKEND>, a: &A, b: &B)
83 where
84 A: GLWECiphertextToRef,
85 B: GLWECiphertextToRef,
86 Module<BACKEND>: VecZnxSub + VecZnxCopy + VecZnxNegateInplace,
87 {
88 #[cfg(debug_assertions)]
89 {
90 assert_eq!(a.n(), self.n());
91 assert_eq!(b.n(), self.n());
92 assert_eq!(a.basek(), b.basek());
93 assert!(self.rank() >= a.rank().max(b.rank()));
94 }
95
96 let min_col: usize = a.rank().min(b.rank()) + 1;
97 let max_col: usize = a.rank().max(b.rank() + 1);
98 let self_col: usize = self.rank() + 1;
99
100 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
101 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
102 let b_ref: &GLWECiphertext<&[u8]> = &b.to_ref();
103
104 (0..min_col).for_each(|i| {
105 module.vec_znx_sub(&mut self_mut.data, i, &a_ref.data, i, &b_ref.data, i);
106 });
107
108 if a.rank() > b.rank() {
109 (min_col..max_col).for_each(|i| {
110 module.vec_znx_copy(&mut self_mut.data, i, &a_ref.data, i);
111 });
112 } else {
113 (min_col..max_col).for_each(|i| {
114 module.vec_znx_copy(&mut self_mut.data, i, &b_ref.data, i);
115 module.vec_znx_negate_inplace(&mut self_mut.data, i);
116 });
117 }
118
119 let size: usize = self_mut.size();
120 (max_col..self_col).for_each(|i| {
121 (0..size).for_each(|j| {
122 self_mut.data.zero_at(i, j);
123 });
124 });
125
126 self.set_basek(a.basek());
127 self.set_k(set_k_binary(self, a, b));
128 }
129
130 fn sub_inplace_ab<A, BACKEND: Backend>(&mut self, module: &Module<BACKEND>, a: &A)
131 where
132 A: GLWECiphertextToRef + Infos,
133 Module<BACKEND>: VecZnxSubABInplace,
134 {
135 #[cfg(debug_assertions)]
136 {
137 assert_eq!(a.n(), self.n());
138 assert_eq!(self.basek(), a.basek());
139 assert!(self.rank() >= a.rank())
140 }
141
142 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
143 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
144
145 (0..a.rank() + 1).for_each(|i| {
146 module.vec_znx_sub_ab_inplace(&mut self_mut.data, i, &a_ref.data, i);
147 });
148
149 self.set_k(set_k_unary(self, a))
150 }
151
152 fn sub_inplace_ba<A, BACKEND: Backend>(&mut self, module: &Module<BACKEND>, a: &A)
153 where
154 A: GLWECiphertextToRef + Infos,
155 Module<BACKEND>: VecZnxSubBAInplace,
156 {
157 #[cfg(debug_assertions)]
158 {
159 assert_eq!(a.n(), self.n());
160 assert_eq!(self.basek(), a.basek());
161 assert!(self.rank() >= a.rank())
162 }
163
164 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
165 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
166
167 (0..a.rank() + 1).for_each(|i| {
168 module.vec_znx_sub_ba_inplace(&mut self_mut.data, i, &a_ref.data, i);
169 });
170
171 self.set_k(set_k_unary(self, a))
172 }
173
174 fn rotate<A, B: Backend>(&mut self, module: &Module<B>, k: i64, a: &A)
175 where
176 A: GLWECiphertextToRef + Infos,
177 Module<B>: VecZnxRotate,
178 {
179 #[cfg(debug_assertions)]
180 {
181 assert_eq!(a.n(), self.n());
182 assert_eq!(self.rank(), a.rank())
183 }
184
185 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
186 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
187
188 (0..a.rank() + 1).for_each(|i| {
189 module.vec_znx_rotate(k, &mut self_mut.data, i, &a_ref.data, i);
190 });
191
192 self.set_basek(a.basek());
193 self.set_k(set_k_unary(self, a))
194 }
195
196 fn rotate_inplace<B: Backend>(&mut self, module: &Module<B>, k: i64)
197 where
198 Module<B>: VecZnxRotateInplace,
199 {
200 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
201
202 (0..self_mut.rank() + 1).for_each(|i| {
203 module.vec_znx_rotate_inplace(k, &mut self_mut.data, i);
204 });
205 }
206
207 fn mul_xp_minus_one<A, B: Backend>(&mut self, module: &Module<B>, k: i64, a: &A)
208 where
209 A: GLWECiphertextToRef + Infos,
210 Module<B>: VecZnxMulXpMinusOne,
211 {
212 #[cfg(debug_assertions)]
213 {
214 assert_eq!(a.n(), self.n());
215 assert_eq!(self.rank(), a.rank())
216 }
217
218 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
219 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
220
221 (0..a.rank() + 1).for_each(|i| {
222 module.vec_znx_mul_xp_minus_one(k, &mut self_mut.data, i, &a_ref.data, i);
223 });
224
225 self.set_basek(a.basek());
226 self.set_k(set_k_unary(self, a))
227 }
228
229 fn mul_xp_minus_one_inplace<B: Backend>(&mut self, module: &Module<B>, k: i64)
230 where
231 Module<B>: VecZnxMulXpMinusOneInplace,
232 {
233 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
234
235 (0..self_mut.rank() + 1).for_each(|i| {
236 module.vec_znx_mul_xp_minus_one_inplace(k, &mut self_mut.data, i);
237 });
238 }
239
240 fn copy<A, B: Backend>(&mut self, module: &Module<B>, a: &A)
241 where
242 A: GLWECiphertextToRef + Infos,
243 Module<B>: VecZnxCopy,
244 {
245 #[cfg(debug_assertions)]
246 {
247 assert_eq!(self.n(), a.n());
248 assert_eq!(self.rank(), a.rank());
249 }
250
251 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
252 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
253
254 (0..self_mut.rank() + 1).for_each(|i| {
255 module.vec_znx_copy(&mut self_mut.data, i, &a_ref.data, i);
256 });
257
258 self.set_k(a.k().min(self.size() * self.basek()));
259 self.set_basek(a.basek());
260 }
261
262 fn rsh<B: Backend>(&mut self, module: &Module<B>, k: usize)
263 where
264 Module<B>: VecZnxRshInplace,
265 {
266 let basek: usize = self.basek();
267 module.vec_znx_rsh_inplace(basek, k, &mut self.to_mut().data);
268 }
269
270 fn normalize<A, B: Backend>(&mut self, module: &Module<B>, a: &A, scratch: &mut Scratch<B>)
271 where
272 A: GLWECiphertextToRef,
273 Module<B>: VecZnxNormalize<B>,
274 {
275 #[cfg(debug_assertions)]
276 {
277 assert_eq!(self.n(), a.n());
278 assert_eq!(self.rank(), a.rank());
279 }
280
281 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
282 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
283
284 (0..self_mut.rank() + 1).for_each(|i| {
285 module.vec_znx_normalize(a.basek(), &mut self_mut.data, i, &a_ref.data, i, scratch);
286 });
287 self.set_basek(a.basek());
288 self.set_k(a.k().min(self.k()));
289 }
290
291 fn normalize_inplace<B: Backend>(&mut self, module: &Module<B>, scratch: &mut Scratch<B>)
292 where
293 Module<B>: VecZnxNormalizeInplace<B>,
294 {
295 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
296 (0..self_mut.rank() + 1).for_each(|i| {
297 module.vec_znx_normalize_inplace(self_mut.basek(), &mut self_mut.data, i, scratch);
298 });
299 }
300}
301
302impl GLWECiphertext<Vec<u8>> {
303 pub fn rsh_scratch_space(n: usize) -> usize {
304 VecZnx::rsh_scratch_space(n)
305 }
306}
307
308// c = op(a, b)
309fn set_k_binary(c: &impl Infos, a: &impl Infos, b: &impl Infos) -> usize {
310 // If either operands is a ciphertext
311 if a.rank() != 0 || b.rank() != 0 {
312 // If a is a plaintext (but b ciphertext)
313 let k = if a.rank() == 0 {
314 b.k()
315 // If b is a plaintext (but a ciphertext)
316 } else if b.rank() == 0 {
317 a.k()
318 // If a & b are both ciphertexts
319 } else {
320 a.k().min(b.k())
321 };
322 k.min(c.k())
323 // If a & b are both plaintexts
324 } else {
325 c.k()
326 }
327}
328
329// a = op(a, b)
330fn set_k_unary(a: &impl Infos, b: &impl Infos) -> usize {
331 if a.rank() != 0 || b.rank() != 0 {
332 a.k().min(b.k())
333 } else {
334 a.k()
335 }
336}