Haskellにインラインアセンブリを書く

Haskellにインラインアセンブリ

先日 inline-c という、Haskellのソースの中にインラインでCのコードを書けるようにするパッケージがリリースされました。これまでの類似のパッケージよりも使いやすい感じで、愚直にFFIを書いたり、ブリッジライブラリを書いたり使ったりするよりやっぱり楽なもんだなあと感心していたんですが、これってもしかしたらインラインのCのインラインにアセンブリ書けば、Haskellに直接インラインでアセンブリを書くこともできるんじゃないか？と、ふと思ったので、やってみたら普通にできましたという話です。

inline-c

inline-c に関しては、GitHubレポジトリに丁寧な README.md があるので、詳しくはこちらを見てくださいというところなんですが、せっかくなので少し試してみましょうか。

Test.hs

{-# LANGUAGE QuasiQuotes #-}

import qualified Language.C.Inline as C

sigma :: C.CInt -> IO C.CInt
sigma x = [C.block| int {
  int i, ret = 0;
  for (i = 1; i <= $(int x); i++)
    ret += i;
  return ret;
}|]

main :: IO ()
main = print =<< sigma 10

inline-cでは、QuasiQuoteでCのコードをHaskellに埋め込むので、LANGUAGEプラグマでこれを使えるようにしないといけません。sigmaという関数は、ご覧のとおり、引数xに対して、1からxまでの和を計算する関数です。基本的にCで普通に書けばいいのですが、$(<型名> <変数名>)などのunquoteによって、Haskell側の値にアクセスできるようになっています。unquoteは単に変数にアクセス出来るだけではなくて、他にもいろいろできて、しかも拡張可能な設計になっているようですが、今回はそこには触れません。

Cの関数をHaskellに埋め込めたのはいいけど、これコンパイルできるのか？っていう話ですが、残念ながら全く普通にというわけにはいかなくて、まずTest.hsをコンパイルすると、コンパイル時IOによって同じディレクトリにTest.cというコードが生成されます。この中に必要なCのコードができ上がるので、これを一緒にリンクする必要があるというわけです。

$ ghc -c Test.hs
$ ls
Test.c  Test.hi  Test.hs  Test.o

普段のHaskellコードのコンパイルでできる.hiファイルと、.oファイルの他に、Test.cというファイルができているのがわかると思います。

Test.c

$ cat Test.c

int inline_c_0_7f1539fa51c0ab8bf8eb3554152827d72733b9bd(int x_inline_c_0) {

  int i, ret = 0;
  for (i = 1; i <= x_inline_c_0; i++)
    ret += i;
  return ret;

}

あとはこれを一緒にリンクしてやればOKです。

$ gcc -c Test.c -o Test.c.o
$ ghc Test.o Test.c.o -o Test
$ ./Test
55

QuasiQuoteは、C.blockの他にもいくつかあって、例えば、インラインのCのコードが単一の式からなるのであれば、

add :: C.CInt -> C.CInt -> IO C.CInt
add x y = [C.exp| int{ $(int x) + $(int y) } |]

のようにreturnを省略できるC.expが使えますし、さらにこの式が「純粋」であるなら、

add' :: C.CInt -> C.CInt -> C.CInt
add' x y = [C.pure| int{ $(int x) + $(int y) } |]

C.pureというのを用いて、IOを含まない純粋な関数を得ることも出来ます。

以上が簡単なinline-cの使い方です。

アセンブリを埋め込む

さて、ではアセンブリを埋め込んでみます。と言っても、別に特別なことはなくて、gcc形式のインラインアセンブリを埋め込んでやるだけです。gcc形式だと環境に制限を与えそうですが、現在のghcはgccに依存していて、Windows版のGHCもgccを同梱しているので、多分どこでも動くと思います。

popcnt :: C.CInt -> IO C.CInt
popcnt x = [C.block| int{
  int ret;
  asm("popcnt %1, %0": "=r"(ret): "r"($(int x)));
  return ret;
}|]

main :: IO ()
main = print =<< popcnt 123

あまりおもしろい例が思いつかなかったので、popcntを呼び出してみました。これはオペランドの立っているビットの数を数える命令です。

$ ghc -c Test.hs -o Test.o
$ gcc -c Test.c -o Test.c.o
$ ghc Test.o Test.c.o -o Test
$ ./Test
6

特に以外でもないのですが、ふつうに動きました。多分これはinline-cが中身のCコードにはあまり深入りせずに、そのままCのコードに吐き出してくれているからだと思います。

$ cat Test.c

int inline_c_3_3833d7deec70767d5514934fe8e980fdad611620(int x_inline_c_0) {

  int ret;
  asm("popcnt %1, %0": "=r"(ret): "r"(x_inline_c_0));
  return ret;

}

使用例

アセンブリ埋め込めて何か嬉しいんか？っていう話なんですが、今の御時世、速度のためにアセンブリを書くなんていう機会もとんとなくなってしまったので、あまり思いつかないんですが、Cからではアクセスしづらいようなものに使えるような気がします。

RDRAND

例えば、最近のIntel CPUがサポートしているハードウェア乱数命令RDRANDなんていうのがあります。

rdrand :: IO C.CInt
rdrand = [C.block| int{
  int ret;
  asm("rdrand %0": "=r"(ret));
  return ret;
}|]

main :: IO ()
main = do
  print =<< rdrand
  print =<< rdrand
  print =<< rdrand

これをコンパイルして実行すると、

$ ./Test
-1932382952
1751702702
-1731380647

それっぽく動いているような気がします。
まあでも、gccのintrinsic関数でもいいような気がします。

{-# LANGUAGE QuasiQuotes     #-}
{-# LANGUAGE TemplateHaskell #-}

import qualified Language.C.Inline as C

C.include "<immintrin.h>"

rdrand' :: IO C.CInt
rdrand' = [C.block| int {
  unsigned int ret;
  _rdrand32_step(&ret);
  return ret;
}|]

C.include で生成するCのコードの先頭にincludeディレクティブを挿入できます。また、これを使うために{-# LANGUAGE TemplateHaskell #-}が必要になっています。

CPUID

もう一つ例として、CPUID命令を呼び出してみます。
この命令は、CPUの種類や、サポートしている命令セットなどを問い合わせたりするのに使います。

eaxレジスタに0を入れて呼び出すと、そのCPUのベンダーIDが12バイトで、ebx, edx, ecxレジスタの順に格納されて返ってきます。

さて、やりたいことはCPUIDを呼び出して、返ってきた値を取得するだけなのですが、3つの値を返さないといけないので、ちょっと面倒です。というのも、C言語はタプルを直接返せないので、関数の返り値としてはマーシャリングできません。普通C言語だと、こういう場合メモリを呼び出し側で確保して、関数にはポインタをわたしてやるのですが、そういうことをHaskellでやります。

withPtrおよびそのバリエーションの関数としてそういうのがinline-cに用意されています。複数のメモリを確保しつつ、そこに書き込んだ値が全体の返り値になるwithPtrs_という関数が便利そうです。

{-# LANGUAGE QuasiQuotes     #-}

import           Data.Bits
import           Data.Char
import qualified Language.C.Inline as C

cpuid :: C.CInt -> IO (C.CUInt, C.CUInt, C.CUInt)
cpuid a =
  C.withPtrs_ $ \(b, d, c) -> [C.block| void{
    asm("movl %0, %%eax;"
        "cpuid;"
        "movl %%ebx, (%1);"
        "movl %%edx, (%2);"
        "movl %%ecx, (%3);"
        :
        : "r"($(int a))
        , "r"($(unsigned int *b)), "r"($(unsigned int *d)), "r"($(unsigned int *c))
        : "%rax", "%rbx", "%rcx", "%rdx"
        );
  }|]

生成されたコードを見てみると、たしかにポインタを受け渡しするコードが生成されているようです。

void inline_c_6_44e2d905032613613a62d5703d30b021155a7f78(int a_inline_c_0, unsigned * b_inline_c_1, unsigned * d_inline_c_2, unsigned * c_inline_c_3) {

    asm("movl %0, %%eax;"
        "cpuid;"
        "movl %%ebx, (%1);"
        "movl %%edx, (%2);"
        "movl %%ecx, (%3);"
        :
        : "r"(a_inline_c_0), "r"(b_inline_c_1), "r"(d_inline_c_2), "r"(c_inline_c_3)
        : "%rax", "%rbx", "%rcx", "%rdx"
        );

}

あとはこれを呼び出して、結果を表示するコードを書きます。

main :: IO ()
main = do
  (ebx, edx, ecx) <- cpuid 0
  let f n = [ chr $ fromIntegral $ n `shiftR` (8*i) .&. 0xff | i <- [0..3]]
  putStrLn $ f ebx ++ f edx ++ f ecx

コンパイルして実行します。

$ ./Test
GenuineIntel

うまく動いたようです。