整系数Discrete Fourier transform技巧

Discrete Fourier transform

向量$\mathbf{a}$的DFT $\mathbf{A}$定义为:

$$Aj=\sum{k=0}^{n-1}{a_k\omega^{jk}}$$

其中$\omega$是primitive root of unity of order $n$,复数时通常取$exp(2i\pi/n)$或$exp(-2i\pi/n)$。

Inverse transform

性质:$\frac{1}{n}\sum_{k=0}^{n-1}{\omega^{jk}} = [j\;mod\;n=0]$

$$\sum_{k=0}^{n-1}{Ak\omega^{-jk}} =
\sum
{k=0}^{n-1}{\sum_{l=0}^{n-1}{al\omega^{kl}}\omega^{-jk}} =
\sum
{k=0}^{n-1}{\sum_{l=0}^{n-1}{a_l\omega^{(l-j)k}}} =
na_j$$

因此

$$aj=\frac{1}{n}\sum{k=0}^{n-1}{A_k\omega^{jk}}$$

Number theoretic transform

DFT可以用于复数以外的其他ring,常用于$\mathbb{Z}/m$。

使用128 bits模数需要高效的u64*u64%u64,其中模数是常数。

硬件除法指令(32 bits、64 bits)

DIV指令性能很低。

1
2
3
4
5
6
extern inline u64 mul_mod(u64 a, u64 b, u64 m)
{
u64 r;
asm("mulq %2\n\tdivq %3" : "=&d"(r), "+a"(a) : "rm"(b), "rm"(m) : "cc");
return r;
}

到AVX-512也没有提供把两个64 bits乘数的积放在一个128 bits寄存器的指令,GCC没有提供用乘法、移位等模拟的u128除以u64的常量除法。

64位mantissa浮点数(32 bits、64 bits)

当模数$P<2^{63}$时可以用64位mantissa浮点数计算u64*u64%u64

由等式$$(a\cdot b\% m) = a\cdot b - \lfloor\frac{a\cdot b}{m}\rfloor\cdot m$$

两边模$2^{64}$,得
\begin{eqnarray}
(a\cdot b\% m)\% 2^{64} &=& (a\cdot b - \lfloor\frac{a\cdot b}{m}\rfloor\cdot m) \% 2^{64} \
&=& ((a\cdot b)\%2^{64} - (\lfloor\frac{a\cdot b}{m}\rfloor\cdot m)\%2^{64}) \% 2^{64}
\end{eqnarray
}

即用u64乘法计算$a\cdot b$的低64位,减去$\lfloor\frac{a\cdot b}{m}\rfloor\cdot m$的低64位。其中$\lfloor\frac{a\cdot b}{m}\rfloor<m$,可以用64位mantissa浮点数(Intel x87 80-bit double-extended precision)计算,再round成u64

round时若向上取整了,减数会大于被减数。若$m<2^{63}$,可以根据差的符号位判断。

1
2
3
4
5
u64 mul_mod(u64 a, u64 b, u64 P)
{
u64 x = a*b, r = x - P*u64((long double)a*b/P+0.5);
return i64(r) < 0 ? r + P : r;
}

存储$P$的倒数$Q=\frac{1}{P}$,用(long double)a*b*Q代替(long double)a*b/P能快些。此时$Q$会引入额外的误差,Matters Computational说适用于$m<2^{62}$,原因不明。

编译器生成的常量除法(32 bits)

对于固定的模$P$,GCC/llvm可以生成u64%u32的高效代码。llvm的lib/Transforms/Utils/IntegerDivision.cpp

Montgomery modular multiplication+Barret reduction(32 bits)

Faster arithmetic for number-theoretic transforms,用乘法和移位代替除法。

快速计算u64%u32,用乘法和移位代替除法。设$2^{m-1}\leq P<2^m$,$\alpha$为大于等于$m$的整数,$\beta$为负整数。

