CSE.hs 10.9 KB
Newer Older
Austin Seipp's avatar
Austin Seipp committed
1 2 3
{-
(c) The AQUA Project, Glasgow University, 1993-1998

4
\section{Common subexpression}
Austin Seipp's avatar
Austin Seipp committed
5
-}
6

7 8
{-# LANGUAGE CPP #-}

9
module CSE (cseProgram) where
10 11 12

#include "HsVersions.h"

13
import CoreSubst
14 15 16
import Var              ( Var )
import Id               ( Id, idType, idInlineActivation, zapIdOccInfo )
import CoreUtils        ( mkAltExpr
17
                        , exprIsTrivial)
18
import Type             ( tyConAppArgs )
19 20
import CoreSyn
import Outputable
21
import BasicTypes       ( isAlwaysActive )
22
import TrieMap
23 24

import Data.List
25

Austin Seipp's avatar
Austin Seipp committed
26
{-
27 28
                        Simple common sub-expression
                        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
29
When we see
30 31
        x1 = C a b
        x2 = C x1 b
32
we build up a reverse mapping:   C a b  -> x1
33
                                 C x1 b -> x2
34 35 36
and apply that to the rest of the program.

When we then see
37 38
        y1 = C a b
        y2 = C y1 b
39
we replace the C a b with x1.  But then we *dont* want to
40
add   x1 -> y1  to the mapping.  Rather, we want the reverse, y1 -> x1
41
so that a subsequent binding
42 43
        y2 = C y1 b
will get transformed to C x1 b, and then to x2.
44

45
So we carry an extra var->var substitution which we apply *before* looking up in the
46 47 48
reverse mapping.


simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
49 50
Note [Shadowing]
~~~~~~~~~~~~~~~~
51
We have to be careful about shadowing.
52
For example, consider
53 54 55
        f = \x -> let y = x+x in
                      h = \x -> x+x
                  in ...
56

57 58 59 60
Here we must *not* do CSE on the inner x+x!  The simplifier used to guarantee no
shadowing, but it doesn't any more (it proved too hard), so we clone as we go.
We can simply add clones to the substitution already described.

61
Note [Case binders 1]
62
~~~~~~~~~~~~~~~~~~~~~~
63 64
Consider

65 66 67
        f = \x -> case x of wild {
                        (a:as) -> case a of wild1 {
                                    (p,q) -> ...(wild1:as)...
68 69 70 71 72 73

Here, (wild1:as) is morally the same as (a:as) and hence equal to wild.
But that's not quite obvious.  In general we want to keep it as (wild1:as),
but for CSE purpose that's a bad idea.

So we add the binding (wild1 -> a) to the extra var->var mapping.
74
Notice this is exactly backwards to what the simplifier does, which is
75
to try to replaces uses of 'a' with uses of 'wild1'
76

77
Note [Case binders 2]
78
~~~~~~~~~~~~~~~~~~~~~~
79
Consider
80
        case (h x) of y -> ...(h x)...
81 82

We'd like to replace (h x) in the alternative, by y.  But because of
83
the preceding [Note: case binders 1], we only want to add the mapping
84
        scrutinee -> case binder
85 86
to the reverse CSE mapping if the scrutinee is a non-trivial expression.
(If the scrutinee is a simple variable we want to add the mapping
87
        case binder -> scrutinee
88
to the substitution
89

90 91
Note [CSE for INLINE and NOINLINE]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
92 93 94
There are some subtle interactions of CSE with functions that the user
has marked as INLINE or NOINLINE. (Examples from Roman Leshchinskiy.)
Consider
95

96
        yes :: Int  {-# NOINLINE yes #-}
97
        yes = undefined
98

99
        no :: Int   {-# NOINLINE no #-}
100
        no = undefined
101

102
        foo :: Int -> Int -> Int  {-# NOINLINE foo #-}
103
        foo m n = n
104

105
        {-# RULES "foo/no" foo no = id #-}
106

107 108
        bar :: Int -> Int
        bar = foo yes
109

110
We do not expect the rule to fire.  But if we do CSE, then we risk
Gabor Greif's avatar
Gabor Greif committed
111
getting yes=no, and the rule does fire.  Actually, it won't because
112
NOINLINE means that 'yes' will never be inlined, not even if we have
Gabor Greif's avatar
Gabor Greif committed
113
yes=no.  So that's fine (now; perhaps in the olden days, yes=no would
114
have substituted even if 'yes' was NOINLINE.
115

116
But we do need to take care.  Consider
117

118 119
        {-# NOINLINE bar #-}
        bar = <rhs>     -- Same rhs as foo
120

121 122
        foo = <rhs>

123
If CSE produces
124
        foo = bar
125 126 127 128 129 130 131 132 133 134 135 136 137 138 139
then foo will never be inlined to <rhs> (when it should be, if <rhs>
is small).  The conclusion here is this:

   We should not add
       <rhs> :-> bar
  to the CSEnv if 'bar' has any constraints on when it can inline;
  that is, if its 'activation' not always active.  Otherwise we
  might replace <rhs> by 'bar', and then later be unable to see that it
  really was <rhs>.

Note that we do not (currently) do CSE on the unfolding stored inside
an Id, even if is a 'stable' unfolding.  That means that when an
unfolding happens, it is always faithful to what the stable unfolding
originally was.

140

141 142 143 144 145 146
Note [CSE for case expressions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
  case f x of y { pat -> ...let y = f x in ... }
Then we can CSE the inner (f x) to y.  In fact 'case' is like a strict
let-binding, and we can use cseRhs for dealing with the scrutinee.
147

Austin Seipp's avatar
Austin Seipp committed
148 149
************************************************************************
*                                                                      *
150
\section{Common subexpression}
Austin Seipp's avatar
Austin Seipp committed
151 152 153
*                                                                      *
************************************************************************
-}
154

155
cseProgram :: CoreProgram -> CoreProgram
156
cseProgram binds = snd (mapAccumL cseBind emptyCSEnv binds)
157 158

cseBind :: CSEnv -> CoreBind -> (CSEnv, CoreBind)
159
cseBind env (NonRec b e)
160 161 162 163 164 165 166 167 168 169 170 171 172 173
  = (env2, NonRec b' e')
  where
    (env1, b') = addBinder env b
    (env2, e') = cseRhs env1 (b',e)

cseBind env (Rec pairs)
  = (env2, Rec (bs' `zip` es'))
  where
    (bs,es) = unzip pairs
    (env1, bs') = addRecBinders env bs
    (env2, es') = mapAccumL cseRhs env1 (bs' `zip` es)

cseRhs :: CSEnv -> (OutBndr, InExpr) -> (CSEnv, OutExpr)
cseRhs env (id',rhs)
174
  = case lookupCSEnv env rhs' of
175 176 177 178 179 180
        Nothing
          | always_active -> (extendCSEnv env rhs' id', rhs')
          | otherwise     -> (env,                      rhs')
        Just id
          | always_active -> (extendCSSubst env id' id, Var id)
          | otherwise     -> (env,                      Var id)
181 182 183 184 185 186 187 188
          -- In the Just case, we have
          --        x = rhs
          --        ...
          --        x' = rhs
          -- We are replacing the second binding with x'=x
          -- and so must record that in the substitution so
          -- that subsequent uses of x' are replaced with x,
          -- See Trac #5996
189
  where
190 191 192 193
    rhs' = cseExpr env rhs

    always_active = isAlwaysActive (idInlineActivation id')
         -- See Note [CSE for INLINE and NOINLINE]
194

195
tryForCSE :: CSEnv -> InExpr -> OutExpr
196
tryForCSE env expr
197
  | exprIsTrivial expr'                   = expr'       -- No point
198
  | Just smaller <- lookupCSEnv env expr' = Var smaller
199 200 201
  | otherwise                             = expr'
  where
    expr' = cseExpr env expr
202

203
cseExpr :: CSEnv -> InExpr -> OutExpr
204 205
cseExpr env (Type t)               = Type (substTy (csEnvSubst env) t)
cseExpr env (Coercion c)           = Coercion (substCo (csEnvSubst env) c)
twanvl's avatar
twanvl committed
206
cseExpr _   (Lit lit)              = Lit lit
207 208
cseExpr env (Var v)                = lookupSubst env v
cseExpr env (App f a)              = App (cseExpr env f) (tryForCSE env a)
209
cseExpr env (Tick t e)             = Tick t (cseExpr env e)
210
cseExpr env (Cast e co)            = Cast (cseExpr env e) (substCo (csEnvSubst env) co)
211 212 213 214
cseExpr env (Lam b e)              = let (env', b') = addBinder env b
                                     in Lam b' (cseExpr env' e)
cseExpr env (Let bind e)           = let (env', bind') = cseBind env bind
                                     in Let bind' (cseExpr env' e)
215
cseExpr env (Case scrut bndr ty alts) = Case scrut' bndr'' ty alts'
216 217 218 219 220 221 222 223 224
                          where
                                alts' = cseAlts env2 scrut' bndr bndr'' alts
                                (env1, bndr') = addBinder env bndr
                                bndr'' = zapIdOccInfo bndr'
                                -- The swizzling from Note [Case binders 2] may
                                -- cause a dead case binder to be alive, so we
                                -- play safe here and bring them all to life
                                (env2, scrut') = cseRhs env1 (bndr'', scrut)
                                -- Note [CSE for case expressions]
225

226
cseAlts :: CSEnv -> OutExpr -> InBndr -> InBndr -> [InAlt] -> [OutAlt]
twanvl's avatar
twanvl committed
227

228
cseAlts env scrut' bndr bndr' alts
229 230
  = map cse_alt alts
  where
231
    (con_target, alt_env)
232
        = case scrut' of
233 234
            Var v' -> (v',     extendCSSubst env bndr v')    -- See Note [Case binders 1]
                                                             -- map: bndr -> v'
235

236 237
            _      ->  (bndr', extendCSEnv env scrut' bndr') -- See Note [Case binders 2]
                                                             -- map: scrut' -> bndr'
238

239
    arg_tys = tyConAppArgs (idType bndr)
240

241
    cse_alt (DataAlt con, args, rhs)
242 243 244 245 246 247 248 249 250
        | not (null args)
                -- Don't try CSE if there are no args; it just increases the number
                -- of live vars.  E.g.
                --      case x of { True -> ....True.... }
                -- Don't replace True by x!
                -- Hence the 'null args', which also deal with literals and DEFAULT
        = (DataAlt con, args', tryForCSE new_env rhs)
        where
          (env', args') = addBinders alt_env args
251 252
          new_env       = extendCSEnv env' con_expr con_target
          con_expr      = mkAltExpr (DataAlt con) args' arg_tys
253 254

    cse_alt (con, args, rhs)
255 256 257
        = (con, args', tryForCSE env' rhs)
        where
          (env', args') = addBinders alt_env args
258

Austin Seipp's avatar
Austin Seipp committed
259 260 261
{-
************************************************************************
*                                                                      *
262
\section{The CSE envt}
Austin Seipp's avatar
Austin Seipp committed
263 264 265
*                                                                      *
************************************************************************
-}
266

267
type InExpr  = CoreExpr         -- Pre-cloning
268 269 270
type InBndr  = CoreBndr
type InAlt   = CoreAlt

271
type OutExpr  = CoreExpr        -- Post-cloning
272 273
type OutBndr  = CoreBndr
type OutAlt   = CoreAlt
274

275
data CSEnv  = CS { cs_map    :: CoreMap (OutExpr, Id)   -- Key, value
276 277 278 279 280
                 , cs_subst  :: Subst }

emptyCSEnv :: CSEnv
emptyCSEnv = CS { cs_map = emptyCoreMap, cs_subst = emptySubst }

281
lookupCSEnv :: CSEnv -> OutExpr -> Maybe Id
282
lookupCSEnv (CS { cs_map = csmap }) expr
283 284 285 286
  = case lookupCoreMap csmap expr of
      Just (_,e) -> Just e
      Nothing    -> Nothing

287 288 289
extendCSEnv :: CSEnv -> OutExpr -> Id -> CSEnv
extendCSEnv cse expr id
  = cse { cs_map = extendCoreMap (cs_map cse) expr (expr,id) }
290 291 292 293

csEnvSubst :: CSEnv -> Subst
csEnvSubst = cs_subst

294
lookupSubst :: CSEnv -> Id -> OutExpr
295
lookupSubst (CS { cs_subst = sub}) x = lookupIdSubst (text "CSE.lookupSubst") sub x
296 297

extendCSSubst :: CSEnv -> Id  -> Id -> CSEnv
298
extendCSSubst cse x y = cse { cs_subst = extendIdSubst (cs_subst cse) x (Var y) }
299 300

addBinder :: CSEnv -> Var -> (CSEnv, Var)
301
addBinder cse v = (cse { cs_subst = sub' }, v')
302 303
                where
                  (sub', v') = substBndr (cs_subst cse) v
304 305

addBinders :: CSEnv -> [Var] -> (CSEnv, [Var])
306
addBinders cse vs = (cse { cs_subst = sub' }, vs')
307 308
                where
                  (sub', vs') = substBndrs (cs_subst cse) vs
309

310
addRecBinders :: CSEnv -> [Id] -> (CSEnv, [Id])
311
addRecBinders cse vs = (cse { cs_subst = sub' }, vs')
312 313
                where
                  (sub', vs') = substRecBndrs (cs_subst cse) vs