Worker-Wrapper transformationについて軽く

Worker-Wrapper transformation(Worker-Wrapper変換)について簡単に説明したいと思います。数年前は調べてもあまり資料見つからなかったのですが、現在は幸いにも簡単に見つかります。

• http://www.cs.nott.ac.uk/~pszgmh/wrapper.pdf
• http://www.cs.nott.ac.uk/~psxjlha/wwmf.pdf

全部を読んではいないんですがどうやらfixを用いたような最適化手法の一種のようです。
まず準備を行います。

fact関数の遅さ

例えばfact関数を単純に再帰を使ってナイーブに書いてみます。

import Data.Function (fix)

fact :: Int -> Int
fact 0 = 1
fact n = n * fact (n-1)

負数の引数はここでは無視します。

さてGHCでは基本すべての値はBoxedな値として扱われます。Unboxedな値も扱えますが、そのためにはMagicHashという拡張を有効にする必要があります。ついでにGHC.PrimとGHC.Typesをimportしておきます。

{-# LANGUAGE MagicHash #-}
module Main where

import Data.Function (fix)
import GHC.Prim
import GHC.Types (Int(I#))

するとIntのコンストラクタを明示的に扱うことができるようになります。I#です。MagicHash拡張はこの#付きのあれこれを扱うために必要です。I#の中にはUnboxedな値が格納されています。Unboxedな値はMagicHash指定時には0#, 1#といった#付きのリテラルで記述することが可能です。これらUnboxedな値の型はInt#とこれまた#付きで記述します。また、Unboxedな型のkindは#です。Unboxedな値を扱う関数群は最後に#が付いています。
つまりはUnboxed関連やprimitives周辺は#をつけることで統一されている様です。

次にfact関数に対して-O0 -ddump-prepをつけてCorePrepを見てみます。

Rec {
-- RHS size: {terms: 25, types: 9, coercions: 0}
Lib.fact [Occ=LoopBreaker] :: GHC.Types.Int -> GHC.Types.Int
[GblId, Arity=1, Str=DmdType, Unf=OtherCon []]
Lib.fact =
  \ (ds_s1Sy [Occ=Once!] :: GHC.Types.Int) ->
    case ds_s1Sy of wild_s1Sz { GHC.Types.I# ds1_s1SA [Occ=Once!] ->
    case ds1_s1SA of _ [Occ=Dead] {
      __DEFAULT ->
        let {
          sat_s1SE [Occ=Once] :: GHC.Types.Int
          [LclId, Str=DmdType]
          sat_s1SE =
            let {
              sat_s1SD [Occ=Once] :: GHC.Types.Int
              [LclId, Str=DmdType]
              sat_s1SD =
                let {
                  sat_s1SC [Occ=Once] :: GHC.Types.Int
                  [LclId, Str=DmdType]
                  sat_s1SC = GHC.Types.I# 1# } in
                GHC.Num.- @ GHC.Types.Int GHC.Num.$fNumInt wild_s1Sz sat_s1SC } in
            Lib.fact sat_s1SD } in
        GHC.Num.* @ GHC.Types.Int GHC.Num.$fNumInt wild_s1Sz sat_s1SE;
      0# -> GHC.Types.I# 1#
    }
    }
end Rec }

GHC.Types.I#がちらほら見えると思います。これは単なるInt型の値コンストラクタなので、caseによるパターンマッチで分解することができます。このコードではI#のパターンマッチを行ったあと、中身の値を0#かどうかでbase caseかどうか判断しています。そしてGHC.Num.-とかGHC.Num.*が現れつつ、再帰を行います。

ここでついでにGHC.Num.-とかGHC.Num.*の実装も見てみましょう。

instance  Num Int  where
    I# x + I# y = I# (x +# y)
    I# x - I# y = I# (x -# y)
    negate (I# x) = I# (negateInt# x)
    I# x * I# y = I# (x *# y)
    abs n  = if n `geInt` 0 then n else negate n

    signum n | n `ltInt` 0 = negate 1
             | n `eqInt` 0 = 0
             | otherwise   = 1

    {-# INLINE fromInteger #-}   -- Just to be sure!
    fromInteger i = I# (integerToInt i)

これによると、例えば3 + 4といった何気ないコードは、まずパターンマッチでそれぞれI#を取り外し、+#というprimitiveな計算を行った後に再度I#をつける、ということをしています。

つまり先のコードは、（最適化をoffにしたコードとはいえ）値コンストラクタをつけたり何度も何度も外したりしながら計算を行っているということを意味しています。もちろんそれによって僕らはUnboxedな値を直に扱うことなくいろいろな恩恵を受けているのでしょうけど、パフォーマンスを考慮している際には邪魔となります。はずしたまま計算すればいいじゃないですか。

Unboxedのまま計算する

Int#に対する演算はGHC.Primで定義されています。その演算子、-#, *#を用いればfact関数は高速に計算できそうです。

さてここで僕らはfix関数を用いてfactは以下のように書き換えることができると知っています。

fact' :: Int -> Int
fact' = fix (\rec n -> if n == 0 then 1 else n * rec (n-1))

このfixに渡している関数は再帰関数ではありません。このコンストラクタを外した版の関数は簡単に定義できます。これをworker関数と呼びます。

fact_worker :: Int# -> Int#
fact_worker = fix (\rec n -> if n == 0# then 1# else n *# rec (n -# 1#))

これに対しfact'と同じ結果を得るには、引数と返り値とを簡単にwrapしてあげれば良さそうです。

fact_wrapper :: (Int# -> Int#) -> Int -> Int
fact_wrapper worker (I# n) = I# (worker n)

これだけです。これをwrapper関数と呼びます。論文のやつとは違いますがまあ。

worker関数とwrapper関数を組み合わせれば目的の関数が得られます。

fact_worker_wrapper :: Int -> Int
fact_worker_wrapper = fact_wrapper fact_worker

この関数は再帰部ではUnboxedな値を用いて高速に計算しつつ、コンストラクタのつけ外しは最初と最後の一回ずつのみに抑えています。

ここでは手作業で行いましたが、似たような変換をコンパイラがやってくれます。この最適化をWorker-Wrapper transformationと呼びます。先のfactは-O2をつけると以下のようなCoreが生成されます。¹

Rec {
-- RHS size: {terms: 18, types: 4, coercions: 0}
Lib.$wfact [InlPrag=[0], Occ=LoopBreaker]
  :: GHC.Prim.Int# -> GHC.Prim.Int#
[GblId,
 Arity=1,
 Caf=NoCafRefs,
 Str=DmdType <S,1*U>,
 Unf=OtherCon []]
Lib.$wfact =
  \ (ww_s2H4 [Occ=Once!] :: GHC.Prim.Int#) ->
    case ww_s2H4 of ds_s2H5 {
      __DEFAULT ->
        case GHC.Prim.-# ds_s2H5 1# of sat_s2H6 { __DEFAULT ->
        case Lib.$wfact sat_s2H6 of ww1_s2H7 { __DEFAULT ->
        GHC.Prim.*# ds_s2H5 ww1_s2H7
        }
        };
      0# -> 1#
    }
end Rec }

-- RHS size: {terms: 10, types: 4, coercions: 0}
Lib.fact [InlPrag=INLINE[0]] :: GHC.Types.Int -> GHC.Types.Int
[GblId,
 Arity=1,
 Caf=NoCafRefs,
 Str=DmdType <S(S),1*U(1*U)>m,
 Unf=OtherCon []]
Lib.fact =
  \ (w_s2H8 [Occ=Once!] :: GHC.Types.Int) ->
    case w_s2H8 of _ [Occ=Dead] { GHC.Types.I# ww1_s2Ha [Occ=Once] ->
    case Lib.$wfact ww1_s2Ha of ww2_s2Hb { __DEFAULT ->
    GHC.Types.I# ww2_s2Hb
    }
    }

$wというprefixはその関数がworker-wrapper変換で生成されたworker関数であることを示しています。
先の手作業と同じようにworker関数はInt# -> Int#の型の関数となり、I#をつけ外しするようなことはなくなりました。

Worker-Wrapper transforamationの適用範囲と制限

任意の再帰関数はfixを用いて書き換えることができます。よってこの最適化の適用範囲は再帰関数全般に渡ると言っていいでしょう大雑把に言って。Worker関数がうまく定義できるかどうかという問題はありますが。そしてBoxedな関数をUnboxedな関数に変換するため、強力な最適化です。変換作業は単純なので、コンパイラが検知できないようなWorker/Wrapperが自明でないコードを自分で書き換えるといったこともできなくはありません。つまりBoxed/Unboxedの変換だけでなく、広範囲に適用することもまあできるでしょう。

しかし今回IntでI#をはずしたように、具体的な型が決まらないとこの最適化はできません。また高階関数でも発動しません。

論文によると、この最適化によって遅くならないことが証明されているようです。すごいですね（読んでない）

他に情報知ってましたらコメントか何かで教えてください。

---

1. -Oでworker-wrapper transformationは発火するようです。 https://twitter.com/kazu_yamamoto/status/804517945808257024 ↩