Intel AVX-512・富岳SVE用ライブラリで高速ループ処理　文字列で与えられた関数から最適化されたSIMDコードを生成

自己紹介

光成滋生氏（以下、光成）：「Intel AVX-512/富岳SVE用SIMDコード生成ライブラリsimdgen」について発表します。

自己紹介です。サイボウズ・ラボで、暗号や最適化に関するR&Dをやっています。2019年3月から、富士通の富岳のAI系のチームの人と共同研究をしています。

「Xbyak」というJITアセンブラの開発者で、富岳用の「Xbyak_aarch64」のコミッターでもあります。暗号系では、2021年に本を書いて、その読書会を毎週木曜日にやっています。興味がある方はご覧ください。

Intel AVX-512・富岳SVE用ライブラリ「simdgen」を開発中

今日やりたいことは、ループ処理の高速化です。難しいことはぜんぜんできていませんが、例えば次のことをします。

（スライドを示して）n個の要素のfloatの配列に対して、例えばexpとlogの値を足して、結果を*dstの各要素に出力することを高速に処理したい。あるいは、n個の要素のfloatの値を取って、それらのcoshを計算して、logを取った結果を全部足した値を求めたい。このような処理の、一般的なことができるライブラリを作っています。

（スライドを示して）「simdgen」というもので、私のGitHubでただいま開発中です。これは、例えば「log(exp(x) + 3)」という文字列を与えると、それに対してIntelのAVX-512のSIMDコード、Windows、Linux、あるいは富岳のSVEというSIMDコードを、関数を実行時に生成して実際にそれを呼び出せるライブラリです。

これはC-APIの例で、SgCreateでコードへのポインタを取り、それに対して関数生成するものを書いて呼び出すと、アドレスが返ってきます。

そのアドレスに対して、自分の好きなfloatの要素とdstの要素を渡してやると、実際にここに書いたコードがSIMD実行され、IntelだったらAVX-512を使ったコード、富岳だったらSVEを使ったコードが生成されて、不要になったら破棄します。

機能としては、カッコ付きの簡単な四則演算、例えば「x + 1 * (2 + x * (3 + x))」みたいなものや、logとexpとcoshとか逆数などを、順次追加していく予定です。

和としては、reduce sumのつもりで、red_sumというのがあります。例えば「red_sum(log(cosh(x)))」とやると、先ほどのこのtargetCと呼ばれるものを計算できるようになります。

ベンチマークの紹介

まず、簡単なベンチマークから紹介します。（スライドを示して）AVXはこういうマシーンでやって、コンパイラはGCCの10で「-Ofast -march-native」です。ぶれないようにTurbo Boostはオフにしています。

SVEは、富士通が出しているFX700で、GCC 11の「-Ofast -march=native」とやります。例えば、今さっき紹介した各要素に対するcoshのlogを取って、全部足すというようなものをCで書いたものと、このsimdgenでやったものをやると、AVX-512でだいたい16倍ぐらい、SVEでは8倍ぐらい速くなりました。

今から、こんなことができるという中身を、ざっくりと紹介していきます。

まず「そもそもSIMDって何？」という人のために一言だけ説明すると、AVX-512や富岳のA64FXというCPUが持っているSVEは、512bit SIMDレジスタを32個持っています。

このSIMDレジスタは何かというと、1個のレジスタで、例えば512bitだったら、32個のfloat16個を同時にまとめて処理できます。

AVX-512でいうと、addps、これはpacked single floatという意味ですが、これに対して、zmmレジスタと、zmm1、zmm2というのをやると、1と2に入った16個同士のfloatの値を各要素ごとに足した値がここに入ります。

同様のことを、富岳のSVEで「fadd(z0.s, z1.s, z2.s)」と書くと、1と2の要素をそれぞれ足してz0に入ります。

今日は、本当に最小限の説明で、これだけしかお話ししませんが、もうちょっと知りたい方は、私の過去のスライドなどをご覧ください。

