インタビュー

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

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

「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を介さないでやれるようにしたいですよね。

※4)
IntelがOSSとして公開しているx86_64機械語のデコード(逆アセンブル)処理用ライブラリのこと。

光成氏と米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自体の適用領域も広がっているので,オープン標準の強みを活かして今後この技術がさまざまな分野で使われるようになることを期待しています。

※5)
富士通研究所とArmはもともと協業関係にあり,定期的にミーティングを行っている。
※6)
IntelはoneDNNをx86専用ライブラリとしていたが,近年はその他のCPUをサポートしていく動きがあり,ARMアーキテクチャを取り入れる動きもある。
※7)
ディープラーニングのさまざまな計算処理を行う前に,データの形式を変換する処理のこと。
※8)
第三者が実装したソースコードを,リポジトリに取り込んでもらうよう依頼すること。
※9)
現状では一部の富岳ユーザーがArm移植版のoneDNNを試用している状況。
※10)
Intel Math Kernel Library。科学,金融分野などにおける数値演算の高速化のためにIntelが独自開発したライブラリのこと。

著者プロフィール

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

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

バックナンバー

2020