この記事は言語実装 Advent Calendar 2016のために書かれました。

元々僕はStandard MLで一からコンパイラを書こうとしていたんですが、ローレベルあんまり詳しくなくてネイティブコードを生成できるようにするのが面倒だったのと、Standard MLのsyntax sugarがあまりにも無糖だったので辛くなってきました。いろいろ最適化を実装して楽しみたいだけなのに自作に固執しなくても良い気がしたので、今度はMinCamlに自作のコンパイラで実装していた最適化を導入して遊ぶことにします。

Common subexpression elimination(以下、CSE)

CSEはプログラム中に同じ部分式が二度以上現れる場合、一度評価した値を使い回す最適化です。例えば式a * b + a * bには部分式a * bが二回現れますが、これにCSEを適用するとlet x = a * b in x + xのようになります。

こうして見ると一般の式に対しての適用が難しいのではないかと思われるかもしれませんが、 MinCamlはK正規形、つまり全ての部分式にletで名前を付けた中間表現を採用しています。従ってa * b + a * bのような式は最適化を施す前に既にlet x = a * b in let y = a * b in x + yのように変換されているため、それほど実装は難しくはありません。既に束縛されている変数の中身を覚えておき、letが現れる度に同じ式に束縛されている変数が無いか調べて適宜置き換えるだけです。

α等価性の判定

同じ式と言っても、fun x -> xとfun y -> yのように束縛変数の名前が違うだけの式も（α）等価な式と見做さなくてはなりません。従ってfun x -> eのような式とfun y -> e'のような式が等価であるかどうか調べる場合、変数名が違っているからといって違う式と決めつけずeと[y→x]e'が等価になるか調べる訳です。

ここで、代入は代入先の式の大きさNに対してO(N)の計算量を要するため、素朴に束縛が現れる度に代入してしまっては計算量が馬鹿になりません。環境を用いて代入を遅延してやれば等価判定を効率的に行える事でしょう。

以上のアイデアを実装したものが次のコードになります。