\begin{eqnarray}
Q &=& \frac{\lfloor\frac{a}{2^{m+\beta}}\rfloor\lfloor\frac{2^{m+\alpha}}{P}\rfloor}{2^{\alpha-\beta}} \
&\leq& \frac{\frac{a}{2^{m+\beta}}\frac{2^{m+\alpha}}{P}}{2^{\alpha-\beta}} \
&=& \frac{a}{P} \
\end{eqnarray
}

$Q\leq\frac{a}{P}$,$a-QP$是$a\%P$的估计值,若大于等于$P$则减去$P$的倍数。

\begin{eqnarray}
Q &>& \frac{(\frac{a}{2^{m+\beta}}-1)(\frac{2^{m+\alpha}}{P}-1)}{2^{\alpha-\beta}}-1 \
&=& \frac{a}{P}-\frac{a}{2^{m+\alpha}}-\frac{2^{m+\beta}}{P}+2^{\beta-\alpha}-1 \
&\geq& \lfloor\frac{a}{P}\rfloor-\frac{a}{2^{m+\alpha}}-\frac{2^{m+\beta}}{P}+2^{\beta-\alpha}-1 \
\end{eqnarray
}

因此

\begin{eqnarray}
0\leq \lfloor\frac{a}{P}\rfloor-Q &\leq& \frac{a}{2^{m+\alpha}}+\frac{2^{m+\beta}}{P}-2^{\beta-\alpha}+1 \
\end{eqnarray
}

令$\alpha=m, \beta=-1$,则$0\leq \lfloor\frac{a}{P}\rfloor-Q \leq 2$,即估计值最多小2,最多两个conditional move指令(if (r >= P) r -= P;)即可修正余数。

令$\alpha=m+1, \beta=-2$,则估计值最多小1。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
extern inline u64 barrett_30_2(u64 a, u64 P, u64 M)
{ // 2^29 < P < 2^30, M = floor(2^61/P)
u64 r = a-((a>>28)*M>>33)*P;
if (r >= P) r -= P;
return r;
}
extern inline u64 barrett_30_1(u64 a, u64 P, u64 M)
{ // 2^29 < P < 2^30, M = floor(2^60/P)
u64 r = a-((a>>29)*M>>31)*P;
if (r >= P) r -= P;
if (r >= P) r -= P;
return r;
}

当模数为常量时,该算法不如编译器生成的常量除法。若模数不固定时可以考虑使用。

Cyclic convolution

$$(a\ast b)j = \sum{k=0}^{n-1}{akb{(j-k)\;mod\;n}}$$

\begin{eqnarray}
(a\ast b)j &=& \sum{p=0}^{n-1}{\sum_{q=0}^{n-1}{[(p+q)\;mod\;n=j]a_pbq}}\
&=& \sum
{p=0}^{n-1}{\sum_{q=0}^{n-1}{[(p+q-j)\;mod\;n=0]a_pbq}}\
&=& \sum
{p=0}^{n-1}{\sum{q=0}^{n-1}{\frac{1}{n}\sum{k=0}^{n-1}{\omega^{pk}\omega^{qk}\omega^{-jk}}a_pbq}}\
&=& \frac{1}{n}\sum
{k=0}^{n-1}{(\sum_{p=0}^{n-1}{ap\omega^{kp}})(\sum{q=0}^{n-1}{bq\omega^{kq}})\omega^{-jk}} \
&=& \frac{1}{n}\sum
{k=0}^{n-1}{(A_kB_k)\omega^{-jk}} \
\end{eqnarray
}

如果要计算linear convolution,比如多项式乘法,可以把$\mathbf{a}$长度补到$2n-1$,求cyclic convolution。

下面讨论用于计算整系数向量卷积的discrete Fourier transform。

精度

使用complex<double>计算convolution,需要保证结果每一项的实数部分在$[-2^{53}-1, 2^{53}-1]$(Number.MIN_SAFE_VALUENumber.MAX_SAFE_INTEGER)范围内,$2^{53}-1$是double能精确表示的最大整数。采取round half to even规则,$2^{53},2^{53}+1$均表示为$2^{53}$,无法区分。

设每项系数的绝对值小于等于$v$,那么convolution结果每一项绝对值小于等于$nv^2$,若$nv^2\leq 2^{53}-1$则可放心使用complex<double>

