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 unsafe fn inverse_butterfly(
88 x: *mut u64,
89 y: *mut u64,
90 w_inv: &u64,
91 w_inv_shoup: &u64,
92 q: &u64,
93 q_twice: &u64,
94) {
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; // TODO (Jay): Hot path
105 *y = w_inv.wrapping_mul(t).wrapping_sub(k.wrapping_mul(*q));
106
107 *x = x_dash;
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 q: u64,
196 q_twice: u64,
197) {
198 debug_assert!(a.len() == psi_inv.len());
199
200 let mut m = a.len();
201 let mut t = 1;
202 while m > 1 {
203 let mut j_1: usize = 0;
204 let h = m >> 1;
205 for i in 0..h {
206 let j_2 = j_1 + t;
207 unsafe {
208 let w_inv = psi_inv.get_unchecked(h + i);
209 let w_inv_shoup = psi_inv_shoup.get_unchecked(h + i);
210
211 for j in j_1..j_2 {
212 let x = a.get_unchecked_mut(j) as *mut u64;
213 let y = a.get_unchecked_mut(j + t) as *mut u64;
214 inverse_butterfly(x, y, w_inv, w_inv_shoup, &q, &q_twice);
215 }
216 }
217 j_1 = j_1 + 2 * t;
218 }
219 t *= 2;
220 m >>= 1;
221 }
222
223 a.iter_mut()
224 .for_each(|a0| *a0 = ((*a0 as u128 * n_inv as u128) % q as u128) as u64);
225}
226
227/// Find n^{th} root of unity in field F_q, if one exists
228///
229/// Note: n^{th} root of unity exists if and only if $q = 1 \mod{n}$
230pub(crate) fn find_primitive_root<R: RngCore>(q: u64, n: u64, rng: &mut R) -> Option<u64> {
231 assert!(n.is_power_of_two(), "{n} is not power of two");
232
233 // n^th root of unity only exists if n|(q-1)
234 assert!(q % n == 1, "{n}^th root of unity in F_{q} does not exists");
235
236 let t = (q - 1) / n;
237
238 for _ in 0..100 {
239 let mut omega = rng.gen::<u64>() % q;
240
241 // \omega = \omega^t. \omega is now n^th root of unity
242 omega = mod_exponent(omega, t, q);
243
244 // We restrict n to be power of 2. Thus checking whether \omega is primitive
245 // n^th root of unity is as simple as checking: \omega^{n/2} != 1
246 if mod_exponent(omega, n >> 1, q) == 1 {
247 continue;
248 } else {
249 return Some(omega);
250 }
251 }
252
253 None
254}
255
256#[derive(Debug)]
257pub struct NttBackendU64 {
258 q: u64,
259 q_twice: u64,
260 n: u64,
261 n_inv: u64,
262 psi_powers_bo: Box<[u64]>,
263 psi_inv_powers_bo: Box<[u64]>,
264 psi_powers_bo_shoup: Box<[u64]>,
265 psi_inv_powers_bo_shoup: Box<[u64]>,
266}
267
268impl NttBackendU64 {
269 fn _new(q: u64, n: usize) -> Self {
270 // \psi = 2n^{th} primitive root of unity in F_q
271 let mut rng = ChaCha8Rng::from_seed([0u8; 32]);
272 let psi = find_primitive_root(q, (n * 2) as u64, &mut rng)
273 .expect("Unable to find 2n^th root of unity");
274 let psi_inv = mod_inverse(psi, q);
275
276 // assert!(
277 // ((psi_inv as u128 * psi as u128) % q as u128) == 1,
278 // "psi:{psi}, psi_inv:{psi_inv}"
279 // );
280
281 let modulus = ModularOpsU64::new(q);
282
283 let mut psi_powers = Vec::with_capacity(n as usize);
284 let mut psi_inv_powers = Vec::with_capacity(n as usize);
285 let mut running_psi = 1;
286 let mut running_psi_inv = 1;
287 for _ in 0..n {
288 psi_powers.push(running_psi);
289 psi_inv_powers.push(running_psi_inv);
290
291 running_psi = modulus.mul(&running_psi, &psi);
292 running_psi_inv = modulus.mul(&running_psi_inv, &psi_inv);
293 }
294
295 // powers stored in bit reversed order
296 let mut psi_powers_bo = vec![0u64; n as usize];
297 let mut psi_inv_powers_bo = vec![0u64; n as usize];
298 let shift_by = n.leading_zeros() + 1;
299 for i in 0..n as usize {
300 // i in bit reversed order
301 let bo_index = i.reverse_bits() >> shift_by;
302
303 psi_powers_bo[bo_index] = psi_powers[i];
304 psi_inv_powers_bo[bo_index] = psi_inv_powers[i];
305 }
306
307 // shoup representation
308 let psi_powers_bo_shoup = psi_powers_bo
309 .iter()
310 .map(|v| shoup_representation_fq(*v, q))
311 .collect_vec();
312 let psi_inv_powers_bo_shoup = psi_inv_powers_bo
313 .iter()
314 .map(|v| shoup_representation_fq(*v, q))
315 .collect_vec();
316
317 // n^{-1} \mod{q}
318 let n_inv = mod_inverse(n as u64, q);
319
320 NttBackendU64 {
321 q,
322 q_twice: 2 * q,
323 n: n as u64,
324 n_inv,
325 psi_powers_bo: psi_powers_bo.into_boxed_slice(),
326 psi_inv_powers_bo: psi_inv_powers_bo.into_boxed_slice(),
327 psi_powers_bo_shoup: psi_powers_bo_shoup.into_boxed_slice(),
328 psi_inv_powers_bo_shoup: psi_inv_powers_bo_shoup.into_boxed_slice(),
329 }
330 }
331}
332
333impl<M: Modulus<Element = u64>> NttInit<M> for NttBackendU64 {
334 fn new(q: &M, n: usize) -> Self {
335 // This NTT does not support native modulus
336 assert!(!q.is_native());
337 NttBackendU64::_new(q.q().unwrap(), n)
338 }
339}
340
341impl NttBackendU64 {
342 fn reduce_from_lazy(&self, a: &mut [u64]) {
343 let q = self.q;
344 a.iter_mut().for_each(|a0| {
345 if *a0 >= q {
346 *a0 = *a0 - q;
347 }
348 });
349 }
350}
351
352impl Ntt for NttBackendU64 {
353 type Element = u64;
354
355 fn forward_lazy(&self, v: &mut [Self::Element]) {
356 ntt_lazy(
357 v,
358 &self.psi_powers_bo,
359 &self.psi_powers_bo_shoup,
360 self.q,
361 self.q_twice,
362 )
363 }
364
365 fn forward(&self, v: &mut [Self::Element]) {
366 ntt(
367 v,
368 &self.psi_powers_bo,
369 &self.psi_powers_bo_shoup,
370 self.q,
371 self.q_twice,
372 );
373 }
374
375 fn backward_lazy(&self, v: &mut [Self::Element]) {
376 ntt_inv_lazy(
377 v,
378 &self.psi_inv_powers_bo,
379 &self.psi_inv_powers_bo_shoup,
380 self.n_inv,
381 self.q,
382 self.q_twice,
383 )
384 }
385
386 fn backward(&self, v: &mut [Self::Element]) {
387 ntt_inv_lazy(
388 v,
389 &self.psi_inv_powers_bo,
390 &self.psi_inv_powers_bo_shoup,
391 self.n_inv,
392 self.q,
393 self.q_twice,
394 );
395 self.reduce_from_lazy(v);
396 }
397}
398
399mod tests {
400 use itertools::Itertools;
401 use rand::{thread_rng, Rng};
402 use rand_distr::Uniform;
403
404 use super::{NttBackendU64, NttInit};
405 use crate::{
406 backend::{ArithmeticOps, ModInit, ModularOpsU64, VectorOps},
407 ntt::Ntt,
408 utils::{generate_prime, negacyclic_mul},
409 };
410
411 const Q_60_BITS: u64 = 1152921504606748673;
412 const N: usize = 1 << 4;
413
414 const K: usize = 128;
415
416 fn random_vec_in_fq(size: usize, q: u64) -> Vec<u64> {
417 thread_rng()
418 .sample_iter(Uniform::new(0, q))
419 .take(size)
420 .collect_vec()
421 }
422
423 #[test]
424 fn native_ntt_backend_works() {
425 // TODO(Jay): Improve tests. Add tests for different primes and ring size.
426 let ntt_backend = NttBackendU64::_new(Q_60_BITS, N);
427 for _ in 0..K {
428 let mut a = random_vec_in_fq(N, Q_60_BITS);
429 let a_clone = a.clone();
430
431 ntt_backend.forward(&mut a);
432 assert_ne!(a, a_clone);
433 ntt_backend.backward(&mut a);
434 assert_eq!(a, a_clone);
435
436 ntt_backend.forward_lazy(&mut a);
437 assert_ne!(a, a_clone);
438 ntt_backend.backward(&mut a);
439 assert_eq!(a, a_clone);
440
441 ntt_backend.forward(&mut a);
442 ntt_backend.backward_lazy(&mut a);
443 // reduce
444 a.iter_mut().for_each(|a0| {
445 if *a0 > Q_60_BITS {
446 *a0 -= *a0 - Q_60_BITS;
447 }
448 });
449 assert_eq!(a, a_clone);
450 }
451 }
452
453 #[test]
454 fn native_ntt_negacylic_mul() {
455 let primes = [25, 40, 50, 60]
456 .iter()
457 .map(|bits| generate_prime(*bits, (2 * N) as u64, 1u64 << bits).unwrap())
458 .collect_vec();
459
460 for p in primes.into_iter() {
461 let ntt_backend = NttBackendU64::_new(p, N);
462 let modulus_backend = ModularOpsU64::new(p);
463 for _ in 0..K {
464 let a = random_vec_in_fq(N, p);
465 let b = random_vec_in_fq(N, p);
466
467 let mut a_clone = a.clone();
468 let mut b_clone = b.clone();
469 ntt_backend.forward_lazy(&mut a_clone);
470 ntt_backend.forward_lazy(&mut b_clone);
471 modulus_backend.elwise_mul_mut(&mut a_clone, &b_clone);
472 ntt_backend.backward(&mut a_clone);
473
474 let mul = |a: &u64, b: &u64| {
475 let tmp = *a as u128 * *b as u128;
476 (tmp % p as u128) as u64
477 };
478 let expected_out = negacyclic_mul(&a, &b, mul, p);
479
480 assert_eq!(a_clone, expected_out);
481 }
482 }
483 }
484}