Haskell/GHCのSIMDについて考える

最近のコンピューターの性能を活用するには、何らかの並列化が必須です。具体的にはSIMDの活用やマルチコア（それとGPU）です。プログラミング言語でこれらを利用できれば「C言語よりも速い」を名乗れます。この記事では並列化技術のうち、SIMDを考えます。

HaskellコンパイラーであるGHCにはSIMDのプリミティブ（データ型と関数）が実装されています。しかし、広く使われているとは言い難いです。

使いづらい要因はいくつか考えられます。使えるバックエンドの少なさが一つで、現状（GHC 9.8時点）では、x86(_64)向けのLLVMバックエンドでしか使えません。AArch64など他のアーキテクチャーや、NCGバックエンドでは使えません。

別の要因として、用意されたプリミティブが少ないというものがあります。例えば、整数のビット演算はできません。

現状

論文

GHCにSIMDプリミティブを実装した人たちの論文が以下です：

  • Geoffrey Mainland, Roman Leshchinskiy and Simon Peyton Jones, Exploiting Vector Instructions with Generalized Stream Fusion, ICFP 2013
    • 元々のタイトルは「Haskell Beats C using Generalized Stream Fusion」だったようです。

論文ではGHCのSIMDプリミティブの他に、vectorパッケージとの連携やData Parallel Haskell (DPH)との連携が考えられていたようですが、現在生き残っているのはGHCのSIMDプリミティブだけです。

ライブラリー

GHCのSIMDプリミティブ型はunboxedなので、普通に使うにはbox化されたラッパーが欲しいです。Hackageを探すとそれっぽいパッケージが2つ見つかります：

設計について

SIMDのAPIをプログラミング言語に露出する際は、アーキテクチャーによる幅の違いをどう吸収するか、あるいはしないのかという問題があります。

GHCでは、128ビット、256ビット、512ビットの幅の型をそれぞれプリミティブ型として提供しています。これらは対応するCPUレジスタにそのまま対応します。つまり、x86系の場合512ビット幅はZMM (AVX-512)に対応し、YMMを2つ使うようなことはしません。

一方、LLVMのvectorの抽象化は、任意の長さのベクトルを書けるようにしており、マシンのSIMD幅よりも長いものは複数の命令・レジスターにコンパイルされます。

現状のGHCの設計では、マシンのSIMD幅の違いを吸収するのはライブラリーの役目となるでしょう。

GHCのSIMDを使いやすくするために何ができるか

ここからは、GHCのSIMDを使いやすくするために何ができるかを考えていきます。

バックエンドを増やす

SIMDに対応するバックエンドが限られていると、なかなかライブラリーでは採用しづらいです。そこで、SIMDに対応するバックエンドを増やすことが考えられます。

x86(_64) NCG

現在、パソコン向けで最もポピュラーなアーキテクチャーはx86系でしょう。よって、x86 NCGに対応することはメリットが大きいと考えられます。

実は既にその作業をした人がいて（2018年のGSoCプロジェクトだったようです）、2019年に一旦本体にmergeされましたが、バグっていたようでrevertされました。

このパッチを現在のGHCを対象に復活させることができたらいいなあ……と思って今いじってみていますが、なんかまだ結構未完成なところもある（要素の型は Float/Double しかサポートされていない）ので、LLVMバックエンドと同等のところへ持っていくには結構作業が必要そうです。

AArch64 LLVM/NCG

x86系の次にポピュラーなアーキテクチャーはAArch64 (Arm64)かと思います。よって、これも対応するべきでしょう。

AArch64のSIMDは128ビットのASIMDと、幅が可変なSVEです。GHCの既存のAPIから使いやすいのはASIMDの方でしょう。SVEは対象CPUが256ビットや512ビットであることを仮定してコンパイルするか、別途専用の型を用意することになるでしょう。

当面はASIMDに絞るとして、LLVMで対応するのはそんなに難しくないはずです。というか、私が試みてマージリクエストを投げました。

NCGでも対応すると良さそうですが、非自明な作業が発生しそうです。

WebAssembly

Wasmにも128 ビットSIMDの拡張があって、今の実行環境ではもう使えるようになっているようです。

issueは立っています：

