Intel AVX-512IFMA52 命令セットによるモンゴメリ乗算の高速化

目次

• 導入
• Intel CPU の AVX-512 命令セットと AVX-512IFMA52 命令サブセット
• モンゴメリ乗算を SIMD 実装する際のデータ配置
• 各種モンゴメリ乗算プログラムのパフォーマンス比較
• 総括
• 参考文献
• 謝辞

導入

Fixstars では暗号アルゴリズムの高速化を行っています。暗号アルゴリズムでは、素体 \(\mathbb{F}_p\) を扱うものが多くあり、その際、\(\mathbb{F}_p\) の上での乗算を実装する必要があります。

この乗算は非負整数の剰余付き乗算 \(z := x \cdot y \pmod p\) となります。素直にこの式に従って計算すると、 \(p\) で割った余りを求めるために除算・剰余算が使用されます。しかし除算・剰余算は、少なくとも現在普及しているアーキテクチャでは効率的ではないため、実装の際は整数除算を実行せずに済むモンゴメリ乗算が採用されます。

本記事では、モンゴメリ乗算のプログラムを 3 種類用意し、パフォーマンス分析を行い、モンゴメリ乗算の SIMD 化の効果を確かめます。

1. AVX-512 命令セットの基本的なサブセット（AVX-512F + AVX-512VL）の組み込み関数を用いた実装

2. 更に AVX-512IFMA52 命令セットの組み込み関数を用いた実装

3. Number Theoretic Library v11.5.1

プログラム 1 および 2 で実装したアルゴリズムは参考文献 [BMSZ2013] のものとなります。

Intel CPU の AVX-512 命令セットと AVX-512IFMA52 命令サブセット

AVX-512 命令セットを搭載したマイクロアーキテクチャは、Knights Landing や Skylake-SP 以降、サーバーやハイエンドデスクトップ向けの CPU に含まれることが多く、ラップトップ向け CPU に関しては Ice Lake 以降となります。とはいえ 2024 年現在、気が付けば手元のラップトップの Intel CPU が Ice Lake 以降のマイクロアーキテクチャだった、ということは珍しくなく、手元でも手軽に AVX-512 命令セット（のサブセット）が使用できる時代になりました。AVX-512 ではマスクレジスタの導入に伴うマスク命令によって、AVX2 よりもかなり SIMD プログラミングがしやすくなったと実感する場面が多いかと思います。

AVX-512 命令セットの新たなサブセットとして AVX-512IFMA52 が Ice Lake マイクロアーキテクチャで実装されました。IFMA は Integer Fused Multiply-Add の略です。Intel のドキュメント等で記載が見られているわけではないですが、倍精度浮動小数点数の仮数部が 52-bit であり、IFMA52 はその回路を利用していると考えられ、参考文献 [GK2016] にもその記述があります。

IFMA52 が使用できない場合、整数乗算の命令は、最長でも vpmuludq 命令によって（32 bit 整数）*（32 bit 整数）=（64 bit 整数）までしか扱えません。一方で IFMA52 を使用できる場合、vpmadd52{h,l}uq 命令によって（52 bit 整数）*（52 bit 整数）=（104 bit 整数）を簡単に扱うことができるようになります。

モンゴメリ乗算をSIMD 実装する際のデータ配置

暗号では数百 ~ 数千ビット長の整数を扱います。すなわち多倍長整数として表現する必要があります。

　１、AVX-512 命令セットの基本的なサブセット（AVX-512F + AVX-512VL）を用いた実装では、整数を 32-bit ずつに分割します。 例えば 150 ビット長の整数を扱う場合、メモリ上の整数の（AoSoA 的な）データ配置は下図のようになります。256-bit ベクトルレジスタを用いた場合、SIMD によって同時に 8 個の整数に対して演算が行われ、8 個の整数はそれぞれ 5 分割されることになります。

　２、更に AVX-512IFMA52 命令セットを用いた実装では、整数を 52-bit ずつに分割します。ただしメモリ上は 64-bit に分割する必要があるため、上位 12-bit 分は 0 をパディングします。例えば 150 ビット長の整数を扱う場合、メモリ上の整数の（AoSoA 的な）データ配置は下図のようになります。256-bit ベクトルレジスタを用いた場合、SIMD によって同時に 4 個の整数に対して演算が行われ、4 個の整数はそれぞれ 3 分割されることになります。

モンゴメリ乗算のアルゴリズムは、この分割数の 2 乗に概ね比例する命令数を要するため、分割数を減らすことで性能向上が期待できます。具体的には、整数を 32-bit ではなく 52-bit ずつに分割することで、アルゴリズムの計算量と各種 SIMD 命令のスループットを考慮すると、モンゴメリ乗算の効率を 85 % 程度高めることができることが見積もれます。

各種モンゴメリ乗算プログラムのパフォーマンス比較

それぞれのモンゴメリ乗算のプログラムをビルドしてパフォーマンスを計測した結果を下図の両対数プロットに記します。AVX-512IFMA52 命令セットを用いた実装では、確かにモンゴメリ乗算が高速となったことが確かめられます。

プロットの見方：

