Cerebras CS-1は、40万もの計算コアをウェハーまるごと1枚に集積した巨大なプロセッサ（WSE - Wafer Scale Engine）をもつスパコンです（写真1）。
写真1 Cerebras社のラボで稼働していたCS-1（2019年12月筆者撮影）
CS-1はAI処理（ディープ・ラーニング）向けのものとして作られました。しかし、CerebrasはSC20[1]において、ニューラルネットワーク分野でなく計算流体力学[2]への適用事例を、そこでの驚異的な処理性能とともに発表しました[3]。
具体的にはNETL（米国エネルギー省の研究所）と協力して、発電所の流動床燃焼（Fluid Bed Combustion - FBC）をモデリングする彼らのソフトウェアの一部をCS-1に移植した際のシステム設計について説明するものです。
計算流体力学と言えば歴史あるスパコン応用分野の1つです。つまりこの発表はCerebrasのアーキテクチャがAI処理以外の、従来的と言っても良いHPC応用分野での新しい可能性を示すものです。
論文を読むのは大変ですが、Cerebras自身によるとても良い解説ブログ[4]が出ています。本記事ではこれ以降、このblogに沿う形で簡単に彼らのアプローチを紹介します。これがCS-1あるいはWSEつまりウェハースケールなプロセッサの今後の展開に興味を持つ人への案内となれば幸いです。
発表の概要
まず、ブログは冒頭で、今回の研究成果について次のようにまとめています。
- 発電所の流動床燃焼をモデリングするNETLのソフトウェアの一部をCS-1に移植したところ、
- NETLで運用されている米国24位、世界82位のスパコンであるJoule 2.0に対して、
- 200倍以上の処理性能を得た
Joule 2.0のシステム構成を見ると、8万6,400コアとあり、なかなか大きなシステムです。Xeon CPUを2つ積んだマシン2,060台と、さらにNVIDIA P100 GPUを2つ積んだマシン100台で構成されます。
対してCS-1は1台、つまり15Uサイズとラック半分以下の大きさですから、ちょっとインパクトがあります。もちろん消費電力も桁違いで、Joule 2.0が450KWなのに対してCS-1は最大20KWと他の比較記事にはあります[5]。
ブログの内容はCS-1のアーキテクチャ概説へと続き、そこではWSEつまりウェハー上にある40万コアの構成について説明されています。計算するコア数ではJoule 2.0に対して数倍程度の優位性しかありませんから、この200倍の性能差がどこから出てきたかが重要です。
CS-1のアーキテクチャ
高い性能が出せた理由を説明する前に、前提知識としてWSEの構造について少し紹介しておきます。といってもおもに2020年1月の筆者自身によるCerebras取材記事「ウェハースケールCPUの誕生――Cerebrasのクレイジーな挑戦」の要約ですので、その内容を覚えておられる方は次の節までスキップしてください。
まずブログにあるWSEの構成図（図1）を見てください。
図1 CS-1ウエハースケールエンジン（WSE）：ブログから引用
図の左に一つの計算コア（tileと呼んでいます）が示されており、それが1つのシリコンダイ領域に51×89個並び、そしてこのダイが12×7個並んでいることが示されています。この12×7ダイが、30cmのシリコンウェハーを丸ごと占領しているようすを図2に示します。
図2 ウェハー上のチップ配置
そして各チップ内部のコアは、図3のような2次元メッシュ構造のネットワークによって相互接続されています。
図3 各チップ内部のコアの相互接続
このタイル状に並べられたコアの群に対して、ニューラルネットワークの層を1つ1つ、図4のように割り付けてデータを流していきます。
図4 ニューラルネットの配置
読者の皆様にはこの構造を頭に置いていただいて、いよいよ計算手法の説明に進みましょう。
関連する計算を局所に閉じ込める
今回の流体計算処理では、空間を600×600×1,500つまり5.6億の小さな直方体（セル）を並べたものとして扱います。600×600を平面に並べ、それぞれ高さ方向に1,500積み上げたと考えてください。そしてこの3Dメッシュ空間を、WSEのプロセッサコアに素直に割り付けます。すばらしくわかりやすい図があるので、先にそれをご覧ください。
図5 3次元の問題を処理要素の2次元ファブリックへマッピング：ブログから引用
シミュレーションでは図の黒い点それぞれが保持するデータ（温度情報など）が、周囲の黒点の状態に影響されて変化していく状況を計算します。赤い矢印（筆者による追加）の先にある点であれば、平面方向に隣接する4つの点と、上下に接する2つの点の情報を使って計算する必要があります。
さて600×600の平面に並ぶ点群は、それぞれWSEのウェハー平面に並ぶ計算コアにそのまま割り当てます。これで40万あるWSEのコアのうちの36万、つまり90％を使うことがわかりますね。高さ方向に積み上げられた1,500セルの計算は、それぞれの位置のコアが担当します。
そして、ブログは各セルの計算に必要なメモリについて、1,500セルでWSEのコアそれぞれが持つ45KBの高速メモリのやはり90%をうまいこと使い切る、と説明しています[6]。
次に計算処理の実体が、疎（スパース）な行列ベクトル演算であることと、各コアが4つの隣接するコアと通信し、その4つの隣接するコアに積み上げられた合計6,000個の値を収集する必要があることが示されています。
WSEのコアはもともと平面方向に隣接する4つのコアと直接通信できるように接続されていることを思い出してください。いかにうまく問題の計算作業をWSEのコアとネットワークシステムに収めたか、感じていただけるでしょうか。
ハードウェアによる支援
この計算モデルで問題となるのは、コア内メモリの読み書き速度と隣接コア間の通信速度です。そしてそこにWSEつまりCS-1の優位性があります。つまりCS-1はこのメモリ読み書きと通信処理に関して、それぞれ特異なハードウェアによって高速計算を実現しているのです。
まずローカルメモリはすべてSRAMで構成されており、1クロックで読み書きすることが可能です[7]。SRAMは通常のCPUがシリコン上に配置する、最速のキャッシュメモリで使われている手法です。WSEも同じくシリコン上のコア、つまり命令実行回路のすぐ隣にメモリを埋め込むために、必然的にSRAMを用いています。当然ですがキャッシュ機構はありません。すべてがキャッシュにある、あるいはキャッシュしかないようなものです。
また隣接コアとの通信もハードウェアによる支援を受けています。これもすばらしくわかりやすい図があるので、先にそれをご覧ください（図6）。
図6 Axのルーティングパターン（blogから引用）
WSEは隣接するコア間の通信メッセージに、ハードウェアに実装された「カラー」と呼ばれる24種のタグを付けることができます。各コアは1サイクルに1ワードずつ、カラーを付けた状態で4つの隣接コアにブロードキャストできます。受信側のコアではこのカラーに紐付けられたスレッドに直接データが受け渡されます[8]。
4つの隣接コアにはそれぞれ異なるカラーを与えてあるため、それぞれの受信に従って即時に適切なスレッドが隣接コアでそれぞれ実行されます。もちろん「即時」に実行できるのはコアが空いている時だけですが、空いていない場合でも完了次第すぐに当該スレッドが実行されます。重要なのはハードウェアがタグを見て、予定されていたスレッドを直接選択し、オーバヘッドなしに実行が始められることです。
これらのCS-1ハードウェアが持つ特別な性質が、クラスタ型スパコンのJoule 2の200倍という性能を出せた理由です。
大量の計算処理をうまくコアに分散配置してコアが持つメモリに収め、必要となる通信を隣接コアに収めることで、1クロックでのメモリ読み書きとオーバヘッドなしの通信（データとソフトウェアの接続）という機構が、コアの計算処理に「待ち」を生じさせないシステムを実現したのです。通常のクラスタ型スパコン、つまり汎用のプロセッサとネットワークシステムで同種のことを実現するのはかなり難しいことがわかるでしょうか。
生まれながらの性質
ところで、ここまで読まれた方は「えらくまた便利な機構がWSEにあったもんだなぁ」と思われたかも知れません。筆者である私も一瞬そう思いました。
しかし考えてみればCS-1は元々データフローマシンです。するとキャッシュなしの低遅延ローカルメモリや、メッセージへのタグ付けと対応するスレッドへの即時受け渡しなどに注力して設計するのは当然とも言えます。つまり高速化に寄与したWSEの各種のハードウェアや仕組みはもともとデータフロー処理を最適化するために作り込まれたもので、今回のNETLの計算流体力学の事例はこの「WSEが生まれながらにもっていた強み」を、既存の問題に「うまくフィットさせた」結果と考えることができます。
こういうアーキテクチャ上の工夫がみごとに収まったのを見ると、ちょっと痺れます。
問題の大きさ、という制約
ここまで良い話ばかり取り上げてきたのですが、もちろんトレードオフがあります。最大の制約は掛けられる問題のサイズでしょう。
まずコアごとのメモリは45Kバイトしかないため、そこに押し込められる数の最大40万倍のセル数にしか空間分割できません。それで必要な精度が得られれば良いですが、それが不足という場合は複数台構成となり、桁違いのオーバヘッドが生じるだろうと想像します[9]。
また今回のように直方体での空間分割にはとてもうまくハマりますが、それ以外の場合はやはり効率が低下するでしょう。そもそも有限体積法以外の何にどの程度うまく適応させられるか自体も確定的なことは言えなさそうです。
一般的なクラスタ型スパコンとの大きな違いはこの点です。つまり普通は「最も効率良く計算が進む領域」から離れるにつれて「計算速度が低下する」ような性能曲線が得られるでしょう。WSEの場合は「とてもうまく行くか、諦めざるを得ないか」に二極化しそうです。
ブログでもその成果について「100万個以上の流体セルをリアルタイムよりも速くシミュレートするのに十分な性能を発揮した世界初のシステムです」としており、あるボリュームに限定した高速性を強調しています。
おわりに
筆者が前掲の取材記事のためにCerebras Systemsのco-founderでありCTOであるGary Lauterbach氏を取材したのは、2019年の末でした（写真2）。
このとき筆者は、これだけ大規模なデータフローマシンは今までなかったのだから、AI以外の新しい応用があり得るのではないかと尋ねています。そのときGaryは「我々はすでに、CS-1にうまくマップできるものでニューラルネットワークではない特定のワークロードを見つけているよ」と答えてくれました。そしてその言葉どおり、一年後に彼らはAI処理ではない、従来的と言っても良い実用的な流体力学計算の問題でインパクトのある結果を出したわけです。
写真2 取材時のGary Lauterbach氏（左）とDirector of ProductのAndy Hock氏（右）（2019年12月筆者撮影）
重要なのは、そのアプローチに特殊な問題設定などがなく、どうやら汎用性があると思えることです。もちろん上に書いたように、万能ではありません。それでも今回の手法に沿う限り、発電所の流動床燃焼に限らず多くの問題に適用できるように思えます。
WSEという新しいチャレンジを応援している筆者としては、元々のAI処理だけでなく、今回のような実用例も増えて欲しいところです。そしてまた異なるアプローチで、しかしWSEの特性がみごとにハマった実装例を出して、もう一度痺れさせて貰いたいものです。
なお、2020年の8月に開催されたHotChipsでは、TSMC 7nmによる 85万コアの第2世代WSE製品がアナウンスされました。これでどんな成果が出てくるか、期待しています。
興味を持たれた方は、ぜひ本記事の下敷きとなったCerebras社のブログをお読みください。
- 参考ブログ：
- AI分野を超えたウェハースケールコンピュータの活用領域: 計算流体力学分野で記録更新
（2020年4月27日追記）
この記事のリリース直後、4/20にCerebrasはコア数を倍増させたWSE-2を搭載したCS-2を発表しました。TSMC 16nmプロセスから7nmに進化し、40万コアから85万コアになりました。その分大規模なモデルを扱えるようになったわけです。
