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

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向けディープラーニングのソフトウェアスタック開発を行っている。

「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の場合は注意しないといけないですよ」っていうのは実は言ってたんですけどね。⁠笑)

「Xbyak_translator_aarch64」でIntel命令からArm命令への変換を実現

小崎:川上さん、Translatorを作るのを思いつくところまでは割とストレートフォワードだと思うんですけど、実現方法がトリッキーでインターネットで話題になりましたよね。

川上:そうなんですか……中の人は知りませんでした。

私自身は去年4月から今の部署に入ってきて、⁠富岳向けにAIを動かすライブラリを作りましょう」という話を最初に聞いて、Xbyakを移植しないといけないというお話も伺っていて。最終的にoneDNNを富岳の上で動かすとなったときに、ただ単にXbyak_aarch64を作るだけでは駄目で、そのうえでoneDNNの移植作業も必要だと考えていました。全部真面目に移植していると大変なことになるので、翻訳機能のようなものを作らないといけないと思っていて、まず昨年の上期はXbyak_aarch64を作って、下期はたぶんTranslatorをやらないといけないと、早い段階で思っていました。そして去年の秋ぐらいに社内で作りますよと宣言をして、Translatorを作り始めました。

光成:最初、その件について川上さんからご相談を受けたときに、Intelの命令一覧がテキストデータで提供されていないか、要するにArmの命令列を自動生成させるための元データはありませんか?  と聞かれて「ないんですよ、私もPDFを見ながら1個1個実装してたんです」と答えたのですが、そこで「Intel XED」※4を使うと割と簡単にできますよというのをご紹介した覚えがあります。

川上:x64の命令の機械語を生成させてからAArch64の命令に変換するという手順を踏んだのは、その手順を踏まないとx64の機械語のエンコードフォーマットを丸々マスターしないといけなくなってしまうので、一旦x64の機械語は出させてしまって、それを光成さんに教えていただいたXEDという逆アセンブルするライブラリを使って人間がわかる情報(たとえば、これはレジスタとレジスタの足し算をする命令である、など)に変換して、そこからAArch64の命令を組み立ててやるっていう、そういう手順を踏んでやると実装が早く進むと思ったからです。

小崎:Arm用にライブラリを作ろうとした場合、Xbyak_aarch64からAArch64向けの命令列をJITで生成するっていうのがストレートフォワードだと思うんですけども。川上さんは、そんなたるいことやってられるかって言い出しまして。どうなったかというオリジナルのXbyakでx64の機械語までは生成できるんだから、それをそのままXEDで逆アセンブルして、それをAArch64の命令列に置き換えるっていうことを、JITコンパイルの過程できちんと全てやれば解決できるだろうっていう、超絶トリッキーなことを言い出して。最初、僕は何を言ってるんだろうって思いました。⁠笑)

川上:Translatorを作れたとなれば、oneDNNのソースコードからIntelの実装思想を読みとってArm向けに実装し直すという作業がいらなくなります。機械的にx64命令と、AArch64命令の対応関係さえ定義できてしまえば、変換対象がoneDNNだろうがなんだろうが変換できてしまう。そんな感じでしょうか。

小崎:あと、AIの世界の独特の事情もあって。デスクトップ向けだとJITコンパイルに1秒かかるだけでユーザーからクレームがつきますが、AIの世界って、1回の計算が3日とかなので、数秒JITやTranslateに時間をかけたところで、その後10%速度が速くなるのであれば余裕で元が取れるみたいな世界観なんですね。このドメイン特有の事情もあって、JITのところに気合を入れるというのが割に合うっていうか。

川上:そうですね。JITコードの生成自体は遅くても、そのあとビッグデータ相手に何百万回も繰り返しそのJITコードが使われるとすると、ペナルティは全然気にならないと。そういう事情はありますね。

光成:将来的にAI以外の分野で使うのであれば、Translate処理を速くするためにXEDを介さないでやれるようにしたいですよね。

光成氏と米Intelとのやりとり

光成:Xbyakに「この機能がほしい、あの機能がほしい」とGitHubのイシューが立っていて、よく見たらIntelの中の人だったということがありましたね。

川上:MKL-DNN自体も4年ぐらいの歴史がありますよね。問い合わせがよく来るようになったのってどれぐらいの時期だったんですか?

光成:2016年ごろにAVX-512を利用できるCPUが発表されましたが、私が使える環境ではなかったので、そんなに急いでやることはないなと思っていたら、GitHubのイシューで「ぜひとも対応してくれ」っていうのが来て、対応を始めたのがきっかけです。Intelが新しいCPU作るたびに、中の人から対応してほしいと依頼が来て、GitHubでやりとりしていました。個人的なメッセージやメールで直接やりとりをしたことはないです。

