fhe-study/arith/src/ntt_u62.rs
arith/src/ntt_u62.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
11const P0: u64 = 8070449433331580929; // max use 1<<62
12const P1: u64 = 8070450532384645121;
13
14// const P0: u64 = 0x80000000080001; // max use 1<<55
15// const P1: u64 = 0x80000000130001;
16
17// const P0: u64 = ((1u128 << 64) - (1u128 << 28) + 1u128) as u64;
18// const P1: u64 = ((1u128 << 64) - (1u128 << 27) + 1u128) as u64;
19// const P0: u64 = ((1u128 << 64) - (1u128 << 28) + 1u128) as u64;
20// const P1: u64 = ((1u128 << 64) - (1u128 << 27) + 1u128) as u64;
21// const P2: u64 = (1 << 60) - (1 << 26) + 1;
22
23#[derive(Debug)]
24pub struct NTT {}
25
26impl NTT {
27 pub fn ntt(
28 n: usize,
29 a: &Vec<u64>,
30 ) -> (
31 Vec<u64>,
32 Vec<u64>,
33 // Vec<u64>,
34 // Vec<u64>,
35 // Vec<u64>,
36 // Vec<u64>,
37 // Vec<u64>,
38 ) {
39 // TODO ensure that: a_i <P0
40
41 // apply modulus p_i
42 let a_0: Vec<u64> = a.iter().map(|a_i| (a_i % P0 + P0) % P0).collect();
43 let a_1: Vec<u64> = a.iter().map(|a_i| (a_i % P1 + P1) % P1).collect();
44 // let a_2: Vec<u64> = a.iter().map(|a_i| (a_i % P2 + P2) % P2).collect();
45 // let a_3: Vec<u64> = a.iter().map(|a_i| (a_i % P3 + P3) % P3).collect();
46 // let a_4: Vec<u64> = a.iter().map(|a_i| (a_i % P4 + P4) % P4).collect();
47 // let a_5: Vec<u64> = a.iter().map(|a_i| (a_i % P5 + P5) % P5).collect();
48 // let a_6: Vec<u64> = a.iter().map(|a_i| (a_i % P6 + P6) % P6).collect();
49
50 let r_0 = NTT_p::ntt(P0, n, &a_0);
51 let r_1 = NTT_p::ntt(P1, n, &a_1);
52 // let r_2 = NTT_p::ntt(P2, n, &a_2);
53 // let r_3 = NTT_p::ntt(P3, n, &a_3);
54 // let r_4 = NTT_p::ntt(P4, n, &a_4);
55 // let r_5 = NTT_p::ntt(P5, n, &a_5);
56 // let r_6 = NTT_p::ntt(P6, n, &a_6);
57
58 (r_0, r_1) //, r_2) //, r_3, r_4) //, r_5, r_6)
59 }
60
61 pub fn intt(
62 n: usize,
63 r: &(
64 Vec<u64>,
65 Vec<u64>,
66 // Vec<u64>,
67 // Vec<u64>,
68 // Vec<u64>,
69 // Vec<u64>,
70 // Vec<u64>,
71 ),
72 ) -> Vec<u64> {
73 let a_0 = NTT_p::intt(P0, n, &r.0);
74 let a_1 = NTT_p::intt(P1, n, &r.1);
75 // let a_2 = NTT_p::intt(P2, n, &r.2);
76 // let a_3 = NTT_p::intt(P3, n, &r.3);
77 // let a_4 = NTT_p::intt(P4, n, &r.4);
78 // let a_5 = NTT_p::intt(P5, n, &r.5);
79 // let a_6 = NTT_p::intt(P6, n, &r.6);
80
81 // Garner CRT for two moduli: combine (r1 mod p1, r2 mod p2) -> Z/(p1*p2)
82 // let inv_p1_mod_p2: u128 = inv_mod_u64(p1 % p2, p2) as u128;
83 // const INV_P1_MOD_P2: u128 = 4895217125691974194;
84
85 reconstruct(a_0, a_1) //, a_2) // , a_3, a_4) //, a_5, a_6)
86 }
87}
88
89fn reconstruct(
90 a0: Vec<u64>,
91 a1: Vec<u64>,
92 // a2: Vec<u64>,
93 // a_3: Vec<u64>,
94 // a_4: Vec<u64>,
95 // a_5: Vec<u64>,
96 // a_6: Vec<u64>,
97) -> Vec<u64> {
98 // let Q = P0 as u128 * P1 as u128;
99
100 // y_i = q/q_i
101 // let y0 = ((u64::MAX as u128 + 1) / P0 as u128);
102 // let y1 = ((u64::MAX as u128 + 1) / P1 as u128);
103 // let y2 = ((u64::MAX as u128 + 1) / P2 as u128) as u64;
104 let y0: u128 = P1 as u128; // N_i =Q/P0 = P1*P2
105 let y1: u128 = P0 as u128;
106 // let y1: u128 = P0 as u128 * P2 as u128;
107 // let y0: u128 = P1 as u128 * P2 as u128; // N_i =Q/P0 = P1*P2
108 // let y1: u128 = P0 as u128 * P2 as u128;
109 // let y2: u128 = P0 as u128 * P1 as u128;
110 // let y0 = (Q / P0 as u128) as u64;
111 // let y1 = (Q / P1 as u128) as u64;
112 // let y2 = ((u64::MAX as u128 + 1) / P2 as u128) as u64;
113 // let y3 = ((u64::MAX as u128 + 1) / P3 as u128) as u64;
114 // let y4 = ((u64::MAX as u128 + 1) / P4 as u128) as u64;
115
116 // y_i^-1 mod q_i = z_i
117 let z0: u128 = inv_mod(P0 as u128, y0); // M_i = N_i^-1 mod q_i
118 let z1: u128 = inv_mod(P1 as u128, y1);
119 // let z2: u128 = inv_mod(P2 as u128, y2);
120 // let y2_inv = inv_mod(P2 as u128, y2);
121 // let y3_inv = inv_mod(P3 as u128, y3);
122 // let y4_inv = inv_mod(P4 as u128, y4);
123
124 // m1 = q1^-1 mod q2
125 // aux = (a2 - a1) * m1 mod q2
126 // a = a1 + (q1 * m1) * aux
127
128 /*
129 let m1 = inv_mod(P1 as u128, P0 as u128) as u64; // P0^-1 mod P1
130 let aux: Vec<u64> = itertools::zip_eq(a0.clone(), a1.clone())
131 .map(|(a0_i, a1_i)| ((a1_i - a0_i) * m1) % P1)
132 .collect();
133 let a: Vec<u64> = itertools::zip_eq(a0, aux)
134 // .map(|(a1_i, aux_i)| a1_i + (P1 * m1) * aux_i)
135 // .map(|(a0_i, aux_i)| a0_i + (P0 * m1) * aux_i)
136 .map(|(a0_i, aux_i)| a0_i + ((P0 * m1) % P1) * aux_i)
137 .collect();
138 a
139 */
140 let p0: u128 = P0 as u128;
141 let p1: u128 = P1 as u128;
142 let a: Vec<u64> = itertools::zip_eq(a0, a1)
143 .map(|(a0_i, a1_i)| a0_i as u128 + ((p0 * z1) % p1) * (((a1_i - a0_i) as u128 * z1) % p1))
144 .map(|v| v as u64)
145 .collect();
146 a
147 // dbg!(a0[0] as u128);
148 // dbg!(a0[0] as u128 * y0);
149 // dbg!(a0[0] as u128 * z0);
150 // dbg!(a0[0] as u128 * y0 * z0);
151 // let a: Vec<u128> = itertools::multizip((a0, a1, a2))
152 // .map(|(a0_i, a1_i, a2_i)| {
153 // a0_i as u128 * y0 * z0 + a1_i as u128 * y1 * z1 + a2_i as u128 * y2 * z2
154 // })
155 // .collect();
156 // dbg!(&a);
157 // let Q = y2 * P2 as u128;
158 // let a: Vec<u128> = a.iter().map(|a_i| a_i % Q).collect();
159 // dbg!(&a);
160 // let q64 = 1_u128 << 64;
161 // let a: Vec<u64> = a.iter().map(|a_i| (a_i % q64) as u64).collect();
162 // a
163
164 /*
165 // x_i*z_i mod q_i
166 let r0: Vec<u64> = a_0.iter().map(|a_i| ((a_i * z0) % P0) * y0).collect();
167 let r1: Vec<u64> = a_1.iter().map(|a_i| ((a_i * z1) % P1) * y1).collect();
168 // let r0: Vec<u64> = a_0.iter().map(|a_i| ((a_i * z0) % P0) * y0).collect();
169 // let r1: Vec<u64> = a_1.iter().map(|a_i| ((a_i * z1) % P1) * y1).collect();
170 // let r2: Vec<u64> = a_2.iter().map(|a_i| ((a_i * y2_inv) % P2) * y2).collect();
171 // let r3: Vec<u64> = a_3.iter().map(|a_i| ((a_i * y3_inv) % P3) * y3).collect();
172 // let r4: Vec<u64> = a_4.iter().map(|a_i| ((a_i * y4_inv) % P4) * y4).collect();
173
174 let r: Vec<u64> = itertools::multizip((r0.iter(), r1.iter()))
175 .map(|(a, b)| a + b)
176 .collect();
177 // let r = r0;
178 //
179 dbg!(&r);
180
181 let p1p2: u128 = (P0 as u128) * (P1 as u128);
182 // let p1p2_inv: u128 = inv_mod((P0 % P1) as u128, P1) as u128;
183 let p1p2_inv: u128 = inv_mod((P0) as u128, P1) as u128;
184 dbg!(&p1p2);
185 dbg!(&p1p2_inv);
186 // let p1p2: u128 = P0 as u128 / 2; // PIHALF
187 let r = r
188 .iter()
189 .map(|c_i_u64| {
190 let c_i = *c_i_u64 as u128;
191 if c_i * 2 >= p1p2 {
192 // if c_i >= p1p2 {
193 c_i.wrapping_sub(p1p2) as u64
194 } else {
195 c_i as u64
196 }
197 })
198 .collect();
199 // let r: Vec<u64> = itertools::multizip((r0.iter(), r1.iter(), r2.iter(), r3.iter(), r4.iter()))
200 // .map(|(a, b, c, d, e)| a + b + c + d + e)
201 // .collect();
202 // let mut r = a_0 + y0_inv + a_1 * y1_inv + a_2 * y2_inv + a_3 * y3_inv + a_4 * y4_inv;
203
204 r
205 */
206}
207
208fn exp_mod(q: u128, x: u128, k: u128) -> u128 {
209 // work on u128 to avoid overflow
210 let mut r = 1u128;
211 let mut x = x.clone();
212 let mut k = k.clone();
213 x = x % q;
214 // exponentiation by square strategy
215 while k > 0 {
216 if k % 2 == 1 {
217 r = (r * x) % q;
218 }
219 x = (x * x) % q;
220 k /= 2;
221 }
222 r
223}
224/// returns x^-1 mod Q
225fn inv_mod(q: u128, x: u128) -> u128 {
226 // by Fermat's Little Theorem, x^-1 mod q \equiv x^{q-2} mod q
227 exp_mod(q, x, q - 2)
228}
229
230#[cfg(test)]
231mod tests {
232 use super::*;
233 use rand_distr::Distribution;
234
235 use anyhow::Result;
236
237 #[test]
238 fn test_dbg() -> Result<()> {
239 println!("{}", 1u128 << 64);
240 let n: usize = 16;
241
242 println!("{}", P0);
243 println!("{}", P1);
244 // let q = 1u128 << 64;
245 // assert!(P0 as u128 * P1 as u128 > (n as u128 * (q * q)) / 2);
246
247 // let a: Vec<u64> = vec![1u64, 2, 3, 4];
248 // let a: Vec<u64> = vec![1u64, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
249 // let a: Vec<u64> = vec![9u64, 8, 7, 6, 0, 9999, 0, 0, 0, 0, 0, 0, 6, 7, 8, 9];
250 use rand::Rng;
251 let mut rng = rand::thread_rng();
252 let a: Vec<u64> = (0..n)
253 // .map(|_| rng.gen_range(0..=(1u64 << 57) - 1) - (1u64 << 56))
254 .map(|_| rng.gen_range(0..(1 << 62)))
255 .collect();
256
257 dbg!(a.len());
258
259 let a_ntt = NTT::ntt(n, &a);
260 dbg!(&a_ntt);
261
262 let a_intt = NTT::intt(n, &a_ntt);
263
264 dbg!(&a_intt);
265 assert_eq!(a_intt, a);
266
267 Ok(())
268 }
269}