PACT 2007 | AA-Sort: A New Parallel Sorting Algorithm for Multi-Core SIMD Processors

古い論文ですが,PACT 2007で発表したSIMDを使うソートアルゴリズムについての解説を書いてみます.これは,VLDB 2015で発表する論文("SIMD- and Cache-Friendly Algorithm for Sorting an Array of Structures") のベースになっており,PACTの論文では整数配列のソートを対象にしていて,VLDBの論文ではそれを構造体配列のソートに拡張しました.PACTの論文で提案した技術のうち,SIMDを使うマージソートは,その後,いろいろな論文で再実装されてCPU上やGPU上でも使われています.

論文とスライドなどはこちらから→http://researcher.ibm.com/researcher/view_person_subpage.php?id=3982

 

論文の概要

やりたいこと

目的は "ソートを高速に実行したい!" の一言につきます.

最近のプロセッサでは,ソートの性能は分岐予測のミスの影響を大きく受けてしまうため,分岐予測の困難な条件分岐をSIMD命令セットにはたいてい含まれているminやmaxという演算命令に置き換えることで高速化できます.ところが,大きな問題として,メモリアクセスへの制約があります.(2007年当時のプロセッサの) SIMD命令は16バイト境界にアラインした,つまり開始アドレスが16の倍数になっている連続した16バイトに対するロードストア以外は性能のオーバーヘッドが大きく,これを避けることが必須でした.(今のCPUやGPUでも,少なくとも不連続アクセスは性能上のオーバーヘッドが大きいので,ストーリーはあまり変わりません.アラインについては昔ほどは気にしなくても良くなっています.) クイックソートなどは,このメモリアクセスの制約の中では高速化することが困難であるため,この論文では2種類の違う特長のあるSIMDに適したソートアルゴリズムを提案して,それらを組み合わせます.

手法の概要

この論文では,データサイズがプロセッサのキャッシュメモリよりも小さい場合に適したアルゴリズムとしてコムソート(comb sort),それよりも大きいデータサイズの場合にマージソートを使い,それぞれをSIMD命令を使って効率よく実行するアルゴリズムを提案しました.全体としては,まず小さなブロックごとにコムソートでソートした後,それをマージしていくという形になります.

コムソートのSIMD

コムソートは,あまり有名ではありませんがバブルソートの簡単な拡張で,計算量をほぼlog(N)になるようにしたものです.挿入ソートとシェルソートの関係がバブルソートとコムソートの関係に似ています.つまりバブルソートでは隣の要素との大小を比較して逆順の場合に入れ替えるという処理を繰り返しますが,コムソートでは隣の要素の代わりに少し離れた要素との比較・交換を行い,イタレーションごとにこの要素間の間隔を縮めていきます.詳細はWikipedia参照 → コムソート - Wikipedia

コムソートはメモリアクセスの順番が入力データに依存しない(ループを回る回数はデータに依存して変わる)ため,クイックソートなどのように入力データに応じて不連続アクセスが必要になる問題を避けることができます.それでも,アラインしないメモリアクセスやループ伝播依存 (loop-carried dependence, 前のイタレーションの結果を次のイタレーションで使うために並列化して同時に実行できない状態)を避ける必要があります.そこで今回の手法では,SIMD命令の幅をS(例えば32-bit整数で128-bit SIMD命令ならS=4)として,N個の入力データをS列 x N/S行の行列と見なして,メモリ内で転置するとアラインや依存の問題なくソートができることを示し,転置順にソートした後で正しい順に直すという手法を提案しました.転置のコストよりもSIMD命令をソートで効率よく使用できる利点の方が大きく,L2キャッシュに納まる範囲のデータに関しては大変高速です.一方で弱点として,イタレーションごとにデータ全体にアクセスするという特徴からキャッシュサイズを超えると極端に性能が劣化するという点があります.

マージソートSIMD

マージソートSIMD命令で効率よく実行する方法としては,bitonicマージソートやodd-evenマージソートが知られていました (例えばGPUTeraSort).bitonicマージソートやodd-evenマージソートは大小比較と交換をデータに依存しない順番で実行することでソートが行えることから,SIMD命令やハードウェア回路での実装に適しています(ソーティングネットワークとか呼ばれます).欠点としては,一段のマージ処理の計算量がO(N log(N)),ソート全体ではO(N log(N)2) になってしまうことがあります.比較を使う普通のマージ処理は,メモリアクセスがデータによっては不連続になってしまうため,SIMD命令には適さない反面,ご存じのように計算量はマージ処理でO(N),それを繰り返すマージソート全体でO(N log(N)) で済みます. 

この論文では,SIMDを使うodd-evenマージ(bitonicマージでも可)と分岐を使う普通のマージ処理のハイブリッド手法を提案しました.擬似コードで書くと以下のようになります.

f:id:inouehrs:20150518132312p:plain

