多次元 FFT と畳み込み

概要

多次元の離散フーリエ変換は、画像処理・偏微分方程式 (PDE) のスペクトル法・相関解析・量子化学など、 多方面で基幹計算となる。本記事は分離可能性 ($e^{-2\pi i (k_x x + k_y y)} = e^{-2\pi i k_x x} \cdot e^{-2\pi i k_y y}$) に基づく row-column 分解、ブロック転置によるキャッシュ最適化、畳み込み定理と ゼロパディング、バッチ FFT を扱う。

API は fft2d / fft3d および simd_fft::batchFFT として公開されている。

2D FFT の row-column 分解

$M \times N$ の 2 次元配列 $x[m, n]$ の DFT は

$$X[k_y, k_x] = \sum_{m=0}^{M-1} \sum_{n=0}^{N-1} x[m, n] \, \omega_M^{k_y m} \, \omega_N^{k_x n}$$

$\omega_M^{k_y m} \cdot \omega_N^{k_x n}$ が分離するため、内側の $n$ について先に和を取ると各行の 1D FFTが、 次に外側の $m$ について和を取ると各列の 1D FFTが得られる。

$$X[k_y, k_x] = \sum_{m=0}^{M-1} \underbrace{\Bigl[ \sum_{n=0}^{N-1} x[m, n] \, \omega_N^{k_x n} \Bigr]}_{\text{行 FFT } Y[m, k_x]} \, \omega_M^{k_y m}$$

合計コストは $M$ 回の長さ $N$ FFT $+$ $N$ 回の長さ $M$ FFT $= O(MN \log(MN))$。 奇数次元・素数次元・ストライド・キャッシュ局所性などの実装上の問題はすべて 1D FFT 側に押し込められる。

転置の高速化 (ブロッキング)

row-major 配置の配列で行方向 FFT は単位ストライドで効率良く動くが、 列方向はストライド $N_x$ でアクセスするためキャッシュラインの再利用ができない。 $N_x = 2048$ で double 複素数 (16 bytes) なら列ステップごとに 32 KB を読み飛ばし、L1 キャッシュ (典型 32–48 KB) に収まらない。

対策として行と列の FFT の間に転置を挟む。 単純な転置は依然ストライドが大きいが、$B \times B$ ($B = 32 \text{ や } 64$) のブロック単位で転置すれば、 ブロック全体が L1 に収まり、書込み側だけストライドが残る。

$$\text{ブロック転置}: \;\; \forall \text{ block } (I, J), \; \forall i, j \in [0, B): \;\; A^T[J B + j, I B + i] = A[I B + i, J B + j]$$

転置 + 連続 1D FFT は、ストライド版 1D FFT より一般に 2〜4 倍速い。 sangi の fft2d はブロック転置を 2 回挟む構成 (転置 → 連続 FFT → 転置で元の配置に戻し → 連続 FFT) を取る。

3D FFT と奥行き方向

3D の場合 $x \to y \to z$ の順に 1D FFT を適用する。row-major 配置 data[z * ny * nx + y * nx + x] では $x$ が単位ストライド、$y$ がストライド $N_x$、$z$ がストライド $N_x N_y$ となる。

奥行き方向 ($z$) のストライドは行列より一段大きく、ブロック転置だけでは L1 に収まらないことがある。 対策として:

• 2 次元スラブ: $z$ 軸方向のサイズを小さく取って一度に処理する。
• 3D ブロック (cubelet) 転置: $B^3$ 立方体の単位で転置。
• 軸ローテーション: $(x, y, z) \to (y, z, x) \to (z, x, y)$ と循環的に順序を変える。

sangi の fft3d は Tensor3D<T> コンテナを使い、各軸方向に 1D FFT を逐次実行する。 中規模 ($\leq 256^3$ 程度) では単純な軸順実装で十分高速で、大規模 GPU 並列化は将来課題である。

fftshift3d は DC 成分 ($k = 0$) をテンソルの中心に移動するためのユーティリティで、可視化や畳み込みでの整列に使う。

畳み込み定理

連続時間信号 $x(t), y(t)$ の畳み込み $(x \ast y)(t) = \int x(\tau) y(t - \tau) d\tau$ のフーリエ変換は

$$\mathcal{F}\{x \ast y\}(\xi) = \mathcal{F}\{x\}(\xi) \cdot \mathcal{F}\{y\}(\xi)$$

離散版でも同じ関係が (周期境界の意味で) 成立する。 従って長さ $N$ の畳み込みは FFT $\to$ 要素ごとの積 $\to$ IFFT で $O(N \log N)$ で計算できる。 直接定義通りの $O(N^2)$ より大幅に高速である。

多次元畳み込みでも同じ枠組みで:

$$\mathcal{F}_{2D}\{X \ast K\}(\xi, \eta) = \mathcal{F}_{2D}\{X\}(\xi, \eta) \cdot \mathcal{F}_{2D}\{K\}(\xi, \eta)$$

