mlton/ssa/inline.fun

   1 (* Copyright (C) 2009 Matthew Fluet.
   2  * Copyright (C) 1999-2007 Henry Cejtin, Matthew Fluet, Suresh
   3  *    Jagannathan, and Stephen Weeks.
   4  * Copyright (C) 1997-2000 NEC Research Institute.
   5  *
   6  * MLton is released under a BSD-style license.
   7  * See the file MLton-LICENSE for details.
   8  *)
   9
  10 functor Inline (S: INLINE_STRUCTS): INLINE =
  11 struct
  12
  13 open S
  14 open Exp Transfer
  15
  16 structure Function =
  17    struct
  18       open Function
  19
  20       fun containsCall (f: Function.t): bool =
  21          Exn.withEscape
  22          (fn escape =>
  23           (Vector.foreach
  24            (Function.blocks f, fn Block.T {transfer, ...} =>
  25             case transfer of
  26                Call _ => escape true
  27              | _ => ())
  28            ; false))
  29       fun containsLoop (f: Function.t): bool =
  30          let
  31             val {get, set, destroy} =
  32                Property.destGetSet (Label.plist, Property.initConst false)
  33          in
  34             Exn.withEscape
  35             (fn escape =>
  36              let
  37                 val _ =
  38                    Function.dfs
  39                    (f, fn (Block.T {label, transfer, ...}) =>
  40                     (set (label, true)
  41                      ; (case transfer of
  42                            Goto {dst, ...} => if get dst then escape true else ()
  43                          | _ => ())
  44                      ; fn () => set (label, false)))
  45              in
  46                 false
  47              end)
  48             before (destroy ())
  49          end
  50    end
  51
  52 local
  53    fun 'a make (dontInlineFunc: Function.t * 'a -> bool)
  54       (Program.T {functions, ...}, a: 'a): Func.t -> bool =
  55       let
  56          val {get = shouldInline: Func.t -> bool,
  57               set = setShouldInline, ...} =
  58             Property.getSetOnce (Func.plist, Property.initConst false)
  59       in
  60          List.foreach
  61          (functions, fn f =>
  62           if not (Function.mayInline f) orelse dontInlineFunc (f, a)
  63              then ()
  64           else setShouldInline (Function.name f, true))
  65          ; Control.diagnostics
  66            (fn display =>
  67             let open Layout
  68             in List.foreach
  69                (functions, fn f =>
  70                 let
  71                    val name = Function.name f
  72                    val shouldInline = shouldInline name
  73                 in
  74                    display
  75                    (seq [Func.layout name, str ": ",
  76                          record [("shouldInline", Bool.layout shouldInline)]])
  77                 end)
  78             end)
  79          ; shouldInline
  80       end
  81 in
  82    val leafOnce = make (fn (f, {size}) =>
  83                         Option.isNone (Function.sizeMax (f, {max = size,
  84                                                              sizeExp = Exp.size,
  85                                                              sizeTransfer =Transfer.size}))
  86                         orelse Function.containsCall f)
  87    val leafOnceNoLoop = make (fn (f, {size}) =>
  88                               Option.isNone (Function.sizeMax (f, {max = size,
  89                                                                    sizeExp = Exp.size,
  90                                                                    sizeTransfer =Transfer.size}))
  91                               orelse Function.containsCall f
  92                               orelse Function.containsLoop f)
  93 end
  94
  95 structure Graph = DirectedGraph
  96 structure Node = Graph.Node
  97
  98 local
  99    fun make (dontInline: Function.t -> bool)
 100       (Program.T {functions, ...}, {size: int option}) =
 101       let
 102          val max = size
 103          type info = {function: Function.t,
 104                       node: unit Node.t,
 105                       shouldInline: bool ref,
 106                       size: int ref}
 107          val {get = funcInfo: Func.t -> info,
 108               set = setFuncInfo, ...} =
 109             Property.getSetOnce
 110             (Func.plist, Property.initRaise ("funcInfo", Func.layout))
 111          val {get = nodeFunc: unit Node.t -> Func.t,
 112               set = setNodeFunc, ...} =
 113             Property.getSetOnce
 114             (Node.plist, Property.initRaise ("nodeFunc", Node.layout))
 115          val graph = Graph.new ()
 116          (* initialize the info for each func *)
 117          val _ =
 118             List.foreach
 119             (functions, fn f =>
 120              let
 121                 val name = Function.name f
 122                 val n = Graph.newNode graph
 123              in
 124                 setNodeFunc (n, name)
 125                 ; setFuncInfo (name, {function = f,
 126                                       node = n,
 127                                       shouldInline = ref false,
 128                                       size = ref 0})
 129              end)
 130          (* Build the call graph. *)
 131          val _ =
 132             List.foreach
 133             (functions, fn f =>
 134              let
 135                 val {name, blocks, ...} = Function.dest f
 136                 val {node, ...} = funcInfo name
 137              in
 138                 Vector.foreach
 139                 (blocks, fn Block.T {transfer, ...} =>
 140                  case transfer of
 141                     Call {func, ...} =>
 142                        (ignore o Graph.addEdge)
 143                        (graph, {from = node, to = #node (funcInfo func)})
 144                   | _ => ())
 145              end)
 146          (* Compute strongly-connected components.
 147           * Then start at the leaves of the call graph and work up.
 148           *)
 149          val _ =
 150             List.foreach
 151             (rev (Graph.stronglyConnectedComponents graph),
 152              fn scc =>
 153              case scc of
 154                 [n] =>
 155                    let
 156                       val {function, shouldInline, size, ...} =
 157                          funcInfo (nodeFunc n)
 158                    in
 159                       if Function.mayInline function
 160                          andalso not (dontInline function)
 161                          then Exn.withEscape
 162                               (fn escape =>
 163                                let
 164                                   val res =
 165                                      Function.sizeMax
 166                                      (function,
 167                                       {max = max,
 168                                        sizeExp = Exp.size,
 169                                        sizeTransfer =
 170                                        fn t =>
 171                                        case t of
 172                                           Call {func, ...} =>
 173                                              let
 174                                                 val {shouldInline, size, ...} =
 175                                                    funcInfo func
 176                                              in
 177                                                 if !shouldInline
 178                                                    then !size
 179                                                    else escape ()
 180                                              end
 181                                         | _ => Transfer.size t})
 182                                in
 183                                   case res of
 184                                      NONE => ()
 185                                    | SOME n => (shouldInline := true
 186                                                 ; size := n)
 187                                end)
 188                       else ()
 189                    end
 190               | _ => ())
 191          val _ =
 192             Control.diagnostics
 193             (fn display =>
 194              let open Layout
 195              in List.foreach
 196                 (functions, fn f =>
 197                  let
 198                     val name = Function.name f
 199                     val {shouldInline, size, ...} = funcInfo name
 200                     val shouldInline = !shouldInline
 201                     val size = !size
 202                  in
 203                     display
 204                     (seq [Func.layout name, str ": ",
 205                           record [("shouldInline", Bool.layout shouldInline),
 206                                   ("size", Int.layout size)]])
 207                  end)
 208              end)
 209       in
 210          ! o #shouldInline o funcInfo
 211       end
 212 in
 213    val leafRepeat = make (fn _ => false)
 214    val leafRepeatNoLoop = make (fn f => Function.containsLoop f)
 215 end
 216
 217 fun nonRecursive (Program.T {functions, ...}, {small: int, product: int}) =
 218    let
 219       type info = {doesCallSelf: bool ref,
 220                    function: Function.t,
 221                    node: unit Node.t,
 222                    numCalls: int ref,
 223                    shouldInline: bool ref,
 224                    size: int ref}
 225       val {get = funcInfo: Func.t -> info,
 226            set = setFuncInfo, ...} =
 227          Property.getSetOnce
 228          (Func.plist, Property.initRaise ("funcInfo", Func.layout))
 229       val {get = nodeFunc: unit Node.t -> Func.t,
 230            set = setNodeFunc, ...} =
 231          Property.getSetOnce
 232          (Node.plist, Property.initRaise ("nodeFunc", Node.layout))
 233       val graph = Graph.new ()
 234       (* initialize the info for each func *)
 235       val _ =
 236          List.foreach
 237          (functions, fn f =>
 238           let
 239              val name = Function.name f
 240              val n = Graph.newNode graph
 241           in
 242              setNodeFunc (n, name)
 243              ; setFuncInfo (name, {doesCallSelf = ref false,
 244                                    function = f,
 245                                    node = n,
 246                                    numCalls = ref 0,
 247                                    shouldInline = ref false,
 248                                    size = ref 0})
 249           end)
 250       (* Update call counts. *)
 251       val _ =
 252          List.foreach
 253          (functions, fn f =>
 254           let
 255              val {name, blocks, ...} = Function.dest f
 256              val {doesCallSelf, ...} = funcInfo name
 257           in
 258              Vector.foreach
 259              (blocks, fn Block.T {transfer, ...} =>
 260               case transfer of
 261                  Call {func, ...} =>
 262                     let
 263                        val {numCalls, ...} = funcInfo func
 264                     in
 265                        if Func.equals (name, func)
 266                           then doesCallSelf := true
 267                        else Int.inc numCalls
 268                     end
 269                | _ => ())
 270           end)
 271       fun mayInline (setSize: bool,
 272                      {function, doesCallSelf, numCalls, size, ...}: info): bool =
 273          Function.mayInline function
 274          andalso not (!doesCallSelf)
 275          andalso let
 276                     val n =
 277                        Function.size
 278                        (function,
 279                         {sizeExp = Exp.size,
 280                          sizeTransfer =
 281                          fn t as Call {func, ...} =>
 282                                let
 283                                   val {shouldInline, size, ...} = funcInfo func
 284                                in
 285                                   if !shouldInline
 286                                      then !size
 287                                      else Transfer.size t
 288                                end
 289                           | t => Transfer.size t})
 290                  in
 291                     if setSize
 292                        then size := n
 293                     else ()
 294                     ; (!numCalls - 1) * (n - small) <= product
 295                  end
 296       (* Build the call graph.  Do not include functions that we already know
 297        * will not be inlined.
 298        *)
 299       val _ =
 300          List.foreach
 301          (functions, fn f =>
 302           let
 303              val {name, blocks, ...} = Function.dest f
 304              val info as {node, ...} = funcInfo name
 305           in
 306              if mayInline (false, info)
 307                 then Vector.foreach
 308                      (blocks, fn Block.T {transfer, ...} =>
 309                       case transfer of
 310                          Call {func, ...} =>
 311                             if Func.equals (name, func)
 312                                then ()
 313                             else (ignore o Graph.addEdge)
 314                                  (graph, {from = node, to = #node (funcInfo func)})
 315                        | _ => ())
 316              else ()
 317           end)
 318       (* Compute strongly-connected components.
 319        * Then start at the leaves of the call graph and work up.
 320        *)
 321       val _ =
 322          List.foreach
 323          (rev (Graph.stronglyConnectedComponents graph),
 324           fn [n] => let val info as {shouldInline, ...} = funcInfo (nodeFunc n)
 325                     in shouldInline := mayInline (true, info)
 326                     end
 327            | _ => ())
 328       val _ =
 329          Control.diagnostics
 330          (fn display =>
 331           let open Layout
 332           in List.foreach
 333              (functions, fn f =>
 334               let
 335                  val name = Function.name f
 336                  val {numCalls, shouldInline, size, ...} = funcInfo name
 337                  val numCalls = !numCalls
 338                  val shouldInline = !shouldInline
 339                  val size = !size
 340               in
 341                  display
 342                  (seq [Func.layout name, str ": ",
 343                        record [("numCalls", Int.layout numCalls),
 344                                ("shouldInline", Bool.layout shouldInline),
 345                                ("size", Int.layout size)]])
 346               end)
 347           end)
 348    in
 349       ! o #shouldInline o funcInfo
 350    end
 351
 352 fun transform {program as Program.T {datatypes, globals, functions, main},
 353                shouldInline: Func.t -> bool,
 354                inlineIntoMain: bool} =
 355    let
 356       val {get = funcInfo: Func.t -> {function: Function.t,
 357                                       isCalledByMain: bool ref},
 358            set = setFuncInfo, ...} =
 359          Property.getSetOnce
 360          (Func.plist, Property.initRaise ("Inline.funcInfo", Func.layout))
 361       val isCalledByMain: Func.t -> bool =
 362          ! o #isCalledByMain o funcInfo
 363       val () = List.foreach (functions, fn f =>
 364                              setFuncInfo (Function.name f,
 365                                           {function = f,
 366                                            isCalledByMain = ref false}))
 367       val () =
 368          Vector.foreach
 369          (#blocks (Function.dest (Program.mainFunction program)),
 370           fn Block.T {transfer, ...} =>
 371           case transfer of
 372              Transfer.Call {func, ...} =>
 373                 #isCalledByMain (funcInfo func) := true
 374            | _ => ())
 375       fun doit (blocks: Block.t vector,
 376                 return: Return.t) : Block.t vector =
 377          let
 378             val newBlocks = ref []
 379             val blocks =
 380                Vector.map
 381                (blocks,
 382                 fn block as Block.T {label, args, statements, transfer} =>
 383                 let
 384                    fun new transfer =
 385                       Block.T {label = label,
 386                                args = args,
 387                                statements = statements,
 388                                transfer = transfer}
 389                 in
 390                   case transfer of
 391                      Call {func, args, return = return'} =>
 392                         let
 393                            val return = Return.compose (return, return')
 394                         in
 395                            if shouldInline func
 396                               then
 397                               let
 398                                  local
 399                                     val {name, args, start, blocks, ...} =
 400                                        (Function.dest o Function.alphaRename)
 401                                        (#function (funcInfo func))
 402                                     val blocks = doit (blocks, return)
 403                                     val _ = List.push (newBlocks, blocks)
 404                                     val name =
 405                                        Label.newString (Func.originalName name)
 406                                     val _ =
 407                                        List.push
 408                                        (newBlocks,
 409                                         Vector.new1
 410                                         (Block.T
 411                                          {label = name,
 412                                           args = args,
 413                                           statements = Vector.new0 (),
 414                                           transfer = Goto {dst = start,
 415                                                            args = Vector.new0 ()}}))
 416                                  in
 417                                     val name = name
 418                                  end
 419                               in
 420                                  new (Goto {dst = name,
 421                                             args = args})
 422                               end
 423                            else new (Call {func = func,
 424                                            args = args,
 425                                            return = return})
 426                         end
 427                    | Raise xs =>
 428                         (case return of
 429                             Return.NonTail
 430                             {handler = Handler.Handle handler, ...} =>
 431                                new (Goto {dst = handler,
 432                                           args = xs})
 433                           | _ => block)
 434                    | Return xs =>
 435                         (case return of
 436                             Return.NonTail {cont, ...} =>
 437                                new (Goto {dst = cont, args = xs})
 438                           | _ => block)
 439                    | _ => block
 440                 end)
 441          in
 442             Vector.concat (blocks::(!newBlocks))
 443          end
 444       val shrink = shrinkFunction {globals = globals}
 445       val functions =
 446          List.fold
 447          (functions, [], fn (f, ac) =>
 448           let
 449              val {args, blocks, mayInline, name, raises, returns, start} =
 450                 Function.dest f
 451              fun keep () =
 452                 let
 453                    val blocks = doit (blocks, Return.Tail)
 454                 in
 455                    shrink (Function.new {args = args,
 456                                          blocks = blocks,
 457                                          mayInline = mayInline,
 458                                          name = name,
 459                                          raises = raises,
 460                                          returns = returns,
 461                                          start = start})
 462                    :: ac
 463                 end
 464           in
 465              if Func.equals (name, main)
 466                 then if inlineIntoMain
 467                         then keep ()
 468                      else f :: ac
 469              else
 470                 if shouldInline name
 471                    then
 472                       if inlineIntoMain
 473                          orelse not (isCalledByMain name)
 474                          then ac
 475                       else keep ()
 476                 else keep ()
 477           end)
 478       val program =
 479          Program.T {datatypes = datatypes,
 480                     globals = globals,
 481                     functions = functions,
 482                     main = main}
 483       val _ = Program.clearTop program
 484    in
 485       program
 486    end
 487
 488 fun inlineLeaf (p, {loops, repeat, size}) =
 489    if size = SOME 0
 490       then p
 491    else transform {program = p,
 492                    shouldInline =
 493                    case (loops, repeat) of
 494                       (false, false) => leafOnce (p, {size = size})
 495                     | (false, true) => leafRepeat (p, {size = size})
 496                     | (true, false) => leafOnceNoLoop (p, {size = size})
 497                     | (true, true) => leafRepeatNoLoop (p, {size = size}),
 498                    inlineIntoMain = true}
 499 fun inlineNonRecursive (p, arg) =
 500    transform {program = p,
 501               shouldInline = nonRecursive (p, arg),
 502               inlineIntoMain = !Control.inlineIntoMain}
 503
 504 end