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