非再帰FFT - Kude’s diary

非再帰 FFTを自分なりに解釈してみた。

FFT
行列
変換を扱いやすくする
変換を変形していく
bit reversalを右に
ACL
転置ver.

FFT

長さ $2^n$ の配列 $A$ は $2^n-1$ 次多項式 $\displaystyle f(x) = \sum_{0 \le k < 2^n} A_k x^k$ と対応付けられる。 $\zeta$ を $1$ の $2^n$ 乗根として、配列 $\left(f(\zeta^0), f(\zeta^1), f(\zeta^2), \ldots, f(\zeta^{2^n-1})\right)$ を、すなわち $x=\zeta^0, \zeta^1, \zeta^2, \ldots, \zeta^{2^n - 1}$ における $f$ の評価値を求めたい。

$f$ の偶数次、奇数次の係数を取ってくることで作られる2つの $2^{n-1}-1$ 次多項式
\begin{align*}
f_0 (x) &= \sum_{0 \le k < 2^{n-1}} A_{2k} x^k\\
f_1 (x) &= \sum_{0 \le k < 2^{n-1}} A_{2k+1} x^k
\end{align*}
を用いることで、 $f(x)$ は
\begin{align*}
f(x) = f_0(x^2) + x f_1(x^2)
\end{align*}
と分けられる。 $f_0$ , $f_1$ に対してFFTをすることで得られる $x=(\zeta^2)^0, (\zeta^2)^1, (\zeta^2)^2, \ldots, (\zeta^2)^{2^{n-1}-1}$ における $f_0$ , $f_1$ の評価値を用いることで、各 $f(\zeta^i)$ は以下の式から計算できる。
\begin{align*}
f(\zeta^i) &=
\begin{cases}
f_0\left((\zeta^2)^i\right) + \zeta^i f_1\left((\zeta^2)^i\right) & (0 \le i < 2^{n - 1})\\
f_0\left((\zeta^2)^{i - 2^{n-1}}\right) - \zeta^{i - 2^{n-1}} f_1\left((\zeta^2)^{i - 2^{n - 1}}\right) & (2^{n - 1} \le i < 2^n)
\end{cases}
\end{align*}

行列

もう少し機械的に扱いたいので、上の操作が何をやっているかを行列の言葉で表してみる。行列 $F_n$ を
\begin{align*}
F_n=(\zeta^{ij})_{(i,j)}=
\begin{pmatrix}
\zeta^0 & \zeta^0 & \zeta^0 & \cdots & \zeta^0 \\
\zeta^0 & \zeta^1 & \zeta^2 & & \zeta^{2^n-1} \\
\zeta^0 & \zeta^2 & \zeta^4 & & \zeta^{2(2^n - 1)}\\
\vdots & & & \ddots & \vdots \\
\zeta^0 & \zeta^{2^n-1} & \zeta^{2(2^n-1)} & \cdots & \zeta^{(2^n-1)(2^n-1)}
\end{pmatrix}
\end{align*}
と定めて $A$ を行ベクトルと見なすと、FFTは $A$ を $A F_n$ に変換する操作と言うことができる。

上で説明したFFTは

$A$ をインデックスの偶奇で2つに分けてそれぞれFFTを行う
それらの結果を組み合わせて答えを求める

