インタビュー

富岳のディープラーニング処理を支えるJITコンパイラ「Xbyak_aarch64」誕生秘話

この記事を読むのに必要な時間:およそ 5 分

TOP500,HPCG,HPL-AI,Graph500での世界1位獲得新型コロナウイルス対策を目的とした試行利用など,話題に事欠かないスーパーコンピュータ「富岳⁠⁠。そのディープラーニング処理を高速化するには,あるOSSの存在が必要不可欠でした。それが,サイボウズ・ラボ⁠株⁠の光成滋生氏が開発したx86/x64向けC++ JITアセンブラ「Xbyak」の設計思想をベースに,光成氏の助言のもと⁠株⁠富士通研究所が開発したArm向けのC++ JITアセンブラ「Xbyak_aarch64」です。Xbyak_aarch64は,富岳上でのディープラーニング処理を実現するキー技術のひとつです。

本記事では,Xbyakの開発者である光成氏を中心に,⁠株)富士通研究所の上席研究員であり,Linuxカーネルへのコアコミッターでもある小崎資広氏を聞き手役として,同研究所シニアリサーチャーの川上健太郎氏,同研究所プロジェクトマネージャーの福本尚人氏を交えて対談を行いました。開発の紆余曲折を赤裸々に語っていただいたほか,今後のOSSアップストリーム方針についても言及いただきました。

※)
本件の技術詳細については⁠株⁠富士通研究所より解説記事が公開されています。併せてご覧ください。

共同開発のきっかけ

光成:きっかけは小崎さんからお声をいただいたことでした。

小崎:僕が富士通研究所の研究員をやりつつサイボウズの技術顧問もしていて,両方に籍があったんです。そして,富士通研究所側の大きなミッションである富岳の可能性を探していくうちに,Xbyakにたどり着いたんです。

福本:富岳に搭載されている「A64FX」など,Arm系CPUでディープラーニングをやろうと思うと,別途ライブラリが必要なことが調査でわかっていました。ちょうどそのときに,IntelがoneDNN※1という名前のライブラリをオープンソースで公開していて,これをベースに始めようという話になりました。

そのoneDNNの中で利用されていたのが光成さんが開発していたXbyakで,ライブラリを作るためには,XbyakをArm向けに移植しなければいけないという話が富士通研究所内で持ち上がっていました。そこで小崎さんから,Xbyakの開発者なら知ってるよということで,つないでもらいました。

小崎:そうですそうです。それで私が光成さんにご相談した結果,快諾していただいて,一緒に共同研究しましょうという枠組みができました。

光成:そういえば,小崎さんの紹介で一緒にお会いするときに福本さんがいきなり凄く細かい最適化の話をはじめて。めちゃめちゃお仲間がいるなあと思いました。

福本:細かい最適化ができないとディープラーニングのライブラリを作るなんてできないだろうと思ったので,こういった素養のある人が集まっていますよ,というアピールでした。⁠笑)

光成 滋生

サイボウズ・ラボ⁠株⁠

さまざまなCPUや特殊プロセッサ向けのmp3やMPEG4などのコーデック開発に関わったのち,2007年サイボウズ・ラボ入社。ペアリングライブラリの世界最速実装を達成したほか,Ethereumなどのブロックチェーン系プロジェクトで利用されているBLS署名ライブラリの開発を行った。Xbyakは入社前から続けている個人開発プロジェクトである。

小崎 資広

⁠株⁠富士通研究所 上席研究員,サイボウズ(株⁠⁠ 技術顧問

しばらく家電の世界で働いた後,2005年に富士通研究所入社。LinuxやRubyのコア開発者として,一貫してオープンソースの世界で働き続けている。2017年にサイボウズ技術顧問就任。

福本 尚人

⁠株⁠富士通研究所 プロジェクトマネージャー

2012年より,富士通研究所でプログラムの高速化に関する研究開発に従事。最新ハードウェア向けに,HPCアプリ,Linpack,行列積などの高速化を行ってきた。Intel CPUを搭載したスーパーコンピュータにおいて,TOP500(スパコンランキング)向けの高速化や大規模分散並列深層学習の高速化に関わり,日本一,世界一性能の達成に貢献している。2019年よりマネージャーとして,Arm向けディープラーニングのソフトウェアスタック開発を行っている。

※1)
Intelが開発したx86向けディープラーニング専用ライブラリ。開発当時はMKL-DNNという名称だった。これを,富岳に搭載されているArm系CPUであるA64FX向けに移植および最適化を行い,世界最高性能を達成するのがプロジェクトの最終目標。

