phantom-zone/src/backend.rs
src/backend.rsRaw
1use std::marker::PhantomData;
2
3use itertools::izip;
4use num_traits::{PrimInt, Signed, ToPrimitive, WrappingAdd, WrappingMul, WrappingSub, Zero};
5
6use crate::{utils::ShoupMul, Matrix, RowMut};
7
8pub trait Modulus {
9 type Element;
10 /// Modulus value if it fits in Element
11 fn q(&self) -> Option<Self::Element>;
12 /// Modulus value as f64 if it fits in f64
13 fn q_as_f64(&self) -> Option<f64>;
14 /// Is modulus native?
15 fn is_native(&self) -> bool;
16 /// -1 in signed representaiton
17 fn neg_one(&self) -> Self::Element;
18 /// Largest unsigned value that fits in the modulus. That is, q - 1.
19 fn largest_unsigned_value(&self) -> Self::Element;
20 /// Smallest unsigned value that fits in the modulus
21 /// Always assmed to be 0.
22 fn smallest_unsigned_value(&self) -> Self::Element;
23 /// Convert unsigned value in signed represetation to i64
24 fn map_element_to_i64(&self, v: &Self::Element) -> i64;
25 /// Convert f64 to signed represented in modulus
26 fn map_element_from_f64(&self, v: f64) -> Self::Element;
27 /// Convert i64 to signed represented in modulus
28 fn map_element_from_i64(&self, v: i64) -> Self::Element;
29}
30
31impl Modulus for u64 {
32 type Element = u64;
33 fn is_native(&self) -> bool {
34 // q of size u64 can never be a naitve modulus
35 false
36 }
37 fn largest_unsigned_value(&self) -> Self::Element {
38 self - 1
39 }
40 fn neg_one(&self) -> Self::Element {
41 self - 1
42 }
43 fn smallest_unsigned_value(&self) -> Self::Element {
44 0
45 }
46 fn map_element_to_i64(&self, v: &Self::Element) -> i64 {
47 assert!(v <= self, "{v} must be <= {self}");
48 if *v >= (self >> 1) {
49 -ToPrimitive::to_i64(&(self - v)).unwrap()
50 } else {
51 ToPrimitive::to_i64(v).unwrap()
52 }
53 }
54 fn map_element_from_f64(&self, v: f64) -> Self::Element {
55 //FIXME (Jay): Before I check whether v is smaller than 0 with `let is_neg =
56 // o.is_sign_negative() && o != 0.0; I'm ocnfused why didn't I simply check <
57 // 0.0?
58 let v = v.round();
59 if v < 0.0 {
60 self - v.abs().to_u64().unwrap()
61 } else {
62 v.to_u64().unwrap()
63 }
64 }
65 fn map_element_from_i64(&self, v: i64) -> Self::Element {
66 if v < 0 {
67 self - v.abs().to_u64().unwrap()
68 } else {
69 v.to_u64().unwrap()
70 }
71 }
72 fn q(&self) -> Option<Self::Element> {
73 Some(*self)
74 }
75 fn q_as_f64(&self) -> Option<f64> {
76 self.to_f64()
77 }
78}
79
80pub trait ModInit {
81 type M;
82 fn new(modulus: Self::M) -> Self;
83}
84
85pub trait GetModulus {
86 type Element;
87 type M: Modulus<Element = Self::Element>;
88 fn modulus(&self) -> &Self::M;
89}
90
91pub trait VectorOps {
92 type Element;
93
94 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element);
95 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]);
96
97 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
98 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
99 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
100 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element);
101 fn elwise_neg_mut(&self, a: &mut [Self::Element]);
102 /// inplace mutates `a`: a = a + b*c
103 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]);
104 fn elwise_fma_scalar_mut(
105 &self,
106 a: &mut [Self::Element],
107 b: &[Self::Element],
108 c: &Self::Element,
109 );
110
111 // fn modulus(&self) -> Self::Element;
112}
113
114pub trait ArithmeticOps {
115 type Element;
116
117 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
118 fn mul_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
119 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
120 fn add_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
121 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
122 fn neg(&self, a: &Self::Element) -> Self::Element;
123
124 // fn modulus(&self) -> Self::Element;
125}
126
127pub struct ModularOpsU64<T> {
128 q: u64,
129 q_twice: u64,
130 logq: usize,
131 barrett_mu: u128,
132 barrett_alpha: usize,
133 modulus: T,
134}
135
136impl<T> ModInit for ModularOpsU64<T>
137where
138 T: Modulus<Element = u64>,
139{
140 type M = T;
141 fn new(modulus: Self::M) -> ModularOpsU64<T> {
142 assert!(!modulus.is_native());
143
144 // largest unsigned value modulus fits is modulus-1
145 let q = modulus.largest_unsigned_value() + 1;
146 let logq = 64 - (q + 1u64).leading_zeros();
147
148 // barrett calculation
149 let mu = (1u128 << (logq * 2 + 3)) / (q as u128);
150 let alpha = logq + 3;
151
152 ModularOpsU64 {
153 q,
154 q_twice: q << 1,
155 logq: logq as usize,
156 barrett_alpha: alpha as usize,
157 barrett_mu: mu,
158 modulus,
159 }
160 }
161}
162
163impl<T> ModularOpsU64<T> {
164 fn add_mod_fast(&self, a: u64, b: u64) -> u64 {
165 debug_assert!(a < self.q);
166 debug_assert!(b < self.q);
167
168 let mut o = a + b;
169 if o >= self.q {
170 o -= self.q;
171 }
172 o
173 }
174
175 fn add_mod_fast_lazy(&self, a: u64, b: u64) -> u64 {
176 debug_assert!(a < self.q_twice);
177 debug_assert!(b < self.q_twice);
178
179 let mut o = a + b;
180 if o >= self.q_twice {
181 o -= self.q_twice;
182 }
183 o
184 }
185
186 fn sub_mod_fast(&self, a: u64, b: u64) -> u64 {
187 debug_assert!(a < self.q);
188 debug_assert!(b < self.q);
189
190 if a >= b {
191 a - b
192 } else {
193 (self.q + a) - b
194 }
195 }
196
197 // returns (a * b) % q
198 ///
199 /// - both a and b must be in range [0, 2q)
200 /// - output is in range [0 , 2q)
201 fn mul_mod_fast_lazy(&self, a: u64, b: u64) -> u64 {
202 debug_assert!(a < 2 * self.q);
203 debug_assert!(b < 2 * self.q);
204
205 let ab = a as u128 * b as u128;
206
207 // ab / (2^{n + \beta})
208 // note: \beta is assumed to -2
209 let tmp = ab >> (self.logq - 2);
210
211 // k = ((ab / (2^{n + \beta})) * \mu) / 2^{\alpha - (-2)}
212 let k = (tmp * self.barrett_mu) >> (self.barrett_alpha + 2);
213
214 // ab - k*p
215 let tmp = k * (self.q as u128);
216
217 (ab - tmp) as u64
218 }
219
220 /// returns (a * b) % q
221 ///
222 /// - both a and b must be in range [0, 2q)
223 /// - output is in range [0 , q)
224 fn mul_mod_fast(&self, a: u64, b: u64) -> u64 {
225 debug_assert!(a < 2 * self.q);
226 debug_assert!(b < 2 * self.q);
227
228 let ab = a as u128 * b as u128;
229
230 // ab / (2^{n + \beta})
231 // note: \beta is assumed to -2
232 let tmp = ab >> (self.logq - 2);
233
234 // k = ((ab / (2^{n + \beta})) * \mu) / 2^{\alpha - (-2)}
235 let k = (tmp * self.barrett_mu) >> (self.barrett_alpha + 2);
236
237 // ab - k*p
238 let tmp = k * (self.q as u128);
239
240 let mut out = (ab - tmp) as u64;
241
242 if out >= self.q {
243 out -= self.q;
244 }
245
246 return out;
247 }
248}
249
250impl<T> ArithmeticOps for ModularOpsU64<T> {
251 type Element = u64;
252
253 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
254 self.add_mod_fast(*a, *b)
255 }
256
257 fn add_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
258 self.add_mod_fast_lazy(*a, *b)
259 }
260
261 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
262 self.mul_mod_fast(*a, *b)
263 }
264
265 fn mul_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
266 self.mul_mod_fast_lazy(*a, *b)
267 }
268
269 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
270 self.sub_mod_fast(*a, *b)
271 }
272
273 fn neg(&self, a: &Self::Element) -> Self::Element {
274 self.q - *a
275 }
276
277 // fn modulus(&self) -> Self::Element {
278 // self.q
279 // }
280}
281
282impl<T> VectorOps for ModularOpsU64<T> {
283 type Element = u64;
284
285 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
286 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
287 *ai = self.add_mod_fast(*ai, *bi);
288 });
289 }
290
291 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
292 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
293 *ai = self.sub_mod_fast(*ai, *bi);
294 });
295 }
296
297 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
298 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
299 *ai = self.mul_mod_fast(*ai, *bi);
300 });
301 }
302
303 fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
304 a.iter_mut().for_each(|ai| *ai = self.q - *ai);
305 }
306
307 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
308 izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
309 *oi = self.mul_mod_fast(*ai, *b);
310 });
311 }
312
313 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
314 izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
315 *oi = self.mul_mod_fast(*ai, *bi);
316 });
317 }
318
319 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
320 a.iter_mut().for_each(|ai| {
321 *ai = self.mul_mod_fast(*ai, *b);
322 });
323 }
324
325 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
326 izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
327 *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *ci));
328 });
329 }
330
331 fn elwise_fma_scalar_mut(
332 &self,
333 a: &mut [Self::Element],
334 b: &[Self::Element],
335 c: &Self::Element,
336 ) {
337 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
338 *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *c));
339 });
340 }
341
342 // fn modulus(&self) -> Self::Element {
343 // self.q
344 // }
345}
346
347impl<T> ModularOpsU64<T> {
348 /// Returns \sum a[i]b[i]
349 pub fn shoup_fma<M: Matrix<MatElement = u64>>(&self, out: &mut M::R, a: &M, a_shoup: &M, b: &M)
350 where
351 M::R: RowMut,
352 {
353 assert!(a.dimension() == a_shoup.dimension());
354 assert!(a.dimension() == b.dimension());
355
356 let q = self.q;
357 let q_twice = self.q << 1;
358
359 // first row (without summation)
360 izip!(
361 out.as_mut().iter_mut(),
362 a.get_row(0),
363 a_shoup.get_row(0),
364 b.get_row(0)
365 )
366 .for_each(|(o, a, a_shoup, b)| {
367 *o = ShoupMul::mul(*b, *a, *a_shoup, q);
368 });
369
370 izip!(
371 a.iter_rows().skip(1),
372 a_shoup.iter_rows().skip(1),
373 b.iter_rows().skip(1)
374 )
375 .for_each(|(a_row, a_shoup_row, b_row)| {
376 izip!(
377 out.as_mut().iter_mut(),
378 a_row.as_ref().iter(),
379 a_shoup_row.as_ref().iter(),
380 b_row.as_ref().iter()
381 )
382 .for_each(|(o, a0, a0_shoup, b0)| {
383 let quotient = ((*a0_shoup as u128 * *b0 as u128) >> 64) as u64;
384 let mut v = (a0.wrapping_mul(b0)).wrapping_add(*o);
385 v = v.wrapping_sub(q.wrapping_mul(quotient));
386
387 if v >= q_twice {
388 v -= q_twice;
389 }
390
391 *o = v;
392 });
393 });
394 }
395}
396impl<T> GetModulus for ModularOpsU64<T>
397where
398 T: Modulus,
399{
400 type Element = T::Element;
401 type M = T;
402 fn modulus(&self) -> &Self::M {
403 &self.modulus
404 }
405}
406
407pub struct WordSizeModulus<T> {
408 modulus: T,
409}
410
411impl<T> ModInit for WordSizeModulus<T>
412where
413 T: Modulus,
414{
415 type M = T;
416 fn new(modulus: T) -> Self {
417 assert!(modulus.is_native());
418 // For now assume ModulusOpsU64 is only used for u64
419 Self { modulus: modulus }
420 }
421}
422
423impl<T> ArithmeticOps for WordSizeModulus<T>
424where
425 T: Modulus,
426 T::Element: WrappingAdd + WrappingSub + WrappingMul + Zero,
427{
428 type Element = T::Element;
429 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
430 T::Element::wrapping_add(a, b)
431 }
432
433 fn add_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
434 self.add(a, b)
435 }
436
437 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
438 T::Element::wrapping_mul(a, b)
439 }
440
441 fn mul_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
442 self.mul(a, b)
443 }
444
445 fn neg(&self, a: &Self::Element) -> Self::Element {
446 T::Element::wrapping_sub(&T::Element::zero(), a)
447 }
448
449 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
450 T::Element::wrapping_sub(a, b)
451 }
452}
453
454impl<T> VectorOps for WordSizeModulus<T>
455where
456 T: Modulus,
457 T::Element: WrappingAdd + WrappingSub + WrappingMul + Zero,
458{
459 type Element = T::Element;
460
461 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
462 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
463 *ai = T::Element::wrapping_add(ai, bi);
464 });
465 }
466
467 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
468 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
469 *ai = T::Element::wrapping_sub(ai, bi);
470 });
471 }
472
473 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
474 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
475 *ai = T::Element::wrapping_mul(ai, bi);
476 });
477 }
478
479 fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
480 a.iter_mut()
481 .for_each(|ai| *ai = T::Element::wrapping_sub(&T::Element::zero(), ai));
482 }
483
484 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
485 izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
486 *oi = T::Element::wrapping_mul(ai, b);
487 });
488 }
489
490 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
491 izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
492 *oi = T::Element::wrapping_mul(ai, bi);
493 });
494 }
495
496 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
497 a.iter_mut().for_each(|ai| {
498 *ai = T::Element::wrapping_mul(ai, b);
499 });
500 }
501
502 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
503 izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
504 *ai = T::Element::wrapping_add(ai, &T::Element::wrapping_mul(bi, ci));
505 });
506 }
507
508 fn elwise_fma_scalar_mut(
509 &self,
510 a: &mut [Self::Element],
511 b: &[Self::Element],
512 c: &Self::Element,
513 ) {
514 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
515 *ai = T::Element::wrapping_add(ai, &T::Element::wrapping_mul(bi, c));
516 });
517 }
518
519 // fn modulus(&self) -> &T {
520 // &self.modulus
521 // }
522}
523
524impl<T> GetModulus for WordSizeModulus<T>
525where
526 T: Modulus,
527{
528 type Element = T::Element;
529 type M = T;
530 fn modulus(&self) -> &Self::M {
531 &self.modulus
532 }
533}