福本:新しいCPUの命令に対応するときって、何かドキュメントと一緒に来たりするんですか? それとも、プルリクエストでコードごと来たりするんですか?

光成:両方あります。最新のIntel CPU用のドキュメントPDFが公開されて、その命令に対応してほしいという依頼が来る場合もありますし、6月ぐらいにIntel Advanced Matrix Extension(AMX)という命令体系が発表されたときには、同時にプルリクエストが来ました。おそらくIntelのほうで内部調整して、仕様公開と同時にプルリクエストを出しているんだと思います。他にはIntelが私のソースコードを取り込んで新命令に対応していたことがあったんですけど、それがバグっていたので、⁠私のほうを使って下さい」とプルリクを出し返したこともありました。

川上:なるほど。福本さん、我々のところに新しいArmの仕様書来ないですよ、まだまだ(Xbyak_aarch64の)リポジトリを盛り上げていかないと。

福本:そうですね。⁠笑)

光成:(お礼にIntelから)何か欲しいなとTwitterになんとなく書いてみたら、Intel日本法人の方のツテで、一昨年に当時一番速いマシンを研究用に貸してくださったこともありましたね。112コアのマシンで、開発研究用途のものがサイボウズ・ラボにやってきました。

福本:「28コアのチップがワンソケットに2つ入ってる凄いやつが来た」ってTwitterで拝見して、うらやましいなと僕も思ってました。

OSSならではのフットワークが活きた

小崎:このチームは、完成度を上げてから公開することが多い富士通の中では、早め早めに公開して、うまくいった面もあるんですよね。そのへんの戦略とかについて福本さんからコメントないですか?

福本:いやー、どうですかね。大きな宣伝をしながらアップロードしたという訳ではなかったと思います。

光成:私からすると、公開されたらすぐにエミュレータとかを使ってテストができて、こちらでもいろいろ実験をしたり、サンプルコードを書いたりできたのはすごく大きくて。ちょっとおもしろい開発スタイルだなと思いました。

福本:確かに、いろいろなところにコードを渡すのに、公開しているほうが非常に楽だったというのは大きなファクターでしたね。たとえば、海外に渡そうと思ったら、輸出系の手続きとかをきちんと踏まなければいけなくて。そういうところが、OSSとして公開すると1回で終わることもありました。

川上:製品じゃなくてOSSなので、多少バグがあることに目をつぶって、早く出す方が価値があるという気楽さはありましたね。

光成:私もバグ修正のプルリクエストとかが気軽に出せたので、共同開発としてはすごくありがたい方法だったと思います。

今後のIntelやArmとの関係性、アップストリーム方針について

福本:先にArmのほうからコンタクトがあって、Armのほうでディープラーニングのソフトウェアスタックをどう作っていくかという意識合わせを、月1ぐらいのペースで定期的にミーティングを開いています[5]⁠。

Intelとの関係について言うと、Arm向けoneDNNを富士通研究所のリポジトリで公開していて、ResNet-50の処理性能向上などの成果を学会などでアピールしていたところ、Intelから本家のoneDNNにアップストリームしませんかと話が来たので、AArch64向けに作ったJITコードなどをアップストリームしています[6]⁠。

具体的なアップストリーム方針としては、まずはXbyak_aarch64と、Xbyak_aarch64で実装したリオーダー処理[7]をoneDNNに組み込み、マージリクエスト[8]を出したところです。我々が書いたソースコードがGitHub上でレビューされ、最近正式に取り込まれました。今後はXbyak_aarch64での実装を関数単位でプルリクエストを出していって、本家のoneDNNがAArch64でも高速に実行できるように持っていく計画です[9]⁠。

光成:個人的な感想で一番びっくりしたのが、IntelにとってライバルであるArmのコードをIntel自身のOSSに入れたことで、すごいな、と思いました。いちユーザーとして驚きです。

小崎:実は富士通社内でもびっくりされてですね。⁠君たちそんなことするの?」と、ハードウェア作ってる部署から主にびっくりされましたね。

光成:ああ、やっぱりそうなんですね。

福本:あと、僕がもうひとつびっくりしたのが、そもそもライブラリをオープンソースで公開しているのがすごいなと思いましたね。Intelは自社製のコンパイラにマスカーネルライブラリ(MKL[10]⁠)を付けたものを、元々商品として売っていたこともあるので。

光成:オープンソースでそのような戦略に舵を取ったから、Xbyakの流れができたのもあるかもしれませんね。XbyakはIntel CPU向けのコンパイラという前提で書いていたものですが、今度はIntel以外のCPUでも動かしたいと、Intelの人から依頼が来るのは面白いなと感じました。

小崎:AI自体の適用領域も広がっているので、オープン標準の強みを活かして今後この技術がさまざまな分野で使われるようになることを期待しています。

おすすめ記事

記事・ニュース一覧