complex<double>还要受到浮点运算误差影响。根据Roundoff Error Analysis of the Fast Fourier Transform,没仔细看,relative error均值为log2(n)*浮点运算精度*变换前系数最大值,对于结果$x$,这个量达到$0.5/x$就很可能出错,增长速度可以看作是$log(n)v^2$,不如$nv^2$。因此通常不必担心浮点运算的误差。

对于模$P$的number theoretic transform,$v\leq P-1$,若$nv^2\leq P$则可放心使用。

998244353, 897581057, 880803841小于$2^{30}$,两倍不超过INT_MAX,且可表示为$k*n+1$,其中$n$为2的幂,适合用作number theoretic transform的模。

设$\mathbf{a},\mathbf{b}$系数取自$[0,v]$的uniform distribution,则$\mathbf{a}\ast\mathbf{b}$系数均值为$nv^2/4$,方差为$nv^4/9$。若把系数平移至$[-v/2, v/2]$,则$\mathbf{a}\ast\mathbf{b}$系数均值$0$,方差为$nv^4/144$。若$\mathbf{a},\mathbf{b}$其中之一independent and identically distributed,则方差会很小。可以用Chebyshev’s inequality等估计系数绝对值,上界$nv^2$可以减小。即使$\mathbf{a},\mathbf{b}$不是independent and identically distributed,也可以用$\mathbf{a}\ast\mathbf{b}=\mathbf{a}\ast(\mathbf{b+c})-\mathbf{a}\ast\mathbf{c}$来计算,$\mathbf{c}$是independent and identically distributed uniform $[-v/2,v/2]$。

方案0:sqrt decomposition(FFT, NTT)

取$M$为接近$\sqrt{v}$的整数,分解$\mathbf{a}=\mathbf{a0}+M\mathbf{a1}$、$\mathbf{a}=\mathbf{a0}+M\mathbf{a1}$,则:

$$\mathbf{a}\ast\mathbf{b} = \mathbf{a0}\ast\mathbf{b0}+M(\mathbf{a0}\ast\mathbf{b1}+\mathbf{a1}\ast\mathbf{b0})+M^2(\mathbf{a1}\ast\mathbf{b1})$$

适当选择$M$可以使$\mathbf{a0},\mathbf{a1}$的系数小于等于$\lfloor\sqrt{v}\rfloor$,convolution结果系数最大值为$n(\lfloor\sqrt{v}\rfloor)^2\approx nv$,比原来的$nv^2$小。

求出$DFT(\mathbf{a0}), DFT(\mathbf{a1}), DFT(\mathbf{b0}), DFT(\mathbf{b1})$后,计算等式右边四个convolution,带权相加即得到原convolution。

需要4次长为$2n$的DFT、1次长为$2n$的inverse DFT。

可以使用Toom-2 (Karatsuba)计算$\mathbf{a0}\ast\mathbf{b0}, \mathbf{a1}\ast\mathbf{b1}, (\mathbf{a0}+\mathbf{a1})\ast(\mathbf{b0}+\mathbf{b1})$,减少为3次DFT、1次inverse DFT。$\mathbf{a0}\ast\mathbf{b1}+\mathbf{a1}\ast\mathbf{b0} = (\mathbf{a0}+\mathbf{a1})\ast(\mathbf{b0}+\mathbf{b1}) - \mathbf{a0}\ast\mathbf{b0} - \mathbf{a1}\ast\mathbf{b1}$。

容易扩展到cube root decomposition等。对于number theoretic transform,分成$k$份需要$2k$个DFT和$2k-1$个IDFT,不如用Chinese remainder theorem。

优化0:imaginary unit用于进位(FFT)

取$S$与$\sqrt(P)$接近且$M=P-S*S%P$尽可能小。

分解$\mathbf{a}=\mathbf{a0}+S\mathbf{a1}$、$\mathbf{b}=\mathbf{b0}+S\mathbf{b1}$。

