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