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