\begin{eqnarray}
IDFT(DFT(a0+i\sqrt{M}a1)\cdot DFT(b0+i\sqrt{M}a1)) &=& (\mathbb{a0}+i\sqrt{M}\mathbb{a1})\ast (\mathbb{b0}+i\sqrt{M}\mathbb{b1}) \
&=& (\mathbf{a0}\ast \mathbf{b0})-M(\mathbf{a1}\ast\mathbf{b1})+i\sqrt{M}((\mathbf{a0}\ast\mathbf{b1})+(\mathbf{a1}\ast\mathbf{b0})) \
\end{eqnarray
}

每一项绝对值变为$\sqrt{1+M^2}$倍了。右边同余于$(\mathbf{a0}\ast \mathbf{b0})+S^2(\mathbf{a1}\ast\mathbf{b1})+i\sqrt{S}(\mathbf{a0}\ast\mathbf{b1}+\mathbf{a1}\ast\mathbf{b0})$,提取虚部与实部即可得到:

$$(\mathbf{a0}\ast \mathbf{b0})+S(\mathbf{a0}\ast\mathbf{b1}+\mathbf{a1}\ast\mathbf{b0})+S^2(\mathbf{a1}\ast\mathbf{b1})$$

需要2次长为$2n$的DFT、1次长为$2n$的inverse DFT。

优化1:正交计算两个实系数向量DFT(FFT)

$\mathbf{a}$的共轭的DFT可由$\mathbf{a}$的DFT求出:

\begin{eqnarray}
DFT(conj(\mathbf{a}))j &=& \sum{k=0}^{n-1}{conj(ak)\omega^{jk}} \
&=& \sum
{k=0}^{n-1}{conj(ak\omega^{-jk})} \
&=& conj(\sum
{k=0}^{n-1}{ak\omega^{-jk}}) \
&=& conj(DFT(\mathbf{a})
{-j\;mod\;n}) \
&=& conj(rev(DFT(\mathbf{a}))_j) \
\end{eqnarray
}

\begin{eqnarray}
DFT(re(\mathbf{a})) &=& (DFT(\mathbf{a})+DFT(conj(\mathbf{a})))/2 = (DFT(\mathbf{a}) + conj(rev(DFT(\mathbf{a}))))/2 \
DFT(im(\mathbf{a})) &=& (DFT(\mathbf{a})-DFT(conj(\mathbf{a})))/(2i) = (DFT(\mathbf{a}) - conj(rev(DFT(\mathbf{a}))))/(2i) \
\end{eqnarray
}

分解$\mathbf{a}=\mathbf{a0}+M\mathbf{a1}$、$\mathbf{b}=\mathbf{b0}+M\mathbf{b1}$后,根据上面的公式,用$DFT(\mathbf{a0}+i\mathbf{a1})$计算$DFT(\mathbf{a0})$和$DFT(\mathbf{a1})$,同法计算$DFT(\mathbf{b0})$和$DFT(\mathbf{b1})$。然后用$IDFT(DFT(\mathbf{a0})\cdot DFT(\mathbf{b0}) + i DFT(\mathbf{a0})\cdot DFT(\mathbf{b1}))$计算出$\mathbf{a0}\ast\mathbf{b0}$与$\mathbf{a0}\ast\mathbf{b1}$,同法计算出$\mathbf{a1}\ast\mathbf{b0}$与$\mathbf{a1}\ast\mathbf{b1}$。

需要2次长为$2n$的DFT、2次长为$2n$的inverse DFT。

奇偶项优化(FFT)

该优化可以和其他方式叠加。

把$\mathbf{a}$看作多项式$A(x)=\sum_{k=0}^{n-1}{a_kx^k}$,同样地,$\mathbf{b}$看作多项式$B(x)$。

奇次项$A0(x)=\sum_{0\leq k<n, k\;\text{is even}}{akx^k}$, 偶次项$A1(x)=\sum{0\leq k<n, k\;\text{is odd}}{a_kx^k}$,同样地,定义$B0(x)$与$B1(x)$。$A0(x)$的系数为$\mathbf{a0},$A1(x)$的系数为$\mathbf{a1}$,令其长为$n$,高位用$0$填充。

