compiler/stranal/WorkWrap.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1993-1998
   3 %
   4 \section[WorkWrap]{Worker/wrapper-generating back-end of strictness analyser}
   5
   6 \begin{code}
   7 module WorkWrap ( wwTopBinds, mkWrapper ) where
   8
   9 import CoreSyn
  10 import CoreUnfold       ( certainlyWillInline, mkInlineUnfolding, mkWwInlineRule )
  11 import CoreUtils        ( exprType, exprIsHNF )
  12 import CoreArity        ( exprArity )
  13 import Var
  14 import Id
  15 import Type             ( Type )
  16 import IdInfo
  17 import Demand
  18 import UniqSupply
  19 import BasicTypes
  20 import VarEnv           ( isEmptyVarEnv )
  21 import Maybes           ( orElse )
  22 import WwLib
  23 import Util             ( lengthIs, notNull )
  24 import Outputable
  25 import MonadUtils
  26
  27 #include "HsVersions.h"
  28 \end{code}
  29
  30 We take Core bindings whose binders have:
  31
  32 \begin{enumerate}
  33
  34 \item Strictness attached (by the front-end of the strictness
  35 analyser), and / or
  36
  37 \item Constructed Product Result information attached by the CPR
  38 analysis pass.
  39
  40 \end{enumerate}
  41
  42 and we return some ``plain'' bindings which have been
  43 worker/wrapper-ified, meaning:
  44
  45 \begin{enumerate}
  46
  47 \item Functions have been split into workers and wrappers where
  48 appropriate.  If a function has both strictness and CPR properties
  49 then only one worker/wrapper doing both transformations is produced;
  50
  51 \item Binders' @IdInfos@ have been updated to reflect the existence of
  52 these workers/wrappers (this is where we get STRICTNESS and CPR pragma
  53 info for exported values).
  54 \end{enumerate}
  55
  56 \begin{code}
  57 wwTopBinds :: UniqSupply -> [CoreBind] -> [CoreBind]
  58
  59 wwTopBinds us top_binds
  60   = initUs_ us $ do
  61     top_binds' <- mapM wwBind top_binds
  62     return (concat top_binds')
  63 \end{code}
  64
  65 %************************************************************************
  66 %*                                                                      *
  67 \subsection[wwBind-wwExpr]{@wwBind@ and @wwExpr@}
  68 %*                                                                      *
  69 %************************************************************************
  70
  71 @wwBind@ works on a binding, trying each \tr{(binder, expr)} pair in
  72 turn.  Non-recursive case first, then recursive...
  73
  74 \begin{code}
  75 wwBind  :: CoreBind
  76         -> UniqSM [CoreBind]    -- returns a WwBinding intermediate form;
  77                                 -- the caller will convert to Expr/Binding,
  78                                 -- as appropriate.
  79
  80 wwBind (NonRec binder rhs) = do
  81     new_rhs <- wwExpr rhs
  82     new_pairs <- tryWW NonRecursive binder new_rhs
  83     return [NonRec b e | (b,e) <- new_pairs]
  84       -- Generated bindings must be non-recursive
  85       -- because the original binding was.
  86
  87 wwBind (Rec pairs)
  88   = return . Rec <$> concatMapM do_one pairs
  89   where
  90     do_one (binder, rhs) = do new_rhs <- wwExpr rhs
  91                               tryWW Recursive binder new_rhs
  92 \end{code}
  93
  94 @wwExpr@ basically just walks the tree, looking for appropriate
  95 annotations that can be used. Remember it is @wwBind@ that does the
  96 matching by looking for strict arguments of the correct type.
  97 @wwExpr@ is a version that just returns the ``Plain'' Tree.
  98
  99 \begin{code}
 100 wwExpr :: CoreExpr -> UniqSM CoreExpr
 101
 102 wwExpr e@(Type {}) = return e
 103 wwExpr e@(Lit  {}) = return e
 104 wwExpr e@(Var  {}) = return e
 105
 106 wwExpr (Lam binder expr)
 107   = Lam binder <$> wwExpr expr
 108
 109 wwExpr (App f a)
 110   = App <$> wwExpr f <*> wwExpr a
 111
 112 wwExpr (Note note expr)
 113   = Note note <$> wwExpr expr
 114
 115 wwExpr (Cast expr co) = do
 116     new_expr <- wwExpr expr
 117     return (Cast new_expr co)
 118
 119 wwExpr (Let bind expr)
 120   = mkLets <$> wwBind bind <*> wwExpr expr
 121
 122 wwExpr (Case expr binder ty alts) = do
 123     new_expr <- wwExpr expr
 124     new_alts <- mapM ww_alt alts
 125     return (Case new_expr binder ty new_alts)
 126   where
 127     ww_alt (con, binders, rhs) = do
 128         new_rhs <- wwExpr rhs
 129         return (con, binders, new_rhs)
 130 \end{code}
 131
 132 %************************************************************************
 133 %*                                                                      *
 134 \subsection[tryWW]{@tryWW@: attempt a worker/wrapper pair}
 135 %*                                                                      *
 136 %************************************************************************
 137
 138 @tryWW@ just accumulates arguments, converts strictness info from the
 139 front-end into the proper form, then calls @mkWwBodies@ to do
 140 the business.
 141
 142 The only reason this is monadised is for the unique supply.
 143
 144 Note [Don't w/w INLINE things]
 145 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 146 It's very important to refrain from w/w-ing an INLINE function (ie one
 147 with an InlineRule) because the wrapper will then overwrite the
 148 InlineRule unfolding.
 149
 150 Furthermore, if the programmer has marked something as INLINE,
 151 we may lose by w/w'ing it.
 152
 153 If the strictness analyser is run twice, this test also prevents
 154 wrappers (which are INLINEd) from being re-done.  (You can end up with
 155 several liked-named Ids bouncing around at the same time---absolute
 156 mischief.)
 157
 158 Notice that we refrain from w/w'ing an INLINE function even if it is
 159 in a recursive group.  It might not be the loop breaker.  (We could
 160 test for loop-breaker-hood, but I'm not sure that ever matters.)
 161
 162 Note [Don't w/w INLINABLE things]
 163 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 164 If we have
 165   {-# INLINABLE f #-}
 166   f x y = ....
 167 then in principle we might get a more efficient loop by w/w'ing f.
 168 But that would make a new unfolding which would overwrite the old
 169 one.  So we leave INLINABLE things alone too.
 170
 171 This is a slight infelicity really, because it means that adding
 172 an INLINABLE pragma could make a program a bit less efficient,
 173 because you lose the worker/wrapper stuff.  But I don't see a way
 174 to avoid that.
 175
 176 Note [Don't w/w inline small non-loop-breker things]
 177 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 178 In general, we refrain from w/w-ing *small* functions, which are not
 179 loop breakers, because they'll inline anyway.  But we must take care:
 180 it may look small now, but get to be big later after other inlining
 181 has happened.  So we take the precaution of adding an INLINE pragma to
 182 any such functions.
 183
 184 I made this change when I observed a big function at the end of
 185 compilation with a useful strictness signature but no w-w.  When
 186 I measured it on nofib, it didn't make much difference; just a few
 187 percent improved allocation on one benchmark (bspt/Euclid.space).
 188 But nothing got worse.
 189
 190 Note [Wrapper activation]
 191 ~~~~~~~~~~~~~~~~~~~~~~~~~
 192 When should the wrapper inlining be active?  It must not be active
 193 earlier than the current Activation of the Id (eg it might have a
 194 NOINLINE pragma).  But in fact strictness analysis happens fairly
 195 late in the pipeline, and we want to prioritise specialisations over
 196 strictness.  Eg if we have
 197   module Foo where
 198     f :: Num a => a -> Int -> a
 199     f n 0 = n              -- Strict in the Int, hence wrapper
 200     f n x = f (n+n) (x-1)
 201
 202     g :: Int -> Int
 203     g x = f x x            -- Provokes a specialisation for f
 204
 205   module Bsr where
 206     import Foo
 207
 208     h :: Int -> Int
 209     h x = f 3 x
 210
 211 Then we want the specialisation for 'f' to kick in before the wrapper does.
 212
 213 Now in fact the 'gentle' simplification pass encourages this, by
 214 having rules on, but inlinings off.  But that's kind of lucky. It seems
 215 more robust to give the wrapper an Activation of (ActiveAfter 0),
 216 so that it becomes active in an importing module at the same time that
 217 it appears in the first place in the defining module.
 218
 219 \begin{code}
 220 tryWW   :: RecFlag
 221         -> Id                           -- The fn binder
 222         -> CoreExpr                     -- The bound rhs; its innards
 223                                         --   are already ww'd
 224         -> UniqSM [(Id, CoreExpr)]      -- either *one* or *two* pairs;
 225                                         -- if one, then no worker (only
 226                                         -- the orig "wrapper" lives on);
 227                                         -- if two, then a worker and a
 228                                         -- wrapper.
 229 tryWW is_rec fn_id rhs
 230   | isNeverActive inline_act
 231         -- No point in worker/wrappering if the thing is never inlined!
 232         -- Because the no-inline prag will prevent the wrapper ever
 233         -- being inlined at a call site.
 234         --
 235         -- Furthermore, don't even expose strictness info
 236   = return [ (fn_id, rhs) ]
 237
 238   | is_thunk && worthSplittingThunk maybe_fn_dmd res_info
 239         -- See Note [Thunk splitting]
 240   = ASSERT2( isNonRec is_rec, ppr new_fn_id )   -- The thunk must be non-recursive
 241     checkSize new_fn_id rhs $
 242     splitThunk new_fn_id rhs
 243
 244   | is_fun && worthSplittingFun wrap_dmds res_info
 245   = checkSize new_fn_id rhs $
 246     splitFun new_fn_id fn_info wrap_dmds res_info rhs
 247
 248   | otherwise
 249   = return [ (new_fn_id, rhs) ]
 250
 251   where
 252     fn_info      = idInfo fn_id
 253     maybe_fn_dmd = demandInfo fn_info
 254     inline_act   = inlinePragmaActivation (inlinePragInfo fn_info)
 255
 256         -- In practice it always will have a strictness
 257         -- signature, even if it's a uninformative one
 258     strict_sig  = strictnessInfo fn_info `orElse` topSig
 259     StrictSig (DmdType env wrap_dmds res_info) = strict_sig
 260
 261         -- new_fn_id has the DmdEnv zapped.
 262         --      (a) it is never used again
 263         --      (b) it wastes space
 264         --      (c) it becomes incorrect as things are cloned, because
 265         --          we don't push the substitution into it
 266     new_fn_id | isEmptyVarEnv env = fn_id
 267               | otherwise         = fn_id `setIdStrictness`
 268                                      StrictSig (mkTopDmdType wrap_dmds res_info)
 269
 270     is_fun    = notNull wrap_dmds
 271     is_thunk  = not is_fun && not (exprIsHNF rhs)
 272
 273 ---------------------
 274 checkSize :: Id -> CoreExpr
 275           -> UniqSM [(Id,CoreExpr)] -> UniqSM [(Id,CoreExpr)]
 276 checkSize fn_id rhs thing_inside
 277   | isStableUnfolding (realIdUnfolding fn_id)
 278   = return [ (fn_id, rhs) ]
 279       -- See Note [Don't w/w INLINABLE things]
 280       -- and Note [Don't w/w INLINABLABLE things]
 281       -- NB: use realIdUnfolding because we want to see the unfolding
 282       --     even if it's a loop breaker!
 283
 284   | certainlyWillInline (idUnfolding fn_id)
 285   = return [ (fn_id `setIdUnfolding` inline_rule, rhs) ]
 286         -- Note [Don't w/w inline small non-loop-breaker things]
 287         -- NB: use idUnfolding because we don't want to apply
 288         --     this criterion to a loop breaker!
 289
 290   | otherwise = thing_inside
 291   where
 292     inline_rule = mkInlineUnfolding Nothing rhs
 293
 294 ---------------------
 295 splitFun :: Id -> IdInfo -> [Demand] -> DmdResult -> Expr Var
 296          -> UniqSM [(Id, CoreExpr)]
 297 splitFun fn_id fn_info wrap_dmds res_info rhs
 298   = WARN( not (wrap_dmds `lengthIs` arity), ppr fn_id <+> (ppr arity $$ ppr wrap_dmds $$ ppr res_info) )
 299     (do {
 300         -- The arity should match the signature
 301       (work_demands, wrap_fn, work_fn) <- mkWwBodies fun_ty wrap_dmds res_info one_shots
 302     ; work_uniq <- getUniqueM
 303     ; let
 304         work_rhs = work_fn rhs
 305         work_id  = mkWorkerId work_uniq fn_id (exprType work_rhs)
 306                         `setIdOccInfo` occInfo fn_info
 307                                 -- Copy over occurrence info from parent
 308                                 -- Notably whether it's a loop breaker
 309                                 -- Doesn't matter much, since we will simplify next, but
 310                                 -- seems right-er to do so
 311
 312                         `setInlineActivation` (inlinePragmaActivation inl_prag)
 313                                 -- Any inline activation (which sets when inlining is active)
 314                                 -- on the original function is duplicated on the worker
 315                                 -- It *matters* that the pragma stays on the wrapper
 316                                 -- It seems sensible to have it on the worker too, although we
 317                                 -- can't think of a compelling reason. (In ptic, INLINE things are
 318                                 -- not w/wd). However, the RuleMatchInfo is not transferred since
 319                                 -- it does not make sense for workers to be constructorlike.
 320
 321                         `setIdStrictness` StrictSig (mkTopDmdType work_demands work_res_info)
 322                                 -- Even though we may not be at top level,
 323                                 -- it's ok to give it an empty DmdEnv
 324
 325                         `setIdArity` (exprArity work_rhs)
 326                                 -- Set the arity so that the Core Lint check that the
 327                                 -- arity is consistent with the demand type goes through
 328
 329         wrap_rhs  = wrap_fn work_id
 330         wrap_prag = InlinePragma { inl_inline = Inline
 331                                  , inl_sat    = Nothing
 332                                  , inl_act    = ActiveAfter 0
 333                                  , inl_rule   = rule_match_info }
 334                 -- See Note [Wrapper activation]
 335                 -- The RuleMatchInfo is (and must be) unaffected
 336                 -- The inl_inline is bound to be False, else we would not be
 337                 --    making a wrapper
 338
 339         wrap_id   = fn_id `setIdUnfolding` mkWwInlineRule work_id wrap_rhs arity
 340                           `setInlinePragma` wrap_prag
 341                           `setIdOccInfo` NoOccInfo
 342                                 -- Zap any loop-breaker-ness, to avoid bleating from Lint
 343                                 -- about a loop breaker with an INLINE rule
 344
 345     ; return ([(work_id, work_rhs), (wrap_id, wrap_rhs)]) })
 346         -- Worker first, because wrapper mentions it
 347         -- mkWwBodies has already built a wrap_rhs with an INLINE pragma wrapped around it
 348   where
 349     fun_ty          = idType fn_id
 350     inl_prag        = inlinePragInfo fn_info
 351     rule_match_info = inlinePragmaRuleMatchInfo inl_prag
 352     arity           = arityInfo fn_info
 353                     -- The arity is set by the simplifier using exprEtaExpandArity
 354                     -- So it may be more than the number of top-level-visible lambdas
 355
 356     work_res_info | isBotRes res_info = BotRes  -- Cpr stuff done by wrapper
 357                   | otherwise         = TopRes
 358
 359     one_shots = get_one_shots rhs
 360
 361 -- If the original function has one-shot arguments, it is important to
 362 -- make the wrapper and worker have corresponding one-shot arguments too.
 363 -- Otherwise we spuriously float stuff out of case-expression join points,
 364 -- which is very annoying.
 365 get_one_shots :: Expr Var -> [Bool]
 366 get_one_shots (Lam b e)
 367   | isId b    = isOneShotLambda b : get_one_shots e
 368   | otherwise = get_one_shots e
 369 get_one_shots (Note _ e) = get_one_shots e
 370 get_one_shots _          = noOneShotInfo
 371 \end{code}
 372
 373 Note [Thunk splitting]
 374 ~~~~~~~~~~~~~~~~~~~~~~
 375 Suppose x is used strictly (never mind whether it has the CPR
 376 property).
 377
 378       let
 379         x* = x-rhs
 380       in body
 381
 382 splitThunk transforms like this:
 383
 384       let
 385         x* = case x-rhs of { I# a -> I# a }
 386       in body
 387
 388 Now simplifier will transform to
 389
 390       case x-rhs of
 391         I# a -> let x* = I# a
 392                 in body
 393
 394 which is what we want. Now suppose x-rhs is itself a case:
 395
 396         x-rhs = case e of { T -> I# a; F -> I# b }
 397
 398 The join point will abstract over a, rather than over (which is
 399 what would have happened before) which is fine.
 400
 401 Notice that x certainly has the CPR property now!
 402
 403 In fact, splitThunk uses the function argument w/w splitting
 404 function, so that if x's demand is deeper (say U(U(L,L),L))
 405 then the splitting will go deeper too.
 406
 407 \begin{code}
 408 -- See Note [Thunk splitting]
 409 -- splitThunk converts the *non-recursive* binding
 410 --      x = e
 411 -- into
 412 --      x = let x = e
 413 --          in case x of
 414 --               I# y -> let x = I# y in x }
 415 -- See comments above. Is it not beautifully short?
 416
 417 splitThunk :: Var -> Expr Var -> UniqSM [(Var, Expr Var)]
 418 splitThunk fn_id rhs = do
 419     (_, wrap_fn, work_fn) <- mkWWstr [fn_id]
 420     return [ (fn_id, Let (NonRec fn_id rhs) (wrap_fn (work_fn (Var fn_id)))) ]
 421 \end{code}
 422
 423
 424 %************************************************************************
 425 %*                                                                      *
 426 \subsection{Functions over Demands}
 427 %*                                                                      *
 428 %************************************************************************
 429
 430 \begin{code}
 431 worthSplittingFun :: [Demand] -> DmdResult -> Bool
 432                 -- True <=> the wrapper would not be an identity function
 433 worthSplittingFun ds res
 434   = any worth_it ds || returnsCPR res
 435         -- worthSplitting returns False for an empty list of demands,
 436         -- and hence do_strict_ww is False if arity is zero and there is no CPR
 437   -- See Note [Worker-wrapper for bottoming functions]
 438   where
 439     worth_it Abs              = True    -- Absent arg
 440     worth_it (Eval (Prod _)) = True     -- Product arg to evaluate
 441     worth_it _                = False
 442
 443 worthSplittingThunk :: Maybe Demand     -- Demand on the thunk
 444                     -> DmdResult        -- CPR info for the thunk
 445                     -> Bool
 446 worthSplittingThunk maybe_dmd res
 447   = worth_it maybe_dmd || returnsCPR res
 448   where
 449         -- Split if the thing is unpacked
 450     worth_it (Just (Eval (Prod ds))) = not (all isAbsent ds)
 451     worth_it _                       = False
 452 \end{code}
 453
 454 Note [Worker-wrapper for bottoming functions]
 455 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 456 We used not to split if the result is bottom.
 457 [Justification:  there's no efficiency to be gained.]
 458
 459 But it's sometimes bad not to make a wrapper.  Consider
 460         fw = \x# -> let x = I# x# in case e of
 461                                         p1 -> error_fn x
 462                                         p2 -> error_fn x
 463                                         p3 -> the real stuff
 464 The re-boxing code won't go away unless error_fn gets a wrapper too.
 465 [We don't do reboxing now, but in general it's better to pass an
 466 unboxed thing to f, and have it reboxed in the error cases....]
 467
 468
 469 %************************************************************************
 470 %*                                                                      *
 471 \subsection{The worker wrapper core}
 472 %*                                                                      *
 473 %************************************************************************
 474
 475 @mkWrapper@ is called when importing a function.  We have the type of
 476 the function and the name of its worker, and we want to make its body (the wrapper).
 477
 478 \begin{code}
 479 mkWrapper :: Type               -- Wrapper type
 480           -> StrictSig          -- Wrapper strictness info
 481           -> UniqSM (Id -> CoreExpr)    -- Wrapper body, missing worker Id
 482
 483 mkWrapper fun_ty (StrictSig (DmdType _ demands res_info)) = do
 484     (_, wrap_fn, _) <- mkWwBodies fun_ty demands res_info noOneShotInfo
 485     return wrap_fn
 486
 487 noOneShotInfo :: [Bool]
 488 noOneShotInfo = repeat False
 489 \end{code}