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