Haskell Advent Calendar 2019 5日目

この冬、神速のサンタクロースがやってくる——

Haskellにおいて、バイト列の表現はByteStringが定番である。ByteStringはPinned領域に直接格納され、空間効率はリストに比べればはるかに良い。しかし、Pinned領域にあるとヒープフラグメンテーションが起こりやすくなるということでもあり、細かい文字列をつなぎ合わせるような使い方はパフォーマンスに悪影響が及ぶ。そのような問題を避けるため、ビルダーと呼ばれる構造が用意されている。 Data.ByteString.Builderは、 word8 42 <> byteString "hello" <> doubleLE 42 のように細かいプリミティブを連結し、 toLazyByteString を呼ぶと最後にByteStringを一気に鋳出せるという仕組みである。ByteStringをちまちま結合するよりも格段に高速であり、waiなどのインターフェイスにも利用されている。

しかし、各パーツをバッファに書き込んでポインタをずらすだけで済む処理にしてはやたら遅い。Builderの実体は、Freeモナドのような再帰的なADTにコンストラクタを重ねる関数であり、構築・分解するオーバーヘッドが大きいのだろう——結局、中間構造を作らないという目的を完全に達成したわけではないのだ。

type BuildStep a = BufferRange -> IO (BuildSignal a) data BuildSignal a = Done {-# UNPACK #-} ! (Ptr Word8) a | BufferFull {-# UNPACK #-} ! Int {-# UNPACK #-} ! (Ptr Word8) (BuildStep a) | InsertChunk {-# UNPACK #-} ! (Ptr Word8) S.ByteString (BuildStep a)

そのパフォーマンスの問題に目を付けたのがfast-builderで、ポインタやRealWorld(IOモナドの実装に使われるステートトークン)をUnboxedタプルに格納するという徹底的な最適化を施した。

newtype Builder = Builder { unBuilder :: DataSink -> BuilderState -> BuilderState } type BuilderState = ( # Addr # , Addr # , State # RealWorld # )

その結果、標準のビルダーの数倍もの高速化を達成した。fast-builderは、バッファに直接データを書き込むという動作を基本とし、3つのモードをサポートしている。

toStrictByteString: 指数関数的に大きくなるバッファにデータを書き込み、StrictなByteStringとして取り出す。

toLazyByteString: スレッドを立ち上げ、バッファが満タンになるたびに新しいチャンクを生成することによって、LazyなByteStringを作り出す。

hPutBuilder: バッファがいっぱいになるたびに、ハンドルにその内容を書き込む。

3つのモードを実現するために、内部のDataSinkというデータ型に必要な情報をまとめているため、拡張性に乏しいという難点がある。ローレベルな内部構造も拡張の難しさに拍車をかけている。

data DataSink = DynamicSink ! (IORef DynamicSink) | GrowingBuffer ! (IORef (ForeignPtr Word8)) | HandleSink ! IO.Handle ! Int ! (IORef Queue)

そんな問題を解決するため、新たにmasonというライブラリを作った。

github.com

masonは、ビルダーを実行する手段として4つの関数を提供する。Strict、LazyなByteStringと、ハンドルおよびソケットへの書き込みだ。StrictなByteStringやソケットに使えるというだけでもオリジナルよりかなり便利になったが、ライブラリに変更を加えずとも新しい使い道を定義することもできる。

toStrictByteString :: BuilderFor GrowingBuffer -> B.ByteString toLazyByteString :: BuilderFor Channel -> BL.ByteString hPutBuilderLen :: Handle -> BuilderFor PutBuilderEnv -> IO Int sendBuilder :: S.Socket -> BuilderFor SocketEnv -> IO Int

そのからくりは型パラメータにある。Builderの中身は、バックエンドに応じた値と、バッファへのポインタを Buffer 型として受け取り、何らかのアクションを実行して返すという至極単純なものだ。

data BuilderFor s = Builder { unBuilder :: s -> Buffer -> IO Buffer } data Buffer = Buffer { bEnd :: {-# UNPACK #-} ! (Ptr Word8) , bCur :: {-# UNPACK #-} ! (Ptr Word8) } type Builder = forall s . Buildable s => BuilderFor s

Buildableクラスを通して、目的に合わせた処理を具現化している。

class Buildable s where byteString :: B.ByteString -> BuilderFor s flush :: BuilderFor s allocate :: Int -> BuilderFor s

StrictなByteStringを生成する場合、 allocate n は現在の倍以上かつn以上の長さのバッファを作るアクションで、flushは何もしない。 byteString は与えられたバイト列をバッファにコピーする関数となる。ハンドルに書き込む場合、 flush は現在のバッファの中身をハンドルに落とし、 byteString の引数が大きい場合はバッファを介さずに直接書き込む——というように、バックエンドに合わせた処理を実装している。

これらのメソッドが与えられると、ensureという関数が定義できる。 ensure n は、バッファにnバイト以上の余裕があることを保証して関数を呼び出し、足りない場合はその前に領域を確保する。実質的にはこれが Builder のスマートコンストラクタであり、新たな部品を作ることを容易にしている。

ensure :: Int -> (Buffer -> IO Buffer) -> Builder ensure mlen cont = Builder $ \ env buf @ (Buffer end ptr) -> if ptr `plusPtr` mlen >= end then do buf' @ (Buffer end' ptr') <- unBuilder flush env buf if mlen <= minusPtr end' ptr' then cont buf' else unBuilder (allocate mlen) env buf' >>= cont else cont buf

このようなチェックをいちいち実行するのは無駄だが、 ensure はモノイドの準同型写像であるため、融合変換によって一つにまとめることができる。最適化が理想的に回れば、可変長の要素がある時だけ空きがチェックされるようにコンパイルされるはずだ。flushやallocateなどが呼ばれるタイミングが変わってしまう半グレ的な最適化だが、まさに必要な振る舞いなので仕方がない。

{-# INLINE[1] ensure #-} {-# RULES "<>/ensure" forall m n f g. ensure m f <> ensure n g = ensure (m + n) (f >=> g) #-}

masonは基本的に最適化が鬼回りする前提で設計されているので、インライン化を怠ったり、中途半端な場所に保存したりすると最高のパフォーマンスが発揮できない点には注意する必要がある。

長々とテクニックをひけらかしたところで、実際遅かったら意味がない。fast-builderのJSONデータを扱うベンチマークを拝借し、Textのエンコード方法などを少し改良して比較してみた*1。

mason/hPutBuilder mean 187.6 μs ( +- 14.02 μs ) fast-builder/hPutBuilder mean 366.5 μs ( +- 50.78 μs ) bytestring/hPutBuilder mean 294.3 μs ( +- 83.62 μs ) mason/toStrictByteString mean 105.1 μs ( +- 23.33 μs ) fast-builder/toStrictByteString mean 223.1 μs ( +- 7.506 μs ) bytestring/toLazyByteString mean 249.3 μs ( +- 16.76 μs ) mason/toLazyByteString mean 101.9 μs ( +- 4.787 μs ) fast-builder/toLazyByteString mean 226.4 μs ( +- 10.56 μs )

圧勝である。bytestringやfast-builderの倍以上と、大変満足のいく結果だ。元のベンチマークのようにTextを一旦 Data.Text.Encoding.encodeUtf8 で変換すると差は若干縮まるが、それでも最速であることに変わりはない。

fast-builder/hPutBuilder mean 372.2 μs ( +- 68.48 μs ) bytestring/hPutBuilder mean 603.6 μs ( +- 641.7 μs ) mason/toStrictByteString mean 215.4 μs ( +- 71.29 μs ) fast-builder/toStrictByteString mean 302.2 μs ( +- 78.57 μs ) bytestring/toLazyByteString mean 498.3 μs ( +- 154.5 μs ) mason/toLazyByteString mean 196.2 μs ( +- 10.76 μs ) fast-builder/toLazyByteString mean 230.7 μs ( +- 17.19 μs ) bytestring/toLazyByteString mean 381.3 μs ( +- 18.35 μs )

fumieval.hatenablog.com

wineryのビルダーをfast-builderからmasonに差し替え、ベンチマークを実行した。

serialise/list/winery/old mean 223.2 μs ( +- 16.28 μs ) serialise/list/winery/new mean 119.8 μs ( +- 21.57 μs ) serialise/list/binary mean 1.597 ms ( +- 394.7 μs ) serialise/list/cereal mean 747.0 μs ( +- 196.5 μs ) serialise/list/serialise mean 494.4 μs ( +- 126.2 μs ) serialise/list/store mean 50.83 μs ( +- 4.517 μs ) serialise/list/aeson mean 7.403 ms ( +- 1.628 ms ) serialise/item/winery/old mean 234.3 ns ( +- 18.57 ns ) serialise/item/winery/new mean 62.61 ns ( +- 4.216 ns ) serialise/item/binary mean 1.738 μs ( +- 383.9 ns ) serialise/item/cereal mean 709.8 ns ( +- 223.7 ns ) serialise/item/serialise mean 583.1 ns ( +- 132.6 ns ) serialise/item/store mean 60.91 ns ( +- 14.96 ns ) serialise/item/aeson mean 8.440 μs ( +- 5.778 μs )

圧倒的に速い！Textをメモリ表現のまま扱うstoreと比較すると、UTF-8でエンコードするwineryには幾らかのハンデがあるが、各要素のエンコードに関してはstoreと同等の記録を出している。

bytestringにも包括的なベンチマークが付属しているので現在取り組んでいる。まだ作業中だが、ぶっちぎりで速いことに疑いの余地はない。

ここからはおまけ。Grisu3というアルゴリズムで浮動小数点数の表示を高速化するプルリクエストがbytestringに送られていたが、長年放置されていた(Reimplement floatDec/doubleDec by winterland1989 · Pull Request #115 · haskell/bytestring · GitHub)のでどのくらい速くなるか試してみた。その差、実に22倍である。

mason/double mean 138.0 ns ( +- 58.45 ns ) fast-builder/double mean 2.379 μs ( +- 231.6 ns ) bytestring/double mean 3.033 μs ( +- 813.4 ns )

また、文字列リテラルをStringを介さずに扱う変更([RFC] Builder: Efficiently handle literal strings by bgamari · Pull Request #132 · haskell/bytestring · GitHub)も放置されていたので、同等のものをこちらで実装した。派手な差はないが確実に速くなっている。

mason/literal mean 599.2 ns ( +- 249.7 ns ) fast-builder/literal mean 759.4 ns ( +- 165.2 ns ) bytestring/literal mean 640.1 ns ( +- 16.93 ns )

まとめ

masonは、高い抽象度を保ちつつも効率がよく、拡張性も確保されている、ある意味Haskellらしいライブラリだ。2020年の世界標準を目指し、さらに磨きをかけていきたい。