• 両対数プロットであることに注意
• 横軸は多倍長整数のビット長
• 縦軸は一秒あたりに達成できたモンゴメリ乗算の回数。この数値が大きいほどモンゴメリ乗算が高速に処理できる
• 図中の各点：
  • 3 種類のプログラムを実行して得られたパフォーマンス。
    1. AVX-512 命令セットの基本的なサブセット（AVX-512F + AVX-512VL）を用いた実装（図中の青点）
    2. 更に AVX-512IFMA52 命令セットを用いた実装（図中の赤点）
    3. Number Theoretic Library v11.5.1（図中の緑点）
  • 各プログラムに対して、2 種類のコンパイラでビルドしたものを用意。
    • Clang 14.0.0 でビルドしたもの（図中の●）
    • GCC 11.4.0 でビルドしたもの（図中の△）
• 図中の各線：
  • 2 種類の SIMD 実装の達成できるパフォーマンスの最大理論値。この理論値は、理論 CPU 周波数（Turbo Boost technology 2.0 における全コア Turbo Boost の場合）、コア数、ソースコードで記述した組み込み関数が対応する各種 SIMD 演算命令のスループットから計算されます。なおここでは、命令のレイテンシやメモリ命令に関しては十分隠蔽できるものと仮定して、計算に含めていません。
  • AVX-512 命令セットの基本的なサブセット（AVX-512F + AVX-512VL）を用いた実装（図中の青線）
  • 更に AVX-512IFMA52 命令セットを用いた実装（図中の赤線）

補足：

• ビット長が 256-bit 以下の時、AVX-512IFMA52 命令セットを用いた実装の Clang ビルドは概ね理論値通りのパフォーマンスを達成しました。
  • なおビット長が 52-bit の時、理論値をわずかに超えたパフォーマンスが得られました。実は 52-bit の時に限っては上図の赤線の更に 4.5% 向上した性能が、より正確な理論性能でした。上図の理論値曲線に用いたアルゴリズムのスループットは「ソースコード内で記述した、SIMD を扱う組み込み関数の呼出回数に対応した命令数」から求めており、「最適化を施されてビルドされた、実際のバイナリのアセンブリ内の命令数」から求めているわけではないため、このような差異が生じます。実際、後者の方法で算出した命令数が、前者の方法で算出した命令数に対して、命令数が1つ減少したことを確認しており、その影響が 4.5% のスループット向上です。
• ビット長が 512-bit 以上の時、今回計算した、達成できるパフォーマンスの最大理論値に対して SIMD 実装のパフォーマンスがやや悪いことが見て取れます。これはベクトルレジスタが不足するようになり、レジスタスピルのための命令が増加したためです。
• 本計測では巨大なデータ列に対するモンゴメリ乗算のパフォーマンスを計測しています。しかしビット長が 32-bit の場合であってもモンゴメリ乗算のアルゴリズムは演算律速であったため、キャッシュミスの影響は微小と言えます。
• 環境情報
  • OS: Ubuntu 22.04.3 LTS
  • CPU: 11th Gen Intel(R) Core(TM) i7-1165G7 @ 2.80GHz

課題：

• ビット長が 256-bit 以下の時、GCC ビルドは非常に低速でした。プロファイラで分析するとほとんどの時間でストールしているので、GCC を使用したい場合はそれに合わせてプログラムの書き方を工夫する必要がありそうです。
• ビット長が 1024-bit 以下の時、Number Theoretic Library は高速ではないことが分かります。一方で 4096-bit の時は今回の SIMD 実装よりわずかに良いパフォーマンスを達成しました。その原因を探るためには Number Theoretic Library のモンゴメリ乗算のアルゴリズムを精査する必要がありそうです。

総括

今回のモンゴメリ乗算の実装では、元々、主に 100 ビット前後の多倍長整数に焦点を当てたものでした。その範囲で SIMD 実装の Clang ビルドは Number Theoretic Library よりも大幅に高速化され、AVX-512IFMA52 を用いた SIMD 実装が AVX-512IFMA52 を用いない SIMD 実装より、理論値通り約 85% 高速となったことが確かめられました。

また、理論値と実測値の差異に関しても突き詰め、その理由を考察して十分に説明可能なものとなりました。

参考文献

• [BMSZ2013] Joppe W Bos, Peter L Montgomery, Daniel Shumow, and Gregory M Zaverucha. Montgomery multiplication using vector instructions. In Selected Areas in Cryptography–SAC 2013: 20th International Conference, Burnaby, BC, Canada, August 14-16, 2013, Revised Selected Papers 20, pages 471–489. Springer, 2013.
• [GK2016] Shay Gueron, Vlad Krasnov. Accelerating Big Integer Arithmetic Using Intel IFMA Extensions. In 2016 IEEE 23rd Symposium on Computer Arithmetic (ARITH), pages 32-38. IEEE, 2016.

謝辞

本記事を完成させるにあたり、インターン生の宮下敦行様からの実装と分析に関する貢献がありましたことを、感謝申し上げます。

このサイトはスパムを低減するために Akismet を使っています。コメントデータの処理方法の詳細はこちらをご覧ください。