POWER (AltiVec)

POWER (PowerPC)にもAltiVecとかいう名前で128ビットのSIMDがあるらしいので、使えるとよさそうです。倍精度には対応していないという噂ですが。

プリミティブを増やす

現在GHCに実装されているSIMDプリミティブは限定的です。よく使われるようなものだと、

  • sqrt
  • abs
  • ビット演算
  • FMA

などが欠けています。

実は、LLVMバックエンドを前提とすれば、現状のGHCにこれらのSIMD命令を出力させることが可能です。unpackしてから要素ごとに演算してpackするような一連の命令列が、LLVMの最適化によって一つの命令に変換されるのです。

sqrtFloatX4# :: FloatX4# -> FloatX4#
sqrtFloatX4# v = case unpackFloatX4# v of
                   (# x0, x1, x2, x3 #) ->
                     packFloatX4# (# sqrtFloat# x0, sqrtFloat# x1, sqrtFloat# x2, sqrtFloat# x3 #)
-- LLVMバックエンドでは1命令にコンパイルされることが期待できる

ただ、いずれNCGにも対応させることを考えると、NCGに同様の最適化を実装するよりは専用のプリミティブ関数を追加したほうが良いと思われます。

このほか、現状では条件分岐に関するプリミティブが欠けています。マスク関連ですね。比較演算とか論理演算、selectみたいなやつです。これはマスクの型をどうするのかとか、検討することがいろいろありそうです。BoolX4# みたいな型を要素数ごとに用意するのか、 FloatX4Mask# みたいな型を要素型と要素数ごとに用意するのか。

プリミティブを追加する際は、現状対応しているx86(_64)以外に、AArch64やPOWER, WebAssemblyの方も見ながら検討していく必要があります。RISC-Vは知らん。

FFIの対応

GHCが対応しないニッチな命令でも、FFI経由で利用できると良さそうです。

さっき試したらLLVMバックエンドではUnliftedFFITypes拡張を有効にすればccallでベクトル型の受け渡しができました。しかし、Windows標準の呼び出し規約だとベクトル型は参照渡しになって効率が悪いかもしれないので、vectorcall呼び出し規約にも対応することが考えられます。

プリミティブを減らす

その辺のSIMD命令セットは整数除算を持たないにも関わらず、GHCのプリミティブ関数にはSIMDな整数除算が含まれます。LLVMはそういうのもいい感じにコンパイルしてくれるのかもしれませんが、NCGでやるのは面倒そうなので、対応する命令がなさそうなやつはプリミティブ関数としては削除するようにした方がいいかもしれません（GHC.Exts には互換性のためにHaskell実装を残しても良い）。

非自明な課題：命令セット拡張をどのように利用するか

x86_64ではSSE2は標準で使えますが、SSE4やAVX2やAVX-512は全てのCPUで使えるわけではありません。ローエンドCPUがAVXに対応してないという話を聞いたことがあります（最近はどうなっているのか知りませんが）。また、Rosetta 2などのエミュレーターはAVXには対応しません。

というわけで、AVXとかを活用するプログラムは、AVX対応CPU専用にビルドするか、あるいは実行時にCPUIDを見て切り替えることになります。現状のGHCでは -mavx2 みたいなフラグを必要とするので、前者のやり方しか対応しないということになります。

特定のCPUだけで動けばいいプログラムであればそれでもいいかもしれませんが、多くのCPUで動いてほしいライブラリーはそれでは困ります。

というわけで、モジュール単位ではなく関数（実行パス？）単位でCPUの機能を有効化する仕組みと、実行時にCPUの情報を見て実行するコードを切り替えるような仕組みが欲しいです。

今やっていること

しばらく前に、各種アーキのSIMD命令の調査メモを書きました。まだx86とASIMDとWebAssemblyしか調べていませんが：

SIMD命令には直交性もクソもないのがわかりますね。

AArch64のASIMDへLLVMバックエンドで対応するやつは、しばらく前にマージリクエストを投げて、レビュー待ちです。

今はx86 NCGへのパッチをいじっています。しかし、ネイティブコード生成に関する私の実力が乏しいため、どこまでできるか、という状況です。

このエントリーをはてなブックマークに追加