2D 畳み込みは 2D FFT × 2 (順) + 要素積 + 2D IFFT × 1 で計算され、空間領域の素朴な $O(M^2 N^2)$ を $O(MN \log(MN))$ に削減する。

線形畳み込みと循環畳み込み

FFT が暗黙に計算するのは長さ $N$ の循環畳み込みであり、入力が $N$ で周期的だと仮定される。 信号処理で多用される線形畳み込みは次のような違いがある:

FFT で線形畳み込みを得るには、両入力を $L \geq M + N - 1$ まで ゼロパディング してから FFT する。 さもないと循環の折り返し (wrap-around) で出力末尾と先頭が汚染される。

overlap-add / overlap-save 法

長い信号 $x$ と短いカーネル $h$ ($N \ll \mathrm{len}(x)$) の畳み込みでは、 $x$ をブロックに分割して各ブロックでフィルタリングし、結果を加算 (overlap-add) または取捨選択 (overlap-save) で結合する。 各ブロック畳み込みは FFT で行い、$x$ 全体の FFT は不要である。 定常 FIR フィルタ実装の標準的手法である。

sangi での運用

• FFT::convolve(a, b) — 内部で長さ $\geq a.size() + b.size() - 1$ の 2 のべき乗までパディングする線形畳み込み。
• RealFFT::convolve(a, b) — 実数版。RealFFT::fft 経由で約 2 倍高速。
• 2D 線形畳み込みは fft2d_real $\to$ 要素積 $\to$ ifft2d を自前で組む (現状 convolve2d 関数は未公開)。

バッチ FFT

同じサイズ $N$ の独立な信号 $K$ 本に FFT を適用する場面はステレオ音声・スペクトログラム・MIMO 通信などで頻出する。 $K$ 本を順次処理するより、$K$ 本を SIMD レーンにインターリーブして同時に処理するほうが速い。

• twiddle テーブル再利用: $K$ 本のバタフライで同じ twiddle を共有でき、L1 キャッシュのヒット率が上がる。
• SIMD レーン詰め: AVX2 では 4 つの float 複素数 (= 8 floats) を 1 つの 256-bit レジスタに収め、4 信号を 1 命令で並列処理。
• 分岐削減: 信号間でループ構造が共通なため、分岐予測ミスがない。

sangi の simd_fft::batchFFT は 4 信号を 1 単位として AVX2 でインターリーブ実行し、端数は通常の simd_fft::fft() で処理する。 回転子テーブルは thread_local でキャッシュされる。

simd_fft::stereoFFT は別系統の最適化で、 インターリーブされたステレオ WAV (L0, R0, L1, R1, ...) を 1 回の複素 FFT で L/R スペクトルに分離する。 共役対称性 ($X[k]$ が実部 L、虚部 R を持つ複素信号の FFT) を利用するため、2 回の実数 FFT より速い。

応用例

画像処理

2D FFT は画像のフィルタリング・特徴抽出・テクスチャ解析の基幹計算である。

• 大カーネル畳み込み: ガウシアンぼかし $\sigma > 10$ 程度ではカーネル幅 $> 60$ となり、空間領域の畳み込みより FFT 経由が速い。
• 位相相関: 二画像の位置ずれを検出する手法。FFT 後に正規化クロススペクトル $X \bar{Y} / |X \bar{Y}|$ を取り、IFFT のピークを画素単位で読む。
• 周波数ドメインフィルタ: バンドパス・ハイパス・notch 等を FFT 領域で乗算 1 回で実装。

PDE スペクトル法

周期境界の偏微分方程式は FFT で空間微分が代数演算に変わる:

$$\frac{\partial}{\partial x} \;\xrightarrow{\;\mathcal{F}\;}\; i k_x, \qquad \nabla^2 \;\xrightarrow{\;\mathcal{F}\;}\; -(k_x^2 + k_y^2 + k_z^2)$$

2D Navier-Stokes、Schrödinger 方程式、Cahn-Hilliard 方程式などのスペクトル法ソルバはこの代数化を利用する。 sangi の fft2d, fft3d は周期境界 PDE のプロトタイピングに直接使える。

相関 (cross-correlation)

畳み込み定理の双対として、相関 $r_{xy}[k] = \sum_n x[n] \, \overline{y[n - k]}$ は $R_{xy}[\xi] = X[\xi] \, \overline{Y[\xi]}$ で計算できる。 テンプレートマッチング・遅延推定・自己相関などで FFT 経由の計算が標準である。

sangi API

参考文献

• Eklundh, J. O. (1972). "A Fast Computer Method for Matrix Transposing". IEEE Trans. Computers, C-21, 801–803.
• Frigo, M. and Johnson, S. G. (2005). "The Design and Implementation of FFTW3". Proc. IEEE, 93, 216–231.
• Boyd, J. P. (2001). Chebyshev and Fourier Spectral Methods, 2nd ed. Dover.
• Oppenheim, A. V. and Schafer, R. W. (2010). Discrete-Time Signal Processing, 3rd ed. Pearson. (overlap-add / overlap-save)