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