競プロ備忘録

競プロerの備忘録

爆速なNTTを実装したい

それぞれ長さ N, Mであるような配列 a, bに対して、 c _ {k} = \sum _ {i+j=k} {a _ {i} b _ {j}}といった形の演算を畳み込みと呼び、多項式の乗算の実装等に使えます。

で、そのライブラリのverify用の問題がyosupo judgeにあります(https://judge.yosupo.jp/problem/convolution_mod)。

私はRust使いなのですが、ちょっと前まではRustの最速コード(私ではないです)は150msec以上かかっていたような記憶があります。C++最速のコード(全体の最速コードでもあります)は50msec弱くらいですので、だいぶ遅いです。

一応RustはCやC++並みのパフォーマンスを出せると謳われているのにこれはちょっと気に入らないですね？というわけで、頑張って高速化を目指したという話です。

タイトルが「爆速な畳み込み」ではないのは、結局畳み込みを実装するために内部で使用するNTTの高速化の話が中心になるからです。細かいことを気にしてはいけません（？）

畳み込みの話

細かい解説はできませんが、畳み込み定理というものがあります。本来はフーリエ変換に関連する定理らしいですが、離散フーリエ変換でも成り立ちます。
冒頭の式で言えば、 cの離散フーリエ変換は、 a, bの離散フーリエ変換を要素ごとに乗算したものと等しいという定理です。

したがって、離散フーリエ変換を高速で計算することができれば、

という手続きによって、2つのデータ列を畳み込むことができます。そんな訳なので、離散フーリエ変換をとにかく高速化しようということです。

DFTの話

DFT(Discrete Fourier Transform, 離散フーリエ変換)とは、フーリエ変換を離散的な時間・周波数領域で適用できるように改良した手法らしいです。
DFTは以下のような式で表されます。

F(ω) = \sum _ {t = 0} ^ {N-1} {f(t) e ^ {- \frac{2 \pi tx}{N}i}}

\piは円周率、 eはネイピア数、iは虚数単位です。 Nは普通、データを格納した配列の長さでしょう。

変換後のデータ長を Mとすれば、各 ω = 0, 1, .., M-1に対して上の式を計算するので、愚直にやれば O(NM)の時間計算量となります。一般的に N \ne Mなことがあるのか知りませんが、以下 N = Mの前提で話を進めます。

DFTにはIDFT(Inverse Discrete Fourier Transform, 逆離散フーリエ変換)という逆変換があり、これは以下のような式で表されます。

f(t) = \frac {1}{N}\sum _ {ω = 0} ^ {N-1} {F(ω) e ^ { \frac{2 \pi tω}{N}i}}

各データに掛けるeなんちゃらが共役な複素数になって、最後に \frac {1}{N}を乗じるだけということです。この \frac{1}{N}については、離散フーリエ変換の導出の過程で、正変換側につくことや、正変換・逆変換両側に \frac{1}{\sqrt {N}}としてつけることもあります。
ただ、畳み込み演算では正変換を2回、逆変換を1回をやるので、逆変換側に押し付けてしまったほうが得な気がします。多くの競プロの記事でも逆変換側に押し付けているので、それに倣うこととます。

(※このパートで1のN乗根について述べていますが、導出の過程を誤っているとTwitterでご指摘をいただきました。そのうち書き直すので読み飛ばしてください。ごめんなさい。)
どちらの変換でも eの肩に複素数を乗せた指数関数をかけていますが、よく見ると e ^ {2 \pi i}が混じっています。みんな大好きオイラーの公式から、 e ^ {i \pi} = -1なので、これは1に等しいです。さらに \frac {tω}{N}乗されているので、これは 1の N乗根を tω乗したものと言えます。

ここで、 1の N乗根 e ^ {- \frac{2 \pi}{N}i}を W _ {N}と表すことで、正変換は以下のようにも表せます。

F(ω) = \sum _ {t = 0} ^ {N-1} {f(t) W _ {N} ^ {tω}}

1の N乗根は、複素数平面上の単位円周上を回転して再度 1に戻ってくるので、回転因子(Twiddle Factor)とも呼ばれます。逆変換における回転因子は正変換のそれと共役な複素数なので、逆方向に回転します。

DFTはこの回転因子を用いて、変換前の次数が N、変換後の次数が Mであれば、 M行 N列の行列によって表現することもできます。例えば、変換前後ともに次数4のDFTの変換行列は以下のようになります。

\begin{pmatrix} W _ {4} ^ {0} & W _ {4} ^ {0} & W _ {4} ^ {0} & W _ {4} ^ {0} \\ W _ {4} ^ {0} & W _ {4} ^ {1} & W _ {4} ^ {2} & W _ {4} ^ {3} \\ W _ {4} ^ {0} & W _ {4} ^ {2} & W _ {4} ^ {4} & W _ {4} ^ {6} \\ W _ {4} ^ {0} & W _ {4} ^ {3} & W _ {4} ^ {6} & W _ {4} ^ {9} \end{pmatrix}

FFTの話

FFT (Fast Fourier Transform, 高速フーリエ変換)はDFTを高速化するためのアルゴリズムです。

有名な手法としてはCooley-Tukey FFTがあり、これは次数を落としながら小さいサイズのDFTを行い、マージして元のサイズでの結果を得るという、分割統治によるアルゴリズムです。この記事でもこれをひたすら突き詰めていきます。

ほかにも、次数が素数であるときに適用できるRader FFTや、互いに素な2数の積であるときに適用できるPrime Factor FFTなどもあるらしいです。そのうち調べてみたいです。

Cooley-Tukey FFT

前の節で述べたように、分割統治によって次数を下げていくことで計算量を落とします。次数は合成数でさえあれば適用できますが、多くの場合2の冪乗で揃えます。以下、 N = 2 ^ l( lは正の整数)とします。

DFTの式について、 tを偶奇で分割してみます。

\sum _ {t = 0} ^ {N-1} {f(t)W _ {N} ^ {tω}} = \sum _ {t=0} ^ {\frac{N}{2}-1} {f(2t) W _ {N} ^ {ω(2t)}} + \sum _ {t = 0} ^ {\frac{N}{2}-1} {f(2t+1) W _{N} ^ {ω(2t+1)}}

ところで、 pq = Nであるとき、 W _ {N} ^ {p} = W _ {q}となります。もともとの定義に立ち返れば、 W _ {N} ^ {p} = e ^ {- \frac {2 \pi p} {N}i} = e ^ {- \frac {2 \pi} {q}i} = W _ {q}となるのが確かめられます。

これを用いて上の式をさらに変形することで、

\sum _ {t = 0} ^ {N-1} {f(t)W _ {N} ^ {tω}} = \sum _ {t=0} ^ {\frac{N}{2}-1} {g(t) W _ {\frac{N}{2}} ^ {tω}} + W _ {N} ^ {ω} \sum _ {t = 0} ^ {\frac{N}{2}-1} {h(t) W _{\frac{N}{2}} ^ {tω}}

となります(ただし、 g(t), h(t)はそれぞれ f(t)の偶数番目、奇数番目を集めた列)。

右辺をよくみると、 g(t), h(t)に対するDFTを行い、 h(t)への適用結果に回転因子を乗じて足す操作になっています。

分解は logN段で終わります。最下段では次数1のDFTが実行されますが、これは何もしないことと同じなので、 O(1)です。結果のマージは各段 O(N)で可能なので、結果的に計算量を O(NlogN)に削減できたことになります。

素朴な実装としては、以下のようになります。愚直なDFTの実行結果と比較していますが、浮動小数点数の比較なので、差が一定値以下のときパスとしています。

Rust Playground

ω \ge \frac{N}{2}のときはどうやってマージするの？と一瞬迷いますが、よくよく考えると、 W _ {\frac{N}{2}}は \frac{N}{2}乗したところで 1に戻りますから、 ω - \frac{N}{2}のときと一致します。よって、 ω - \frac{N}{2}の結果を再利用できます。 W _ {N}についても、 W _ {N} ^ {k + \frac{N}{2}} = -W _ {N} ^ {k}ですので、同様に ω - \frac{N}{2}の結果を利用できます。

NTTの話

これまで複素数でDFTを計算してきましたが、浮動小数点数の演算が必須のため、誤差が出ます。誤差が出るのは嬉しくないのでどうにかしたいですが、どうやらDFTは、1のN乗根が存在するような数のグループであれば適用できるようです。

そこで突然ですが、正の整数を素数 pで割った余りで構成される数のグループに適用することを考えます。これがNTT(Number Theoretic Transform, 数論変換)です。

このようなグループは必ず、 p-1乗して初めて 1と合同になる数( pの原始根)を持つので、これを使えます。

しかし素数なら何でもよいわけではなくて、FFTを適用するには 2 ^ l乗して 1になる数が必要なので、 p-1の素因数になるべく多く 2を含むような素数が望ましいです。このような素数をNTT Friendly素数と呼び、競プロでよく使われる998244353などが該当します。

上で述べたFFTのアルゴリズムをそのまま用いてNTTを実装すると、以下のようになります。逆変換であるINTTも実装して、愚直な畳み込み演算と結果を比較しています。
Rust Playground

非再帰化

よく知られているように再帰関数は遅いことが多いです。これまでNTTは再帰関数として実装していたので、非再帰化します。

頑張って再帰関数の構造を紐解いて無理やり非再帰化しても良いですが、FFTのシグナルフローダイアグラムと呼ばれるものを用いれば、容易に非再帰のルーチンを導くことができます。フローダイアグラムが蝶の羽のような形をしているので、導かれたルーチンはバタフライ演算と呼ばれます。

フローダイアグラムはググればたくさん出てきますが、フローの最初と最後のどちらが大きいバタフライかで2種類のバリエーションがあるはずです。
上で導いたFFTのアルゴリズムは tを偶奇で分割する方法でしたが、これは時間間引き(Decimation In Time, DIT)というバリエーションにあたり、対応するフローダイアグラムは最初のバタフライが小さいほうです。

フローを流す前のデータの添え字を見てみるとヘンテコな並びをしていますが、これはビット反転順序という並びになっています。たとえば N = 8であれば、

(0, 1, 2, 3, 4, 5, 6, 7) \\→ (000, 001, 010, 011, 100, 101, 110, 111) \\→ (000, 100, 010, 110, 001, 101, 011, 111) \\→ (0, 4, 2, 6, 1, 5, 3, 7)

というわけです。なので、正しい結果を得るためには、ルーチンの前にビット反転並べ替えをする必要があります。

実装は以下のような形になります。
Rust Playground

再帰関数で実装していた際には、データをインデックスの偶奇で分割するために作業用のメモリを必要としていましたが、この実装では元のデータ列をそのまま書き換えて結果を得ることができます。このような演算はin-place演算と言われます。
再帰実装の際はin-placeでない代わりに並べ替えは必要としませんでした。このように、作業用メモリを必要とする代わりに自然な並びで結果を出力するアルゴリズムとしてStockham FFTというものがあるようです。名前を聞いたことがあるくらいで詳しくは知らないので、これもそのうちちゃんと調べたいです。

ビット反転並べ替えのためにだいぶ非効率なことをしていますが、これは簡単に書ける方法をとったに過ぎません。実際にはより効率的な方法があるので、それを調べて使うのが良いでしょう。
大浦さんという方のFFTの解説ページに載っていたものが、今のところ試した中では最も速かったです。「大浦　FFT」などとググれば出てきます。

時間間引き・周波数間引き

先ほど「時間間引き」というCooley Tukey FFTのバリエーションが出てきましたが、これは tがDFTにおける時間変数を表し、これを分割して導出するためにそのような名前がついています。

DFTによって変換された関数は周波数関数になりますが、周波数変数 ωによって分割するバリエーションもあります。これが周波数間引き(Decimation In Frequency, DIF)であり、フローダイアグラムはフローの最初のほうが大きなバタフライになります。

こちらはフローを流しきった後のデータがビット反転順序となるので、ルーチンの最後でビット反転並べ替えをすればよいです。

これを実装した結果が以下です。時間間引き・周波数間引きのNTTの結果を比較して、一致することを確かめています。
Rust Playground

ちなみに、時間間引き・周波数間引きのルーチンは、互いに転置写像の関係にあり、機械的に導く方法があるようです。もっとちゃんと調べてみたいです。
(ネタ元：FFT の転置写像 – 37zigenのHP)

ビット反転並べ替えの削除

これまでみてきたように、DITはルーチンの冒頭、DIFはルーチンの末尾でビット反転並べ替えが必要でした。

ところで、畳み込み演算の用途であれば、DFTとIDFTは常に対で使います。また、ビット反転並べ替えは2回行えば元通りです。では、DFTにはDIFのルーチンを、IDFTにはDITのルーチンを用いれば、ビット反転並べ替えを削除できそうではないでしょうか？

これは実際、正しい結果を得ることができます。実装は以下の通りです。愚直な畳み込みの結果と比較しています。
Rust Playground

基数4のアルゴリズム

※Twitterでご指摘いただきましたが、以下の乗算回数の考察には誤りがあります。
複素数でのFFTの際には W _ {4}の乗算が実部・虚部の入れ替えに相当するため、乗算が減ります。しかしNTTの際には W _ {4}の乗算は必要です。
下記のコード例でいうところのroot, root2, root3の乗算を含めて考察できていないのですが、これを含めると基数4のアルゴリズムであっても乗算の回数は減っていません。
従って、最外側のループの段数が減ること以上の効果はないかもしれません。
一応元の記載も残しておきますが、NTTの実装の改善には役立たない可能性があります。ごめんなさい。
以下、元の記載------------------

ルーチンの最内側のループの内部に注目すると、これはサイズ2のNTTです。これをもっと大きなサイズにできたら速くなりそうじゃないでしょうか？

サイズ2のNTTでは乗算1回、加算2回が必要ですから、サイズ4のNTTをサイズ2に分割して実行する場合は乗算4回、加算8回が必要です。

一方、分割しない場合は乗算1回、加算8回でできます。これはDFTの変換行列を眺めるとわかります。

\begin{pmatrix} W _ {4} ^ {0} & W _ {4} ^ {0} & W _ {4} ^ {0} & W _ {4} ^ {0} \\ W _ {4} ^ {0} & W _ {4} ^ {1} & W _ {4} ^ {2} & W _ {4} ^ {3} \\ W _ {4} ^ {0} & W _ {4} ^ {2} & W _ {4} ^ {4} & W _ {4} ^ {6} \\ W _ {4} ^ {0} & W _ {4} ^ {3} & W _ {4} ^ {6} & W _ {4} ^ {9} \end{pmatrix} = \begin{pmatrix} W _ {4} ^ {0} & W _ {4} ^ {0} & W _ {4} ^ {0} & W _ {4} ^ {0} \\ W _ {4} ^ {0} & W _ {4} ^ {1} & W _ {4} ^ {2} & W _ {4} ^ {3} \\ W _ {4} ^ {0} & W _ {4} ^ {2} & W _ {4} ^ {0} & W _ {4} ^ {2} \\ W _ {4} ^ {0} & W _ {4} ^ {3} & W _ {4} ^ {2} & W _ {4} ^ {1} \end{pmatrix} = \begin{pmatrix} 1 & 1 & 1 & 1 \\ 1 & W _ {4} & -1 & -W _ {4} \\ 1 & -1 & 1 & -1 \\ 1 & -W _ {4} & -1 & W _ {4} \end{pmatrix}

元のデータ列が (x _ {0}, x _ {1}, x _ {2}, x _ {3})であれば、 x _ {0} + x _ {2}, x _ {0} - x _ {2}, x _ {1} + x _ {3}, (x _ {1} - x _ {3})W _ {4}を前計算してから4回加減算すればよいので、結果、乗算1回、加算8回となります。

ルーチンの最内側のNTTの次数によって基数Nのアルゴリズムなどと言ったりするので、これは基数4のアルゴリズム、今までの実装は基数2のアルゴリズムとなります。

基数は大きくとるほど乗算の回数を減らすことが可能ですが、1度にロードしないといけない配列要素が増えるためキャッシュ効率的にデメリットが出てきたり、実装がやたら複雑な割に効果が薄くなったりしがちなので、基数4のアルゴリズムを使うことが多いようです。

また、基数2,4のアルゴリズムを組み合わせてフローダイアグラムを変則的に進めることで乗算回数を最適化するSplit-radix FFTというアルゴリズムもあります。残念ながら私の力では実装できなかったので、興味ある方は自力で調べてみてください。

以下、基数4のアルゴリズムの実装です。
Rust Playground

データ長が 4^lで表せない場合(つまり N = 2 ^ {2l+1}の形になる場合)、基数2のNTTで帳尻合わせをする必要があります。データ長を調整しても良いですが、メモリ効率的に嬉しくないですし、1のN乗根が取れない長さになるとそもそも正しい結果が得られません。

ルーチンの最内側の形をみると、NTTのほうはDFT行列を単純に乗じた形ではないし、INTTのほうはrootを変なかけ方をしているしで、微妙に頭が混乱しますが、基数2のアルゴリズムを力ずくで頑張って紐解くと、この形にできます。

回転因子をキャッシュする

ルーチンの中で使用する回転因子は、データ長 Nに対して N個しかありません。なので、キャッシュしておけば回転因子を求めるための掛け算を減らすことができます。

頑張ってやるだけなので、実装は省略します。

回転因子が 1であるときの乗算を削除する

見ればわかることですが、回転因子は最初は必ず 1です。
1を掛け算するのは無駄なので、削除します。

ちゃんと数えてはいないですが、フローダイアグラムをみても 1の掛け算はそれなりの数がありますから、効果が期待できます。
-1であるときも削除しても良いかもしれません。

やるだけなので、実装は省略します。以後の実装例でも無駄に実装が複雑になるので、1や-1の乗算を場合分けして削除することはしていません。

フローダイアグラムの添え字をすべてビット反転順序にする

先ほどビット反転並べ替えの文脈で登場した大浦さんの解説サイトで紹介されていた方法で、AtCoder Libraryで採用されている方法でもあります。

この方法では回転因子の指数の順序もビット反転順序になりますが、その代わり最内側のループではずっと同じ回転因子を乗じることになるため、最内側の回転因子の更新のための乗算がごっそりなくなります。

では、外側のループではどうやって回転因子を更新すればよいでしょうか？

大浦さんのサイトでは、単純に今欲しい指数のビット反転を行って回転因子を計算していますが、AtCoder Libraryでは(少なくとも私からすると)非常に天才的な方法で計算します。

AtCoder Libraryの実装を見ると、fft_info()という構造体があり、その中でrate2, rate3, irate2, irate3という配列を定義しているのがわかります。それぞれ、基数2のNTT、基数4のNTT、基数2のINTT、基数4のINTTで利用します。

例えばrate2の要素は、以下のような値となっています。

rate2[i] = W _ {2 ^ 2} ^ {-1}W _ {2 ^ 3} ^ {-1}W _ {2 ^ 4} ^ {-1}...W _ {2 ^ {i+2}}

これを使って、例えばサイズ8のNTTの1段目に必要な回転因子を計算するとします。その場合、 W _ {8} ^ {0}, W _ {8} ^ {2}, W _ {8} ^ {1}, W _ {8} ^ {3}が欲しいですが、以下のような手続きでこれを得られます。

  • 初期値 wを 1 (= W _ {8} ^ {0})とし、 0番目の回転因子とする
  • インデックス iを 0とする
  • w = w * rate2[i.trailing\_ones()]で更新する
  • 更新後の wは、 i+1番目の回転因子である
  • iを 2まで増やしながら上記の手続きを繰り返し、 0から 3番目の回転因子を順に得る

実際、

i = 0, w = 1 \\ i = 1, w = 1 * rate[0] = W _ {4} = W _ {8} ^ {2} \\ i = 2, w = W _ {8} ^ {2} * rate[1] = W _ {8} ^ {2} * W _ {4} ^ {-1} * W _ {8} = W _ {8} ^ {2} * W _ {8} ^ {-2} * W _ {8} = W _ {8} ^ {1} \\ i = 3, w = W _ {8} * rate[0] = W _ {8} * W _ {4} = W _ {8} * W _ {8} ^ {2} = W _ {8} ^ {3}

となっており、欲しい回転因子の列が得られます。
凄いですね。競プロの典型か何かなのでしょうか？全く理解不能です…

実装例は以下のようになります…といっても、AtCoder Libraryと同じ実装なので、AtCoder Libraryを見ると良いでしょう。
Rust Playground

Rustに限らずC++でも、コンパイラオプションなどで拡張しない限り、コンパイル時計算できる演算量は限られています。しかしこの手法ではキャッシュがたかだか数十要素なので、コンパイル時計算も容易にできます。

多項式乗算の実装ではキャッシュを何度も使いまわしたい場面が多いでしょうから、とても嬉しいですね。
ただ、現状のAtCoderのRustバージョンでは、const fnの中でif分岐、whileループすらできないため、数十要素のキャッシュも簡単ではありませんが…

また、この手法ではフローダイアグラムの形が変わり、DIT、DIFともに向きが逆転します。ビット反転並べ替えが必要なパートも変わるため、NTTにはDITを、INTTにはDIFを使う必要があります。

ところで、フローダイアグラムの最初が大きなバタフライである(素朴なアルゴリズムではDIFに適用する)バタフライ演算を、Gentleman Sandeバラフライ、そうでないほうをCooley Tukeyバタフライと呼びます。
この名前は、バタフライの形に紐づいているのか、それともDIT、DIFに紐づいているのか、どちらなのでしょうか？いまだに分からないままです。
この節のアルゴリズムでは、フローダイアグラムの形とDIT・DIFの紐づきが変わっているので、どちらのバタフライ演算をCooley Tukey, Gentleman Sandeと呼べばいいのか分からないのです…ご存じの方、教えてください。

Six-step FFT

ルーチンの核となるループに注目すると、特に大きいサイズのNTTでは、かなり離れた位置にある要素を何度も取得しなければならない場面があり、キャッシュ効率的に嬉しくなさそうです。
キャッシュ効率を向上させるアルゴリズムとして、Six-step FFTというものがあります。この導出のためには、以下のようにDFTの式を変形します。

F(kN _ {1} + l) \\ = \sum _ {pN _ {0} + q = 0} ^ {N - 1} {f(pN _ {0} + q)W _ {N} ^ {(kN _ {1} + l)(pN _ {0} + q)}} \\ = \sum _ {q = 0} ^ {N _ {0} - 1} {\{(\sum _ {p = 0} ^ {N _ {1} - 1} {f (pN _ {0} + q)W _ {N _ {1}} ^ {lp}})W _ {N} ^ {lq}\}W _ {N _ {0}} ^ {kq}}

ただし、 N = N _ {0}N _ {1}です。
入力データ列を N _ {0}行 N _ {1}列の行列として解釈すると、この式は以下のような手続きによる計算を意味します。

  • データ行列を転置する
  • 各行 N _ {1}サイズのFFTを実行する
  • W _ {N} ^ {lq}をすべての要素に乗じる
  • データ行列を転置する
  • 各行 N _ {0}サイズのFFTを実行する
  • データ行列を転置する

N _ {0}, N _ {1}を \sqrt {N}に近くなるようにとれば、FFTの実行サイズは非常に小さくなり、キャッシュ効率が向上することが見込めます。
N = 2 ^ {l}なので、一般に N _ {0} \le N _ {1}とすれば、 N _ {0} = N _ {1}か 2N _ {0} = N _ {1}かの2通りがあります。

前者の場合は正方行列ですから、容易にin-placeでの行列転置が可能です。
後者の場合は正方行列ではなく、一般的に正方行列でない行列のin-place転置は容易ではありません。色々悩んだのですが、私は結局以下の資料からアルゴリズムを拝借しました。
[2011.11524] Speeding up decimal multiplication
要するには、 N \times 2N行列の場合、 N \times Nの部分行列2つを転置し、行を並び替えることで全体として転置できるというわけです。正方行列については上記の通りin-place転置可能、行の並び替えは O(N)サイズの作業メモリがあれば可能です。

また、FFTの結果はビット反転順序で返ってくるので、普通に実装すればFFTの実行ごとにビット反転並べ替えをする必要があります。
これを省くには、先に4ステップ目の行列転置を行い、 lをビット反転した指数を持つ回転因子を順に乗じればよいです。ここでも、前節の方法を応用できます。
ビット反転並べ替えを省いた場合、やはり全体としてビット反転順序で結果が返ってきます。

さらに、DFT・IDFTを対で使う場合は、それぞれ6ステップ目・1ステップ目の行列転置を省けます。転置の転置は元の行列なので、それは確かにそうですね。

さてここまで述べておいてなんですが、私が実装した結果では、Six-step FFTでの高速化は達成できませんでした。
真の理由はわかりませんが、高速な行列転置が実装できなかったことや、競プロで扱うようなデータサイズでは小さすぎてキャッシュ効率の向上度合いが大きくないというのが理由だと思っています。
2ステップ目、5ステップ目のFFTはそれぞれ、各行の中以外でのデータの依存関係はありませんから、並列実行が容易な気がします。もし並列実行が効果を発揮するような環境であれば、また結果は変わるのかもしれません。

実装は省略します。

ベクトル命令の使用

最終手段みたいなものですが、AVX2までのベクトル命令を用いた並列化を実装します。
なんでAVX512じゃないの？といえば、RustではstableでAVX512命令を使えないですし、そもそもAVX512をサポートするCPUを載せたPCやサーバなんて持っていないからです。エミュレータ的なものがあると聞きますが、コンテスト中にそんなの使うのは現実的ではないので、AVX2までで我慢します。

AVX2の場合、32bit整数8個の演算を同時に行えます。最内側のループのサイズが8以上であればレジスタへのロードは容易にできるので、効果は高そうです。

難点は、AVX系命令にもSSE系命令にも、整数除算・整数剰余算命令がないことです。
そこで、Montgomery剰余乗算を実装します。Montgomery剰余乗算のアルゴリズムを用いれば、加減乗算のみで乗算剰余を計算することが可能です。
Montgomery剰余乗算については既存の優れた記事がたくさんあるので、調べてみてください。(優れてはいないですが、私の既存記事もあります：Montgomery剰余乗算の学習メモ - 競プロ備忘録)

やるだけなので説明することもあまりありません。実装は以下の通りです。
Rust Playground
INTTでの \frac{1}{N}の乗算、データ列のマージも、 N \ge 8であれば、ベクトル化して高速化しています。入出力がある際は、Montgomery表現との間の変換・復元も8要素まとめて行うことで高速化できます。

上の実装はちょっと手抜きしていますが、本気でやれば N < 8の場合を除いて、最内側のループサイズが8未満の場合も含めてすべてのルーチンをベクトル化することができます。
この場合は単純に連続要素をロード・ストアすることができませんが、要素のロードはgather命令1つで可能です。計算後の要素のストアにはscatter命令を使いたくなりますが、AVX2までにはscatter命令が何故かありません。なので、shuffle, blend, pack, unpackなどを駆使して値を並べ替え、連続要素としてstoreする必要があります。
私のライブラリの実装ではそのあたりも頑張ってベクトル化しているので、暇な人は探してみてください。

あんまり適当な実装をし過ぎると、ベクトルレジスタが枯渇して意図せぬところにロード・ストアが挟まれ、速度が低下することがあります。その場合は計算の順序を入れ替えたり、ブロックを切ってライフタイムを制限したりといった工夫が必要になるかもしれません。

結果・まとめ

以上のような高速化の工夫を行った私のライブラリの実装で、冒頭に紹介したジャッジに投げてみたところ、最も高速なときで69msecとなりました。
C++の最速実装にはさすがに及びませんでしたが、Rustでの提出コード中では最速の実装となりました。爆速といえるかは微妙ですが、高速とはいってもよさそうでしょうか？
実際には入出力もそこそこ重いですから、その高速化が達成できないと、C++の実装には勝てない気がしています。

問題点としては、ちょっとコードが長すぎです。AtCoder Library Practice Contestにも全く同じ問題があるため、そちらのジャッジにも投げてみたのですが、cargo equipのminify機能を利用してもなお72000Byteです。こどふぉのジャッジではコード長制限オーバーです。
冗長なコードや無駄なコードもまだ多いので、コード削減は今後の課題になりそうです。

ちなみに、簡潔で高速なRust実装はyosupo judgeや上述のAtCoderのジャッジに存在していて、toomerさん(HBitさんと同一人物らしいです)の実装やtosさんの実装は、SIMDの実装がないにも関わらず非常に高速でかつコード長が短いです。
コンパクトなライブラリを志向する方は、そちらの実装を参考にするのが良いと思います。

理解しきれていないアルゴリズムもあります。「フローダイアグラムの添え字をすべてビット反転順序にする」の章については、フローダイアグラムの形から直感的には理解しているのですが、原理を論理的に示せてはいません。これからも既存の論文や記事を探すなどの努力が必要そうです。

参考にした資料

小野測器さんのコラムやる夫で学ぶディジタル信号処理FFT (高速フーリエ・コサイン・サイン変換) の概略と設計法OTFFT: High Speed FFT Library「Six-Step FFTとEight-Step FFT」Speeding up decimal multiplicationその他多数、参考にさせていただきました。