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