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