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