という操作をしていた。例えば $n=3$ の場合、これは次のような行列の分解に対応する。
\begin{align*}
F_3 &=
\begin{pmatrix}
\zeta^{0} & \zeta^{0} & \zeta^{0} & \zeta^{0} & \zeta^{0} & \zeta^{0} & \zeta^{0} & \zeta^{0} \\
\zeta^{0} & \zeta^{1} & \zeta^{2} & \zeta^{3} & \zeta^{4} & \zeta^{5} & \zeta^{6} & \zeta^{7} \\
\zeta^{0} & \zeta^{2} & \zeta^{4} & \zeta^{6} & \zeta^{8} & \zeta^{10} & \zeta^{12} & \zeta^{14} \\
\zeta^{0} & \zeta^{3} & \zeta^{6} & \zeta^{9} & \zeta^{12} & \zeta^{15} & \zeta^{18} & \zeta^{21} \\
\zeta^{0} & \zeta^{4} & \zeta^{8} & \zeta^{12} & \zeta^{16} & \zeta^{20} & \zeta^{24} & \zeta^{28} \\
\zeta^{0} & \zeta^{5} & \zeta^{10} & \zeta^{15} & \zeta^{20} & \zeta^{25} & \zeta^{30} & \zeta^{35} \\
\zeta^{0} & \zeta^{6} & \zeta^{12} & \zeta^{18} & \zeta^{24} & \zeta^{30} & \zeta^{36} & \zeta^{42} \\
\zeta^{0} & \zeta^{7} & \zeta^{14} & \zeta^{21} & \zeta^{28} & \zeta^{35} & \zeta^{42} & \zeta^{49}
\end{pmatrix} \\
&=
\left(
\begin{pmatrix}
(\zeta^2)^{0} & & (\zeta^2)^{0} & & (\zeta^2)^{0} & & (\zeta^2)^{0} & \\
& & & & & & & \\
(\zeta^2)^{0} & & (\zeta^2)^{1} & & (\zeta^2)^{2} & & (\zeta^2)^{3} & \\
& & & & & & & \\
(\zeta^2)^{0} & & (\zeta^2)^{2} & & (\zeta^2)^{4} & & (\zeta^2)^{6} & \\
& & & & & & & \\
(\zeta^2)^{0} & & (\zeta^2)^{3} & & (\zeta^2)^{6} & & (\zeta^2)^{9} & \\
& & & & & & &
\end{pmatrix}
+
\begin{pmatrix}
& & & & & & & \\
& (\zeta^2)^{0} & & (\zeta^2)^{0} & & (\zeta^2)^{0} & & (\zeta^2)^{0} \\
& & & & & & & \\
& (\zeta^2)^{0} & & (\zeta^2)^{1} & & (\zeta^2)^{2} & & (\zeta^2)^{3} \\
& & & & & & & \\
& (\zeta^2)^{0} & & (\zeta^2)^{2} & & (\zeta^2)^{4} & & (\zeta^2)^{6} \\
& & & & & & & \\
& (\zeta^2)^{0} & & (\zeta^2)^{3} & & (\zeta^2)^{6} & & (\zeta^2)^{9}
\end{pmatrix}
\right) \\
& \qquad \begin{pmatrix}
1&&&& 1&&& \\
\zeta^0&&&& -\zeta^0&&& \\
&1&&& &1&& \\
&\zeta^1&&& &-\zeta^1&& \\
&&1&& &&1& \\
&&\zeta^2&& &&-\zeta^2& \\
&&&1& &&&1 \\
&&&\zeta^3& &&&-\zeta^3 \\
\end{pmatrix}
\end{align*}
このまま色々弄っていくのはつらいので、もうちょっと分かりやすくこれらの行列を表してみたい。

変換を扱いやすくする

まず以下のように記法を定めておく。

$k$ bit整数を考えるとき左方をLSBとする。例えば $6$ を表す $5$ bit整数は $01100$ となる。
$k$ bit整数 $i$ のbit $0$ を $i_0$ 、bit $[1, k)$ を $i_{[1, k)}$ などと書くことにする。
適宜bitの連結をそのままつなげることで表すことにする。例えば $k$ bit整数 $i$ に対し $0i$ と書いたらbit $0$ が $0$ 、bit $[1, k+1)$ が $i$ である $k+1$ bit整数を表すとする。
$i$ のbit reversalを $\overline{i}$ と書くことにする。例えば $\overline{01100} = 00110$ となる。
$1$ の $2^n$ 乗根を $1^{0.\overbrace{\text{$00 \cdots 0$}}^{n - 1}1}$ と書くことにする。例えば $1$ を表す $3$ bit整数 $i=100$ を用いて $1$ の $8$ 乗根は $1^{0.\overline{i}}$ と書ける。

（操作1）、（操作2）を表す行列を次のように書くことにする。

（操作1）について

配列 $A$ のサイズは $2^n$ なので各インデックスは $n$ bit整数と見なせる。また配列 $A$ はペア $(i, A_i)$ の集合と見なせる。「偶奇で分けてFFTを行う」というのは「bit $0$ の値ごとにグループ分けし、各グループ内でbit $[1, n)$ をインデックスに用いてFFTを行う」ことであることから、その行列を

