Rewrite Rulesについて軽く

Rewrite Rules(書き換え規則)について軽く説明します。
rewrite rulesはコードを書き換える最適化手法です。コードの書き換えに関してはユーザ自身が指定できるところがポイントです。
モジュラリティとパフォーマンスを両立するために使えたりします。

命令型言語で書き換え規則

javascriptを用いてfor文でArrayの変換を複数回実行することを考えます。

var doSomething1 = function (elem) {
    return elem + 2;
};
var doSomething2 = function (elem) {
    return elem * elem;
};

var initArr = [1,2,3,4,5];
var intermediate = [];
for (var i = 0; i < initArr.length; i++) {
    var elem = initArr[i];
    intermediate[i] = doSomething1(elem);
}

var result = [];
for (var i = 0; i < intermediate.length; i++) {
    var elem = intermediate[i];
    result[i] = doSomething2(elem);
}

console.log(result);

これはforが二回分回ってcostlyなので以下のように書き換えることができます。

var initArr = [1,2,3,4,5];
var result = [];
for (var i = 0; i < initArr.length; i++) {
    var elem = initArr[i];
    result[i] = doSomething2(doSomething1(elem));
}

まあこの書き換えはいいと思います。arrayは一回だけ走査するようになって、intermediateとかいう中間の配列も排除できました。

この書き換えをコンパイラにやらせるというのがよくあるrewrite rulesのサンプルです。

RULESから始まるGHC特殊コメントに記述するのですが、=のLHSをRHSに書き換えてくれるというものです。

{-# RULES
"map/map"    forall f g xs.  map f (map g xs) = map (f.g) xs
  #-}

map f (map g xs) を map (f . g) xsに書き換えてくれるというこの書き換え規則"map/map"は、先ほどJSの手作業でやったことをやってくれます。

これはストリームフュージョンの一例としても出てきます。ストリームフュージョンはGHCではrewrite rulesを用いて達成しています。

このListに対するmapはFunctorのfmapへ一般化できるので、

{-# RULES
"fmap/fmap"    forall f g xs.  fmap f (fmap g xs) = fmap (f.g) xs
  #-}

とか書けます。よーしじゃあどんどん汎用的な書き換えルールにしていこうねーという話はあるようで、興味がある型はVectorとかPipesあたり見ると面白いかもしれません。項書き換え系がどうのこうのという話もありましたね（知らない）

簡単な例

IntをStringへエンコード/デコードする例を考えます。

encodeInt :: Int -> String
encodeInt n = unsafePerformIO $ do
    threadDelay $ 3 * 1000 * 1000 -- 3sec
    return $ show n

decodeInt :: String -> Int
decodeInt s = unsafePerformIO $ do
    threadDelay $ 3 * 1000 * 1000 -- 3sec
    readIO s

decode失敗は無視します。
このエンコード/デコードは非常に重い処理で、3秒くらいかかってしまうようです。なので可能なら省略したいと思うのは当然の事でしょう。
なので以下の様なルールを書きます。

{-# RULES
"encodeInt/decodeInt" forall n. decodeInt (encodeInt n) = n
"decodeInt/encodeInt" forall n. encodeInt (decodeInt n) = n
  #-}

encodeIntの直後にdecodeIntがあったら、もしくはその逆があったら、それらを省略して引数の値をそのまま返す、というルールです。
繰り返しですがデコード失敗は無視しています。無視しない場合は片方使えません。

そしてコンパイルするとどうなるかというと、これだけではこのルールはうまく発火しない可能性があります。

どういうことかというと、この書き換え規則の前にinliningが走るケースです。
inliningが先に走ってしまうと、encodeInt, decodeIntという関数自体が呼び出し側から消え去ってしまうので、その後はルールが適用できなくなってしまいます。
ここで必要になってくるのがPhase Controlです。

Phase Control

GHCのSimplifierは複数回interations(default=4)を回しますが、そのiterationはいくつかのphaseに分割されており、phase numberがふってあります。このphase numberによって制御が可能です(default=2)。
phase numberは初めに"gentle"があり、その後 2, 1, 0 と下がっていきます。

Phase Controlは、[k]と書くとphase kもしくはその後を意味し、[~k]と書くとk以前(kを含まない)を意味します。

たとえば、以下です。

{-# RULES
"encodeInt/decodeInt" [~2] forall n. decodeInt (encodeInt n) = n
"decodeInt/encodeInt" [~2] forall n. encodeInt (decodeInt n) = n
  #-}

{-# INLINE [2] encodeInt #-}
{-# INLINE [2] decodeInt #-}

こうするとphase 2より前にRulesを発火し、 phase 2、もしくはより後にINLINEを実行することが明示できます。

とはいえphase "gentle"ではinlineをdisableにしてrulesを発火させるようにしているようなので、表層的な単純なrulesならおそらくphase control無しでも発火すると思われます。

rewrite rulesの発火具合は
-ddump-simpl-statsとか-ddump-rule-rewritesとかつけると出力されます。
以下は-ddump-simpl-statsの出力の一部です。
decodeInt/encodeInt, encodeInt/decodeInt といったルール名で発火したことが確認できます。

44 RuleFired
  9 Class op >>
  8 *#
  8 Class op *
  7 Class op show
  3 Class op >>=
  2 Class op readsPrec
  2 Class op return
  1 SPEC $fShow[]
  1 decodeInt/encodeInt
  1 encodeInt/decodeInt
  1 unpack
  1 unpack-list

以下は-ddump-rule-rewritesの出力の一部です。Before/Afterが確認できます。

Rule fired
    Rule: encodeInt/decodeInt
    Before: Main.decodeInt ValArg Main.encodeInt (GHC.Types.I# 12345#)
    After:  (\ (n_a18V [Occ=Once] :: GHC.Types.Int) -> n_a18V)
              (GHC.Types.I# 12345#)
    Cont:   Stop[RuleArgCtxt] GHC.Types.Int

次の様にINLINEをわざと先に発火させることで、rulesが発火しなくなることも確認できるでしょう。

{-# RULES
"encodeInt/decodeInt" [1] forall n. decodeInt (encodeInt n) = n
"decodeInt/encodeInt" [1] forall n. encodeInt (decodeInt n) = n
  #-}

{-# INLINE [~1] encodeInt #-}
{-# INLINE [~1] decodeInt #-}

こうやって3sec + 3secで6secくらいかかるようなencodeInt/decodeIntの組み合わせが一瞬で返ってくる様になりました。

main :: IO ()
main = do
    getCurrentTime >>= print
    print $ decodeInt (encodeInt 12345)
    print $ encodeInt (decodeInt "12345")
    getCurrentTime >>= print

Rewrite Rulesの制限と意義

Rewrite rulesは便利ですが、書き換え前と書き換え後の意味が同じであることはコンパイラではチェックしてくれません。型チェックはしてくれます。

Rewrite rulesの意義の一つはモジュラリティと速度の両立でしょう。
モジュラーにしようと別々に定義したものが、組み合わせると個別に入ったオーバーヘッドによって全体の効率が落ちてしまう、というのはモジュラリティを意識して開発するとわりと至るところで起こりうる問題です。
モジュラリティはほしい、速度も欲しい、APIは増やしたくない、という時にこういった書き換え規則は役にたつでしょう。