\begin{eqnarray}
A(x)B(x) &=& (A0(x^2)+x A1(x^2))(B0(x^2)+x B1(x^2)) \
&=& A0(x^2)B0(x^2)+x^2A1(x^2)B1(x^2) + x(A0(x^2)B1(x^2)+A1(x^2)B0(x^2)) \
\end{eqnarray
}

用正交计算两个实系数向量DFT的方式,用2次长度为$n$(之前都是 $2n$)的DFT计算$DFT(\mathbf{a0}), DFT(\mathbf{a1})$, $DFT(\mathbf{b0}), DFT(\mathbf{b1})$。$\mathbf{a1}$循环右移1位的DFT的第$j$项等于$DFT(\mathbf{a1})_j\omega^j$,因此根据$A1(x^2)$的DFT的系数可以得到$x^2A1(x^2)$的DFT的系数。

构造长为$n$的向量$\mathbf{c}$:
$$c_j = (a0_j b0_j + a1_j b1_j \omega^j) + i(a0_j b1_j + a1_j b0_j)$$

$DFT(\mathbf{c})$的实部为结果的偶次项系数,虚部为结果的奇次项系数。

需要2次长为$n$的DFT、1次长为$n$的inverse DFT。

方案1:Chinese remainder theorem(NTT)

取$t$个可用于number theoretic transform的质数$P_0,P1,\ldots,P{t-1}$,使$M = \prod_{0\leq j<t}{P_j}\geq nv^2$,计算$t$个NTT,之后用Chinese remainder theorem合并。

求$x \equiv v_j\quad(\text{mod}\;M/P_j),\quad 0\leq j<t$有至少两种算法。

经典算法(Gauss’s algorithm)

Gauss之前也有很多人提出。

对于每个$P_j$用Blankinship’s algorithm计算$P_j p_j \equiv 1\quad(\text{mod}\;M/P_j)$。

$$x = (\sum_{j=0}^{t-1}{v_jP_jp_j}) \% M$$

注意$M\geq nv^2$可能超出机器single-precision表示范围,该算法不适合求$x\% P$。

Garner’s algorithm

定义
\begin{eqnarray}
c_1 &=& inv(P_0, P_1) \
c_2 &=& inv(P_0P_1, P_2) \
c_3 &=& inv(P_0P_1P_2, P_3) \
\ldots \
\end{eqnarray
}

\begin{eqnarray}
y_0 &=& v_0 \% P_0 & x_0 &=& y_0 \
y_1 &=& (v_1-x_0)c_1 \% P_1 & x_1 &=& x_0 + y_1P_0 \
y_2 &=& (v_2-x_1)c_2 \% P_2 & x_2 &=& x_1 + y_2P_0P_1 \
y_3 &=& (v_3-x_2)c_3 \% P_3 & x_3 &=& x_2 + y_3P_0P_1P_2 \
\ldots \
\end{eqnarray
}

$xj$满足前$j+1$个方程,$x=x{t-1}$满足所有方程。

稍加变形可用于求$x\% P$。

\begin{eqnarray}
y_0 &=& v_0 \% P_0 & x_0 &=& y_0 \% P \
y_1 &=& (v_1-(y_0)\% P_1)c_1 \% P_1 & x_1 &=& (x_0 + y_1P_0) \% P \
y_2 &=& (v_2-(y_0+y_1P_0)\% P_2)c_2 \% P_2 & x_2 &=& (x_1 + y_2P_0P_1) \% P \
y_3 &=& (v_3-(y_0+y_1P_0+y_2P_0P_1)\% P_3)c_3 \% P_3 & x_3 &=& (x_2 + y_3P_0P_1P_2) \% P \
\ldots \
\end{eqnarray
}

原来的每个$x_j$是精确值,现在只有$x_j%P$的结果,因此计算$yj$时之前的$x{0.dotsj-1}$都不可复用,需要重新计算。时间复杂度上升为$O(t^2)$。

测试

https://gist.github.com/MaskRay/fac2042058dd5d9e59953f18f3f3978a

