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