compiler/simplCore/FloatIn.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   3 %
   4 %************************************************************************
   5 %*                                                                      *
   6 \section[FloatIn]{Floating Inwards pass}
   7 %*                                                                      *
   8 %************************************************************************
   9
  10 The main purpose of @floatInwards@ is floating into branches of a
  11 case, so that we don't allocate things, save them on the stack, and
  12 then discover that they aren't needed in the chosen branch.
  13
  14 \begin{code}
  15 module FloatIn ( floatInwards ) where
  16
  17 #include "HsVersions.h"
  18
  19 import CoreSyn
  20 import CoreUtils        ( exprIsHNF, exprIsDupable )
  21 import CoreFVs          ( CoreExprWithFVs, freeVars, freeVarsOf, idRuleVars )
  22 import Id               ( isOneShotBndr, idType )
  23 import Var
  24 import Type             ( isUnLiftedType )
  25 import VarSet
  26 import Util             ( zipEqual, zipWithEqual, count )
  27 import UniqFM
  28 import Outputable
  29 \end{code}
  30
  31 Top-level interface function, @floatInwards@.  Note that we do not
  32 actually float any bindings downwards from the top-level.
  33
  34 \begin{code}
  35 floatInwards :: [CoreBind] -> [CoreBind]
  36 floatInwards = map fi_top_bind
  37   where
  38     fi_top_bind (NonRec binder rhs)
  39       = NonRec binder (fiExpr [] (freeVars rhs))
  40     fi_top_bind (Rec pairs)
  41       = Rec [ (b, fiExpr [] (freeVars rhs)) | (b, rhs) <- pairs ]
  42 \end{code}
  43
  44 %************************************************************************
  45 %*                                                                      *
  46 \subsection{Mail from Andr\'e [edited]}
  47 %*                                                                      *
  48 %************************************************************************
  49
  50 {\em Will wrote: What??? I thought the idea was to float as far
  51 inwards as possible, no matter what.  This is dropping all bindings
  52 every time it sees a lambda of any kind.  Help! }
  53
  54 You are assuming we DO DO full laziness AFTER floating inwards!  We
  55 have to [not float inside lambdas] if we don't.
  56
  57 If we indeed do full laziness after the floating inwards (we could
  58 check the compilation flags for that) then I agree we could be more
  59 aggressive and do float inwards past lambdas.
  60
  61 Actually we are not doing a proper full laziness (see below), which
  62 was another reason for not floating inwards past a lambda.
  63
  64 This can easily be fixed.  The problem is that we float lets outwards,
  65 but there are a few expressions which are not let bound, like case
  66 scrutinees and case alternatives.  After floating inwards the
  67 simplifier could decide to inline the let and the laziness would be
  68 lost, e.g.
  69
  70 \begin{verbatim}
  71 let a = expensive             ==> \b -> case expensive of ...
  72 in \ b -> case a of ...
  73 \end{verbatim}
  74 The fix is
  75 \begin{enumerate}
  76 \item
  77 to let bind the algebraic case scrutinees (done, I think) and
  78 the case alternatives (except the ones with an
  79 unboxed type)(not done, I think). This is best done in the
  80 SetLevels.lhs module, which tags things with their level numbers.
  81 \item
  82 do the full laziness pass (floating lets outwards).
  83 \item
  84 simplify. The simplifier inlines the (trivial) lets that were
  85  created but were not floated outwards.
  86 \end{enumerate}
  87
  88 With the fix I think Will's suggestion that we can gain even more from
  89 strictness by floating inwards past lambdas makes sense.
  90
  91 We still gain even without going past lambdas, as things may be
  92 strict in the (new) context of a branch (where it was floated to) or
  93 of a let rhs, e.g.
  94 \begin{verbatim}
  95 let a = something            case x of
  96 in case x of                   alt1 -> case something of a -> a + a
  97      alt1 -> a + a      ==>    alt2 -> b
  98      alt2 -> b
  99
 100 let a = something           let b = case something of a -> a + a
 101 in let b = a + a        ==> in (b,b)
 102 in (b,b)
 103 \end{verbatim}
 104 Also, even if a is not found to be strict in the new context and is
 105 still left as a let, if the branch is not taken (or b is not entered)
 106 the closure for a is not built.
 107
 108 %************************************************************************
 109 %*                                                                      *
 110 \subsection{Main floating-inwards code}
 111 %*                                                                      *
 112 %************************************************************************
 113
 114 \begin{code}
 115 type FreeVarsSet   = IdSet
 116
 117 type FloatingBinds = [(CoreBind, FreeVarsSet)]
 118         -- In reverse dependency order (innermost binder first)
 119
 120         -- The FreeVarsSet is the free variables of the binding.  In the case
 121         -- of recursive bindings, the set doesn't include the bound
 122         -- variables.
 123
 124 fiExpr :: FloatingBinds         -- Binds we're trying to drop
 125                                 -- as far "inwards" as possible
 126        -> CoreExprWithFVs       -- Input expr
 127        -> CoreExpr              -- Result
 128
 129 fiExpr to_drop (_, AnnVar v) = mkCoLets' to_drop (Var v)
 130
 131 fiExpr to_drop (_, AnnType ty) = ASSERT( null to_drop )
 132                                  Type ty
 133 fiExpr to_drop (_, AnnCast expr co)
 134   = Cast (fiExpr to_drop expr) co       -- Just float in past coercion
 135
 136 fiExpr _ (_, AnnLit lit) = Lit lit
 137 \end{code}
 138
 139 Applications: we do float inside applications, mainly because we
 140 need to get at all the arguments.  The next simplifier run will
 141 pull out any silly ones.
 142
 143 \begin{code}
 144 fiExpr to_drop (_,AnnApp fun arg)
 145   = mkCoLets' drop_here (App (fiExpr fun_drop fun) (fiExpr arg_drop arg))
 146   where
 147     [drop_here, fun_drop, arg_drop] = sepBindsByDropPoint False [freeVarsOf fun, freeVarsOf arg] to_drop
 148 \end{code}
 149
 150 Note [Floating in past a lambda group]
 151 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 152 * We must be careful about floating inside inside a value lambda.
 153   That risks losing laziness.
 154   The float-out pass might rescue us, but then again it might not.
 155
 156 * We must be careful about type lambdas too.  At one time we did, and
 157   there is no risk of duplicating work thereby, but we do need to be
 158   careful.  In particular, here is a bad case (it happened in the
 159   cichelli benchmark:
 160         let v = ...
 161         in let f = /\t -> \a -> ...
 162            ==>
 163         let f = /\t -> let v = ... in \a -> ...
 164   This is bad as now f is an updatable closure (update PAP)
 165   and has arity 0.
 166
 167 * Hack alert!  We only float in through one-shot lambdas,
 168   not (as you might guess) through lone big lambdas.
 169   Reason: we float *out* past big lambdas (see the test in the Lam
 170   case of FloatOut.floatExpr) and we don't want to float straight
 171   back in again.
 172
 173   It *is* important to float into one-shot lambdas, however;
 174   see the remarks with noFloatIntoRhs.
 175
 176 So we treat lambda in groups, using the following rule:
 177
 178  Float in if (a) there is at least one Id,
 179          and (b) there are no non-one-shot Ids
 180
 181  Otherwise drop all the bindings outside the group.
 182
 183 This is what the 'go' function in the AnnLam case is doing.
 184
 185 Urk! if all are tyvars, and we don't float in, we may miss an
 186       opportunity to float inside a nested case branch
 187
 188 \begin{code}
 189 fiExpr to_drop lam@(_, AnnLam _ _)
 190   | go False bndrs      -- Float in
 191   = mkLams bndrs (fiExpr to_drop body)
 192
 193   | otherwise           -- Dump it all here
 194   = mkCoLets' to_drop (mkLams bndrs (fiExpr [] body))
 195
 196   where
 197     (bndrs, body) = collectAnnBndrs lam
 198
 199     go seen_one_shot_id [] = seen_one_shot_id
 200     go seen_one_shot_id (b:bs)
 201       | isTyVar       b = go seen_one_shot_id bs
 202       | isOneShotBndr b = go True bs
 203       | otherwise       = False  -- Give up at a non-one-shot Id
 204 \end{code}
 205
 206 We don't float lets inwards past an SCC.
 207         ToDo: keep info on current cc, and when passing
 208         one, if it is not the same, annotate all lets in binds with current
 209         cc, change current cc to the new one and float binds into expr.
 210
 211 \begin{code}
 212 fiExpr to_drop (_, AnnNote note@(SCC _) expr)
 213   =     -- Wimp out for now
 214     mkCoLets' to_drop (Note note (fiExpr [] expr))
 215
 216 fiExpr to_drop (_, AnnNote note@(CoreNote _) expr)
 217   = Note note (fiExpr to_drop expr)
 218 \end{code}
 219
 220 For @Lets@, the possible ``drop points'' for the \tr{to_drop}
 221 bindings are: (a)~in the body, (b1)~in the RHS of a NonRec binding,
 222 or~(b2), in each of the RHSs of the pairs of a @Rec@.
 223
 224 Note that we do {\em weird things} with this let's binding.  Consider:
 225 \begin{verbatim}
 226 let
 227     w = ...
 228 in {
 229     let v = ... w ...
 230     in ... v .. w ...
 231 }
 232 \end{verbatim}
 233 Look at the inner \tr{let}.  As \tr{w} is used in both the bind and
 234 body of the inner let, we could panic and leave \tr{w}'s binding where
 235 it is.  But \tr{v} is floatable further into the body of the inner let, and
 236 {\em then} \tr{w} will also be only in the body of that inner let.
 237
 238 So: rather than drop \tr{w}'s binding here, we add it onto the list of
 239 things to drop in the outer let's body, and let nature take its
 240 course.
 241
 242 Note [extra_fvs (1): avoid floating into RHS]
 243 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 244 Consdider let x=\y....t... in body.  We do not necessarily want to float
 245 a binding for t into the RHS, because it'll immediately be floated out
 246 again.  (It won't go inside the lambda else we risk losing work.)
 247 In letrec, we need to be more careful still. We don't want to transform
 248         let x# = y# +# 1#
 249         in
 250         letrec f = \z. ...x#...f...
 251         in ...
 252 into
 253         letrec f = let x# = y# +# 1# in \z. ...x#...f... in ...
 254 because now we can't float the let out again, because a letrec
 255 can't have unboxed bindings.
 256
 257 So we make "extra_fvs" which is the rhs_fvs of such bindings, and
 258 arrange to dump bindings that bind extra_fvs before the entire let.
 259
 260 Note [extra_fvs (s): free variables of rules]
 261 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 262 Consider let x{rule mentioning y} = rhs in body
 263 Here y is not free in rhs or body; but we still want to dump bindings
 264 that bind y outside the let.  So we augment extra_fvs with the
 265 idRuleVars of x.
 266
 267
 268 \begin{code}
 269 fiExpr to_drop (_,AnnLet (AnnNonRec id rhs@(rhs_fvs, ann_rhs)) body)
 270   = fiExpr new_to_drop body
 271   where
 272     body_fvs = freeVarsOf body
 273
 274     rule_fvs = idRuleVars id    -- See Note [extra_fvs (2): free variables of rules]
 275     extra_fvs | noFloatIntoRhs ann_rhs
 276               || isUnLiftedType (idType id) = rule_fvs `unionVarSet` rhs_fvs
 277               | otherwise                   = rule_fvs
 278         -- See Note [extra_fvs (2): avoid floating into RHS]
 279         -- No point in floating in only to float straight out again
 280         -- Ditto ok-for-speculation unlifted RHSs
 281
 282     [shared_binds, extra_binds, rhs_binds, body_binds]
 283         = sepBindsByDropPoint False [extra_fvs, rhs_fvs, body_fvs] to_drop
 284
 285     new_to_drop = body_binds ++                         -- the bindings used only in the body
 286                   [(NonRec id rhs', rhs_fvs')] ++       -- the new binding itself
 287                   extra_binds ++                        -- bindings from extra_fvs
 288                   shared_binds                          -- the bindings used both in rhs and body
 289
 290         -- Push rhs_binds into the right hand side of the binding
 291     rhs'     = fiExpr rhs_binds rhs
 292     rhs_fvs' = rhs_fvs `unionVarSet` floatedBindsFVs rhs_binds `unionVarSet` rule_fvs
 293                         -- Don't forget the rule_fvs; the binding mentions them!
 294
 295 fiExpr to_drop (_,AnnLet (AnnRec bindings) body)
 296   = fiExpr new_to_drop body
 297   where
 298     (ids, rhss) = unzip bindings
 299     rhss_fvs = map freeVarsOf rhss
 300     body_fvs = freeVarsOf body
 301
 302         -- See Note [extra_fvs (1,2)]
 303     rule_fvs = foldr (unionVarSet . idRuleVars) emptyVarSet ids
 304     extra_fvs = rule_fvs `unionVarSet`
 305                 unionVarSets [ fvs | (fvs, rhs) <- rhss
 306                              , noFloatIntoRhs rhs ]
 307
 308     (shared_binds:extra_binds:body_binds:rhss_binds)
 309         = sepBindsByDropPoint False (extra_fvs:body_fvs:rhss_fvs) to_drop
 310
 311     new_to_drop = body_binds ++         -- the bindings used only in the body
 312                   [(Rec (fi_bind rhss_binds bindings), rhs_fvs')] ++
 313                                         -- The new binding itself
 314                   extra_binds ++        -- Note [extra_fvs (1,2)]
 315                   shared_binds          -- Used in more than one place
 316
 317     rhs_fvs' = unionVarSets rhss_fvs `unionVarSet`
 318                unionVarSets (map floatedBindsFVs rhss_binds) `unionVarSet`
 319                rule_fvs         -- Don't forget the rule variables!
 320
 321     -- Push rhs_binds into the right hand side of the binding
 322     fi_bind :: [FloatingBinds]      -- one per "drop pt" conjured w/ fvs_of_rhss
 323             -> [(Id, CoreExprWithFVs)]
 324             -> [(Id, CoreExpr)]
 325
 326     fi_bind to_drops pairs
 327       = [ (binder, fiExpr to_drop rhs)
 328         | ((binder, rhs), to_drop) <- zipEqual "fi_bind" pairs to_drops ]
 329 \end{code}
 330
 331 For @Case@, the possible ``drop points'' for the \tr{to_drop}
 332 bindings are: (a)~inside the scrutinee, (b)~inside one of the
 333 alternatives/default [default FVs always {\em first}!].
 334
 335 \begin{code}
 336 fiExpr to_drop (_, AnnCase scrut case_bndr ty alts)
 337   = mkCoLets' drop_here1 $
 338     mkCoLets' drop_here2 $
 339     Case (fiExpr scrut_drops scrut) case_bndr ty
 340          (zipWith fi_alt alts_drops_s alts)
 341   where
 342         -- Float into the scrut and alts-considered-together just like App
 343     [drop_here1, scrut_drops, alts_drops] = sepBindsByDropPoint False [scrut_fvs, all_alts_fvs] to_drop
 344
 345         -- Float into the alts with the is_case flag set
 346     (drop_here2 : alts_drops_s)           = sepBindsByDropPoint True alts_fvs alts_drops
 347
 348     scrut_fvs    = freeVarsOf scrut
 349     alts_fvs     = map alt_fvs alts
 350     all_alts_fvs = unionVarSets alts_fvs
 351     alt_fvs (_con, args, rhs) = foldl delVarSet (freeVarsOf rhs) (case_bndr:args)
 352                                 -- Delete case_bndr and args from free vars of rhs
 353                                 -- to get free vars of alt
 354
 355     fi_alt to_drop (con, args, rhs) = (con, args, fiExpr to_drop rhs)
 356
 357 noFloatIntoRhs :: AnnExpr' Var (UniqFM Var) -> Bool
 358 noFloatIntoRhs (AnnLam b _) = not (is_one_shot b)
 359         -- IMPORTANT: don't say 'True' for a RHS with a one-shot lambda at the top.
 360         -- This makes a big difference for things like
 361         --      f x# = let x = I# x#
 362         --             in let j = \() -> ...x...
 363         --                in if <condition> then normal-path else j ()
 364         -- If x is used only in the error case join point, j, we must float the
 365         -- boxing constructor into it, else we box it every time which is very bad
 366         -- news indeed.
 367
 368 noFloatIntoRhs rhs = exprIsHNF (deAnnotate' rhs)        -- We'd just float right back out again...
 369
 370 is_one_shot :: Var -> Bool
 371 is_one_shot b = isIdVar b && isOneShotBndr b
 372 \end{code}
 373
 374
 375 %************************************************************************
 376 %*                                                                      *
 377 \subsection{@sepBindsByDropPoint@}
 378 %*                                                                      *
 379 %************************************************************************
 380
 381 This is the crucial function.  The idea is: We have a wad of bindings
 382 that we'd like to distribute inside a collection of {\em drop points};
 383 insides the alternatives of a \tr{case} would be one example of some
 384 drop points; the RHS and body of a non-recursive \tr{let} binding
 385 would be another (2-element) collection.
 386
 387 So: We're given a list of sets-of-free-variables, one per drop point,
 388 and a list of floating-inwards bindings.  If a binding can go into
 389 only one drop point (without suddenly making something out-of-scope),
 390 in it goes.  If a binding is used inside {\em multiple} drop points,
 391 then it has to go in a you-must-drop-it-above-all-these-drop-points
 392 point.
 393
 394 We have to maintain the order on these drop-point-related lists.
 395
 396 \begin{code}
 397 sepBindsByDropPoint
 398     :: Bool                 -- True <=> is case expression
 399     -> [FreeVarsSet]        -- One set of FVs per drop point
 400     -> FloatingBinds        -- Candidate floaters
 401     -> [FloatingBinds]      -- FIRST one is bindings which must not be floated
 402                             -- inside any drop point; the rest correspond
 403                             -- one-to-one with the input list of FV sets
 404
 405 -- Every input floater is returned somewhere in the result;
 406 -- none are dropped, not even ones which don't seem to be
 407 -- free in *any* of the drop-point fvs.  Why?  Because, for example,
 408 -- a binding (let x = E in B) might have a specialised version of
 409 -- x (say x') stored inside x, but x' isn't free in E or B.
 410
 411 type DropBox = (FreeVarsSet, FloatingBinds)
 412
 413 sepBindsByDropPoint _is_case drop_pts []
 414   = [] : [[] | _ <- drop_pts]   -- cut to the chase scene; it happens
 415
 416 sepBindsByDropPoint is_case drop_pts floaters
 417   = go floaters (map (\fvs -> (fvs, [])) (emptyVarSet : drop_pts))
 418   where
 419     go :: FloatingBinds -> [DropBox] -> [FloatingBinds]
 420         -- The *first* one in the argument list is the drop_here set
 421         -- The FloatingBinds in the lists are in the reverse of
 422         -- the normal FloatingBinds order; that is, they are the right way round!
 423
 424     go [] drop_boxes = map (reverse . snd) drop_boxes
 425
 426     go (bind_w_fvs@(bind, bind_fvs) : binds) drop_boxes@(here_box : fork_boxes)
 427         = go binds new_boxes
 428         where
 429           -- "here" means the group of bindings dropped at the top of the fork
 430
 431           (used_here : used_in_flags) = [ any (`elemVarSet` fvs) (bindersOf bind)
 432                                         | (fvs, _) <- drop_boxes]
 433
 434           drop_here = used_here || not can_push
 435
 436                 -- For case expressions we duplicate the binding if it is
 437                 -- reasonably small, and if it is not used in all the RHSs
 438                 -- This is good for situations like
 439                 --      let x = I# y in
 440                 --      case e of
 441                 --        C -> error x
 442                 --        D -> error x
 443                 --        E -> ...not mentioning x...
 444
 445           n_alts      = length used_in_flags
 446           n_used_alts = count id used_in_flags -- returns number of Trues in list.
 447
 448           can_push = n_used_alts == 1           -- Used in just one branch
 449                    || (is_case &&               -- We are looking at case alternatives
 450                        n_used_alts > 1 &&       -- It's used in more than one
 451                        n_used_alts < n_alts &&  -- ...but not all
 452                        bindIsDupable bind)      -- and we can duplicate the binding
 453
 454           new_boxes | drop_here = (insert here_box : fork_boxes)
 455                     | otherwise = (here_box : new_fork_boxes)
 456
 457           new_fork_boxes = zipWithEqual "FloatIn.sepBinds" insert_maybe fork_boxes used_in_flags
 458
 459           insert :: DropBox -> DropBox
 460           insert (fvs,drops) = (fvs `unionVarSet` bind_fvs, bind_w_fvs:drops)
 461
 462           insert_maybe box True  = insert box
 463           insert_maybe box False = box
 464
 465     go _ _ = panic "sepBindsByDropPoint/go"
 466
 467
 468 floatedBindsFVs :: FloatingBinds -> FreeVarsSet
 469 floatedBindsFVs binds = unionVarSets (map snd binds)
 470
 471 mkCoLets' :: FloatingBinds -> CoreExpr -> CoreExpr
 472 mkCoLets' to_drop e = foldl (flip (Let . fst)) e to_drop
 473         -- Remember to_drop is in *reverse* dependency order
 474
 475 bindIsDupable :: Bind CoreBndr -> Bool
 476 bindIsDupable (Rec prs)    = all (exprIsDupable . snd) prs
 477 bindIsDupable (NonRec _ r) = exprIsDupable r
 478 \end{code}