SIMD命令の幅をSとすると(図ではS=4),2つの入力列 (vavb) からそれぞれS個ずつのデータをvectorレジスタ (図ではvMinとvMax) に読み込みます.この2つのvectorレジスタSIMDで実装したodd-evenマージでマージします (図ではvector_mergeメソッド).このodd-evenマージはmin/maxとpermutation(並べ替え)命令を繰り返すだけで,条件分岐無しで実装できます.これにより,もともとはvMinとvMaxのそれぞれの中でS個の値がソート済みだった物が,vMinとvMaxを合わせた2*S個の値としてソート済みになります.ここで小さい側の半分(vMin)をメモリに書き出します.どちらかの入力列のポインタを進めるために,ふたつの入力列の"次の要素"の値をスカラ値として比較し,条件分岐を使って小さかった方からS個のデータをvMinに読み込み次のイタレーションに進みます.これを繰り返すことでマージ処理を行います.各イタレーションで片側からしか読まなくても平気なのかと不思議に思うかもしれませんが,vMaxに含まれているS個のデータは,次の要素の大きい方よりも必ず小さくなります(各入力列はソート済みですから!).そのため,この値は次のS個の出力に入ることは無いため次のイタレーションではまだ不要です.

マージ処理にSIMDを使うことの利点としては,マージ処理がデータ並列でまとめてできるという点に加えて,ポインタを進めるための条件分岐の数が1/Sに減るという点があります.この条件分岐はランダムな入力に対しては,分岐予測が効かないため,分岐予測ミスで大きな性能ペナルティになっていますので,その数を減らすことは性能のためには重要です.(ちなみにSIMD版コムソートでは,データ依存の条件分岐がほぼ0になります)

実際の実装ではメモリアクセスの回数を減らすために,普通の (2-way) マージソートの代わりに4-wayマージソートSIMDで高速化して使用しています.4-wayマージソートは,一度に4つの入力列を1つの出力列にマージすることで,メモリアクセスの回数を減らして,特に複数のコアを使って並列ソートを行うときの性能を向上させます.(このmultiwayマージが後にVLDBの論文の伏線になるとは,当時は思ってもみませんでしたが)

性能

PACTの論文では,このアルゴリズムPowerPC 970MP (MacでいうG5)とCellプロセッサ (プレステ3で使ってたプロセッサ) で実装して評価しました.

PPC970の上では,STLのstd::sortの約3.3倍,IBMのライブラリ製品ESSLのソートの1.8倍の性能になりました.また4コアで2.7倍の性能向上が得られました.一方,SIMDで実装したGPUTeraSort(Bitonicマージソートの改良)は計算量が大きいので,シングルスレッドでの実行時間が長い上に,メモリアクセスの頻度が高く,4コアで2.1倍の性能向上にとどまりました.

Cellの上では,自分のアルゴリズムは16コアで12倍以上の性能向上が得られましたが,GPUTeraSortは16コアで7.1倍の性能向上でした(コア間での直接データ転送とか実装でだいぶ頑張ったのに).

ちなみに,当時は2.4GHzのCell 16コアで32M個の32bit整数のソートで実行時間が0.158秒だったのが,VLDBに使った実装を2.9GHzのSandyBridge 16コアで実行して測ってみた所,0.045秒とかでハードウェアの進歩を感じました.当時はCellの性能が爆速という感じだったんですが・・・.SandyBridge上の並列版std::sortに対してならCellの結果でまだ勝ってます.

論文-その後

自分はPPCとCellでしか実装しなかったのですが,その後Intelの人たちがマージソート部分をx86上で実装してくれて,いろいろな(主にSIGMODとかVLDBとかデータベース系の)論文で使われています.Intelの人たちの重要な発見としては,vectorレジスタ複数組み合わせて,より大きなbitonicマージを実装することで,命令レベルの並列度が上がって,性能がより向上するという点があります.(自分はなぜこれに気づけなかったのか反省します・・)

また,SSEではpermutationに制約が多いので,vectorレジスタのマージにodd-evenマージをうまく実装できないのでbitonicマージを使っています(それでもPPCよりはだいぶ命令数多いですが).自分のVLDBの論文ではmin/maxとrotateを繰り返すだけというとても単純な実装がintel上ではbitonicマージよりも速かったということで,そちらを使っています.(PACTの論文にも,permuteが使えない時の対策として軽く記述はありますが)

Brasov

PACT 2007はルーマニアの奥のほうにあるBrasov(ブラショフ)という町での開催でした.(チェアいわく)東ヨーロッパでCopmuter Scienceのメジャーな国際会議をやるのが初めてとかで,驚くほど気合が入っていたのが大変に印象に残りました.3日間の会議で3回宴会があり,特に2日目は午後を潰してドラキュラ城とかあちこち観光して回りました.それ以降,この時ほど気合の入った国際会議は見たことがありません.ちなみに,この論文は4回落ちて5回目での採録でしたが,折れずに頑張ったかいがあったという感じでした. 

アルゴリズムの名前

AA-sortという名前は実はあまり自分では気に入っていません.査読者に,アルゴリズムには必ず名前をつけろと言われてつけたという経緯があります.Aligned-Accessの略でAAですが,最初はA2と表記していて,C3radix sortあたりのネーミングに影響を受けています.