富岳版XbyakがIntelの深層学習ライブラリoneDNNにmergeされる

初めに

サイボウズ・ラボの光成です。

このたび、Intelの公式深層学習ライブラリoneDNNに、富士通が開発しているスーパーコンピュータ富岳向けのPull Requestがmergeされました。 その開発に関わることになった経緯を紹介します。

目次

  • 概要
  • Xbyakとは
  • 動機
  • Intelとの関わり
  • 富士通との関わり

概要

富士通研究所はスーパーコンピュータ富岳で深層学習(ディープラーニング)を高速に処理するためのソフトウェアを開発してます。 そのためにIntelが開発している深層学習ライブラリoneDNNを富岳に移植して改良しています。 このたび、その成果の一部が本家のoneDNNに取り込まれました。

富岳はA64FXというArm v8-Aにベクトル演算機能SVEが追加されたCPUを持ちます。 oneDNNを富岳に移植するには、私が開発しているXbyakのA64FX用Xbyak_aarch64が必要でした。詳しい話は富士通研究所の方々がblogで紹介されるのでここではXbyakがどう関わっているかの紹介をします。

Xbyakとは

Xbyakとは私が14年近く前から開発しているx86/x64向けのJITアセンブラです。 C++のヘッダファイルで提供され、使い勝手としてはC++の文法が使える高機能なアセンブラです。

JIT(Just In Time)とはプログラム実行時にコード生成するという意味です。 通常のC++プログラミングはコンパイル時にコードが静的に決まり、実行時にはその通りに動作します。 JITによって実行時に決まるパラメータに応じた柔軟な最適化が可能になります。 Javaのランタイム実行エンジンやモダンブラウザのJavaScript実行エンジン、最近ではPHP 8などが与えられたプログラムに応じたJITを行っています。 Xbyakは、そのような機能の開発を自分でやりたい場合に支援します。

動機

私はコーデックや暗号系のライブラリの開発をすることが多かったのですが、最適化のためにアセンブリ言語(以下asmと表記)を利用するとき不便に感じることがよくありました。

一つはC/C++の構造体やクラスのメンバー変数にアクセスするためにそれらのメンバー変数が構造体の先頭から何バイト目にあるか、一つ一つasm側でオフセットを設定しないといけない点です。 コード情報が重複するので、メンバー変数を入れ換えたり追加したときにasm側も修正が必要で更新し忘れや間違いが発生しやすいです。大規模な開発ではスクリプトやツールを使ってC/C++のコードからオフセットを自動設定することもしていましたが汎用性は高くありません(templateの構造体のパースなどコストが高すぎて無理)。

struct X {
  int x;
  int y;
  int a[4];
};
// nasmの場合
struc X
 .x resd 1
 .y resd 1
 .a resd 4
endstruc

templateでoffsetが変わる構造体の扱いは難しい。

template<size_t N>
struct X {
  uint64_t x[N];
  uint64_t y[N];
};

もう一つは、IntelのCPUはしばしば新しい命令が追加されるのですが古いCPUでも動作可能にしつつ、新しい命令にも対応しようとするとasm側のコードが煩雑になる点です。 またアセンブラの擬似命令(ifやloop、マクロといった機能)は文法が独特で分かりにくいことが多いというのもあります。

Xbyakの記述は通常のC++の文法に従うのでこの2点を解決します。 個人的には「C++の文法でasmを記述できる」というのが最も便利に感じます。

AVX-512のbfloat16が使えるときとそうでないときでコード生成を変える例(jit_avx512_core_amx_conv_kernel.cppより引用 )

if (is_bf16) {
    vmovdqu16(zmm_src_tmp, ptr[reg_src + offset]);
    vpermw(zmm_dst, zmm_idx, zmm_src);
    vmovdqu16(ptr[reg_dst + offset], zmm_dst);
} else {
    vmovdqu8(zmm_src_tmp, ptr[reg_src + offset]);
    vpermb(zmm_dst, zmm_idx, zmm_src);
    vmovdqu8(ptr[reg_dst + offset], zmm_dst);
}

より詳しい説明は blog.cybozu.io をごらんください。

