fhe-study/arith/src/ntt_u64.rs
arith/src/ntt_u64.rsRaw
1//! This file implements the wrapper on top of the ntt.rs to be able to compute
2//! the NTT for non-prime modulus, specifically for modulus 2^64 (for u64).
3
4use crate::ntt::NTT as NTT_p;
5
6// const P0: u64 = 17293822569241362433;
7
8// const P0: u64 = 4611686018427387905;
9// const P1: u64 = 4611686018326724609;
10
11// const P0: u64 = 8070449433331580929; // max use 1<<62 -1
12// const P1: u64 = 8070450532384645121;
13
14// const P0: u64 = 0x80000000080001; // max use 1<<55 -1
15// const P1: u64 = 0x80000000130001;
16
17const BITS: u128 = 64;
18const P0: u64 = ((1u128 << BITS) - (1u128 << 28) + 1u128) as u64;
19const P1: u64 = ((1u128 << BITS) - (1u128 << 27) + 1u128) as u64;
20const P2: u64 = ((1u128 << BITS) - (1u128 << 26) + 1u128) as u64;
21// const P0: u64 = ((1u128 << BITS) - (1u128 << 18) + 1u128) as u64;
22// const P1: u64 = ((1u128 << BITS) - (1u128 << 17) + 1u128) as u64;
23// const P2: u64 = ((1u128 << BITS) - (1u128 << 16) + 1u128) as u64;
24// const P0: u64 = 0x0FFFFFFF0000001; // 56 bits each P_i
25// const P1: u64 = 0x0FFFFFFE8000001;
26// const P2: u64 = 0x0FFFFFFE4000001;
27
28#[derive(Debug)]
29pub struct NTT {}
30
31impl NTT {
32 pub fn ntt(n: usize, a: &Vec<u64>) -> (Vec<u64>, Vec<u64>, Vec<u64>) {
33 // apply modulus p_i
34 let a_0: Vec<u64> = a.iter().map(|a_i| a_i % P0).collect();
35 let a_1: Vec<u64> = a.iter().map(|a_i| a_i % P1).collect();
36 let a_2: Vec<u64> = a.iter().map(|a_i| a_i % P2).collect();
37 // let a_0: Vec<u64> = a.iter().map(|a_i| (a_i % P0 + P0) % P0).collect();
38 // let a_1: Vec<u64> = a.iter().map(|a_i| (a_i % P1 + P1) % P1).collect();
39 // let a_2: Vec<u64> = a.iter().map(|a_i| (a_i % P2 + P2) % P2).collect();
40
41 dbg!(&a_0, &a_1, &a_2);
42 let r_0 = NTT_p::ntt(P0, n, &a_0);
43 let r_1 = NTT_p::ntt(P1, n, &a_1);
44 let r_2 = NTT_p::ntt(P2, n, &a_2);
45 dbg!(&r_0, &r_1, &r_2);
46
47 (r_0, r_1, r_2)
48 }
49
50 pub fn intt(n: usize, r: &(Vec<u64>, Vec<u64>, Vec<u64>)) -> Vec<u64> {
51 let a_0 = NTT_p::intt(P0, n, &r.0);
52 let a_1 = NTT_p::intt(P1, n, &r.1);
53 let a_2 = NTT_p::intt(P2, n, &r.2);
54 dbg!(&a_0, &a_1, &a_2);
55 // TODO WIP
56 // let a_0: Vec<u64> = a_0.iter().map(|a_i| a_i % P0).collect();
57 // let a_1: Vec<u64> = a_1.iter().map(|a_i| a_i % P1).collect();
58 // let a_2: Vec<u64> = a_2.iter().map(|a_i| a_i % P2).collect();
59
60 reconstruct(a_0, a_1, a_2)
61 }
62}
63
64/// applies CRT to reconstruct the composite original value
65fn reconstruct(a0: Vec<u64>, a1: Vec<u64>, a2: Vec<u64>) -> Vec<u64> {
66 let p0: u128 = P0 as u128;
67 let p1: u128 = P1 as u128;
68 let p2: u128 = P2 as u128;
69 // y_i = q/q_i
70 let y0: u128 = p1 * p2; // N_i =Q/P0 = P1*P2
71 let y1: u128 = p0 * p2;
72 let y2: u128 = p0 * p1;
73
74 // y_i^-1 mod q_i = z_i
75 dbg!(P0, y0);
76 let z0: u128 = inv_mod(p0, y0); // M_i = N_i^-1 mod q_i
77 let z1: u128 = inv_mod(p1, y1);
78 let z2: u128 = inv_mod(p2, y2);
79
80 // m1 = q1^-1 mod q2
81 // aux = (a2 - a1) * m1 mod q2
82 // a = a1 + (q1 * m1) * aux
83
84 // m1 == z1
85 // let m1 = inv_mod(P1 as u128, P0 as u128); // P0^-1 mod P1
86 // let aux: Vec<u64> = itertools::zip_eq(a0.clone(), a1.clone())
87 // .map(|(a0_i, a1_i)| ((a1_i - a0_i) * m1) % P1)
88 // .collect();
89 // let a: Vec<u64> = itertools::zip_eq(a0, a1)
90 // // .map(|(a1_i, aux_i)| a1_i + (P1 * m1) * aux_i)
91 // // .map(|(a0_i, aux_i)| a0_i + (P0 * m1) * aux_i)
92 // .map(|(a0_i, a1_i)| a0_i as u128 + ((p0 * z1) % p1) * (((a1_i - a0_i) as u128 * z1) % p1))
93 // .map(|v| v as u64)
94 // // let a: Vec<u64> = itertools::zip_eq(a0, a1)
95 // // .map(|(a0_i, a1_i)| {
96 // // ((((a0_i as u128 * z0) % P0 as u128) * y0) + (a1_i as u128 * z1 % P1 as u128) * y1)
97 // // as u64
98 // // })
99 // .collect();
100 // a
101
102 // dbg!(a0[0] as u128);
103 // dbg!(a0[0] as u128 * y0);
104 // dbg!(a0[0] as u128 * z0);
105 // dbg!(a0[0] as u128 * y0 * z0);
106 dbg!(&y0, &y1, &y2);
107 dbg!(&z0, &z1, &z2);
108 // WIP, using BigUint to use Q with 192 bits (product of 3 u64)
109 let Q = BigUint::from_u64(P0).unwrap()
110 * BigUint::from_u64(P1).unwrap()
111 * BigUint::from_u64(P2).unwrap();
112 let max_u64 = BigUint::from_u128(1_u128 << 64).unwrap();
113 // let Q = P0 as u128 * P1 as u128 * P2 as u128;
114 let a: Vec<u64> = itertools::multizip((a0, a1, a2))
115 .map(|(a0_i, a1_i, a2_i)| {
116 // (a0_i as u128 * y0 * z0)// % Q
117 // + (a1_i as u128 * y1 * z1)// % Q
118 // + (a2_i as u128 * y2 * z2) // % Q
119 mul_3_big(a0_i as u128, y0, z0, &Q) // TODO rm %Q, since it returns a BigUint, and %Q
120 // is done later at the end of the additions
121 + mul_3_big(a1_i as u128, y1, z1, &Q)
122 + mul_3_big(a2_i as u128, y2, z2, &Q)
123 })
124 .map(|v| {
125 // WIP, using BigUint to use Q with 192 bits (product of 3 u64)
126 // ((BigUint::from_u128(v).unwrap() % Q.clone()) % max_u64.clone())
127 ((v % Q.clone()) % max_u64.clone()).to_u64().unwrap()
128 })
129 // .map(|v| v % (1 << 63))
130 // .map(|v| v % Q)
131 // .map(|v| v as u32)
132 // .map(|v| v as u64)
133 .collect();
134 a
135 // dbg!(&a);
136 // let Q = y2 * P2 as u128;
137 // let a: Vec<u128> = a.iter().map(|a_i| a_i % Q).collect();
138 // dbg!(&a);
139 // let q64 = 1_u128 << 64;
140 // let a: Vec<u64> = a.iter().map(|a_i| (a_i % q64) as u64).collect();
141 // a
142}
143
144fn exp_mod(q: u128, x: u128, k: u128) -> u128 {
145 // work on u128 to avoid overflow
146 let mut r = 1u128;
147 let mut x = x.clone();
148 let mut k = k.clone();
149 x = x % q;
150 // exponentiation by square strategy
151 while k > 0 {
152 if k % 2 == 1 {
153 r = (r * x) % q;
154 }
155 x = (x * x) % q;
156 k /= 2;
157 }
158 r
159}
160/// returns x^-1 mod Q, assuming x and Q are coprime, generally Q is prime
161fn inv_mod_new(q: u128, x: u128) -> u128 {
162 // by Fermat's Little Theorem, x^-1 mod q \equiv x^{q-2} mod q
163 // exp_mod(q, x, q - 2)
164 exp_mod(q, x, q - 2)
165}
166
167fn inv_mod(m: u128, a: u128) -> u128 {
168 if m == 1 {
169 panic!("m==1");
170 }
171
172 let mut m = m.clone();
173 let mut a = a.clone();
174 let m0 = m.clone();
175 let mut x0: i128 = 0;
176 let mut x1: i128 = 1;
177
178 while a > 1 {
179 let q = a / m;
180 let t = m.clone();
181
182 m = a % m;
183 a = t.clone();
184
185 let t = x0;
186 x0 = x1 - (q as i128) * x0;
187 x1 = t;
188 }
189
190 if x1 < 0 {
191 x1 += m0 as i128;
192 }
193
194 x1 as u128
195}
196
197use num_bigint::BigUint;
198use num_traits::{FromPrimitive, ToPrimitive};
199fn mul_3_big(a: u128, b: u128, c: u128, Q: &BigUint) -> BigUint {
200 let r = (BigUint::from_u128(a).unwrap()
201 * BigUint::from_u128(b).unwrap()
202 * BigUint::from_u128(c).unwrap());
203 // % Q;
204 dbg!(&r);
205 let r = r % Q;
206 dbg!(&r);
207 r
208 // let max_u64 = BigUint::from_u128(1_u128 << 64).unwrap();
209 // dbg!(&r % &max_u64);
210 // (r % max_u64).to_u128().unwrap()
211}
212
213#[cfg(test)]
214mod tests {
215 use super::*;
216 use rand_distr::Distribution;
217
218 use anyhow::Result;
219
220 #[test]
221 fn test_inv_mod() -> Result<()> {
222 // test vectors used in this test generated in SageMath
223
224 let p: u64 = 0x0FFFFFFF0000001;
225
226 let x = 3;
227 let x_inv = inv_mod(p as u128, x);
228 assert_eq!(x_inv, 48038395846328321);
229
230 let x = (1_u128 << 64) - 1;
231 let x_inv = inv_mod(p as u128, x);
232 assert_eq!(x_inv, 25433122709808565);
233
234 let x = 1_u128 << 120;
235 let x_inv = inv_mod(p as u128, x);
236 assert_eq!(x_inv, 281474976710656);
237
238 // other prime
239 let p: u64 = ((1u128 << 64) - (1u128 << 28) + 1u128) as u64;
240
241 let x = 3;
242 let x_inv = inv_mod(p as u128, x);
243 assert_eq!(x_inv, 12297829382294077441);
244
245 let x = (1_u128 << 64) - 1;
246 let x_inv = inv_mod(p as u128, x);
247 assert_eq!(x_inv, 10530692608818076599);
248
249 let x = 1_u128 << 120;
250 let x_inv = inv_mod(p as u128, x);
251 assert_eq!(x_inv, 18374686616574755070);
252
253 Ok(())
254 }
255
256 #[test]
257 fn test_reconstruct() -> Result<()> {
258 let n: usize = 16;
259
260 use num_bigint::BigUint;
261 use num_traits::{FromPrimitive, ToPrimitive};
262 let Q = BigUint::from_u64(P0).unwrap()
263 * BigUint::from_u64(P1).unwrap()
264 * BigUint::from_u64(P2).unwrap();
265 let q = BigUint::from_u128(1_u128 << BITS).unwrap();
266 let N = BigUint::from_usize(n).unwrap();
267 let big2 = BigUint::from_u64(2).unwrap();
268 assert!(Q > (N * (&q * &q)) / big2); // sanity check
269
270 use rand::Rng;
271 let mut rng = rand::thread_rng();
272 let a: Vec<u64> = (0..n)
273 .map(|_| rng.gen_range(0..((1u128 << 64) - 1) as u64))
274 .collect();
275 // let a = vec![14713100818624219214];
276 dbg!(&a);
277
278 let a_0: Vec<u64> = a.iter().map(|a_i| a_i % P0).collect();
279 let a_1: Vec<u64> = a.iter().map(|a_i| a_i % P1).collect();
280 let a_2: Vec<u64> = a.iter().map(|a_i| a_i % P2).collect();
281 // let a_0: Vec<u64> = a.iter().map(|a_i| (a_i % P0 + P0) % P0).collect();
282 // let a_1: Vec<u64> = a.iter().map(|a_i| (a_i % P1 + P1) % P1).collect();
283 // let a_2: Vec<u64> = a.iter().map(|a_i| (a_i % P2 + P2) % P2).collect();
284 dbg!(&a_0, &a_1, &a_2);
285 let a_reconstructed = reconstruct(a_0, a_1, a_2);
286
287 dbg!(&a_reconstructed);
288 assert_eq!(a_reconstructed, a);
289
290 Ok(())
291 }
292
293 #[test]
294 fn test_ntt() -> Result<()> {
295 // println!("{}", 1u128 << 64);
296 let n: usize = 1;
297
298 // println!("{}", P0);
299 // println!("{}", P1);
300 // let q = 1u128 << 64;
301 // assert!(P0 as u128 * P1 as u128 > (n as u128 * (q * q)) / 2);
302
303 // let a: Vec<u64> = vec![1u64, 2, 3, 4];
304 // let a: Vec<u64> = vec![1u64, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
305 // let a: Vec<u64> = vec![9u64, 8, 7, 6, 0, 9999, 0, 0, 0, 0, 0, 0, 6, 7, 8, 9];
306 use rand::Rng;
307 let mut rng = rand::thread_rng();
308 let a: Vec<u64> = (0..n)
309 .map(|_| rng.gen_range(0..((1u128 << 64) - 1) as u64))
310 .collect();
311
312 dbg!(&a);
313
314 let a_ntt = NTT::ntt(n, &a);
315 dbg!(&a_ntt);
316
317 let a_intt = NTT::intt(n, &a_ntt);
318
319 dbg!(&a_intt);
320 assert_eq!(a_intt, a);
321
322 Ok(())
323 }
324}