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