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