こんにちは．Jubatus Advent Calendar 6日目のこの記事ではJubatus近傍探索の高速化をしてみたお話をしようと思います．

きっかけ

Jubatus Casual Talks #4に遊びに行った際に，ヱヂリウム株式会社渡邉さんの「まだCPUで消耗してるの？Jubatusによる近傍探索のGPUを利用した高速化」という発表をみました．
とても面白い発表でJubatusの近傍探索(=ハッシュ探索)はメモリ帯域がボトルネックだからGPUで探索させれば高速だぜっ！という内容でした．

しかしCPUはメモリコントローラを内蔵しており，マルチソケットCPUならメモリはそれぞれのCPUにぶら下がっているため，NUMAを意識した探索を行えば，GPU並のメモリ帯域を実現できるはずです．GPUと異なり搭載メモリサイズもとても安価に大容量にすることが出来ます．

NUMA概要

今回は利用できるマルチソケットマシンがAMDのOpteron 6300シリーズしか無かったため，以下のようなトポロジになっています．

1つのソケットに2つのCPUが載っているちょっと特殊なCPUです．それぞれのCPU(node)間はHyperTransportという12.8GB/sの帯域の出るインターコネクトで接続されています．
1つのCPUあたり4本しかHyperTransportリンクを出せないため，全てのCPUとフルメッシュ接続出来ず，他のCPUコアを経由しないと遠いメモリにはアクセスできないといった構成です．

NUMAとは全てのCPUが同じメモリ空間を共有するのですが，メモリへのアクセスコストが不均等なシステムのことです．今回の環境もNUMAですのでCPU0(node0)からCPU7(node7)に接続されているメモリにアクセスするためには，一度CPU1(node1)やCPU6(node6)を経由する必要があるため，アクセスコストが不均等になっています．

図の吹き出しの通りCPUにぶら下がっているメモリの帯域(25.6GB/s)よりもCPU同士を繋ぐインターコネクトの帯域のほうが狭い(12.8GB/s)ので，遠くのメモリにアクセスした場合のペナルティは結構大きくなりそうです．

GPUとのメモリ帯域の比較

メモリは今回使った環境ではDDR3-1600の2ch接続でそれぞれのCPUと接続されています．DDR3-1600 2chの帯域は25.6GB/sらしいので，全メモリ帯域は204.8GB/s(25.6GB/s*8)になります．
最新のビデオカードだとnVIDIA GeForce GTX 1080で320GB/s，nVIDIA Tesla P100だと732GB/sなのでGPUには届きませんが，GTX1080に近いメモリ帯域ですのでGPUほどアクセスパターンを気にせずとも性能の出るCPUに勝機はありそうです．

現在のJubatusをNUMA環境で動かすと...

Jubatusはクライアントから受け取ったDatumをハッシュ化してモデルに書き込むのですが，NUMA環境では遠くのメモリにアクセスするペナルティが大きいので一番近いメモリに空き容量がある限り，そのメモリに書き込もうとします．Jubatusの近傍探索は最近マルチスレッドに対応しましたので，マルチスレッド動作した場合以下のようにそれぞれのCPU上で動作しているJubatusは最も近いメモリにモデルを書き込みます．

モデル構築時は上記の絵のようにメモリ帯域をフルに活用することが出来ます．

では，探索時はどうでしょう？
Jubatusの近傍探索ではクライアントが指定したDatumと，モデルに書き込まれている全てのハッシュ値との距離を計算する必要があるため，全てのモデルをスキャンする必要があります．マルチスレッドで動作した場合下の図のようにそれぞれのCPU上で動作しているJubatusはいろんな場所にあるメモリにアクセスしてしまうため，NUMA環境の場合，非常に効率の悪いメモリアクセスパターンになってしまいます．これはJubatusはNUMAを考慮しておらず，メモリとの距離を気にせずに適当に割り当てられたメモリ空間に記録されたハッシュ値との距離計算を実施してしまうためです．

今のJubatusで，渡邉さんが目標とする2048bitハッシュ2000万件探索100ms以内を達成できるか評価してみると，195ms となってしまいました．実効メモリ帯域に換算すると，24.5GB/s(2048[bit]/8*2000[万件]/0.195[s]/2^30)となり，理論値である204.8GB/sに大きく届きません．

近傍探索のNUMA対応

ではJubatusの近傍探索をNUMAに対応させるにはどうしたら良いでしょうか？
いろんな方法があるとは思いますが，今回は実装が楽な以下の方法をとってみました．

Jubatusにモデルのメモリ配置を仮想メモリ空間のアドレスで均等分割する機能を追加し，学習後やモデルロード後にその機能を呼び出すことで，仮想メモリ空間と物理メモリの場所を固定します．そして探索時は仮想メモリ空間と物理メモリの場所の対応が分かっているので，それに合わせて探索領域を各スレッドに割り当てる．という方法です．
パッチはこれです．250行ぐらいの修正で済みました．

上記の修正を加えることにより，探索時もモデル構築時と同じように以下のようにメモリ帯域をフルに活用した探索が出来るはずです．

実際に実行してみた結果，以下のグラフのようになりました．

NUMA対応させることで2048bit 2000万件の探索が45msと4倍ほど高速になり渡邉さんの目標の半分以下の時間で探索が出来るようになりました．ただ，メモリ帯域に換算すると106GB/sとなり理論値の204GBの半分という結果でした．(もしかして204GB/sって双方向帯域の合計なのかな？)

おわりに

上記の評価はDDR3という1世代前の4プロセッサマシンで実施しましたが，現世代の構成ではDDR4+8ソケットが(お手頃じゃないかもしれませんが)可能です．この場合の合計メモリ帯域は816GB/sとなりnVIDIA Tesla P100のHBM2を上回る帯域が手に入ります．
メモリ帯域あたりの価格はGPUの方が安そうですが(P100は100万円ぐらい，8ソケットXeonは1000万円以上？)，プログラミングしやすさとかメモリ搭載量(P100は16GBに対して8ソケットXeonは数TB)とか勘案すると…　いや，両方共高すぎでちょっと手が届きませんね．

今回，GPUでまだ消耗しているの？という感じの記事を書きましたが，僕自身はお手軽なGPGPUが大好き派で，仮想メモリ空間共有できるAMDのHSA構想に数年前から期待していますし，nVIDIAもPascal(?)あたりから仮想メモリ空間共有が出来るようになったので，手を出してみたいなと思ってます．
AMDには早くHBM2が載ったZen APUを出して欲しいですね！
なんか終わり方がよくわからない感じになってしまいましたが，この辺で．