[commit: ghc] wip/T14152: Do some ad-hoc CSE in Exitification (d169843)

[git at git.haskell.org]

Repository : ssh://git@git.haskell.org/ghc

On branch  : wip/T14152
Link       : http://ghc.haskell.org/trac/ghc/changeset/d169843c69af695aa4531148a57c8b325c52e4dc/ghc

>---------------------------------------------------------------

commit d169843c69af695aa4531148a57c8b325c52e4dc
Author: Joachim Breitner <mail at joachim-breitner.de>
Date:   Wed Sep 6 15:17:12 2017 +0100

    Do some ad-hoc CSE in Exitification
    
    So that code like this
    
        safe :: Int -> [Int] -> [[Int]]
        safe x2 ds6
          = case ds6 of
              [] -> lvl6;
              q : l | x2     == q -> h x2 ds5
                    | x2 + 1 == q -> h x2 ds5
                    | x2 + 2 == q -> h x2 ds5
                    | x2 + 3 == q -> h x2 ds5
                    | otherwise   -> safe (x+10) l
    
    only gets one exit join point, instead of 4 identically.
    
    (There is currently no CSE for join points, and even if there were, it
    is cheap to do it here.)
    
    This does happen at least in GHC somewhere (exhibited by a core lint
    error when this patch was buggy). Nofib does not show any differences,
    not even in code size, which makes me believe that a later pass would remove the
    duplication (common block elimination on the Cmm level?) and the end
    result is not affected.
    
    But in the interest of compilation time it still seems useful to avoid
    this duplication from the get-go.


>---------------------------------------------------------------

d169843c69af695aa4531148a57c8b325c52e4dc
 compiler/simplCore/Exitify.hs | 96 +++++++++++++++++++++++++++++++++----------
 1 file changed, 74 insertions(+), 22 deletions(-)

diff --git a/compiler/simplCore/Exitify.hs b/compiler/simplCore/Exitify.hs
index 51dd675..0ceaf82 100644
--- a/compiler/simplCore/Exitify.hs
+++ b/compiler/simplCore/Exitify.hs
@@ -47,6 +47,7 @@ import VarSet
 import VarEnv
 import CoreFVs
 import FastString
+import TrieMap
 import Type
 
 import Data.Bifunctor
@@ -95,7 +96,7 @@ exitifyProgram binds = map goTopLvl binds
 --   join-points outside the joinrec.
 exitify :: InScopeSet -> [(Var,CoreExpr)] -> (CoreExpr -> CoreExpr)
 exitify in_scope pairs =
