世界最大かつ最高速を実現した小脳モデル

Shoubu上での人工小脳の開発

4月上旬の時点ではShoubu全体の約80%にあたる1008 PEZY-SCが使える状態で、単精度のピーク演算性能は2.6PFlopsであった。そこに、ネコ一匹分に相当する10億ニューロンの小脳を実装した。この小脳を平面に広げると、62mm×64mmの面積となる。これは世界最大規模で、これを時間刻み1msで計算して実時間シミュレーションを行うのは世界最高速で、最も精緻な小脳モデルであるという。

シミュレーションを行うに当たって、この小脳を平らに広げて、1008個のPEZY-SCチップに対応させる。この時、1チップに2mm角の小脳を分担させることになる。

1チップには、1M個の顆粒細胞(GR)1024個のゴルジ細胞(GO)、32個のプルキンエ細胞(PKJ)、32個の分子層介在細胞(BS)、1個の下オリーブ核(IO)、1個の前庭神経核(VN)が含まれることになる。

プログラムとしてはシナプスのコンダクタンスを計算するループと膜電位を計算してスパイクを発生するループがあり、これを全細胞について繰り返す。そして、コンダクタンス計算、膜電位の更新、スパイク生成の部分は、PEZY-SC上で動作するカーネルを呼び出して計算する構造になっている。

プログラムの実装は、ホストコードはC99、PEZY-SCのコードはPZCL(OpenCLの主要機能をサポートするPEZY版の言語)で記述している。このとき、先ず、元のCUDAプログラムをOpenCLに書き直し、それをさらにPZCLに書き直すという2ステップの移植を行うとスムーズにできたとのことである。

とりあえずPEZY-SC上で、8192スレッド並列で動作させるのは容易で、敷居は低かったとのことである。ただし、その版では性能は低く、実時間のシミュレーションには遠く及ばず、性能向上との戦いになったという。

PEZY-SCは次の図のように、4個のPEを含むVillage、4個のVillageからなるCity、16個のCityからなるPrefectureという階層構造で出来ており、4つのPrefectureがPEZY-SCチップを構成している。それぞれの階層にキャッシュが付いており、性能向上のためには、これらをうまく使うことが重要になる。

PEZY-SCチップ1個に実装する顆粒細胞の総数は1Mであるが、これを32×32クラスタに分割し、各クラスタは1024細胞とする。そして1PEには1クラスタ(1024細胞)を割り当てる。クラスタ内の細胞は必ず同時にアクセスされるので、この割り当てはプログラムの処理上も有利である。

また、PEZY-SCのPEは16KiBのローカルメモリを持っているので、これを積極的に利用したいが、1024個の細胞のV、gex、ginhという3つの変数をローカルメモリに置くと、それだけで12KiBを占めてしまう。そうなると、4KiBしかスタックに使えるローカルメモリが残らない。1PEで8スレッドを動かせるが、そうするとスレッドあたり512バイトのスタックで、さずがにこれではスタック領域が足りないので4スレッドだけを使うという実装にしている。さらに、スタックの気持ちになってコードを書き、なんとか4KiBに抑えたという。