綿貫 晃雅

Cerebras CS-3上のGNN学習・性能分析

シリコンウェハ全体を一つのプロセッサとして使うWafer-Scale Engine（WSE）を搭載したCerebras CS-3上に、代表的なグラフニューラルネットワークであるGraphSAGEの学習ワークフローを実装しました。学部卒業研究として始め、現在は性能分析と実装改善の両面から発展させています。

GPU向けの機械学習基盤は成熟していますが、GNN学習は、その基盤が得意な規則的な密行列演算だけでは完結しません。頂点ごとに異なる数の近傍をたどり、疎で不規則な集約を行うためです。先行研究 ↗では、大規模グラフを用いる典型的なDGL学習で、GPU利用率が約10%にとどまった例も報告されています。

そこで、WSEのウェハスケールな計算資源をGNN学習に生かせるかを、CS-3とNVIDIA H100の比較を通じて検証しました。同じデータセットで同程度の検証精度に到達した一方、現行の高水準ワークフローではCS-3の学習スループットが大きく下回りました。分析からは、このワークフローがWSE上のデータ配置や疎集約を十分に制御・観測できていない可能性も見えました。

成果をSWoPP 2026で発表し、現在は性能差の原因をハードウェア側とソフトウェアスタック側へ切り分けながら、実装を改善しています。疎で不規則なワークロードを新しいアーキテクチャで効率よく実行する条件を明らかにし、ハードウェアとプログラミング環境の双方の改善につなげることを目指しています。

LLVM ProjectにおけるClangIR・MLIR開発

研究とは別に、個人のOSS活動としてLLVM Projectの開発に参加しています。

OpenCL Cのようなアクセラレータ向け言語では、型、アドレス空間、カーネル引数などの情報が変換途中で失われると、後段のコード生成や実行時処理に利用できなくなります。

ClangIR（CIR）は、Clangが扱うC/C++系のソース表現と、機械語生成に近いLLVM IRの間に置かれ、ソース言語の意味をより多く保つ中間表現です。CIRの基盤であるMLIRは、用途やハードウェアごとに中間表現を定義し、それらを段階的に変換する仕組みを提供します。

私はこの基盤上で、OpenCLカーネル引数のメタデータをCIRで表現・生成し、MLIR上でLLVM IRを表すLLVM dialectを経て、LLVM IRまで伝える機構を実装しました。これらの変更はLLVM Projectへマージされています。また、不正な入力によるクラッシュを診断へ置き換える修正や、回帰テストの拡充にも継続して取り組んでいます。

この段階的な変換基盤により、複数のアクセラレータ向け実装を一つの仕組みで保守しながら、変換される処理の意味と付随情報も追跡できます。将来は、この知見をランタイムやハードウェアが残す実行記録と結び、異種アクセラレータ上のAI処理を検証する基盤へつなげたいと考えています。