ベクトル化されパフォーマンスポータブルなQuicksort
Vectorized and performance-portable Quicksort (opensource.googleblog.com)
要約
Googleは、従来のstd::sortよりも約10倍高速で、最新のCPUアーキテクチャ全体でポータブルなQuicksortの実装をオープンソースで公開しました。この高速化は、SIMD（Single Instruction, Multiple Data）命令を活用し、特にカラム型データベースで重要となるデータレイアウトに最適化することで達成されました。この新しい実装は、多様なデータ型と命令セットをサポートし、既存のアーキテクチャ固有のアルゴリズムをも凌駕するパフォーマンスを示しています。
全文翻訳
本日、C++のstd::sortよりも約10倍高速に配列をソートでき、最先端のアーキテクチャ固有のアルゴリズムをも凌駕しつつ、すべてのモダンCPUアーキテクチャでポータブルなオープンソースコードを公開します。ここでは、これをどのように達成したかを説明します。
まず、背景を説明します。最近のトレンドとして、レコードや「行」のすべてのフィールドを格納するのではなく、特定の列のすべての値を連続して格納するカラム型データベースがあります。これは、SQLクエリの主要な構成要素であるフィルタリングやソートを高速化できるため、このデータレイアウトに焦点を当てます。
ソートは広く研究されてきた分野ですが、どのようにして10倍の高速化を見つけることができたのでしょうか？その答えはSIMD/ベクトル命令にあります。これらは、単一の命令で複数の独立した要素に対して操作を実行します。例えば、AVX-512命令セットを使用する場合、一度に16個のfloat32を操作したり、Arm NEONでは4個を操作したりします。
SIMDについて既に精通している方は、スーパーコンピュータ、機械学習アプリケーションの線形代数、ビデオ処理、またはJPEG XLのような画像コーデックで使用されているのを聞いたことがあるかもしれません。しかし、SIMD操作は独立した要素のみを扱う場合、隣接する配列要素を並べ替えるソートをどのように実行できるのでしょうか？
例えば、256要素の配列をソートする特別な方法があると想像してください。その場合、より大きな配列をソートするためのクイックソートアルゴリズムは、それを2つのサブ配列に分割することからなります。1つは「ピボット」値（理想的には中央値）より小さい要素、もう1つはそれ以外の要素です。その後、サブ配列が最大256要素になるまで再帰し、それらをソートするために特別な方法を使用します。パーティショニングはCPU時間の大部分を占めるため、SIMDを使用して高速化できれば、高速なソートが可能になります。
幸いなことに、最新の命令セット（Arm SVE、RISC-V V、x86 AVX-512）には、パーティショニングに適した特別な命令が含まれています。要素がピボットより小さいかどうかを示す、別途用意されたyes/no値の入力が与えられると、この「compress-store」命令は、対応する入力が「yes」である要素のみを連続したメモリに格納します。その後、yes/no値を論理的に否定し、命令を再度適用することで、要素を別のパーティションに書き込むことができます。この戦略は、AVX-512固有のクイックソートで使用されてきました。しかし、compress-storeを持たないAVX2のような他の命令セットはどうでしょうか？以前の研究では、permute命令を使用してこの命令をエミュレートする方法が示されています。
私たちはこれらの技術を基盤として、3つのアーキテクチャにわたる6つの命令セットにポータブルで、実際には以前のアーキテクチャ固有のソートをも凌駕する、初のベクトル化クイックソートを実現しました。私たちの実装はHighwayのポータブルSIMD関数を使用しているため、プラットフォームごとに約3,000行のC++を再実装する必要はありません。Highwayは、compress-storeが利用可能な場合はそれを使用し、それ以外の場合は同等のpermute命令を使用します。以前の最先端技術（これも32ビット整数に特化していました）とは異なり、16〜128ビットの入力範囲全体をサポートします。
私たちの単一のポータブル実装にもかかわらず、AVX2、AVX-512（Intel Skylake）、およびArm NEON（Apple M1）の両方で記録的な速度を達成しています。100万個の32/64/128ビットの数値に対して、Apple M1で実行される私たちのコードは、499/471/466 MB/sの速度でソート済み出力を生成できます。3 GHzのSkylakeでAVX-512を使用した場合、速度は1123/1119/1120 MB/sです。興味深いことに、AVX-512はAVX2よりも1.4〜1.6倍高速です。これは、追加の労力なしで得られる価値のある速度向上です（HighwayはCPUで利用可能な命令を確認し、利用可能な最良のものを利用します）。AVX2で実行する場合、798 MB/sを測定しますが、AVX2に最適化された以前の最先端技術は699 MB/sしか管理できません。比較すると、標準ライブラリは同じCPUで58/128/117 MB/sしか到達しないため、数値の種類に応じて9〜19倍の高速化を達成しました。
以前は、ソートは高コストと考えられていました。単一のCPUコアで1 GB/sでソートできるようになったことで、どのような新しいアプリケーションや機能が解き放たれるかを見るのが楽しみです。Apache2ライセンスのソースコードはGithubで利用可能であり（質問やコメントがあればissueを開いてください）、私たちの論文は、実装の詳細な説明と評価（256要素の特殊ケースを含む）を提供しています。
Jan Wassenberg – Brain Computer Architecture Research