内部実装の説明

（スライドを示して）内部実装について、大まかな流れを紹介します。まず、与えられた文字列を構文解析して、簡単な構文木を作ります。それを元に、どのようにSIMDレジスタを使うかという割り振りを決定して、その後、SIMDコードを生成します。

ここからトークンリストみたいなものを作るまでは、富岳とIntelは同一です。最後、構文解析したあとの構文木から、AVX用コード生成のところでXbyakを使って、SVE用にはXbyak_aarch64を使っています。

（スライドを示して）入力例です。1+x*(2+x*3+1/x)を文字列として入れてみると、ご存じの方が多いと思うのですが、逆ポーランド記法で解釈すると、これは定数1に変数x、var{0}はこの場合、変数xに相当しますが、定数1に、変数0に、定数2に、変数xに、float{3}を掛けたものを足して、1を変数xで割ったものというのが、（スライドを示して）ここですね。ここを足す。これがここです。その後、掛けるのはxです。最後に、足す、で、ここの"+"に相当しますという出力をします。

（スライドを示して）ここまでやるのは、Intelとaarch64で一緒で、その後はSIMD生成で、AVX-512の場合は、例えばこのようなコードが生成されます。若干のレジスタ割り当てと定数移動の最適化をしているので、このようなコードになります。

具体例で見ると、まずソースから16個分のfloatレジスタを読み込みます。このループの前に、この式に現れる定数をレジスタにあらかじめ入れておきます。

まず、ここがどういう意味かというと、今呼んだxに3という定数をmulで掛けて、4に入れます。だからこれはzm=x*3という意味です。

そのあとはvaddpsで、定数2を、今やった4に足します。だからzm4というのは2+x*3になります。

今度は、vdivpsで、zm5、zm1、zm0となると、zm5は定数1をxという、ここで読み込んだ値で割ったものになって、足して、以下同様のようにして、このような計算をして、最後に結果を書き込むようなコードが実行時に生成されます。

（スライドを示して）SVEでも同様で、同じように文字列の構文解析を終わった後、事前に各定数が読まれて、ld1wでz0にソースから値が読まれて掛け算して、足し算して、割り算して、同様に最後まで書き込む処理が行われます。

「simdgen」ライブラリの特徴

このライブラリの特徴は、関数呼び出しは全部インライン展開されて、ループ中にコールは一切呼ばれないということです。

Intelコンパイラでは、exp_psやlog_psと呼ばれるSIMDレジスタ用の組み込み関数的なものが用意されていて、今日紹介したようなループは、そういうものを組み合わせて作ることが多いのですが、それらは、事前にコンパイルされたライブラリの関数を呼び出すことが多いです。

それと、それぞれの関数ごとにレジスタ退避・復元がされますが、このライブラリは、JIT生成する時に全部ガッチャンコされて行われます。

それぞれの関数計算で必要なレジスタも、事前にまとめて全部ループの外で設定されます。（スライドを示して）例えば「exp(x3)+1」をSVEで実行すると、この3や1といったexpの計算に必要な定数がループの前に設定されたあと、ここでメモリからz0に入れて、それに、ここの場合、z1に3という定数が入っているので、掛けてz9という値が計算されます。ここがx3に相当します。

そして、ここから下が、このexpの計算になって、最後にここでz9にその出力が入って、足す1がこの行で計算されて出力されます。このようになっています。

ループのベンチマークについて、もう1度確認すると、先ほど（AVXの）16倍と（SVEの）8倍です。そもそもSVEを使った場合、8倍、10倍近く遅く、ちょっと残念な感じがしていたので、もう少しなんとかならないかと考えました。

原因は、A64FXがIntelのアーキテクチャに比べてレイテンシがかなり大きいからです。そのため、ループアンロールをしないと隠蔽できず、CPUの実行ユニットが空いてしまいます。

