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
112pub trait ArithmeticOps {
113 type Element;
114 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
115 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
116 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
117 fn neg(&self, a: &Self::Element) -> Self::Element;
118}
119
120pub trait ArithmeticLazyOps {
121 type Element;
122 fn mul_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
123 fn add_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
124}
125
126pub trait ShoupMatrixFMA<M: Matrix>
127where
128 M::R: RowMut,
129{
130 /// Returns summation of row-wise product of matrix a and b.
131 fn shoup_matrix_fma(&self, out: &mut M::R, a: &M, a_shoup: &M, b: &M);
132}
133
134pub struct ModularOpsU64<T> {
135 q: u64,
136 q_twice: u64,
137 logq: usize,
138 barrett_mu: u128,
139 barrett_alpha: usize,
140 modulus: T,
141}
142
143impl<T> ModInit for ModularOpsU64<T>
144where
145 T: Modulus<Element = u64>,
146{
147 type M = T;
148 fn new(modulus: Self::M) -> ModularOpsU64<T> {
149 assert!(!modulus.is_native());
150
151 // largest unsigned value modulus fits is modulus-1
152 let q = modulus.largest_unsigned_value() + 1;
153 let logq = 64 - (q + 1u64).leading_zeros();
154
155 // barrett calculation
156 let mu = (1u128 << (logq * 2 + 3)) / (q as u128);
157 let alpha = logq + 3;
158
159 ModularOpsU64 {
160 q,
161 q_twice: q << 1,
162 logq: logq as usize,
163 barrett_alpha: alpha as usize,
164 barrett_mu: mu,
165 modulus,
166 }
167 }
168}
169
170impl<T> ModularOpsU64<T> {
171 fn add_mod_fast(&self, a: u64, b: u64) -> u64 {
172 debug_assert!(a < self.q);
173 debug_assert!(b < self.q);
174
175 let mut o = a + b;
176 if o >= self.q {
177 o -= self.q;
178 }
179 o
180 }
181
182 fn add_mod_fast_lazy(&self, a: u64, b: u64) -> u64 {
183 debug_assert!(a < self.q_twice);
184 debug_assert!(b < self.q_twice);
185
186 let mut o = a + b;
187 if o >= self.q_twice {
188 o -= self.q_twice;
189 }
190 o
191 }
192
193 fn sub_mod_fast(&self, a: u64, b: u64) -> u64 {
194 debug_assert!(a < self.q);
195 debug_assert!(b < self.q);
196
197 if a >= b {
198 a - b
199 } else {
200 (self.q + a) - b
201 }
202 }
203
204 // returns (a * b) % q
205 ///
206 /// - both a and b must be in range [0, 2q)
207 /// - output is in range [0 , 2q)
208 fn mul_mod_fast_lazy(&self, a: u64, b: u64) -> u64 {
209 debug_assert!(a < 2 * self.q);
210 debug_assert!(b < 2 * self.q);
211
212 let ab = a as u128 * b as u128;
213
214 // ab / (2^{n + \beta})
215 // note: \beta is assumed to -2
216 let tmp = ab >> (self.logq - 2);
217
218 // k = ((ab / (2^{n + \beta})) * \mu) / 2^{\alpha - (-2)}
219 let k = (tmp * self.barrett_mu) >> (self.barrett_alpha + 2);
220
221 // ab - k*p
222 let tmp = k * (self.q as u128);
223
224 (ab - tmp) as u64
225 }
226
227 /// returns (a * b) % q
228 ///
229 /// - both a and b must be in range [0, 2q)
230 /// - output is in range [0 , q)
231 fn mul_mod_fast(&self, a: u64, b: u64) -> u64 {
232 debug_assert!(a < 2 * self.q);
233 debug_assert!(b < 2 * self.q);
234
235 let ab = a as u128 * b as u128;
236
237 // ab / (2^{n + \beta})
238 // note: \beta is assumed to -2
239 let tmp = ab >> (self.logq - 2);
240
241 // k = ((ab / (2^{n + \beta})) * \mu) / 2^{\alpha - (-2)}
242 let k = (tmp * self.barrett_mu) >> (self.barrett_alpha + 2);
243
244 // ab - k*p
245 let tmp = k * (self.q as u128);
246
247 let mut out = (ab - tmp) as u64;
248
249 if out >= self.q {
250 out -= self.q;
251 }
252
253 return out;
254 }
255}
256
257impl<T> ArithmeticOps for ModularOpsU64<T> {
258 type Element = u64;
259
260 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
261 self.add_mod_fast(*a, *b)
262 }
263
264 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
265 self.mul_mod_fast(*a, *b)
266 }
267
268 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
269 self.sub_mod_fast(*a, *b)
270 }
271
272 fn neg(&self, a: &Self::Element) -> Self::Element {
273 self.q - *a
274 }
275
276 // fn modulus(&self) -> Self::Element {
277 // self.q
278 // }
279}
280
281impl<T> ArithmeticLazyOps for ModularOpsU64<T> {
282 type Element = u64;
283 fn add_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
284 self.add_mod_fast_lazy(*a, *b)
285 }
286 fn mul_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
287 self.mul_mod_fast_lazy(*a, *b)
288 }
289}
290
291impl<T> VectorOps for ModularOpsU64<T> {
292 type Element = u64;
293
294 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
295 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
296 *ai = self.add_mod_fast(*ai, *bi);
297 });
298 }
299
300 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
301 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
302 *ai = self.sub_mod_fast(*ai, *bi);
303 });
304 }
305
306 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
307 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
308 *ai = self.mul_mod_fast(*ai, *bi);
309 });
310 }
311
312 fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
313 a.iter_mut().for_each(|ai| *ai = self.q - *ai);
314 }
315
316 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
317 izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
318 *oi = self.mul_mod_fast(*ai, *b);
319 });
320 }
321
322 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
323 izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
324 *oi = self.mul_mod_fast(*ai, *bi);
325 });
326 }
327
328 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
329 a.iter_mut().for_each(|ai| {
330 *ai = self.mul_mod_fast(*ai, *b);
331 });
332 }
333
334 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
335 izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
336 *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *ci));
337 });
338 }
339
340 fn elwise_fma_scalar_mut(
341 &self,
342 a: &mut [Self::Element],
343 b: &[Self::Element],
344 c: &Self::Element,
345 ) {
346 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
347 *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *c));
348 });
349 }
350
351 // fn modulus(&self) -> Self::Element {
352 // self.q
353 // }
354}
355
356impl<M: Matrix<MatElement = u64>, T> ShoupMatrixFMA<M> for ModularOpsU64<T>
357where
358 M::R: RowMut,
359{
360 fn shoup_matrix_fma(&self, out: &mut <M as Matrix>::R, a: &M, a_shoup: &M, b: &M) {
361 assert!(a.dimension() == a_shoup.dimension());
362 assert!(a.dimension() == b.dimension());
363
364 let q = self.q;
365 let q_twice = self.q << 1;
366
367 // first row (without summation)
368 izip!(
369 out.as_mut().iter_mut(),
370 a.get_row(0),
371 a_shoup.get_row(0),
372 b.get_row(0)
373 )
374 .for_each(|(o, a, a_shoup, b)| {
375 *o = ShoupMul::mul(*b, *a, *a_shoup, q);
376 });
377
378 izip!(
379 a.iter_rows().skip(1),
380 a_shoup.iter_rows().skip(1),
381 b.iter_rows().skip(1)
382 )
383 .for_each(|(a_row, a_shoup_row, b_row)| {
384 izip!(
385 out.as_mut().iter_mut(),
386 a_row.as_ref().iter(),
387 a_shoup_row.as_ref().iter(),
388 b_row.as_ref().iter()
389 )
390 .for_each(|(o, a0, a0_shoup, b0)| {
391 let quotient = ((*a0_shoup as u128 * *b0 as u128) >> 64) as u64;
392 let mut v = (a0.wrapping_mul(b0)).wrapping_add(*o);
393 v = v.wrapping_sub(q.wrapping_mul(quotient));
394
395 if v >= q_twice {
396 v -= q_twice;
397 }
398
399 *o = v;
400 });
401 });
402 }
403}
404
405impl<T> GetModulus for ModularOpsU64<T>
406where
407 T: Modulus,
408{
409 type Element = T::Element;
410 type M = T;
411 fn modulus(&self) -> &Self::M {
412 &self.modulus
413 }
414}
415
416pub struct WordSizeModulus<T> {
417 modulus: T,
418}
419
420impl<T> ModInit for WordSizeModulus<T>
421where
422 T: Modulus,
423{
424 type M = T;
425 fn new(modulus: T) -> Self {
426 assert!(modulus.is_native());
427 // For now assume ModulusOpsU64 is only used for u64
428 Self { modulus: modulus }
429 }
430}
431
432impl<T> ArithmeticOps for WordSizeModulus<T>
433where
434 T: Modulus,
435 T::Element: WrappingAdd + WrappingSub + WrappingMul + Zero,
436{
437 type Element = T::Element;
438 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
439 T::Element::wrapping_add(a, b)
440 }
441
442 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
443 T::Element::wrapping_mul(a, b)
444 }
445
446 fn neg(&self, a: &Self::Element) -> Self::Element {
447 T::Element::wrapping_sub(&T::Element::zero(), a)
448 }
449
450 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
451 T::Element::wrapping_sub(a, b)
452 }
453}
454
455impl<T> VectorOps for WordSizeModulus<T>
456where
457 T: Modulus,
458 T::Element: WrappingAdd + WrappingSub + WrappingMul + Zero,
459{
460 type Element = T::Element;
461
462 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
463 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
464 *ai = T::Element::wrapping_add(ai, bi);
465 });
466 }
467
468 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
469 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
470 *ai = T::Element::wrapping_sub(ai, bi);
471 });
472 }
473
474 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
475 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
476 *ai = T::Element::wrapping_mul(ai, bi);
477 });
478 }
479
480 fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
481 a.iter_mut()
482 .for_each(|ai| *ai = T::Element::wrapping_sub(&T::Element::zero(), ai));
483 }
484
485 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
486 izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
487 *oi = T::Element::wrapping_mul(ai, b);
488 });
489 }
490
491 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
492 izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
493 *oi = T::Element::wrapping_mul(ai, bi);
494 });
495 }
496
497 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
498 a.iter_mut().for_each(|ai| {
499 *ai = T::Element::wrapping_mul(ai, b);
500 });
501 }
502
503 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
504 izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
505 *ai = T::Element::wrapping_add(ai, &T::Element::wrapping_mul(bi, ci));
506 });
507 }
508
509 fn elwise_fma_scalar_mut(
510 &self,
511 a: &mut [Self::Element],
512 b: &[Self::Element],
513 c: &Self::Element,
514 ) {
515 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
516 *ai = T::Element::wrapping_add(ai, &T::Element::wrapping_mul(bi, c));
517 });
518 }
519
520 // fn modulus(&self) -> &T {
521 // &self.modulus
522 // }
523}
524
525impl<T> GetModulus for WordSizeModulus<T>
526where
527 T: Modulus,
528{
529 type Element = T::Element;
530 type M = T;
531 fn modulus(&self) -> &Self::M {
532 &self.modulus
533 }
534}
535
536#[cfg(test)]
537mod tests {
538 use super::*;
539 use itertools::Itertools;
540 use rand::{thread_rng, Rng};
541 use rand_distr::Uniform;
542
543 #[test]
544 fn fma() {
545 let mut rng = thread_rng();
546 let prime = 36028797017456641;
547 let ring_size = 1 << 3;
548
549 let dist = Uniform::new(0, prime);
550 let d = 2;
551 let a0_matrix = (0..d)
552 .into_iter()
553 .map(|_| (&mut rng).sample_iter(dist).take(ring_size).collect_vec())
554 .collect_vec();
555 // a0 in shoup representation
556 let a0_shoup_matrix = a0_matrix
557 .iter()
558 .map(|r| {
559 r.iter()
560 .map(|v| {
561 // $(v * 2^{\beta}) / p$
562 ((*v as u128 * (1u128 << 64)) / prime as u128) as u64
563 })
564 .collect_vec()
565 })
566 .collect_vec();
567 let a1_matrix = (0..d)
568 .into_iter()
569 .map(|_| (&mut rng).sample_iter(dist).take(ring_size).collect_vec())
570 .collect_vec();
571
572 let modop = ModularOpsU64::new(prime);
573
574 let mut out_shoup_fma_lazy = vec![0u64; ring_size];
575 modop.shoup_matrix_fma(
576 &mut out_shoup_fma_lazy,
577 &a0_matrix,
578 &a0_shoup_matrix,
579 &a1_matrix,
580 );
581 let out_shoup_fma = out_shoup_fma_lazy
582 .iter()
583 .map(|v| if *v >= prime { v - prime } else { *v })
584 .collect_vec();
585
586 // expected
587 let mut out_expected = vec![0u64; ring_size];
588 izip!(a0_matrix.iter(), a1_matrix.iter()).for_each(|(a_r, b_r)| {
589 izip!(out_expected.iter_mut(), a_r.iter(), b_r.iter()).for_each(|(o, a0, a1)| {
590 *o = (*o + ((*a0 as u128 * *a1 as u128) % prime as u128) as u64) % prime;
591 });
592 });
593
594 assert_eq!(out_expected, out_shoup_fma);
595 }
596}