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