poulpy/rns/src/dft/ntt.rs
rns/src/dft/ntt.rsRaw
1use crate::dft::DFT;
2use crate::modulus::barrett::Barrett;
3use crate::modulus::montgomery::Montgomery;
4use crate::modulus::prime::Prime;
5use crate::modulus::ReduceOnce;
6use crate::modulus::WordOps;
7use crate::modulus::{BARRETT, NONE, ONCE};
8use itertools::izip;
9
10#[allow(dead_code)]
11pub struct Table<O> {
12 prime: Prime<O>,
13 psi: O,
14 psi_forward_rev: Vec<Barrett<u64>>,
15 psi_backward_rev: Vec<Barrett<u64>>,
16 q: O,
17 two_q: O,
18 four_q: O,
19}
20
21impl Table<u64> {
22 pub fn new(prime: Prime<u64>, nth_root: usize) -> Table<u64> {
23 assert!(
24 nth_root & (nth_root - 1) == 0,
25 "invalid argument: nth_root = {} is not a power of two",
26 nth_root
27 );
28
29 let psi: u64 = prime.primitive_nth_root(nth_root as u64);
30
31 let psi_mont: Montgomery<u64> = prime.montgomery.prepare::<ONCE>(psi);
32 let psi_inv_mont: Montgomery<u64> = prime.montgomery.pow(psi_mont, prime.phi - 1);
33
34 let mut psi_forward_rev: Vec<Barrett<u64>> = vec![Barrett(0, 0); (nth_root >> 1) as usize];
35 let mut psi_backward_rev: Vec<Barrett<u64>> = vec![Barrett(0, 0); (nth_root >> 1) as usize];
36
37 psi_forward_rev[0] = prime.barrett.prepare(1);
38 psi_backward_rev[0] = prime.barrett.prepare(1);
39
40 let log_nth_root_half: u32 = (nth_root >> 1).log2() as _;
41
42 let mut powers_forward: u64 = 1u64;
43 let mut powers_backward: u64 = 1u64;
44
45 for i in 1..(nth_root >> 1) as usize {
46 let i_rev: usize = i.reverse_bits_msb(log_nth_root_half);
47
48 prime
49 .montgomery
50 .mul_external_assign::<ONCE>(psi_mont, &mut powers_forward);
51 prime
52 .montgomery
53 .mul_external_assign::<ONCE>(psi_inv_mont, &mut powers_backward);
54
55 psi_forward_rev[i_rev] = prime.barrett.prepare(powers_forward);
56 psi_backward_rev[i_rev] = prime.barrett.prepare(powers_backward);
57 }
58
59 let q: u64 = prime.q();
60
61 Self {
62 prime: prime,
63 psi: psi,
64 psi_forward_rev: psi_forward_rev,
65 psi_backward_rev: psi_backward_rev,
66 q: q,
67 two_q: q << 1,
68 four_q: q << 2,
69 }
70 }
71}
72
73impl DFT<u64> for Table<u64> {
74 fn forward_inplace(&self, a: &mut [u64]) {
75 self.forward_inplace::<false>(a)
76 }
77
78 fn forward_inplace_lazy(&self, a: &mut [u64]) {
79 self.forward_inplace::<true>(a)
80 }
81
82 fn backward_inplace(&self, a: &mut [u64]) {
83 self.backward_inplace::<false>(a)
84 }
85
86 fn backward_inplace_lazy(&self, a: &mut [u64]) {
87 self.backward_inplace::<true>(a)
88 }
89}
90
91impl Table<u64> {
92 pub fn forward_inplace<const LAZY: bool>(&self, a: &mut [u64]) {
93 self.forward_inplace_core::<LAZY, 0, 0>(a);
94 }
95
96 pub fn forward_inplace_core<const LAZY: bool, const SKIPSTART: u8, const SKIPEND: u8>(
97 &self,
98 a: &mut [u64],
99 ) {
100 let n: usize = a.len();
101 assert!(
102 n & n - 1 == 0,
103 "invalid a.len()={} must be a power of two",
104 n
105 );
106
107 assert!(
108 n <= self.psi_forward_rev.len(),
109 "invalid a.len()={} > psi_forward_rev.len()={}",
110 n,
111 self.psi_forward_rev.len()
112 );
113
114 let log_n: u32 = usize::BITS - ((n as usize) - 1).leading_zeros();
115
116 let start: u32 = SKIPSTART as u32;
117 let end: u32 = log_n - (SKIPEND as u32);
118
119 for layer in start..end {
120 let (m, size) = (1 << layer, 1 << (log_n - layer - 1));
121 let t: usize = 2 * size;
122 if layer == log_n - 1 {
123 if LAZY {
124 izip!(a.chunks_exact_mut(t), &self.psi_forward_rev[m..]).for_each(
125 |(a, psi)| {
126 let (a, b) = a.split_at_mut(size);
127 self.dit_inplace::<false>(&mut a[0], &mut b[0], psi);
128 debug_assert!(
129 a[0] < self.two_q,
130 "forward_inplace_core::<LAZY=true> output {} > {} (2q-1)",
131 a[0],
132 self.two_q - 1
133 );
134 debug_assert!(
135 b[0] < self.two_q,
136 "forward_inplace_core::<LAZY=true> output {} > {} (2q-1)",
137 b[0],
138 self.two_q - 1
139 );
140 },
141 );
142 } else {
143 izip!(a.chunks_exact_mut(t), &self.psi_forward_rev[m..]).for_each(
144 |(a, psi)| {
145 let (a, b) = a.split_at_mut(size);
146 self.dit_inplace::<true>(&mut a[0], &mut b[0], psi);
147 self.prime.barrett.reduce_assign::<BARRETT>(&mut a[0]);
148 self.prime.barrett.reduce_assign::<BARRETT>(&mut b[0]);
149 debug_assert!(
150 a[0] < self.q,
151 "forward_inplace_core::<LAZY=false> output {} > {} (q-1)",
152 a[0],
153 self.q - 1
154 );
155 debug_assert!(
156 b[0] < self.q,
157 "forward_inplace_core::<LAZY=false> output {} > {} (q-1)",
158 b[0],
159 self.q - 1
160 );
161 },
162 );
163 }
164 } else if t >= 16 {
165 izip!(a.chunks_exact_mut(t), &self.psi_forward_rev[m..]).for_each(|(a, psi)| {
166 let (a, b) = a.split_at_mut(size);
167 izip!(a.chunks_exact_mut(8), b.chunks_exact_mut(8)).for_each(|(a, b)| {
168 self.dit_inplace::<true>(&mut a[0], &mut b[0], psi);
169 self.dit_inplace::<true>(&mut a[1], &mut b[1], psi);
170 self.dit_inplace::<true>(&mut a[2], &mut b[2], psi);
171 self.dit_inplace::<true>(&mut a[3], &mut b[3], psi);
172 self.dit_inplace::<true>(&mut a[4], &mut b[4], psi);
173 self.dit_inplace::<true>(&mut a[5], &mut b[5], psi);
174 self.dit_inplace::<true>(&mut a[6], &mut b[6], psi);
175 self.dit_inplace::<true>(&mut a[7], &mut b[7], psi);
176 });
177 });
178 } else {
179 izip!(a.chunks_exact_mut(t), &self.psi_forward_rev[m..]).for_each(|(a, psi)| {
180 let (a, b) = a.split_at_mut(size);
181 izip!(a, b).for_each(|(a, b)| self.dit_inplace::<true>(a, b, psi));
182 });
183 }
184 }
185 }
186
187 #[inline(always)]
188 fn dit_inplace<const LAZY: bool>(&self, a: &mut u64, b: &mut u64, t: &Barrett<u64>) {
189 debug_assert!(*a < self.four_q, "a:{} q:{}", a, self.four_q);
190 debug_assert!(*b < self.four_q, "b:{} q:{}", b, self.four_q);
191 a.reduce_once_assign(self.two_q);
192 let bt: u64 = self.prime.barrett.mul_external::<NONE>(t, b);
193 *b = *a + self.two_q - bt;
194 *a += bt;
195 if !LAZY {
196 a.reduce_once_assign(self.two_q);
197 b.reduce_once_assign(self.two_q);
198 }
199 }
200
201 pub fn backward_inplace<const LAZY: bool>(&self, a: &mut [u64]) {
202 self.backward_inplace_core::<LAZY, 0, 0>(a);
203 }
204
205 pub fn backward_inplace_core<const LAZY: bool, const SKIPSTART: u8, const SKIPEND: u8>(
206 &self,
207 a: &mut [u64],
208 ) {
209 let n: usize = a.len();
210 assert!(
211 n & n - 1 == 0,
212 "invalid x.len()= {} must be a power of two",
213 n
214 );
215
216 assert!(
217 n <= self.psi_backward_rev.len(),
218 "invalid a.len()={} > psi_backward_rev.len()={}",
219 n,
220 self.psi_backward_rev.len()
221 );
222
223 let log_n = usize::BITS - ((n as usize) - 1).leading_zeros();
224
225 let start: u32 = SKIPEND as u32;
226 let end: u32 = log_n - (SKIPSTART as u32);
227
228 for layer in (start..end).rev() {
229 let (m, size) = (1 << layer, 1 << (log_n - layer - 1));
230 let t: usize = 2 * size;
231 if layer == 0 {
232 let n_inv: Barrett<u64> = self.prime.barrett.prepare(self.prime.inv(n as u64));
233 let psi: Barrett<u64> = self.prime.barrett.prepare(
234 self.prime
235 .barrett
236 .mul_external::<ONCE>(&n_inv, &self.psi_backward_rev[1].0),
237 );
238
239 izip!(a.chunks_exact_mut(2 * size)).for_each(|a| {
240 let (a, b) = a.split_at_mut(size);
241 izip!(a.chunks_exact_mut(8), b.chunks_exact_mut(8)).for_each(|(a, b)| {
242 self.dif_last_inplace::<LAZY>(&mut a[0], &mut b[0], &psi, &n_inv);
243 self.dif_last_inplace::<LAZY>(&mut a[1], &mut b[1], &psi, &n_inv);
244 self.dif_last_inplace::<LAZY>(&mut a[2], &mut b[2], &psi, &n_inv);
245 self.dif_last_inplace::<LAZY>(&mut a[3], &mut b[3], &psi, &n_inv);
246 self.dif_last_inplace::<LAZY>(&mut a[4], &mut b[4], &psi, &n_inv);
247 self.dif_last_inplace::<LAZY>(&mut a[5], &mut b[5], &psi, &n_inv);
248 self.dif_last_inplace::<LAZY>(&mut a[6], &mut b[6], &psi, &n_inv);
249 self.dif_last_inplace::<LAZY>(&mut a[7], &mut b[7], &psi, &n_inv);
250 });
251 });
252 } else if t >= 16 {
253 izip!(a.chunks_exact_mut(t), &self.psi_backward_rev[m..]).for_each(|(a, psi)| {
254 let (a, b) = a.split_at_mut(size);
255 izip!(a.chunks_exact_mut(8), b.chunks_exact_mut(8)).for_each(|(a, b)| {
256 self.dif_inplace::<true>(&mut a[0], &mut b[0], psi);
257 self.dif_inplace::<true>(&mut a[1], &mut b[1], psi);
258 self.dif_inplace::<true>(&mut a[2], &mut b[2], psi);
259 self.dif_inplace::<true>(&mut a[3], &mut b[3], psi);
260 self.dif_inplace::<true>(&mut a[4], &mut b[4], psi);
261 self.dif_inplace::<true>(&mut a[5], &mut b[5], psi);
262 self.dif_inplace::<true>(&mut a[6], &mut b[6], psi);
263 self.dif_inplace::<true>(&mut a[7], &mut b[7], psi);
264 });
265 });
266 } else {
267 izip!(a.chunks_exact_mut(2 * size), &self.psi_backward_rev[m..]).for_each(
268 |(a, psi)| {
269 let (a, b) = a.split_at_mut(size);
270 izip!(a, b).for_each(|(a, b)| self.dif_inplace::<true>(a, b, psi));
271 },
272 );
273 }
274 }
275 }
276
277 #[inline(always)]
278 fn dif_inplace<const LAZY: bool>(&self, a: &mut u64, b: &mut u64, t: &Barrett<u64>) {
279 debug_assert!(*a < self.two_q, "a:{} q:{}", a, self.two_q);
280 debug_assert!(*b < self.two_q, "b:{} q:{}", b, self.two_q);
281 let d: u64 = self
282 .prime
283 .barrett
284 .mul_external::<NONE>(t, &(*a + self.two_q - *b));
285 *a = *a + *b;
286 a.reduce_once_assign(self.two_q);
287 *b = d;
288 if !LAZY {
289 a.reduce_once_assign(self.q);
290 b.reduce_once_assign(self.q);
291 }
292 }
293
294 fn dif_last_inplace<const LAZY: bool>(
295 &self,
296 a: &mut u64,
297 b: &mut u64,
298 psi: &Barrett<u64>,
299 n_inv: &Barrett<u64>,
300 ) {
301 debug_assert!(*a < self.two_q);
302 debug_assert!(*b < self.two_q);
303 if LAZY {
304 let d: u64 = self
305 .prime
306 .barrett
307 .mul_external::<NONE>(psi, &(*a + self.two_q - *b));
308 *a = self.prime.barrett.mul_external::<NONE>(n_inv, &(*a + *b));
309 *b = d;
310 } else {
311 let d: u64 = self
312 .prime
313 .barrett
314 .mul_external::<ONCE>(psi, &(*a + self.two_q - *b));
315 *a = self.prime.barrett.mul_external::<ONCE>(n_inv, &(*a + *b));
316 *b = d;
317 }
318 }
319}
320
321#[cfg(test)]
322mod tests {
323 use super::*;
324
325 #[test]
326 fn test_ntt() {
327 let q_base: u64 = 0x800000000004001;
328 let q_power: usize = 1;
329 let prime_instance: Prime<u64> = Prime::<u64>::new(q_base, q_power);
330 let n: usize = 32;
331 let two_nth_root: usize = n << 1;
332 let ntt_table: Table<u64> = Table::<u64>::new(prime_instance, two_nth_root);
333 let mut a: Vec<u64> = vec![0; n as usize];
334 for i in 0..a.len() {
335 a[i] = i as u64;
336 }
337
338 let b: Vec<u64> = a.clone();
339 ntt_table.forward_inplace::<false>(&mut a);
340 ntt_table.backward_inplace::<false>(&mut a);
341 assert!(a == b);
342 }
343}