というわけで、ループアンロール回数を指定できるようにしました。これは単にループアンロールするだけなので、大きくするとレジスタが足りなくてエラーになってしまいます。

（スライドを示して）先ほどのコードの例で、unroll=2で生成してみると、1個を読み出して処理していたのが、隣のメモリも同時に読むコードになります。2個ずつ、同じだけどレジスタ番号だけ異なるものが生成されているのがわかると思います。これがアンロールしたものがそれぞれ出ていっている感じです。

アンロールをした結果と高速化を実現するための最適化

（スライドを示して）アンロールをした時の評価をしてみました。expという関数とlogとinvとcoshに、Cで実装した場合がこんな感じだったのが、アンロールなしの最初のものだとこのような評価で、アンロールを2ですると、全部がだいたい半分ぐらいに、きれいに速くなっています。

3回アンロールしてもざっと3分の1ぐらいになっているので、SVEはかなりレイテンシが大きいことがわかります。4回アンロールしてもまだ速くなっています。logは、logの計算の定数に必要なものが多いので、アンロールできなくなってエラーになってしまいました。逆数でもやはり速くなっています。

AVX-512で同じコードをやってみたんですが、SVEほどには速くはなっていません。多少速くはなっていますが、アンロールの効果がここではあまりありませんでした。これはレイテンシがそこまで大きくないので、1回のループでも、unroll=1でも大丈夫という感じです。

（スライドを示して）ループアンロールの効果があるのでもっとしたいということで、今回速くしたいコードでやってみました。このcoshをexpの内部で呼び出すのですが、それとlogの組み合わせで、ぜんぜんアンロールできなくて、10数倍しか速くなっていないという状況です。

これをなんとかしたいということです。ループアンロールに回すレジスタが足りないのなら、最初は定数レジスタを全部ループの外に追いやるという最適化をしていたのですが、逆にそれを元に戻して、必要なところで都度計算というか設定したらどうなるかを試してみました。

そうすると、ループアンロール用に回すレジスタがたくさん使えるようになって、AVXでもそれなりに速くなってきました。これだけ複雑な計算をすると、レイテンシが隠れないところがあったわけです。そうすると、ループアンロールするごとに速くなっていきました。SVEのほうは劇的に効果がありました。そんなことができるようになっています。

細かい最適化としては、これはこの例の、特有の例なんですが、logというのは、xが1の時はすごく桁落ちしやすいです。ここは精度計算じゃないですが、近似計算する時のセオリーで、x=1の時の付近は別パスで計算しないといけないというルールがあります。

しかし、最後に全部足している場合は、1つでもxが1ではないところがあると、結果は0ではなくなってしまうので、そうすると、がんばって0付近の高精度計算をする必要はありません。

なので、logの細かい計算をオフにするオプションを付けてみました。それでやってみると、1回計算する時でも若干速くなって、加えて、それ用にレジスタの必要がなくなったので、ループアンロールも5回までできるようになり、これぐらい速くなりました。

というわけで結果的に、一番最初のスライドで見せた16倍だったものが26倍ぐらいまで上がって、SVEのほうは、8倍しか速くならなかったのが31倍と、ざっくり4倍ぐらい速くなりました。元のCの関数から比べたら2、30倍速くなった感じです。

Pythonからライブラリを呼び出す方法

このライブラリをPythonから呼べるようにもしています。実際にどういうやり方をするか、簡単に紹介します。

PythonからCのライブラリを呼ぶには、Pythonの「ctypes」というライブラリを使います。そしてcdll.LoadLibraryというものを使ってshared libraryを読み込みます。

その後、Python側でCの関数を呼び出すために、普通の配列ではなく、「c_float * n」というので、「float src[n]」というCの関数相当のをPython側で準備します。

そのあとに、文字列を渡して、実際にsimdgenのAPIを呼べば、JIT生成された関数ポインタが返ってくるのですが、その関数ポインタをCの関数で呼び出すためにキャストしないといけません。そこで、restypeというctypesの変数の型に注釈を入れます。（スライドを示して）このあたりになってくるとかなり、へんちくりんなのですが、これもPythonのコードです。