「Xbyak_aarch64」開発のハマリどころ

小崎:まず我々の大きなディシジョンとして,Xbyakを拡張するんじゃなくて,思想だけを受け継いだ別のソフトウェア(Xbyak_aarch64)を作ったわけじゃないですか。どういう経緯でそう決まったのかとか,コメントを貰ってもいいですか。

川上:XbyakとXbyak_aarch64では対象とする命令セットがまったく違うので,Xbyakを直接いじることは考えなかったです。ただ,Xbyakは非常に作りがわかりやすいので,Xbyak_aarch64側で新しく何か設計をすることはほとんどなくて,いきなり実装に着手できました。ただ,何もないところからAArch64向けの命令を生成できるようにしないといけないので,基本的な命令でも何千もあって,物量的には非常に多かったですね。光成さんも長い歴史の中でちょっとずつXbyakを拡張されてきたと思うのですが,トータルだと工数って非常にかかってるんじゃないかと想像してます。

光成:そうですね。

川上:それを,Xbyak_aarch64を作るだけじゃなくて,最終目標としてoneDNNを富岳の上で動かすところまで考えると,あまり時間はとれなかったので集中して開発しました。あと,もしXbyak_aarch64自体にバグがあって,間違った機械語を生成してしまうと,後々Xbyak_aarch64を使ってアプリを作ったときに,バグの発見って多分非常に難しくなると思うので,検証は非常に入念にやったつもりでいます。

私自身,元々LSI※2の設計を生業としていたので,ハードウェアの設計でミスが起こって作り直しになると,時間的にも製造コスト的にも非常に影響が大きいので,きっちり検証をやって製造するっていう工程を踏むんですけれども。そのときの知見が活きたと思います。

実を言うとバグは1個か2個くらいはあったんですが,光成さんに直してもらいました。何千個もある命令の中で1個2個なので,割合としては非常に少ないと思います。後々アプリを実際に開発しているときに何か問題があって,というような手戻りはなかったのは非常に良かったかなと思っています。

川上 健太郎

⁠株⁠富士通研究所 シニアリサーチャー

2007年,富士通研究所入社。画像コーデックLSIやセンサノードの研究開発に携わる。回路,アルゴリズム,ソフトウェアなどさまざまな階層を横断的に最適化することにより,システム全体の高効率化・低電力化を実現する技術開発に強みを持つ。2019年よりArm HPC環境向けのディープラーニング処理ソフトウェア開発に従事。

小崎:要約すると,数の暴力で死にそうになったってことですね。それ以外に,何というか質的な面倒くささとかあったんですか? たとえばXbyakって,アセンブラっぽいものをC++上に素直に書けるとかいっても,C++と構文がバッティングしてしまうケースとかありそうじゃないですか。

川上:そこはたぶん,最初にXbyakが弾除けになってくれてたのかな。

光成:先ほどすぐに実装に入られたとおっしゃいましたけど,はじめの5月~6月は設計について議論していた気がします。私がXbyakの開発でハマったところとか,クラスやメモリアドレッシングをこうしたらいいんじゃないか,とかいうのを共有しました。

川上:そうですね,思い出してきました。最初のほうでいろいろアドバイスをいただいていたから,ハマることはなかったかな。

福本:そういえばXbyak_aarch64開発の後期では,IntelとArmのCPUアーキテクチャの差で,キャッシュかなにかのクリアをちゃんとしてあげないと,Armだと動かない※3という話もありましたよね。

川上:ああ,ありましたね。Iキャッシュ(命令キャッシュ)をフラッシュしてやらないと,Dキャッシュ(データキャッシュ)に実行コードを書いてもIキャッシュに反映されない……そんな感じですかね。

光成:それも最初のディスカッションのときに「AArch64の場合は注意しないといけないですよ」っていうのは実は言ってたんですけどね。⁠笑)

※2)
Large-Scale Integration,大規模集積回路のこと。
※3)
Intelの場合は8080やMS-DOSの時代から,後方互換性を満足するため,命令コードを自己書き換えした場合にこれを検出して自動で命令キャッシュをリフレッシュする挙動になっている。富岳に使われているA64FXなど,命令キャッシュとデータキャッシュの搭載が当たり前になった時代の後に設計されたCPUは,JITコードを生成した後,CPUの命令キャッシュを意図的にリフレッシュしないと,生成したJITコードを実行できない場合がある。

著者プロフィール

宮島幸太(みやじまこうた)

Software Design編集部所属。技術評論社2018年入社。

バックナンバー

2020