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
6pub trait Modulus {
7 type Element;
8 /// Modulus value if it fits in Element
9 fn q(&self) -> Option<Self::Element>;
10 /// Modulus value as f64 if it fits in f64
11 fn q_as_f64(&self) -> Option<f64>;
12 /// Is modulus native?
13 fn is_native(&self) -> bool;
14 /// -1 in signed representaiton
15 fn neg_one(&self) -> Self::Element;
16 /// Largest unsigned value that fits in the modulus. That is, q - 1.
17 fn largest_unsigned_value(&self) -> Self::Element;
18 /// Smallest unsigned value that fits in the modulus
19 /// Always assmed to be 0.
20 fn smallest_unsigned_value(&self) -> Self::Element;
21 /// Convert unsigned value in signed represetation to i64
22 fn map_element_to_i64(&self, v: &Self::Element) -> i64;
23 /// Convert f64 to signed represented in modulus
24 fn map_element_from_f64(&self, v: f64) -> Self::Element;
25 /// Convert i64 to signed represented in modulus
26 fn map_element_from_i64(&self, v: i64) -> Self::Element;
27}
28
29impl Modulus for u64 {
30 type Element = u64;
31 fn is_native(&self) -> bool {
32 // q of size u64 can never be a naitve modulus
33 false
34 }
35 fn largest_unsigned_value(&self) -> Self::Element {
36 self - 1
37 }
38 fn neg_one(&self) -> Self::Element {
39 self - 1
40 }
41 fn smallest_unsigned_value(&self) -> Self::Element {
42 0
43 }
44 fn map_element_to_i64(&self, v: &Self::Element) -> i64 {
45 assert!(v < self);
46
47 if *v > (self >> 1) {
48 ToPrimitive::to_i64(&(self - v)).unwrap()
49 } else {
50 ToPrimitive::to_i64(v).unwrap()
51 }
52 }
53 fn map_element_from_f64(&self, v: f64) -> Self::Element {
54 //FIXME (Jay): Before I check whether v is smaller than 0 with `let is_neg =
55 // o.is_sign_negative() && o != 0.0; I'm ocnfused why didn't I simply check <
56 // 0.0?
57 let v = v.round();
58 if v < 0.0 {
59 self - v.to_u64().unwrap()
60 } else {
61 v.to_u64().unwrap()
62 }
63 }
64 fn map_element_from_i64(&self, v: i64) -> Self::Element {
65 if v < 0 {
66 self - v.to_u64().unwrap()
67 } else {
68 v.to_u64().unwrap()
69 }
70 }
71 fn q(&self) -> Option<Self::Element> {
72 Some(*self)
73 }
74 fn q_as_f64(&self) -> Option<f64> {
75 self.to_f64()
76 }
77}
78
79pub trait ModInit {
80 type M;
81 fn new(modulus: Self::M) -> Self;
82}
83
84pub trait GetModulus {
85 type Element;
86 type M: Modulus<Element = Self::Element>;
87 fn modulus(&self) -> &Self::M;
88}
89
90pub trait VectorOps {
91 type Element;
92
93 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element);
94 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]);
95
96 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
97 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
98 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
99 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element);
100 fn elwise_neg_mut(&self, a: &mut [Self::Element]);
101 /// inplace mutates `a`: a = a + b*c
102 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]);
103 fn elwise_fma_scalar_mut(
104 &self,
105 a: &mut [Self::Element],
106 b: &[Self::Element],
107 c: &Self::Element,
108 );
109
110 // fn modulus(&self) -> Self::Element;
111}
112
113pub trait ArithmeticOps {
114 type Element;
115
116 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
117 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
118 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
119 fn neg(&self, a: &Self::Element) -> Self::Element;
120
121 // fn modulus(&self) -> Self::Element;
122}
123
124pub struct ModularOpsU64<T> {
125 q: u64,
126 logq: usize,
127 barrett_mu: u128,
128 barrett_alpha: usize,
129 modulus: T,
130}
131
132impl<T> ModInit for ModularOpsU64<T>
133where
134 T: Modulus<Element = u64>,
135{
136 type M = T;
137 fn new(modulus: Self::M) -> ModularOpsU64<T> {
138 assert!(!modulus.is_native());
139
140 // largest unsigned value modulus fits is modulus-1
141 let q = modulus.largest_unsigned_value() + 1;
142 let logq = 64 - (q + 1u64).leading_zeros();
143
144 // barrett calculation
145 let mu = (1u128 << (logq * 2 + 3)) / (q as u128);
146 let alpha = logq + 3;
147
148 ModularOpsU64 {
149 q,
150 logq: logq as usize,
151 barrett_alpha: alpha as usize,
152 barrett_mu: mu,
153 modulus,
154 }
155 }
156}
157
158impl<T> ModularOpsU64<T> {
159 fn add_mod_fast(&self, a: u64, b: u64) -> u64 {
160 debug_assert!(a < self.q);
161 debug_assert!(b < self.q);
162
163 let mut o = a + b;
164 if o >= self.q {
165 o -= self.q;
166 }
167 o
168 }
169
170 fn sub_mod_fast(&self, a: u64, b: u64) -> u64 {
171 debug_assert!(a < self.q);
172 debug_assert!(b < self.q);
173
174 if a > b {
175 a - b
176 } else {
177 (self.q + a) - b
178 }
179 }
180
181 /// returns (a * b) % q
182 ///
183 /// - both a and b must be in range [0, 2q)
184 /// - output is in range [0 , q)
185 fn mul_mod_fast(&self, a: u64, b: u64) -> u64 {
186 debug_assert!(a < 2 * self.q);
187 debug_assert!(b < 2 * self.q);
188
189 let ab = a as u128 * b as u128;
190
191 // ab / (2^{n + \beta})
192 // note: \beta is assumed to -2
193 let tmp = ab >> (self.logq - 2);
194
195 // k = ((ab / (2^{n + \beta})) * \mu) / 2^{\alpha - (-2)}
196 let k = (tmp * self.barrett_mu) >> (self.barrett_alpha + 2);
197
198 // ab - k*p
199 let tmp = k * (self.q as u128);
200
201 let mut out = (ab - tmp) as u64;
202
203 if out >= self.q {
204 out -= self.q;
205 }
206
207 return out;
208 }
209}
210
211impl<T> ArithmeticOps for ModularOpsU64<T> {
212 type Element = u64;
213
214 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
215 self.add_mod_fast(*a, *b)
216 }
217
218 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
219 self.mul_mod_fast(*a, *b)
220 }
221
222 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
223 self.sub_mod_fast(*a, *b)
224 }
225
226 fn neg(&self, a: &Self::Element) -> Self::Element {
227 self.q - *a
228 }
229
230 // fn modulus(&self) -> Self::Element {
231 // self.q
232 // }
233}
234
235impl<T> VectorOps for ModularOpsU64<T> {
236 type Element = u64;
237
238 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
239 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
240 *ai = self.add_mod_fast(*ai, *bi);
241 });
242 }
243
244 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
245 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
246 *ai = self.sub_mod_fast(*ai, *bi);
247 });
248 }
249
250 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
251 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
252 *ai = self.mul_mod_fast(*ai, *bi);
253 });
254 }
255
256 fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
257 a.iter_mut().for_each(|ai| *ai = self.q - *ai);
258 }
259
260 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
261 izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
262 *oi = self.mul_mod_fast(*ai, *b);
263 });
264 }
265
266 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
267 izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
268 *oi = self.mul_mod_fast(*ai, *bi);
269 });
270 }
271
272 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
273 a.iter_mut().for_each(|ai| {
274 *ai = self.mul_mod_fast(*ai, *b);
275 });
276 }
277
278 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
279 izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
280 *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *ci));
281 });
282 }
283
284 fn elwise_fma_scalar_mut(
285 &self,
286 a: &mut [Self::Element],
287 b: &[Self::Element],
288 c: &Self::Element,
289 ) {
290 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
291 *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *c));
292 });
293 }
294
295 // fn modulus(&self) -> Self::Element {
296 // self.q
297 // }
298}
299
300impl<T> GetModulus for ModularOpsU64<T>
301where
302 T: Modulus,
303{
304 type Element = T::Element;
305 type M = T;
306 fn modulus(&self) -> &Self::M {
307 &self.modulus
308 }
309}
310
311pub struct WordSizeModulus<T> {
312 modulus: T,
313}
314
315impl<T> ModInit for WordSizeModulus<T>
316where
317 T: Modulus,
318{
319 type M = T;
320 fn new(modulus: T) -> Self {
321 assert!(modulus.is_native());
322 // For now assume ModulusOpsU64 is only used for u64
323 Self { modulus: modulus }
324 }
325}
326
327impl<T> ArithmeticOps for WordSizeModulus<T>
328where
329 T: Modulus,
330 T::Element: WrappingAdd + WrappingSub + WrappingMul + Zero,
331{
332 type Element = T::Element;
333 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
334 T::Element::wrapping_add(a, b)
335 }
336
337 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
338 T::Element::wrapping_mul(a, b)
339 }
340
341 fn neg(&self, a: &Self::Element) -> Self::Element {
342 T::Element::wrapping_sub(&T::Element::zero(), a)
343 }
344
345 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
346 T::Element::wrapping_sub(a, b)
347 }
348
349 // fn modulus(&self) -> &T {
350 // &self.modulus
351 // }
352}
353
354impl<T> VectorOps for WordSizeModulus<T>
355where
356 T: Modulus,
357 T::Element: WrappingAdd + WrappingSub + WrappingMul + Zero,
358{
359 type Element = T::Element;
360
361 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
362 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
363 *ai = T::Element::wrapping_add(ai, bi);
364 });
365 }
366
367 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
368 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
369 *ai = T::Element::wrapping_sub(ai, bi);
370 });
371 }
372
373 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
374 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
375 *ai = T::Element::wrapping_mul(ai, bi);
376 });
377 }
378
379 fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
380 a.iter_mut()
381 .for_each(|ai| *ai = T::Element::wrapping_sub(&T::Element::zero(), ai));
382 }
383
384 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
385 izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
386 *oi = T::Element::wrapping_mul(ai, b);
387 });
388 }
389
390 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
391 izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
392 *oi = T::Element::wrapping_mul(ai, bi);
393 });
394 }
395
396 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
397 a.iter_mut().for_each(|ai| {
398 *ai = T::Element::wrapping_mul(ai, b);
399 });
400 }
401
402 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
403 izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
404 *ai = T::Element::wrapping_add(ai, &T::Element::wrapping_mul(bi, ci));
405 });
406 }
407
408 fn elwise_fma_scalar_mut(
409 &self,
410 a: &mut [Self::Element],
411 b: &[Self::Element],
412 c: &Self::Element,
413 ) {
414 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
415 *ai = T::Element::wrapping_add(ai, &T::Element::wrapping_mul(bi, c));
416 });
417 }
418
419 // fn modulus(&self) -> &T {
420 // &self.modulus
421 // }
422}