NTT int使用小于$2^{31}$的素数,编译器用乘法模拟常数除法。
NTT long,编译器无法优化常数除法,性能很低,使用浮点mul_mod会略快于128位除以64位的DIV指令。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
n microseconds algorithm
131072 3860 NTT dit2 int
131072 6104 FFT dif2
131072 6712 FFT dit2
131072 6912 NTT dif2 int
131072 6936 Montgomery+Barrett NTT dif2 int
131072 9592 NTT dif2 long non-constant P
131072 10122 NTT dit2 long
131072 13169 NTT dif2 int non-constant P
131072 15419 NTT dif2 long
262144 8993 NTT dif2 int
262144 9036 NTT dit2 int
262144 9670 Montgomery+Barrett NTT dif2 int
262144 15484 FFT dit2
262144 17601 FFT dif2
262144 19731 NTT dit2 long
262144 20527 NTT dif2 long non-constant P
262144 21910 NTT dif2 long
262144 29457 NTT dif2 int non-constant P
524288 18502 NTT dif2 int
524288 20110 Montgomery+Barrett NTT dif2 int
524288 23156 NTT dit2 int
524288 39890 FFT dif2
524288 39904 FFT dit2
524288 44145 NTT dif2 long non-constant P
524288 45038 NTT dit2 long
524288 46334 NTT dif2 long
524288 65265 NTT dif2 int non-constant P
1048576 43648 NTT dit2 int
1048576 45704 NTT dif2 int
1048576 46167 Montgomery+Barrett NTT dif2 int
1048576 104362 NTT dit2 long
1048576 107571 NTT dif2 long non-constant P
1048576 119743 FFT dif2
1048576 122029 NTT dif2 long
1048576 122174 FFT dit2
1048576 144370 NTT dif2 int non-constant P
2097152 122989 Montgomery+Barrett NTT dif2 int
2097152 137276 NTT dif2 int
2097152 143955 NTT dit2 int
2097152 293222 FFT dif2
2097152 338580 FFT dit2
2097152 352833 NTT dif2 int non-constant P
2097152 360372 NTT dif2 long non-constant P
2097152 422108 NTT dit2 long
2097152 423817 NTT dif2 long
4194304 455859 NTT dit2 int
4194304 467340 NTT dif2 int
4194304 490114 Montgomery+Barrett NTT dif2 int
4194304 779945 FFT dif2
4194304 839698 FFT dit2
4194304 904096 NTT dit2 long
4194304 956174 NTT dif2 long
4194304 969572 NTT dif2 long non-constant P
4194304 1074858 NTT dif2 int non-constant P
8388608 1052072 NTT dit2 int
8388608 1138089 NTT dif2 int
8388608 1189775 Montgomery+Barrett NTT dif2 int
8388608 1737166 FFT dif2
8388608 1839095 FFT dit2
8388608 2053195 NTT dif2 long
8388608 2072172 NTT dif2 long non-constant P
8388608 2186451 NTT dit2 long
8388608 2893584 NTT dif2 int non-constant P

花哨的Montgomery+Barrett不如常量除法的NTT int,好处是一份代码可以适用于多个模数,而NTT int得用template或其他方式为各个模数生成不同代码。

不受到Level 3 cache制约时,Montgomery NTT只需要NTT int 60%的时间,此时每次重新计算unit root代替lookup table会快些。

一般来说,decimation in frequency(Sande-Tukey,从较大的n计算到较小的n)优于decimation in time(Cooley-Tukey,从较小的n计算到较大的n),可能是因为decimation in frequency的butterfly数据依赖小些。

FFT有超过10%时间花在计算unit roots上,而NTT只有5%。考虑到FFT往往能正交计算两个序列,而NTT只能计算一个,且double有53位精度而NTT int只能使用$2^{32}$以下的素数(当前代码只能处理$2^{31}$以下的),FFT通常优于NTT。

References

感谢ftiasch老师教导。

uwi,https://www.hackerrank.com/rest/contests/w23/challenges/sasha-and-swaps-ii/hackers/uwi/download_solution:Modified Montgomery+Barrett变体+Garner’s algorithm: