最強にして最速のビルダー、mason

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)

-- | 'BuildSignal's abstract signals to the caller of a 'BuildStep'. There are
-- three signals: 'done', 'bufferFull', or 'insertChunks signals
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というデータ型に必要な情報をまとめているため、拡張性に乏しいという難点がある。ローレベルな内部構造も拡張の難しさに拍車をかけている。

-- | Specifies where bytes generated by a builder go.
data DataSink
  = DynamicSink !(IORef DynamicSink)
    -- ^ The destination of data changes while the builder is running.
  | GrowingBuffer !(IORef (ForeignPtr Word8))
    -- ^ Bytes are accumulated in a contiguous buffer.
  | HandleSink !IO.Handle !Int{-next buffer size-} !(IORef Queue)
    -- ^ Bytes are first accumulated in the 'Queue', then flushed to the
    -- 'IO.Handle'.

そんな問題を解決するため、新たに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) -- ^ end of the buffer (next to the last byte)
  , bCur :: {-# UNPACK #-} !(Ptr Word8) -- ^ current position
  }

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年の世界標準を目指し、さらに磨きをかけていきたい。

hackage.haskell.org