そして、実際にCのAPI、SgGetFuncAddrをctypesのPythonライブラリを使って呼び出すことによって、実行できるようになっています。

関数を呼び出すところでは、float配列に対して、byrefでCの生のポインタを取り出して、POINTER(c_float)にキャストして呼び出してやるとできます。

Pythonには「Numpy」という便利なライブラリがあるので、それからも呼び出すためには、Numpyのndarrayを、さらにPythonのctypes.arrayに変換してからもう1度渡すというかたちです。

このあたりの変換は、ポインタのすげ替えをするだけで、実際にある配列をいじるわけではないので、それほどコストは高くないです。

私はNumpyをぜんぜん使ったことがないので、これが速いコードなのかよくわからないのですが、このコードを書いたものと比較してみると、ざっくり70倍ぐらい速かったです。

NumpyはIntel CPU向けにはけっこう最適化されているらしくて、このリファレンスを見ると、SIMDの最適化していますとドキュメントで書かれていました。

まとめ

まとめます。ループアンロールが重要ということがわかりました。特にSVEでは重要です。最終的にはsimdgenはできるだけ自動的にアンロールできるようにしています。

所感ですが、今回このsimdgenを作っている段階で、内部用のフレームワーク的な部分がだいぶ固まってきたので、関数を追加する部分がだいぶやりやすくなってきました。あと、IntelがPythonのNumpy専用のものを出していて、かなりすごいらしいのであとで調べます。今後、M1 MacやAVX-2にも対応できたらと思っています。

発表は以上です。

質疑応答

司会者：ありがとうございます。光成さんといえばSIMDライブラリの人と、Kernel/VM界隈ではすっかり定着してきたのではないかという印象です。今回また、新しい生成ライブラリを作ったのですね。

光成：そうですね。

司会者：アーキテクチャで差がこんなにあるんだなと、コメントがありました。

アーキテクチャが同じでも、世代によって、レイテンシがもう少し改善されているとかが起きると、このあたりのハックや、前提条件が変わってくるのは、ありそうだということですよね。

光成：そうですね。今回のライブラリだと実行時に生成するので、新しいIntelのCPUで新しい命令が追加されていれば、その分使いやすいので、いろいろ実験はできるかなと思います。

司会者：RISC-Vの勉強会で、ISAは同じでも、内部の仕組みというかそのあたりのデコーダがよくなっているのか、世代ごとで性能がぜんぜん変わるから、こういうハックは、実際に測定しないとなかなかわからなくて、そのハードでしか発揮しないことは多いと聞きました。「やはり、今もこういうのがあるんだな」と聞きながら思いました。でも、これだけ時間が変わると、本当にどんどん手入れをしたくなりますよね。この数字はとてつもないですよね。

光成：そうですね、けっこう。SVEは今回は普通のGCCを使っていますが、富士通が提供している富岳用のコンパイラを使うと、これよりはもっといい数値が出るんです。ライセンスの関係もあって、今回は紹介していないです。

司会者：なるほど、富岳のSVEとIntelのAVX、直接比較するには変態味の強い命令、そうなんですね。まとめてドンと処理をする。ジェネリックな命令ではないと言ったらいいのか、ニュアンスはわかります。

光成：ちょっとだけ、時間がかかってしまいますが、いいですか?

司会者：お願いします。

光成：（スライドを示して）実験的なコードで、例えば、x+3というコードをやって、環境変数で設定できるのですが、dump=codeという設定で生成されたバイナリ……あ、codeというファイルにダンプされるので、こういうふうにして、例えばrsiから読み込んで、足し算して、ストアするという命令が今生成されたのですが、このアンロールを例えば3にすると、3つ分ずつやっているのが、見てわかるかなと思います。