(* check alpha-equivalence *)
let rec equiv sub e1 e2 =
  match e1, e2 with
  | IfEq (x1, y1, e11, e12), IfEq (x2, y2, e21, e22)
  | IfLE (x1, y1, e11, e12), IfLE (x2, y2, e21, e22) ->
      equiv sub (Var x1) (Var x2) &&
      equiv sub (Var y1) (Var y2) &&
      equiv sub e11 e21 &&
      equiv sub e12 e22
  | Let ((x1, _), e11, e12), Let ((x2, _), e21, e22) ->
      equiv sub e11 e21 &&
      equiv (M.add x1 x2 sub) e12 e22
  | LetRec ({ name = (x1, _); args = yt1s; body = e11 }, e12),
    LetRec ({ name = (x2, _); args = yt2s; body = e21 }, e22) ->
      let sub' = M.add x1 x2 sub in
      equiv
        (M.add_list2
          (List.map fst yt1s)
          (List.map fst yt2s) sub') e11 e21 &&
      equiv sub' e12 e22
  | LetTuple (xt1s, y1, e1), LetTuple (xt2s, y2, e2) ->
      equiv sub (Var y1) (Var y2) &&
      equiv
        (M.add_list2
          (List.map fst xt1s)
          (List.map fst xt2s) sub) e1 e2
  | _, _ -> Alpha.g sub e1 = Alpha.g sub e2 (* tenuki *)

let equiv sub e1 e2 =
  try equiv sub e1 e2 with
  (* inconsistent arity *)
  | Invalid_argument _ -> false

MinCamlの中間表現において変数名はユニークなものが振り直されているため、captureを恐れずに環境を使い回しています。
if x <= y then e1 else e2と1のように、構文の形だけを見れば代入を行う前に比較を打ち切れる式もありますが、そいつらに対応するとコードの量が洒落にならなくなるのでちょっと妥協した実装にもなっています。早すぎる最適化ってやつですね。
また、MinCamlは複数引数の関数を許すのでfun (x, y) -> eみたいな式とfun (x', y', z') -> e'みたいな式を比較する事も考えられますが、そのような場合はM.add_list2がInvalid_argumentを投げるのでちゃんとcatchしてあげます。

De Bruijn indexやlocally namelessなんかを採用していればα等価な式は全て同じASTになってたんでしょうが、こいつらはこいつらでつらみがあるのでちょっと…

CSEの実装

前に書いた通り、既に束縛されている変数の中身を覚えておいて、同じ式が現れたら置き換えるだけです。その際にprint_endline "hoge"のような副作用がある式が評価後の値に置き換えられるとまずいので、副作用の有無を判定しなければなりません。また、let x = y in eみたいな形の式が現れるのは気持ち悪いので、こちらも環境で遅延しつつ代入してやります。

以上のアイデアを実装すると次のようになります。

(* common subexpression elimination *)
(* variable names must be unique *)
let rec g env fenv tpenv sub = function
  | IfEq (x, y, e1, e2) ->
      IfEq (find x sub, find y sub, g env fenv tpenv sub e1, g env fenv tpenv sub e2)
  | IfLE (x, y, e1, e2) ->
      IfLE (find x sub, find y sub, g env fenv tpenv sub e1, g env fenv tpenv sub e2)
  | Let ((x, t), e1, e2) ->
      let e1' = g env fenv tpenv sub e1 in
      if effect e1' then
        Let ((x, t), e1', g env fenv tpenv sub e2)
      else
        (try
           (* linear search *)
           let (x', _) = List.find (fun (_, e) -> equiv sub e e1') env in
           Format.eprintf "removing common subexpression %s@." x;
           g env fenv tpenv (M.add x x' sub) e2
         with
         | Not_found ->
             Let ((x, t), e1', g ((x, e1') :: env) fenv tpenv sub e2))
  | LetRec ({ name = (x, _); args = yts; body = e1 } as fundef, e2) ->
      let e1' = g env fenv tpenv sub e1 in
      (try
         (* linear search *)
         let (x', yts', e1') = List.find (fun (x', yts', e1') ->
           try
             equiv
               (M.add_list2
                 (List.map fst yts)
                 (List.map fst yts')
                 (M.add x x' sub)) e1 e1'
           with
           (* inconsistent arity *)
           | Invalid_argument _ -> false) fenv in
         Format.eprintf "removing common subexpression %s@." x;
         (* function declarations have no side-effect *)
         g env fenv tpenv (M.add x x' sub) e2
      with
      | Not_found ->
          LetRec ({ fundef with body = e1' }, g env ((x, yts, e1) :: fenv) tpenv sub e2))
  | LetTuple (xts, y, e) ->
      let y' = find y sub in
      (match M.find y' tpenv with
       | xts' ->
           Format.eprintf "removing common subexpression %s@."
             (List.fold_right ( ^ ) (List.map (( ^ ) " ") @@ List.map fst xts) "");
           (* projections have no side-effect *)
           (* exception Invarid_argument must not be raised if well-typed *)
           g env fenv tpenv
             (M.add_list2
               (List.map fst xts)
               (List.map fst xts') sub) e
       | exception Not_found ->
           LetTuple (xts, y', g env fenv (M.add y' xts tpenv) sub e))
  | e -> Alpha.g sub e

let f = g [] [] M.empty M.empty

MinCamlの中間言語には3種類のletが存在するため、苦肉の策として変数の中身を覚えておく環境も3つ用意しています。等価な式を探す際に線形探索しているのはご愛嬌。多分式の大きさとかで二分探索してから等価性を判定する関数を使うようにすればもっとコンパイルが速くなるんでしょうが、それほどコンパイル速度に困ってないのでこの辺で止めておきます。 ~~コンパイル時間を速くしたければ良いPCを買え~~

結び

僕が手を加えたMinCamlはこちらで公開しています。必修の講義で何故か自作CPUのアセンブリを吐くコンパイラを書く羽目になった人とかいかがですか。

fetburner.core

コアダンプ

MinCamlに簡単な最適化を追加する

Common subexpression elimination(以下、CSE)

α等価性の判定

CSEの実装

結び