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