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`関数を単純に再帰を使ってナイーブに書いてみます。

```hs
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しておきます。

```hs
{-# 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を見てみます。

```hs
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.*`の実装も見てみましょう。

```hs
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`は以下のように書き換えることができると知っています。

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

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

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

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

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

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

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

```hs
fact_worker_wrapper :: Int -> Int
fact_worker_wrapper = fact_wrapper fact_worker
```

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

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

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

```hs
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#`をはずしたように、具体的な型が決まらないとこの最適化はできません。また高階関数でも発動しません。

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

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