整数環FFT

適当な条件を満たす環や体では DFT を利用した時と同様の計算方法で 畳み込み乗算を行うことができる。 その際にはもちろん FFT と同様の式変換を行うことで $O(n\log n)$ の計算量で演算することが可能であり、 そのような方法を「整数環 FFT」もしくは 「数論変換 (Number-Theoretical Transform;NTT)」と呼ぶ。 特に、いくつか特殊な環や体で行う整数環 FFT については DGT (離散ガロア変換)、 DFT (離散フェルマー変換) といったように名前がつけられていることがある。

概要

適当な整数 $n$ に対して

  • 乗算における $n$ の逆元 $n^{-1}$ が存在する
  • 1 の原始 $n$ 乗根 $\omega$ が存在する

という条件を満たす環においては \[ A_k = \sum_{j=0}^{n-1} a_j \omega^{jk},\quad a_j = n^{-1} \sum_{k=0}^{n-1} A_k \omega^{-jk} \]

という相互変換が成り立つ。 これを用いて畳み込み乗算と同じ計算手順を踏むと同様に

\[ c_l = n^{-1}\sum_{k=0}^{n-1} \sum_{j=0}^{n-1} \sum_{j'=0}^{n-1} a_j b_{j'} \omega^{(j+j'-l)k} \] \[ \sum_{k=0}^{n-1}\omega^{jk} = \begin{cases} 0 & (j \bmod n \neq 0)\\ n & (j \bmod n = 0) \end{cases} \]

の 2 つが成り立つので、同様に畳み込み乗算を計算することができる。

64 bit 整数で行える剰余環 FFT

[JT03] では剰余環 $\mathbb{Z}/m\mathbb{Z}$ での計算やその見つけ方が紹介されている。 概要としてはまず $n$ を決め、$m=an+1$ として条件を満たすように 適当に決めた範囲で $a$ と $\omega$ を探索する。

ここでは $\log_2m$ が 25、51、64 付近になる組み合わせを探索した結果を 表1に示した。 $\log m\simeq 25, 51$ のものは倍精度浮動小数点数を用いて、 $\log m\simeq 64$ のものは 64 ビット整数を使って計算しやすい。 ただし $\log m\simeq 51, 64$ のものについては $a\cdot b \bmod m$ の計算を誤差なく行うために特殊な命令、もしくはコード上の工夫が必要となるので注意してほしい。

この表より小さな $n'$ で計算したい場合は倍数となる $n$ を用いて $w'=w^{n/n'}$ とすれば計算できる。

表1. 整数環FFTに使える数の例
$n$$\omega$$m$$\log_2m$
$2^{21}$38$11n+1$24.459
$3\times2^{20}$19$9n+1$24.755
$5\times2^{19}$19$10n+1$24.644
$2^{46}$67$27n+1$50.755
$3\times2^{45}$8$18n+1$50.755
$5\times2^{45}$13$7n+1$50.129
$2^{59}$87$27n+1$63.755
$3\times2^{58}$2$18n+1$63.755
$5\times2^{57}$51$19n+1$63.570

mod $2^{64}-2^{32}+1$

素数 $p=2^{64}-2^{32}+1$ を用いた剰余環 $\mathbb{Z}/p\mathbb{Z}$ は体となる。 そして 7 が原始 $p-1=2^{32}\cdot3\cdot5\cdot17\cdot257\cdot65537$ 乗根となるため $\omega=7^{(p-1)/n}$ を使うことで $n=2^{32}$ までの長さなら 2 基底 FFT ルーチンで対応できる。 さらに 1 回ずつまでであれば 3 基底や 5 基底ルーチンも適用させることでより長い $n$ でも計算することができる。

また、$7^{5(p-1)/192}=2$ が成り立つため、$n \lt 2^{26}$ に対して $\{a_j\}=\{(7^{5(p-1)/192/n})^j\}$ を利用することで整数演算を用いながら IBDWT と同じ原理を適用することも可能になる。

この体での計算では

\[ \begin{cases} 2^{64} = 2^{32} -1\\ 2^{96} = -1\\ 2^{128} = -2^{32}\\ 2^{192} = 1 \end{cases} \]

が成り立つので、$\bmod p$ の演算が容易になる。 具体的には一旦 128 bit になる計算結果を 32 bit 毎に区切って下の桁から順に $x_0$, $x_1$, $x_2$, $x_3$ とすると

\[ \begin{eqnarray} && x_0 + 2^{32}x_1 + 2^{64}x_2 + 2^{96}x_3\\ &=& 2^{64}x_2 + 2^{32}x_1 + (x_0-x_3) & (2^{96} = -1 \ {\rm を適用})\\ &=& 2^{32}(x_1+x_2) + x_0 - x_3 - x_2 & (2^{64} = 2^{32} - 1 \ {\rm を適用}) \end{eqnarray} \]

という形になり、演算命令として mod が必要なくなる。

Schönhage-Strassen Algorithm

FFT を利用した多倍長乗算の難点の 1 つに、1 変数に代入できる数の大きさがレジスタサイズで抑えられることがあげられる。 つまり倍精度浮動小数点数を使った FFT を利用した畳み込み演算では途中どんなに正確に計算できたとしても結果が $2^{53}$ を超えると誤差が出る。 そのため大きな桁数の多倍長乗算をするためにはより多くの数に分割し、 1 つ 1 つの変数に入れる数の桁数を減らすしかない。 これは通常の NTT を利用する場合も同様で、畳み込み演算の結果が $m$ 以上になる演算はできない。

そこで NTT を行う各要素も多倍長数にすればその欠点を克服することができるのではないかと考えられる。 ただ素直に $m$ を適当な長さの多倍長整数にするとその条件を回避することはできるが、NTT の中の計算を $\bmod m$ の環で行うことになるため、その計算そのものが大変になる。 具体的な例としては そもそも適切な $\omega$ を求めること、$a\cdot b\bmod m$ の計算が多倍長剰余演算であること、$w^{jk}$ のテーブルが大きすぎること、などが挙げられるだろう。 もちろんいずれも計算が不可能になるような問題ではないが、NTT そのものの時間、空間コストが高くなってしまうのは自明なことだろう。

この様な問題をも上手く回避したアルゴリズムが Schönhage-Strassen Algorithm (SSA) である。 SSA では要素数 $n=2^N$ に対し $m=2^{n/2}+1$ を設定し NTT を計算する。 すると $\omega=2$ が原始 $n$ 乗根となるのでビットシフトで $\omega^{jk}$ との乗算を行うことができる。 逆変換で必要な $n^{-1}$ との乗算も $2^{n-N}$ で計算できるので難しくない。 $\bmod m$ の計算は、数値を 2 進数の形で格納している前提であれば $N$ ビットごとに区切ることで $\bmod 2^{64}-2^{32}+1$ の例と同様の形で処理することができる他、 負巡回畳込み演算を利用する方法もある。 一方で 2 進数でのデータ格納が必須になるのでデータ入出力を行う際には 基数変換をしなければならない。

計算例

1 つ簡単な計算例として、長さ 8 の 2 数列 \[ \begin{eqnarray} x &=& \{3,1,4,1,5,9,2,6\}\\ y &=& \{5,3,5,8,9,7,9,3\} \end{eqnarray} \]

の畳み込み乗算 $x \circledast y$ を計算してみよう。 ここで巡回畳込みになることを回避する ためにどちらも 0 を 8 要素ずつ後ろに追加し $n=16$、すなわち $m=2^{n/2}+1=257$ なので $\pmod {257}$ の剰余環で SSA を計算することにする。 まず $\omega^{jk}=2^{jk}$ なので順方向の変換は \[ A=\begin{pmatrix} 2^0&2^0&2^0&\cdots &2^0\\ 2^0&2^1&2^2&\cdots &2^{15}\\ \vdots\\ 2^0&2^{15}&2^{30}&\cdots &2^{225} \end{pmatrix} = \begin{pmatrix} 2^0&2^0&2^0&\cdots &2^0\\ 2^0&2^1&2^2&\cdots &2^{15}\\ \vdots\\ 2^0&2^{15}&2^{14}&\cdots &2^1 \end{pmatrix} \] を使ってそれぞれ \[ A \begin{pmatrix}3\\1\\4\\1\\5\\9\\2\\6\\0\\0\\0\\0\\0\\0\\0\\0\end{pmatrix} = \begin{pmatrix}31\\8\\192\\4\\50\\130\\205\\160\\254\\189\\125\\113\\211\\5\\241\\186\end{pmatrix} , \quad A \begin{pmatrix}5\\3\\5\\8\\9\\7\\9\\3\\0\\0\\0\\0\\0\\0\\0\\0\end{pmatrix} = \begin{pmatrix}49\\138\\236\\196\\241\\177\\81\\207\\7\\67\\142\\238\\16\\214\\39\\88\end{pmatrix} \]

と計算できる。ちなみにコンピュータで実装する際には 2 基底 FFT アルゴリズムを適用させることを前提としている。 次に同じインデックスの要素同士をかけ合わせて SSA の逆演算をする。 \[ A^{-1}=n^{-1} \begin{pmatrix} 2^0&2^0&2^0&\cdots &2^0\\ 2^0&2^{-1}&2^{-2}&\cdots &2^{-15}\\ \vdots\\ 2^0&2^{-15}&2^{-30}&\cdots &2^{-225} \end{pmatrix} = 2^{12} \begin{pmatrix} 2^0&2^0&2^0&\cdots &2^0\\ 2^0&2^{15}&2^{14}&\cdots &2^1\\ \vdots\\ 2^0&2^1&2^2&\cdots &2^15 \end{pmatrix} \]

ここでは数式上、元の SSA の係数行列 $A$ の逆行列 $A^{-1}$ と書いているが、 プログラムとしては $\{a_jk\}=2^{-jk}$ として SSA と同様に計算させ、最後に $1/n=2^{-4}$ をかければよい。 この $1/n$ の乗算は上記の数式のように $2^{12}$ に変換 ($2^n=1$ なので $2^{-4}=2^{n-4}=2^{12}$) してシフト計算させても構わないが、$\pmod {2^{n/2}+1}$ の方を利用して 下位 $\log_2n=4$ ビットを 0 にするように $m$ を足してビットシフトする方法もある。 \[ z = A^{-1} \begin{pmatrix}31*49\\8*138\\192*236\\4*196\\50*241\\130*177\\205*81\\160*207\\254*7\\189*67\\125*142\\113*238\\211*16\\5*214\\241*39\\186*88\end{pmatrix} = \begin{pmatrix}15\\14\\38\\46\\83\\127\\140\\176\\191\\183\\177\\164\\87\\60\\18\\0\end{pmatrix} \]

いずれにせよ上記のような結果が求まり、$x\circledast y$ と等しいことが確認できる。 なおこの計算例ではたまたま大丈夫だったが SSA は $\pmod m$ での計算するため畳み込み乗算結果が $m$ 以上になると通常の畳み込み乗算と結果が異なってしまうので注意してほしい。 なので多倍長乗算の一部として使う場合、畳み込み乗算の理論最大値が $m$ 以上にならないように 分割数 $n$ などを調整する必要がある。

計算量の概算

ここで SSA を利用した多倍長乗算の計算量 $M_{\rm SSA}(n)$ について考えてみよう。 サイズ $m$ の多倍長数を $n$ 個利用して計算するとすれば FFT と同様に考えることで \[ M_{\rm SSA}(n) = O(n\log n \cdot M(m)) \]

となる。 ここで再帰的に多倍長演算で SSA を利用することを考えると、2 段階の再帰計算で $O(n\log n\log^2n)$、3 段階の再帰計算で $O(n\log n\log^2n\log^3n)$ と再帰回数が増える度に $\log^kn$ が増えていく。 ただし実用的には 1 京桁程度の計算でも 2 段階の再帰計算で十分であることと $O(\log^3n)$ くらいからは実質的に定数と差がない程度に小さいため SSA の計算量は $O(n\log n\log^2n)$ とすることもある。