-    \body ->mkExitLets exits (mkLetRec pairs' body)
+    \body -> mkExitLets exits (mkLetRec pairs' body)
   where
     mkExitLets ((exitId, exitRhs):exits') = mkLetNonRec exitId exitRhs . mkExitLets exits'
     mkExitLets [] = id
@@ -107,7 +108,7 @@ exitify in_scope pairs =
     -- Which are the recursive calls?
     recursive_calls = mkVarSet $ map fst pairs
 
-    (pairs',exits) = (`runState` []) $ do
+    (pairs',exits) = runExitifyState in_scope $
         forM ann_pairs $ \(x,rhs) -> do
             -- go past the lambdas of the join point
             let (args, body) = collectNAnnBndrs (idJoinArity x) rhs
@@ -122,7 +123,7 @@ exitify in_scope pairs =
     -- It uses a state monad to keep track of floated binds
     go :: [Var]           -- ^ variables to abstract over
        -> CoreExprWithFVs -- ^ current expression in tail position
-       -> State [(Id, CoreExpr)] CoreExpr
+       -> ExitifyM CoreExpr
 
     go captured ann_e
         -- Do not touch an expression that is already a join call with no free
@@ -144,10 +145,10 @@ exitify in_scope pairs =
         | is_exit = do
             -- Assemble the RHS of the exit join point
             let rhs = mkLams args e
-                ty = exprType rhs
-            let avoid = in_scope `extendInScopeSetList` captured
+            -- Remember what is in scope here
+            nowInScope captured
             -- Remember this binding under a suitable name
-            v <- addExit avoid ty (length args) rhs
+            v <- addExit (length args) rhs
             -- And call it from here
             return $ mkVarApps (Var v) args
       where
@@ -208,16 +209,41 @@ exitify in_scope pairs =
     go _ ann_e = return (deAnnotate ann_e)
 
 
+type ExitifyM = State ExitifyState
+data ExitifyState = ExitifyState
+        { es_in_scope_acc :: InScopeSet -- ^ combined in_scope_set of all call sites
+        , es_in_scope     :: InScopeSet -- ^ final in_scope_set
+        , es_joins        :: [(JoinId, CoreExpr)] -- ^ exit join points
+        , es_map          :: CoreMap JoinId
+            -- ^ reverse lookup map, see Note [Avoid duplicate exit points]
+        }
+
+-- Runs the ExitifyM monad, and feeds in the final es_in_scope_acc as the
+-- es_in_scope to use
+runExitifyState :: InScopeSet -> ExitifyM a -> (a, [(JoinId, CoreExpr)])
+runExitifyState in_scope_init f = (res, es_joins state)
+  where
+    (res, state) = runState f (ExitifyState in_scope_init in_scope [] emptyTM)
+    in_scope = es_in_scope_acc state
+
+-- Keeps track of what is in scope at all the various positions where
+-- we want to call an exit join point
+nowInScope :: [Var] -> ExitifyM ()
+nowInScope captured = do
+    st <- get
+    put (st { es_in_scope_acc = es_in_scope_acc st `extendInScopeSetList` captured})
+
 -- Picks a new unique, which is disjoint from
 --  * the free variables of the whole joinrec
 --  * any bound variables (captured)
 --  * any exit join points created so far.
-mkExitJoinId :: InScopeSet -> Type -> JoinArity -> ExitifyM JoinId
-mkExitJoinId in_scope ty join_arity = do
-    fs <- get
-    let avoid = in_scope `extendInScopeSetList` (map fst fs)
-                         `extendInScopeSet` exit_id_tmpl -- just cosmetics
-    return (uniqAway avoid exit_id_tmpl)
+mkExitJoinId :: Type -> JoinArity -> ExitifyM JoinId
+mkExitJoinId ty join_arity = do
+    st <- get
+    let in_scope = es_in_scope st `extendInScopeSet` exit_id_tmpl -- cosmetic only
+    let v = uniqAway in_scope exit_id_tmpl
+    put (st { es_in_scope = es_in_scope st `extendInScopeSet` v})
+    return v
   where
     exit_id_tmpl = mkSysLocal (fsLit "exit") initExitJoinUnique ty
                     `asJoinId` join_arity
@@ -230,16 +256,22 @@ mkExitJoinId in_scope ty join_arity = do
                , occ_int_cxt = False
                , occ_tail = AlwaysTailCalled join_arity }
 
-addExit :: InScopeSet -> Type -> JoinArity -> CoreExpr -> ExitifyM JoinId
-addExit in_scope ty join_arity rhs = do
-    -- Pick a suitable name
-    v <- mkExitJoinId in_scope ty join_arity
-    fs <- get
-    put ((v,rhs):fs)
-    return v
-
-
-type ExitifyM =  State [(JoinId, CoreExpr)]
+-- Adds a new exit join point
+-- (or re-uses an existing one)
+addExit :: JoinArity -> CoreExpr -> ExitifyM JoinId
+addExit join_arity rhs = do
+    st <- get
+    -- See Note [Avoid duplicate exit points]
+    case lookupTM rhs (es_map st) of
+        Just v -> return v
+        Nothing -> do
+            -- Pick a suitable name
+            v <- mkExitJoinId (exprType rhs) join_arity
+            st <- get
+            put (st { es_joins = (v,rhs) : es_joins st
+                    , es_map = insertTM rhs v (es_map st)
+                    })
+            return v
 
 {-
 Note [Interesting expression]
@@ -318,4 +350,24 @@ For `postInlineUnconditionally` and unfolding-based inlining, the function
 `simplLetUnfolding` simply gives exit join points no unfolding, which prevents
 this kind of inlining.
 
+Note [Avoid duplicate exit points]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+If we have
+
+  joinrec go 0     x y = t (x*x)
+          go 10    x y = t (x*x)
+          go (n-1) x y = call go (n-1) (x+y)
+  in …
+
+we want to create only _one_ exit join point:
+
+  join exit x = t (x*x)
+  joinrec go 0     x y = call exit x
+          go 10    x y = call exit x
+          go (n-1) x y = call go (n-1) (x+y)
+  in …
+
+we do so by keeping a `CoreMap JoinId` around, and `addExit` checks for
+if we can re-use an already created exit join point.
 -}