phantom-zone/src/ntt.rs
src/ntt.rsRaw
1use itertools::{izip, Itertools};
2use rand::{thread_rng, Rng, RngCore, SeedableRng};
3use rand_chacha::{rand_core::le, ChaCha8Rng};
4
5use crate::{
6 backend::{ArithmeticOps, ModInit, ModularOpsU64, Modulus},
7 utils::{mod_exponent, mod_inverse, shoup_representation_fq},
8};
9
10pub trait NttInit<M> {
11 /// Ntt istance must be compatible across different instances with same `q`
12 /// and `n`
13 fn new(q: &M, n: usize) -> Self;
14}
15
16pub trait Ntt {
17 type Element;
18 fn forward_lazy(&self, v: &mut [Self::Element]);
19 fn forward(&self, v: &mut [Self::Element]);
20 fn backward_lazy(&self, v: &mut [Self::Element]);
21 fn backward(&self, v: &mut [Self::Element]);
22}
23
24/// Forward butterfly routine for Number theoretic transform. Given inputs `x <
25/// 4q` and `y < 4q` mutates x and y in place to equal x' and y' where
26/// x' = x + wy
27/// y' = x - wy
28/// and both x' and y' are \in [0, 4q)
29///
30/// Implements Algorithm 4 of [FASTER ARITHMETIC FOR NUMBER-THEORETIC TRANSFORMS](https://arxiv.org/pdf/1205.2926.pdf)
31pub fn forward_butterly_0_to_4q(
32 mut x: u64,
33 y: u64,
34 w: u64,
35 w_shoup: u64,
36 q: u64,
37 q_twice: u64,
38) -> (u64, u64) {
39 debug_assert!(x < q * 4, "{} >= (4q){}", x, 4 * q);
40 debug_assert!(y < q * 4, "{} >= (4q){}", y, 4 * q);
41
42 if x >= q_twice {
43 x = x - q_twice;
44 }
45
46 // TODO (Jay): Hot path expected. How expensive is it?
47 let k = ((w_shoup as u128 * y as u128) >> 64) as u64;
48 let t = w.wrapping_mul(y).wrapping_sub(k.wrapping_mul(q));
49
50 (x + t, x + q_twice - t)
51}
52
53pub fn forward_butterly_0_to_2q(
54 mut x: u64,
55 mut y: u64,
56 w: u64,
57 w_shoup: u64,
58 q: u64,
59 q_twice: u64,
60) -> (u64, u64) {
61 debug_assert!(x < q * 4, "{} >= (4q){}", x, 4 * q);
62 debug_assert!(y < q * 4, "{} >= (4q){}", y, 4 * q);
63
64 if x >= q_twice {
65 x = x - q_twice;
66 }
67
68 let k = ((w_shoup as u128 * y as u128) >> 64) as u64;
69 let t = w.wrapping_mul(y).wrapping_sub(k.wrapping_mul(q));
70
71 let ox = x.wrapping_add(t);
72 let oy = x.wrapping_sub(t);
73
74 (
75 (ox).min(ox.wrapping_sub(q_twice)),
76 oy.min(oy.wrapping_add(q_twice)),
77 )
78}
79
80/// Inverse butterfly routine of Inverse Number theoretic transform. Given
81/// inputs `x < 2q` and `y < 2q` mutates x and y to equal x' and y' where
82/// x'= x + y
83/// y' = w(x - y)
84/// and both x' and y' are \in [0, 2q)
85///
86/// Implements Algorithm 3 of [FASTER ARITHMETIC FOR NUMBER-THEORETIC TRANSFORMS](https://arxiv.org/pdf/1205.2926.pdf)
87pub fn inverse_butterfly(
88 x: u64,
89 y: u64,
90 w_inv: u64,
91 w_inv_shoup: u64,
92 q: u64,
93 q_twice: u64,
94) -> (u64, u64) {
95 debug_assert!(x < q_twice, "{} >= (2q){q_twice}", x);
96 debug_assert!(y < q_twice, "{} >= (2q){q_twice}", y);
97
98 let mut x_dash = x + y;
99 if x_dash >= q_twice {
100 x_dash -= q_twice
101 }
102
103 let t = x + q_twice - y;
104 let k = ((w_inv_shoup as u128 * t as u128) >> 64) as u64;
105 let y = w_inv.wrapping_mul(t).wrapping_sub(k.wrapping_mul(q));
106
107 (x_dash, y)
108}
109
110/// Number theoretic transform of vector `a` where each element can be in range
111/// [0, 2q). Outputs NTT(a) where each element is in range [0,2q)
112///
113/// Implements Cooley-tukey based forward NTT as given in Algorithm 1 of https://eprint.iacr.org/2016/504.pdf.
114pub fn ntt_lazy(a: &mut [u64], psi: &[u64], psi_shoup: &[u64], q: u64, q_twice: u64) {
115 assert!(a.len() == psi.len());
116
117 let n = a.len();
118 let mut t = n;
119
120 let mut m = 1;
121 while m < n {
122 t >>= 1;
123 let w = &psi[m..];
124 let w_shoup = &psi_shoup[m..];
125
126 if t == 1 {
127 for (a, w, w_shoup) in izip!(a.chunks_mut(2), w.iter(), w_shoup.iter()) {
128 let (ox, oy) = forward_butterly_0_to_2q(a[0], a[1], *w, *w_shoup, q, q_twice);
129 a[0] = ox;
130 a[1] = oy;
131 }
132 } else {
133 for i in 0..m {
134 let a = &mut a[2 * i * t..(2 * (i + 1) * t)];
135 let (left, right) = a.split_at_mut(t);
136
137 for (x, y) in izip!(left.iter_mut(), right.iter_mut()) {
138 let (ox, oy) = forward_butterly_0_to_4q(*x, *y, w[i], w_shoup[i], q, q_twice);
139 *x = ox;
140 *y = oy;
141 }
142 }
143 }
144
145 m <<= 1;
146 }
147}
148
149/// Same as `ntt_lazy` with output in range [0, q)
150pub fn ntt(a: &mut [u64], psi: &[u64], psi_shoup: &[u64], q: u64, q_twice: u64) {
151 assert!(a.len() == psi.len());
152
153 let n = a.len();
154 let mut t = n;
155
156 let mut m = 1;
157 while m < n {
158 t >>= 1;
159 let w = &psi[m..];
160 let w_shoup = &psi_shoup[m..];
161
162 if t == 1 {
163 for (a, w, w_shoup) in izip!(a.chunks_mut(2), w.iter(), w_shoup.iter()) {
164 let (ox, oy) = forward_butterly_0_to_2q(a[0], a[1], *w, *w_shoup, q, q_twice);
165 a[0] = ox.min(ox.wrapping_sub(q_twice));
166 a[1] = oy.min(oy.wrapping_sub(q_twice));
167 }
168 } else {
169 for i in 0..m {
170 let a = &mut a[2 * i * t..(2 * (i + 1) * t)];
171 let (left, right) = a.split_at_mut(t);
172
173 for (x, y) in izip!(left.iter_mut(), right.iter_mut()) {
174 let (ox, oy) = forward_butterly_0_to_4q(*x, *y, w[i], w_shoup[i], q, q_twice);
175 *x = ox;
176 *y = oy;
177 }
178 }
179 }
180
181 m <<= 1;
182 }
183}
184
185/// Inverse number theoretic transform of input vector `a` with each element can
186/// be in range [0, 2q). Outputs vector INTT(a) with each element in range [0,
187/// 2q)
188///
189/// Implements backward number theorectic transform using GS algorithm as given in Algorithm 2 of https://eprint.iacr.org/2016/504.pdf
190pub fn ntt_inv_lazy(
191 a: &mut [u64],
192 psi_inv: &[u64],
193 psi_inv_shoup: &[u64],
194 n_inv: u64,
195 n_inv_shoup: u64,
196 q: u64,
197 q_twice: u64,
198) {
199 assert!(a.len() == psi_inv.len());
200
201 let mut m = a.len() >> 1;
202 let mut t = 1;
203
204 while m > 0 {
205 let w_inv = &psi_inv[m..];
206 let w_inv_shoup = &psi_inv_shoup[m..];
207
208 if m == 1 {
209 let (left, right) = a.split_at_mut(t);
210
211 for (x, y) in izip!(left.iter_mut(), right.iter_mut()) {
212 let (ox, oy) = inverse_butterfly(*x, *y, w_inv[0], w_inv_shoup[0], q, q_twice);
213
214 *x = (n_inv.wrapping_mul(ox)).wrapping_sub(
215 q.wrapping_mul(((ox as u128 * n_inv_shoup as u128) >> 64) as u64),
216 );
217 *y = (n_inv.wrapping_mul(oy)).wrapping_sub(
218 q.wrapping_mul(((oy as u128 * n_inv_shoup as u128) >> 64) as u64),
219 );
220 }
221 } else {
222 for i in 0..m {
223 let a = &mut a[2 * i * t..2 * (i + 1) * t];
224 let (left, right) = a.split_at_mut(t);
225
226 for (x, y) in izip!(left.iter_mut(), right.iter_mut()) {
227 let (ox, oy) = inverse_butterfly(*x, *y, w_inv[i], w_inv_shoup[i], q, q_twice);
228 *x = ox;
229 *y = oy;
230 }
231 }
232 }
233
234 // for i in 0..m {
235 // let a = &mut a[2 * i * t..2 * (i + 1) * t];
236 // let (left, right) = a.split_at_mut(t);
237
238 // for (x, y) in izip!(left.iter_mut(), right.iter_mut()) {
239 // let (ox, oy) = inverse_butterfly(*x, *y, w_inv[i], w_inv_shoup[i], q,
240 // q_twice); *x = ox;
241 // *y = oy;
242 // }
243 // }
244
245 // for i in 0..h {
246 // let j_2 = j_1 + t;
247 // unsafe {
248 // let w_inv = psi_inv.get_unchecked(h + i);
249 // let w_inv_shoup = psi_inv_shoup.get_unchecked(h + i);
250
251 // for j in j_1..j_2 {
252 // let x = a.get_unchecked_mut(j) as *mut u64;
253 // let y = a.get_unchecked_mut(j + t) as *mut u64;
254 // inverse_butterfly(x, y, w_inv, w_inv_shoup, &q, &q_twice);
255 // }
256 // }
257 // j_1 = j_1 + 2 * t;
258 // }
259 t *= 2;
260 m >>= 1;
261 }
262
263 // a.iter_mut().for_each(|a0| {
264 // *a0 = (n_inv.wrapping_mul(*a0))
265 // .wrapping_sub(((*a0 as u128 * n_inv_shoup as u128) >> 64) as u64)
266 // });
267}
268
269/// Find n^{th} root of unity in field F_q, if one exists
270///
271/// Note: n^{th} root of unity exists if and only if $q = 1 \mod{n}$
272pub(crate) fn find_primitive_root<R: RngCore>(q: u64, n: u64, rng: &mut R) -> Option<u64> {
273 assert!(n.is_power_of_two(), "{n} is not power of two");
274
275 // n^th root of unity only exists if n|(q-1)
276 assert!(q % n == 1, "{n}^th root of unity in F_{q} does not exists");
277
278 let t = (q - 1) / n;
279
280 for _ in 0..100 {
281 let mut omega = rng.gen::<u64>() % q;
282
283 // \omega = \omega^t. \omega is now n^th root of unity
284 omega = mod_exponent(omega, t, q);
285
286 // We restrict n to be power of 2. Thus checking whether \omega is primitive
287 // n^th root of unity is as simple as checking: \omega^{n/2} != 1
288 if mod_exponent(omega, n >> 1, q) == 1 {
289 continue;
290 } else {
291 return Some(omega);
292 }
293 }
294
295 None
296}
297
298#[derive(Debug)]
299pub struct NttBackendU64 {
300 q: u64,
301 q_twice: u64,
302 n: u64,
303 n_inv: u64,
304 n_inv_shoup: u64,
305 psi_powers_bo: Box<[u64]>,
306 psi_inv_powers_bo: Box<[u64]>,
307 psi_powers_bo_shoup: Box<[u64]>,
308 psi_inv_powers_bo_shoup: Box<[u64]>,
309}
310
311impl NttBackendU64 {
312 fn _new(q: u64, n: usize) -> Self {
313 // \psi = 2n^{th} primitive root of unity in F_q
314 let mut rng = ChaCha8Rng::from_seed([0u8; 32]);
315 let psi = find_primitive_root(q, (n * 2) as u64, &mut rng)
316 .expect("Unable to find 2n^th root of unity");
317 let psi_inv = mod_inverse(psi, q);
318
319 // assert!(
320 // ((psi_inv as u128 * psi as u128) % q as u128) == 1,
321 // "psi:{psi}, psi_inv:{psi_inv}"
322 // );
323
324 let modulus = ModularOpsU64::new(q);
325
326 let mut psi_powers = Vec::with_capacity(n as usize);
327 let mut psi_inv_powers = Vec::with_capacity(n as usize);
328 let mut running_psi = 1;
329 let mut running_psi_inv = 1;
330 for _ in 0..n {
331 psi_powers.push(running_psi);
332 psi_inv_powers.push(running_psi_inv);
333
334 running_psi = modulus.mul(&running_psi, &psi);
335 running_psi_inv = modulus.mul(&running_psi_inv, &psi_inv);
336 }
337
338 // powers stored in bit reversed order
339 let mut psi_powers_bo = vec![0u64; n as usize];
340 let mut psi_inv_powers_bo = vec![0u64; n as usize];
341 let shift_by = n.leading_zeros() + 1;
342 for i in 0..n as usize {
343 // i in bit reversed order
344 let bo_index = i.reverse_bits() >> shift_by;
345
346 psi_powers_bo[bo_index] = psi_powers[i];
347 psi_inv_powers_bo[bo_index] = psi_inv_powers[i];
348 }
349
350 // shoup representation
351 let psi_powers_bo_shoup = psi_powers_bo
352 .iter()
353 .map(|v| shoup_representation_fq(*v, q))
354 .collect_vec();
355 let psi_inv_powers_bo_shoup = psi_inv_powers_bo
356 .iter()
357 .map(|v| shoup_representation_fq(*v, q))
358 .collect_vec();
359
360 // n^{-1} \mod{q}
361 let n_inv = mod_inverse(n as u64, q);
362
363 NttBackendU64 {
364 q,
365 q_twice: 2 * q,
366 n: n as u64,
367 n_inv,
368 n_inv_shoup: shoup_representation_fq(n_inv, q),
369 psi_powers_bo: psi_powers_bo.into_boxed_slice(),
370 psi_inv_powers_bo: psi_inv_powers_bo.into_boxed_slice(),
371 psi_powers_bo_shoup: psi_powers_bo_shoup.into_boxed_slice(),
372 psi_inv_powers_bo_shoup: psi_inv_powers_bo_shoup.into_boxed_slice(),
373 }
374 }
375}
376
377impl<M: Modulus<Element = u64>> NttInit<M> for NttBackendU64 {
378 fn new(q: &M, n: usize) -> Self {
379 // This NTT does not support native modulus
380 assert!(!q.is_native());
381 NttBackendU64::_new(q.q().unwrap(), n)
382 }
383}
384
385impl NttBackendU64 {
386 fn reduce_from_lazy(&self, a: &mut [u64]) {
387 let q = self.q;
388 a.iter_mut().for_each(|a0| {
389 if *a0 >= q {
390 *a0 = *a0 - q;
391 }
392 });
393 }
394}
395
396impl Ntt for NttBackendU64 {
397 type Element = u64;
398
399 fn forward_lazy(&self, v: &mut [Self::Element]) {
400 ntt_lazy(
401 v,
402 &self.psi_powers_bo,
403 &self.psi_powers_bo_shoup,
404 self.q,
405 self.q_twice,
406 )
407 }
408
409 fn forward(&self, v: &mut [Self::Element]) {
410 ntt(
411 v,
412 &self.psi_powers_bo,
413 &self.psi_powers_bo_shoup,
414 self.q,
415 self.q_twice,
416 );
417 }
418
419 fn backward_lazy(&self, v: &mut [Self::Element]) {
420 ntt_inv_lazy(
421 v,
422 &self.psi_inv_powers_bo,
423 &self.psi_inv_powers_bo_shoup,
424 self.n_inv,
425 self.n_inv_shoup,
426 self.q,
427 self.q_twice,
428 )
429 }
430
431 fn backward(&self, v: &mut [Self::Element]) {
432 ntt_inv_lazy(
433 v,
434 &self.psi_inv_powers_bo,
435 &self.psi_inv_powers_bo_shoup,
436 self.n_inv,
437 self.n_inv_shoup,
438 self.q,
439 self.q_twice,
440 );
441 self.reduce_from_lazy(v);
442 }
443}
444
445mod tests {
446 use itertools::Itertools;
447 use rand::{thread_rng, Rng};
448 use rand_distr::Uniform;
449
450 use super::{NttBackendU64, NttInit};
451 use crate::{
452 backend::{ArithmeticOps, ModInit, ModularOpsU64, VectorOps},
453 ntt::Ntt,
454 utils::{generate_prime, negacyclic_mul},
455 };
456
457 const Q_60_BITS: u64 = 1152921504606748673;
458 const N: usize = 1 << 4;
459
460 const K: usize = 128;
461
462 fn random_vec_in_fq(size: usize, q: u64) -> Vec<u64> {
463 thread_rng()
464 .sample_iter(Uniform::new(0, q))
465 .take(size)
466 .collect_vec()
467 }
468
469 #[test]
470 fn native_ntt_backend_works() {
471 // TODO(Jay): Improve tests. Add tests for different primes and ring size.
472 let ntt_backend = NttBackendU64::_new(Q_60_BITS, N);
473 for _ in 0..K {
474 let mut a = random_vec_in_fq(N, Q_60_BITS);
475 let a_clone = a.clone();
476
477 ntt_backend.forward(&mut a);
478 assert_ne!(a, a_clone);
479 ntt_backend.backward(&mut a);
480 assert_eq!(a, a_clone);
481
482 ntt_backend.forward_lazy(&mut a);
483 assert_ne!(a, a_clone);
484 ntt_backend.backward(&mut a);
485 assert_eq!(a, a_clone);
486
487 ntt_backend.forward(&mut a);
488 ntt_backend.backward_lazy(&mut a);
489 // reduce
490 a.iter_mut().for_each(|a0| {
491 if *a0 > Q_60_BITS {
492 *a0 -= *a0 - Q_60_BITS;
493 }
494 });
495 assert_eq!(a, a_clone);
496 }
497 }
498
499 #[test]
500 fn native_ntt_negacylic_mul() {
501 let primes = [25, 40, 50, 60]
502 .iter()
503 .map(|bits| generate_prime(*bits, (2 * N) as u64, 1u64 << bits).unwrap())
504 .collect_vec();
505
506 for p in primes.into_iter() {
507 let ntt_backend = NttBackendU64::_new(p, N);
508 let modulus_backend = ModularOpsU64::new(p);
509 for _ in 0..K {
510 let a = random_vec_in_fq(N, p);
511 let b = random_vec_in_fq(N, p);
512
513 let mut a_clone = a.clone();
514 let mut b_clone = b.clone();
515 ntt_backend.forward_lazy(&mut a_clone);
516 ntt_backend.forward_lazy(&mut b_clone);
517 modulus_backend.elwise_mul_mut(&mut a_clone, &b_clone);
518 ntt_backend.backward(&mut a_clone);
519
520 let mul = |a: &u64, b: &u64| {
521 let tmp = *a as u128 * *b as u128;
522 (tmp % p as u128) as u64
523 };
524 let expected_out = negacyclic_mul(&a, &b, mul, p);
525
526 assert_eq!(a_clone, expected_out);
527 }
528 }
529 }
530}