現在はXbyakを使ってペアリング演算の実装mclや準同型暗号の論文を書いたり、Ethereum 2などブロックチェーン系プロジェクトで利用されている署名ライブラリblsを開発したりしています。

Intelとの関わり

前置きが長くなりました。 主に自分のためのツールとして開発をしていたのですが2016年、GitHubのissueでAVX-512に対応してという依頼が来ます。 当時AVX-512が使えるKnights Landingはまだ発売されてなくて、そんなに急いで対応しなくていいんじゃない?と聞いたのですがぜひとも欲しいというので対応することにしました。

とは言えAVX-512用のマニュアルは千ページ以上あり、エンコーディングも複雑で新設される命令も非常に多く、かなり苦労しました。マニュアルの表記が間違っていたので指摘したこともあります。 どうにか対応すると、BVLCが開発していた深層学習フレームワークCaffeのIntel向け最適化バージョンに組み込まれました。

それがうまくいったようで、その後Intelが新しく開発したMKL DNN(現oneDNN)のCPU向け最適化コードで幅広く利用されるようになります。 深層学習では計算パラメータが非常に多く、ユーザが与えたパラメータや、実行時のCPUのキャッシュサイズなどに応じて行列のサイズやデータの形式が変わることがあり、それらに対応するのがJITだと都合がよいのです。 また前節で述べたように、より新しい命令に対応したコード生成をしやすいという利点もあります。

oneDNNは機械学習フレームワークとして有名なTensorFlowやPyTorchなどで利用されています。

新しい命令が発表され、IntelのCPUの公式マニュアルが更新されると、Xbyakをその命令に対応させます。 発表された直後は、まだgccやclangが対応していないこともあり、動作確認が難しいです。GitHubで「実装したけどちゃんと動いているか分からない。」と書くと、Intelの人に「大丈夫。未公開ツールで動作確認した。」という返事が来たことがあります。

2020年の6月にIntelがIntel AMXという新しい命令体系を発表したときには、ほぼ同時にXbyakに対応pull requestがやってくるという、個人ではなかなか得難い体験もしました(cf. Intel Begins Volleying Open-Source Patches Around Intel AMX)。

Xbyakは先月10月に登場したAVX向けの機械学習命令VNNIにも対応し、oneDNNでは既にAMXやこれらの命令を利用した開発が進められています。

富士通との関わり

最初に述べたように、富岳のCPU A64FXはAArch64というスマートフォンに載っているCPUと同じ系統の汎用CPUで、SVEというAVX-512に似たSIMD命令を利用できます。

SVEについてはたとえば8月に開催されたARM入門勉強会で私が発表したIntro to SVE 富岳のA64FXを触ってみた(動画 : https://youtu.be/3pfjrIyjhBc?t=12462)を紹介しておきます。

www.slideshare.net

youtu.be

富岳で機械学習まわりの開発を容易にするためにoneDNNの移植計画が立ち上がり、そこでoneDNNのエンジンで利用されているXbyakの富岳版を開発するための共同研究が始まったのが2019年の4月でした。 移植に関わる詳細な情報は富士通研究所のblogや、技術評論社による対談記事がgihyo.jpで公開されています。合わせてご参照ください。

Xbyak_aarch64は私もコントリビュータとしてバグ取りや改良をしています。 そしてXbyak_aarch64を使った富岳用のoneDNNの開発が進み、その成果の一部を富士通が本家Intelに統合してもらうよう依頼し、了承されました。個人的には富士通がIntelに出したPull Request

github.com

の中で、富士通の人が書いてくれた謝辞に私の名前があるのをIntelの人がみつけて、

Acknowledgment We thank S. Mitsunari (Cybozu Labs, Inc.), the developer of the original Xbyak.

Yeah, we greatly thank Mitsunari-san for the Xbyak on x86 too! 🎉 His work helped us a lot to make oneDNN development faster and simpler.

コメントを下さったのがとてもうれしかったです。 今後も開発に関わり、よりよいものにしたいです。

まとめ

XbyakというJITアセンブラの紹介とそれがIntelの深層学習ライブラリに利用され、富岳でも使われるようになった経緯を紹介しました。