GHCのSIMDサポートの進捗：整数ベクトルの演算

ブログ等にちょいちょい書いていますが、私はこの数ヶ月ほどGHCのSIMDサポートの改善を進めています。

この3ヶ月間はx86 NCGで整数のSIMD演算を使えるようにするパッチを書いており、数日前にマージされました。一つのマイルストーンだと思うので記事にしておきます。

やったこと

128ビット幅の整数ベクトルに対する演算（足し算、掛け算等）とシャッフルをx86 NCGに実装しました。

差分は +7991行、-161行 と、かなり大物になりました。SSEがカオスなのと、テストケースが膨大なのが行数に効いていると思います。

LLVMバックエンドでは以前から整数ベクトルを使えていましたが、今回、LLVMに頼らなくてもコンパイルできるようになるのが進歩です。

x86のSIMD命令セットはSSE、AVX、AVX-512と拡張されてきましたが、まずは基本（古いCPUでも使える）ということでSSEを使います。ただ、SSEは正直カオスで、例えば整数の乗算が一部の幅にしか用意されていなかったり、64ビット整数からなるベクトルの比較にSSE4.2が必要だったりします。

シャッフルっぽいことができる命令は色々ありますが、特定のパターンの命令に対応するシャッフルはその特定の命令にコンパイルするようにしたいです。例えば、

shuffleInt32X4# a b (# 0#, 4#, 1#, 5# #)

というコードは PUNPCKLDQ という命令（C言語の組み込み関数だと _mm_unpacklo_epi32）が行うシャッフルと同じことをするので、PUNPCKLDQ 命令1個（それに加えて、必要なら MOV 系の命令）にコンパイルされてほしいです。そういう賢い命令選択をやるには、どうしても場合分けが多くなります。テストケースでもそれらをカバーしたいので、各パターンに対して数通りのシャッフルを乱数で生成してテストケースとしました。

シャッフルは実装もコード量も多いので、別のmerge requestに分けた方が良かったかもしれません。

ともかく、これで128ビット幅の整数ベクトルはx86上で不自由なく使えるようになったと思います。

実装の過程で他のバグを見つけたりしたので、単に新機能の実装ではなく、GHCのSIMDサポートの安定化に寄与していると思います。

次のステップ

でかいことを成し遂げたので、一旦立ち止まって今後のことを考えたいです。今後やることは色々考えられます。

GHCのSIMDプリミティブの実装を進めているわけですが、一般ユーザーが使うには「プリミティブ」は不親切です。そこで、SIMDを使いやすくするライブラリーを実装したいです。すでにある程度実装を進めており、ベクトルの内積を fold' (+) 0 (zipWith (*) xs ys) という風に書くといい感じにstream fusionされるようになりました。ただ、もう少し面白い応用も用意したいです。

今回はx86 NCGのサポートを実装しましたが、私が普段使っているマシンはApple Silicon、つまりAArch64なので、Arm NEON (ASIMD) もNCGで扱えるようにしたいです。AArch64は後発のアーキテクチャーなのでx86ほどカオスではないと思いますが、64ビット整数の乗算がなかったり、TBL命令などの一部の命令が「複数の連続した番号のレジスター」を扱うので、そこら辺に注意が必要かと思います。

現在のGHCによるSIMDプリミティブは演算の種類が限定的で、足りないものが色々あります。abs, sqrt, ビット演算あたりは、議論の余地が少なく、大抵のアーキテクチャーには対応する命令があると思うので、追加しても良いと思います。一方で、マスクに関する命令（Bool のSIMD版。比較や選択など）は議論の余地があると思うので、実装の順序としては後回しにしたいです。

SIMDに限らない（スカラーにも関わってくる）話ですが、IEEE 754-2019で規定されたmin/max系の演算を追加したいです。関連issue：#25350: Floating-point min/max primops should have consistent behavior across platforms · Issues · Glasgow Haskell Compiler / GHC · GitLab

最近のCPUには、半精度浮動小数点数にネイティブ対応しているものがあります。IntelのSapphire Rapidsや、Apple Siliconです。なので、GHCからそれらに対応する命令を出力できると良いかもしれません。ただ、呼び出し規約の関係でLLVMの改修も必要そうです。関連issue：#15983: Built-in support for half-floats · Issues · Glasgow Haskell Compiler / GHC · GitLab

Windowsのvectorcall呼び出し規約に対応すると良いかもしれません。が、WindowsでSIMDをガンガン使う人がどのくらいいるか謎なので、優先度は低いです。

現在のGHCは -mavx512f みたいなオプションで一部のAVX-512の機能を有効化できます。NCGではあまり使っていないと思いますが、LLVMバックエンドで利用されるかもしれません。ただ、VL拡張やBW、DQ拡張などのよく使われそうなやつが未サポートなので、それらを使えるようにするフラグを追加すると良いかもしれません。

前に「コンパイル時定数しか受け付けない引数」という記事を書きましたが、現状のGHCのSIMDのシャッフルに対応する組み込み関数は、(# Int#, Int#, Int#, Int# #) のような型のコンパイル時定数しか受け付けません。これはラップしづらいので、記事に書いたように、型クラスで

type Int32X4Shuffle :: (Natural, Natural, Natural, Natural) -> Constraint
class Int32X4Shuffle mask where
  shuffleInt32X4# :: Int32X4# -> Int32X4# -> Int32X4#
-- Int32X4Shuffle のインスタンスはcompiler magicで作る

-- ラッパー
shuffleInt32X4 :: X4 Int32 -> X4 Int32 -> forall mask -> Int32X4Shuffle mask => X4 Int32

という形で使えるようにしたいです。これの実装にはcompiler magicが必要ですが、GHCに手を加えなくても、コンパイラープラグインの形でできるかもしれません（未検証）。このアイディアを検証したいです。

次のステップの前に

HaskellでのSIMDについてやりたいことは色々ありますが、「HaskellでSIMDを使えるようにするプロジェクト」に対する私の貢献は余暇の時間でやっており、しかも私自身はSIMDを活用して解決しなければならない課題を抱えているわけではない（「SIMDを使えるようにすること」自体が楽しくてやっている）ので、進み具合は私の時間の使い方に左右されます。

具体的には、これからしばらくの間は、6月の技術書典18に向けた執筆活動を優先させたいです。ただ、「SIMDを扱うためのライブラリー」は早めに公開できるようにしたいところです。

Spread the love