と書くことにする。

（操作2）について

このパートはさらに2つの変換に分解できる。

各 $n-1$ bit整数 $i$ に対し、 $A_{0i}$ と $A_{1i}$ を $A_{0i} + 1^{0.0\overline{i}}A_{1i}$ と $A_{0i} - 1^{0.0\overline{i}}A_{1i}$ に写す。
インデックスを1個分左巡回シフトする。つまり $A$ をペア $(i, A_i)$ の集合と見た時、各 $(i_0 i_{[1,n)}, A_i)$ を $(i_{[1,n)} i_0, A_i)$ に写した集合に変える。

（操作2-1）を表す行列を

、（操作2-2）を表す行列を

と書くことにする。

補足

はクロネッカー積を用いて $F_{n-1} \otimes I_2$ とも書ける。ただし

はある行列 $X, T_{n-1}$ を用いて $T_{n-1} \otimes X$ と書けるという主張ではなく、インデックスのbit $0$ を $X$ の箇所が担当しているという気持ちで書いている。

変換を変形していく

これらを用いると、 $F_n$ は

と表せる。`スワップ'みたいなことを繰り返して

を前方に持ってくることを考えてみる。

はインデックスのbit列を左巡回シフトする変換であることを踏まえると、 $F_n$ は以下のように変形できる。

例えば $n=4$ の場合、これを再帰的に用いると次のように変形できる。

ここで、

はbit reversalを表す変換になっている。これを

と表すことにすると、

と表せる。

のような変換がバタフライ演算と呼ばれる。

bit reversalを右に

以上からbit reversal + バタフライ演算を行うことで非再帰 FFTが実装できるが、ここでbit reversalを後ろに持っていくことを考えてみる。

がインデックスのbit列を逆順にする変換であることを踏まえると、 $F_4$ は以下のように変形できる。

ただし、 $\overline{T_i}$ は $T_i$ とは逆順にインデックスのbit列を扱うようにしたものとする。さっきは変換のときに $1^{0.0\overline{i}}$ を掛けていたが、今度は逆順に扱うようになったので $1^{0.0i}$ が掛けられるようになる。例えば

は、 $A_{1010}$ と $A_{1110}$ を $A_{1010} + 1^{0.010}A_{1110}$ と $A_{1010} - 1^{0.010}A_{1110}$ に写す。

一番最後に置かれたbit reversalは評価値を正しい順番に並び替えているだけなので、その操作を無視しても畳み込み定理は依然成り立つ。したがって畳み込みの実装において無視してしまって良い。ただし原子根を逆元に取り替えるだけでは逆変換とはならなくなる。逆変換をするには素直に巻き戻しをすれば良い。つまり

の表す変換をすれば良い。

またこのとき、変換後にインデックス $i$ に来る値は $f$ の $x=1^{0.i}$ における評価値になる。これを踏まえての意味を考えてみる。 $k$ bit整数 $i$ に対し、 $2^{n-k}-1$ 次多項式 $f_i$ を
\begin{align*}
f_i(x) = \sum_{j: n - k \text{ bit整数}} A_{ij}x^j
\end{align*}
と定める*1。例としてインデックス $1011$ における値の変遷に注目してみると、 $f_{1011}(1^{0.}) \mapsto f_{101}(1^{0.1}) \mapsto f_{10}(1^{0.11}) \mapsto f_{1}(1^{0.011}) \mapsto f(1^{0.1011})$ と変換されていき、 $f$ の添字が $1$ の肩に移っていく感じになっている。

逆にこのように変換を進めていこうという気持ちになることで、この非再帰 FFTを直接導ける。例えば $f$ の添字のbit長を $2$ から $1$ にするステップにおいて $f_{10}(1^{0.11})$ と $f_{11}(1^{0.11})$ から $f_{1}(1^{0.011})$ と $f_{1}(1^{0.111})$ を求めるには、今 $f_{1}(x) = f_{10}(x^2) + x f_{11}(x^2)$ が成り立つので、
\begin{align*}
f_{1}(1^{0.011}) &= f_{10}(1^{0.11}) + 1^{0.011} f_{11}(1^{0.11}) \\
f_{1}(1^{0.111}) &= f_{10}(1^{0.11}) - 1^{0.011} f_{11}(1^{0.11})
\end{align*}
のようにすれば良い。

ACL

ACLのbutterflyもこのような変換を行っている。高速化のために可能であれば2 bit分ずつ処理するようになっているが、ここでは1 bitバージョンに注目する。

int p = 1 << (h - len - 1);
mint rot = 1;
for (int s = 0; s < (1 << len); s++) {
    int offset = s << (h - len);
    for (int i = 0; i < p; i++) {
        auto l = a[i + offset];
        auto r = a[i + offset + p] * rot;
        a[i + offset] = l + r;
        a[i + offset + p] = l - r;
    }
    if (s + 1 != (1 << len))
        rot *= info.rate2[bsf(~(unsigned int)(s))];
}

h-lenが現在の添字のbit長、rotが $1^{0.0s}$ になっている。

例えば $\verb|h| = 8, \verb|len| = 3$ のとき、 $s$ は $000, 100, 010, 110, 001, 101, 011, 111$ の順に回る。そのループ内において各4 bit整数 $i$ に対し $l = f_{i0}(1^{0.s})$ と $r = f_{i1}(1^{0.s})$ を $f_i(1^{0.0s})$ と $f_i(1^{0.1s})$ に写している。

rot *= info.rate2[bsf(~(unsigned int)(s))];

のパートが複雑に見えるが、これは単に $1^{0.0s}$ から $1^{0.0(s+1)}$ を計算している。
例えば $s = 110$ のとき、 $1^{0.0s}$ から $1^{0.0(s+1)}$ を、つまり $1^{0.0110}$ から $1^{0.0001}$ を得るには、 $1^{0.0110}$ に対して $1^{0.011}$ の逆元と $1^{0.0001}$ の積を掛けてやればよい。繰り上がりが止まるbitの位置がbsf(~(unsigned int)(s))によってbit 2と分かる。info.rate2[2]には $1^{0.011}$ の逆元と $1^{0.0001}$ の積が前計算されて入っており、それが参照される。

2 bitバージョンについては、1 bitバージョン2回と考えるほかに、 $f_{00}(1^{0.s})$ と $f_{10}(1^{0.s})$ と $f_{01}(1^{0.s})$ と $f_{11}(1^{0.s})$ を $f(1^{0.00s})$ と $f(1^{0.10s})$ と $f(1^{0.01s})$ と $f(1^{0.11s})$ に写すようなことをしていると考えることもできる。

転置ver.

$F_n$ は対称行列なので転置しても同じ変換を表す。例えば

の転置を取ることで

を得る。ここで、

は

の転置行列とする。転置を取ることは重みをそのままに「入口」と「出口」を交換することと見なせる。

は $A_{i0}$ と $A_{i1}$ を $A_{i0} + 1^{0.0\overline{i}}A_{i1}$ と $A_{i0} - 1^{0.0\overline{i}}A_{i1}$ に写す変換なので、これは以下のような回路を表す。

よってこれの転置となる回路は

となり、

は $A_{i0}$ と $A_{i1}$ を $A_{i0} + A_{i1}$ と $1^{0.0\overline{i}}(A_{i0} - A_{i1})$ に写す変換になっている。

これもbit reversalが最後に来ているので無視できる。こっちの方針で実装すると以下のようになる。

int p = 1 << (len - 1);
const mint root_len = info.root[len];
for (int s = 0; s < (1 << (h - len)); s++) {
    mint rot = 1;
    int offset = s << len;
    for (int i = 0; i < p; i++) {
        auto l = a[i + offset];
        auto r = a[i + offset + p];
        a[i + offset] = l + r;
        a[i + offset + p] =
            (unsigned long long)(mint::mod() + l.val() - r.val()) *
            rot.val();
        rot *= root_len;
    }
}

空間的局所性を高めるために2重ループの内側で下方bitを回すようにしている。そのため二重ループ内でrotを毎回更新している。手元でACLの1 bitバージョンとこれを比較したが、長さ $2^{20}$ のときに前者は $28066383 \,\mathrm{ns}$ 、後者は $37628898 \,\mathrm{ns}$ になって、ACLの方が速い。

*1: $ij$ はbitの連結の意味