Haskell
assembly
ghc

【低レベルHaskell】Haskell (GHC) でもインラインアセンブリに肉薄したい!

C言語やC++にはインラインアセンブリやintrinsics(組み込み関数)があり、最近のCPUで追加された命令を活用したコードを簡単に書くことができます。

一方、Haskell (GHC) にはそのような仕組みはありません1。しかしそれで諦めるのは早い。なんとかしてHaskellからインラインアセンブリっぽいことをしてみましょう。

筆者の環境はx86_64なMacで動かしているGHC 8.6.5です。他の環境では状況が異なる可能性があります。C言語の処理系はClangを使いますが、GCCでも多分同じです。


題材:64ビット整数の積の上位と下位

2つの64ビット整数の積を計算し、上位64ビットと下位64ビット(合わせて128ビット)をそれぞれ計算することを考えます。

C言語やHaskellにあるような通常の乗算 (*) :: Word64 -> Word64 -> Word64 では下位64ビットしか計算できません。一方、x86のアセンブリ言語のレベルでは、乗算の際に上位64ビットも一緒に計算されます。

こういう「機械語レベルでは容易だが、C言語やHaskellのレベルでは非自明な処理」には、インラインアセンブリ(や組み込み関数)を使いたいところです。

(実はGHCには timesWord2# :: Word# -> Word# -> (# Word#, Word# #) という組み込み関数があるので、これを使えば一発です。それでもこの題材を選んだのは、「GHCの組み込み関数と比較した場合にどの程度遅くなるか」を検証できるようにするためです。)


C言語では

GCC/Clangには __int128 型があるので、それを使えば一発です。インラインアセンブリもintrinsicsも必要なかった(ぇ)

unsigned __int128 wideningMul(uint64_t a, uint64_t b)

{
return (unsigned __int128)a * (unsigned __int128)b;
}

あえてインラインアセンブリで書くとすれば、こんな感じでしょうか:

uint64_t wideningMul_inlasm(uint64_t a, uint64_t b, uint64_t *outHigh)

{
uint64_t lo, hi;
asm("movq %2, %%rax;"
// mulq は %rax とオペランド(この場合は %3)の積を計算し、
// 上位64ビットを %rdx, 下位64ビットを %rax に入れる
"mulq %3;"
"movq %%rax, %0;"
"movq %%rdx, %1;"
: "=r"(lo), "=r"(hi)
: "r"(a), "r"(b)
: "%rax", "%rdx");
*outHigh = hi;
return lo;
}

さて、この処理では2つの uint64_t を受け取って、2つの uint64_t を返します。C言語の構文には多値返却がないので、値の返却方法として次のいずれかを選択することになります:


  • 構造体の値返し:struct uint128 { uint64_t lo, hi; } のような構造体を定義して値で返す



    • unsigned __int128 を値で返すのも内部的にはこれっぽい(詳しくは x86_64 のABIを確認。Unix系の場合は「System V ABI x86_64」とでもググればよろしい)



  • ポインタを受け取って渡す:2番目以降の返却値の受け渡しする場所をポインターとして引数に取る(例:さっき書いた wideningMul_inlasm 関数)

前者の例として、C標準の div, ldiv, lldiv 関数は ホニャララdiv_t という構造体を値で返しています。

構造体を値で返すメリットは、構造体が小さければ値をレジスタに載せたまま返すことができることです。

一方、構造体の値返しのデメリットは、他の言語のFFIが対応していない場合があることです。実際、現在のGHCのC FFIは構造体の値での受け渡すことに対応していません。

一応、Cの構造体をFFIで値渡しできるようにしよう、という提案 (proposal) がありますが、動きがありません:


C FFIを使う(ポインタ使用)


Safe FFI

Haskellにできないことは他の言語の力を借りよう!ということで、HaskellにはFFIがあります。これを使えばC言語で書いた関数を呼び出すことができます。

早速やってみましょう:

#include <stdint.h>


extern uint64_t wideningMul_with_ptr(uint64_t a, uint64_t b, uint64_t *outHigh)
{
unsigned __int128 result = (unsigned __int128)a * (unsigned __int128)b;
*outHigh = (uint64_t)(result >> 64);
return (uint64_t)result;
}

先ほど書いたように、現状のGHCのFFIでは構造体を値渡しできないので、返り値の1つはポインタを経由することにします。

Haskell側のコードはこんな感じです:

foreign import ccall "wideningMul_with_ptr"

c_wideningMul_with_ptr :: Word64 -> Word64 -> Ptr Word64 -> IO Word64

wideningMulWithPtr :: Word64 -> Word64 -> Word128
wideningMulWithPtr !a !b = unsafePerformIO $ do
Marshal.alloca $ \outHigh -> do
lo <- c_wideningMul_with_ptr a b outHigh
hi <- peek outHigh
return $ Word128 hi lo

ポインタを扱うということは、領域を確保したり値を読みだしたりするために IO を行うということです。しかし「64ビット整数の乗算」という処理全体としては純粋だと考えられるので、 unsafePerformIO を使って純粋な関数として書いています2

(ちなみに、 Word128 型は wide-word パッケージの Data.WideWord.Word128 型を使いました。)

試してみると、

> wideningMulWithPtr 123 456

56088
> 123 * 456
56088
> wideningMulWithPtr (2^63) (2^62) -- CPUの命令を使う
42535295865117307932921825928971026432
> 2^125 -- 多倍長整数で計算する
42535295865117307932921825928971026432

という風に正しく計算できているようです。

C言語のコードを別のファイルに書くのが面倒であれば、

のように inline-c のようなパッケージを使うのも一つの手段かもしれません。


Unsafe FFI

さて、HaskellのFFIにはsafety levelという概念があります。デフォルトではsafeで、これは「呼び出した外部のコードからHaskellの関数をコールバックしても安全」であることを意味します。

safeの逆はunsafeで、「呼び出した外部のコードからHaskellにコールバックしたら何が起きるかわからんよ」という意味です。

unsafeなFFIはリスクがある分、オーバーヘッドが小さいことが期待できます。

safety levelの指定は、 foreign import 宣言の呼び出し規約の直後に safe または unsafe を書きます:

foreign import ccall unsafe "wideningMul_with_ptr"

c_wideningMul_with_ptr :: Word64 -> Word64 -> Ptr Word64 -> IO Word64

inline-cパッケージを使う場合は、quasiquoter (exp, pure, block) として Language.C.Inline.Unsafe にあるものを使います。


C FFIを使う(2回呼び出す)

先のセクションでは、Cで書いた関数から複数の値を返すためにポインタを経由しました。

しかし、ポインタで値を受け渡しするのは、レジスタ渡しに比べて遅い気がします[要出典]。ポインタを使わずに値を受け渡しできれば、それに越したことはありません。

幸い、今回の題材は「64ビット整数の乗算」という、コストが低い演算です。これだったら、上位64ビット、下位64ビットのそれぞれに対して同じ計算をするのもアリではないでしょうか?

// 下位64ビットを計算する

extern uint64_t wideningMul_lo(uint64_t a, uint64_t b)
{
unsigned __int128 result = (unsigned __int128)a * (unsigned __int128)b;
return (uint64_t)result;
}

// 上位64ビットを計算する
extern uint64_t wideningMul_hi(uint64_t a, uint64_t b)
{
unsigned __int128 result = (unsigned __int128)a * (unsigned __int128)b;
return (uint64_t)(result >> 64);
}

foreign import ccall unsafe "wideningMul_lo"

c_wideningMul_lo :: Word64 -> Word64 -> Word64

foreign import ccall unsafe "wideningMul_hi"
c_wideningMul_hi :: Word64 -> Word64 -> Word64

wideningMul2 :: Word64 -> Word64 -> Word128
wideningMul2 !a !b = Word128 (c_wideningMul_hi a b) (c_wideningMul_lo a b)

実際のところ「2回呼び出してレジスタ渡しで完結させる」方針が「ポインタで渡す」方針に比べてどうなのかは、後で比較します。


黒魔術:foreign import prim

注意:unsafePerformIOやUnsafe FFIを見て「unsafeこわ」と思った健全な方は、今のうちに帰っパ3しておくことをお勧めします。


GHCのPrimOpsについて

さて、たかだかCPUの命令1つのためにFFIでC言語で書いた関数を呼び出す、というのはどうしてもコストがかかる[要出典]印象が拭えません。

FFIのオーバーヘッドをもっと削減し、GHCの組み込み関数に近いパフォーマンスを出すようにはできないでしょうか?

GHC Wikiのページ (Primitive Operations (PrimOps)) によると、GHCの組み込み関数には、inline PrimOps, out-of-line PrimOps, foreign out-of-line PrimOps (foreign import prim) の3種類があるようです。


  • inline PrimOps: その場で命令列に展開される。GHCにハードコードされている。

  • out-of-line PrimOps: 専用の呼び出し規約に従う。GHCにハードコードされている。

  • foreign out-of-line PrimOps: 専用の呼び出し規約に従う。ライブラリー開発者が foreign import prim によって定義できる。GHC本体に付属するライブラリー向け。

この中では(というか、CPUの特定の命令を叩くためにHaskellで使えるあらゆる方法の中で)inline PrimOpsが一番低コストだと思われますが、命令一つを使いたいがためにGHC本体に手を入れるのはやりすぎ4です。というわけで、必然的に foreign out-of-line PrimOps を使うことになります。


foreign import prim を使う

foreign import prim を使うHaskellコードの例は次のようになります:

{-# LANGUAGE GHCForeignImportPrim, UnliftedFFITypes, MagicHash, UnboxedTuples #-}

foreign import prim "wideningMul_prim"
wideningMul_prim# :: Word# -> Word# -> (# Word#, Word# #)

wideningMul :: Word64 -> Word64 -> Word128
wideningMul (W64# a) (W64# b)
= case wideningMul_prim# a b of
(# lo, hi #) -> Word128 (W64# hi) (W64# lo)

一見すると foreign import の呼び出し規約が ccall から prim に変わっただけですね。安心した!(型名やタプルに # がついているのは低レベルHaskellではよくあることなので今更驚きません。ccallでも # を使いまくるFFIはできますし)

使うGHC拡張ですが、 foreign import prim を使うには GHCForeignImportPrim 拡張5が必要です。また、引数と返り値は基本的にunboxed typesしか使えないので、 UnliftedFFITypes 拡張も必要になります。MagicHash, UnboxedTuples は言わずもがなですね。


foreign out-of-line PrimOpsを実装する

Haskell側でオレオレPrimOpsを使う準備はできたので、次は実装を用意します。

GHCのout-of-line PrimOpsはCmmで書くことが想定されているようですが、


  • Cmmはよくわからん!

  • Cmmでインラインアセンブリや環境依存の組み込み関数を使えるかわからん!

ので、アセンブリで直接書くことにします(後述しますが、Cmm, アセンブリベタ書きの他にLLVM IRをいじる、という第3の道もあります)。

GHCのPrimOps用の呼び出し規約は、x86_64の場合は


  • 引数と返り値は全てレジスタで渡す。最初が %rbx, 2番目が %r14, 3番目が %rsi, 4番目が %rdi, ...

  • 呼び出し側に戻る時は jmp *(%rbp) する

という感じのようです。レジスタの使い方に関しては MachRegs.h を参照してください。

実際のアセンブリコードは次のようになります:

    .globl _wideningMul_prim

_wideningMul_prim:
## 第一引数:%rbx
## 第二引数:%r14
movq %rbx, %rax
mulq %r14 ## %rax と %r14の積を計算して、上位を %rdx, 下位を %rax に入れる
movq %rax, %rbx ## %rbx に最初の返却値 (lo) を入れる
movq %rdx, %r14 ## %r14 に第二の返却値 (hi) を入れる
jmp *(%rbp)

先ほどリンクを貼った MachRegs.h のコメントによると %rax%rdx はcaller-saveのようなので好きに使って大丈夫です。たぶん。

なお、どうしてもC言語で書きたかったら、「特定のプロトタイプを持つ関数として書いてClangでコンパイル、-S -emit-llvm で出力されたLLVM IRをいじって呼び出し規約を cc106 にする」という方法があるみたいです。詳しくは下記リンクを参照してください。

参考になるかもしれないリンク集:


ベンチマーク

「64ビット整数2つの乗算を128ビット整数として得る」ための実装方法を色々考えました。というわけで比較です。比較対象は以下です:


  • Haskellの多倍長演算(Integer 型)を使う

  • オペランドを wide-word パッケージの Word128 型に変換してから乗算する



    • wide-word パッケージは内部的には timesWord2# を使っている




  • GHC.PrimtimesWord2# を使う


    • GHCの内部的には WordMul2 と呼ばれる inline PrimOps の一種



  • C FFIを使う


    • Safe FFI / Unsafe FFI

    • ポインタ(unsafePerformIO / unsafeDupablePerformIO / ****PerformIO) / 2回呼び出し




  • foreign import prim を使う

64ビット演算2つの乗算で128ビットを得る演算としてはすでにGHC組み込みの timesWord2# があるので、「inline PrimOpsと(foreign) out-of-line PrimOpsでどのくらい差があるのか」の比較も行えます。

どれが一番早いか予想しておくと、inline PrimOpsである timesWord2# が一番早く、foreign out-of-line PrimOps (foreign import prim) が続くと考えられます。Integer の多倍長演算は一番遅い気がします。

実際のベンチマーク結果は以下です:

benchmarked Haskell/via Integer

time 87.65 ns (86.21 ns .. 89.81 ns)
0.995 R² (0.990 R² .. 0.999 R²)
mean 88.75 ns (88.00 ns .. 90.05 ns)
std dev 3.287 ns (2.321 ns .. 4.606 ns)
variance introduced by outliers: 18% (moderately inflated)

benchmarked Haskell/via Word128
time 10.81 ns (10.48 ns .. 11.05 ns)
0.994 R² (0.988 R² .. 0.997 R²)
mean 11.22 ns (11.04 ns .. 11.51 ns)
std dev 750.1 ps (536.4 ps .. 995.9 ps)
variance introduced by outliers: 41% (moderately inflated)

benchmarked Haskell/timesWord2#
time 11.38 ns (11.08 ns .. 11.88 ns)
0.983 R² (0.956 R² .. 0.998 R²)
mean 11.29 ns (11.13 ns .. 11.92 ns)
std dev 785.5 ps (406.9 ps .. 1.636 ns)
variance introduced by outliers: 42% (moderately inflated)

benchmarked safe FFI/pointer/unsafePerformIO
time 200.2 ns (197.9 ns .. 203.6 ns)
0.997 R² (0.994 R² .. 0.999 R²)
mean 206.5 ns (204.3 ns .. 210.3 ns)
std dev 9.844 ns (6.456 ns .. 15.06 ns)
variance introduced by outliers: 27% (moderately inflated)

benchmarked safe FFI/pointer/unsafeDupablePerformIO
time 204.5 ns (202.5 ns .. 207.1 ns)
0.999 R² (0.998 R² .. 1.000 R²)
mean 205.6 ns (204.5 ns .. 206.9 ns)
std dev 3.975 ns (3.177 ns .. 5.495 ns)

benchmarked safe FFI/pointer/accursedUnutterablePerformIO
time 205.2 ns (200.7 ns .. 209.3 ns)
0.997 R² (0.995 R² .. 0.998 R²)
mean 212.9 ns (209.3 ns .. 218.4 ns)
std dev 14.79 ns (10.53 ns .. 21.35 ns)
variance introduced by outliers: 45% (moderately inflated)

benchmarked safe FFI/calling twice
time 217.6 ns (213.2 ns .. 223.7 ns)
0.988 R² (0.969 R² .. 0.997 R²)
mean 234.1 ns (228.2 ns .. 244.2 ns)
std dev 25.88 ns (18.20 ns .. 35.56 ns)
variance introduced by outliers: 67% (severely inflated)

benchmarked unsafe FFI/pointer/unsafePerformIO
time 32.70 ns (32.24 ns .. 33.14 ns)
0.999 R² (0.997 R² .. 0.999 R²)
mean 32.92 ns (32.67 ns .. 33.36 ns)
std dev 1.131 ns (556.8 ps .. 2.215 ns)
variance introduced by outliers: 15% (moderately inflated)

benchmarked unsafe FFI/pointer/unsafeDupablePerformIO
time 32.11 ns (30.94 ns .. 33.40 ns)
0.991 R² (0.985 R² .. 0.997 R²)
mean 33.29 ns (32.65 ns .. 34.33 ns)
std dev 2.810 ns (2.055 ns .. 4.352 ns)
variance introduced by outliers: 55% (severely inflated)

benchmarked unsafe FFI/pointer/accursedUnutterablePerformIO
time 34.13 ns (32.00 ns .. 36.37 ns)
0.981 R² (0.967 R² .. 0.997 R²)
mean 31.47 ns (31.08 ns .. 32.25 ns)
std dev 1.780 ns (1.032 ns .. 3.210 ns)
variance introduced by outliers: 34% (moderately inflated)

benchmarked unsafe FFI/calling twice
time 13.67 ns (13.46 ns .. 13.85 ns)
0.999 R² (0.998 R² .. 0.999 R²)
mean 13.64 ns (13.54 ns .. 13.75 ns)
std dev 355.1 ps (295.4 ps .. 449.9 ps)
variance introduced by outliers: 11% (moderately inflated)

benchmarked foreign import prim
time 13.66 ns (12.23 ns .. 15.10 ns)
0.961 R² (0.941 R² .. 0.990 R²)
mean 12.36 ns (12.11 ns .. 12.83 ns)
std dev 1.132 ns (674.7 ps .. 1.689 ns)
variance introduced by outliers: 58% (severely inflated)

まとめると


  • 最速はGHC組み込みの timesWord2# 関数で、11ns程度。内部的に timesWord2# を使う Word128 も同等のパフォーマンスが出ている。

  • 次点は foreign import prim で、13ns程度。「unsafe FFI + ポインタを使わず2回呼び出す」もほぼ同等のパフォーマンス。

  • その次は「unsafe FFI + ポインタ渡し」で、32ns程度。unsafePerformIOの代わりにunsafeDupablePerformIOや名状し難い冒涜的なやつを使ってもほぼ同じ。

  • その次が Integer の多倍長計算を使うやつで、87ns程度。

  • 最下位は「safe FFI」で、200ns程度。unsafePerformIOの代わりにunsafeDupablePerformIOや名状し難い冒涜的なやつを使ってもほぼ同じ。「ポインタ渡し」よりも「2回呼び出し」の方が若干遅いのは、safe FFIのコストがポインタ渡しのコストを上回ったからでしょう。

となります。

使ったソースコードはそのうち公開します。


雑感・まとめ

GHC組み込みでインラインの命令列に置き換えられるものが最速なのは当然として、 foreign import prim でそれに肉薄するパフォーマンスが出ることがわかりました。

意外なのは「unsafe FFI + ポインタを使わず2回呼び出し」が foreign import prim とほぼ互角だったことです。CのFFIでもやり方次第では foreign import prim に匹敵するということです。パフォーマンスの差がわずかであれば、 foreign import prim の黒魔術よりもCのFFIを選択するという判断はアリでしょう。

逆に、CのFFIを使っていても、不必要にsafeレベルを使っていると、「CPUの便利な命令を使わずに愚直にHaskellで書いた」バージョンよりも遅くなってしまいました。


  • 高速化のためにCのFFIを使う場合はsafety levelが不必要にsafeになっていないか確認しろ

というのは結構大事な教訓ではないでしょうか。





  1. 一応、SIMD命令はある程度使えるようですが。 



  2. Unsafe Haskellに精通した方なら「この場合なら unsafeDupablePerformIO が使えるし、なんなら 〈検閲削除〉PerformIO の方が…」と思われるかもしれません。安心してください、これらの実験結果も最後に載せておきます。 



  3. 「帰ってパイソンしとこ……」の略[嘘] 



  4. GHC本体に取り込んでもらえるなら苦労の甲斐もあるかもしれませんが、ニッチな命令に対応する組み込み関数をGHC本体に取り込んでもらえる見込みは低いでしょう。GHCにAES命令を入れようというIssueが過去にあったようですが、won't fixでcloseされています。  



  5. いかにも「GHC専用!外部のライブラリーは使うな!」感がある名前ですね。実際、GHC User's Guideにこの拡張の名前は載っていません(foreign import prim 自体は載っています)。 



  6. cc10というのはGHCの呼び出し規約のLLVM内での名称です(→LLVMの呼び出し規約のドキュメント)。GHCのLLVMバックエンドが使っていると思われます。