ghc/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 #include "HsVersions.h"
  10
  11 import CoreSyn
  12 import CoreUnfold       ( Unfolding, certainlySmallEnoughToInline, calcUnfoldingGuidance )
  13 import CmdLineOpts      ( opt_UF_CreationThreshold , opt_D_verbose_core2core,
  14                           opt_D_dump_worker_wrapper
  15                         )
  16 import CoreLint         ( beginPass, endPass )
  17 import CoreUtils        ( coreExprType, exprEtaExpandArity )
  18 import Const            ( Con(..) )
  19 import DataCon          ( DataCon )
  20 import MkId             ( mkWorkerId )
  21 import Id               ( Id, idType, getIdStrictness, setIdArity, isOneShotLambda,
  22                           setIdStrictness, getIdDemandInfo, getInlinePragma,
  23                           setIdWorkerInfo, getIdCprInfo )
  24 import VarSet
  25 import Type             ( Type, isNewType, splitForAllTys, splitFunTys )
  26 import IdInfo           ( mkStrictnessInfo, noStrictnessInfo, StrictnessInfo(..),
  27                           CprInfo(..), exactArity, InlinePragInfo(..)
  28                         )
  29 import Demand           ( Demand, wwLazy )
  30 import SaLib
  31 import UniqSupply       ( UniqSupply, initUs_, returnUs, thenUs, mapUs, getUniqueUs, UniqSM )
  32 import UniqSet
  33 import WwLib
  34 import Outputable
  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
  65 wwTopBinds :: UniqSupply
  66              -> [CoreBind]
  67              -> IO [CoreBind]
  68
  69 wwTopBinds us binds
  70   = do {
  71         beginPass "Worker Wrapper binds";
  72
  73         -- Create worker/wrappers, and mark binders with their
  74         -- "strictness info" [which encodes their worker/wrapper-ness]
  75         let { binds' = workersAndWrappers us binds };
  76
  77         endPass "Worker Wrapper binds" (opt_D_dump_worker_wrapper ||
  78                                         opt_D_verbose_core2core) binds'
  79     }
  80 \end{code}
  81
  82
  83 \begin{code}
  84 workersAndWrappers :: UniqSupply -> [CoreBind] -> [CoreBind]
  85
  86 workersAndWrappers us top_binds
  87   = initUs_ us $
  88     mapUs wwBind top_binds `thenUs` \ top_binds' ->
  89     returnUs (concat top_binds')
  90 \end{code}
  91
  92 %************************************************************************
  93 %*                                                                      *
  94 \subsection[wwBind-wwExpr]{@wwBind@ and @wwExpr@}
  95 %*                                                                      *
  96 %************************************************************************
  97
  98 @wwBind@ works on a binding, trying each \tr{(binder, expr)} pair in
  99 turn.  Non-recursive case first, then recursive...
 100
 101 \begin{code}
 102 wwBind  :: CoreBind
 103         -> UniqSM [CoreBind]    -- returns a WwBinding intermediate form;
 104                                 -- the caller will convert to Expr/Binding,
 105                                 -- as appropriate.
 106
 107 wwBind (NonRec binder rhs)
 108   = wwExpr rhs                                          `thenUs` \ new_rhs ->
 109     tryWW True {- non-recursive -} binder new_rhs       `thenUs` \ new_pairs ->
 110     returnUs [NonRec b e | (b,e) <- new_pairs]
 111       -- Generated bindings must be non-recursive
 112       -- because the original binding was.
 113
 114 ------------------------------
 115
 116 wwBind (Rec pairs)
 117   = mapUs do_one pairs          `thenUs` \ new_pairs ->
 118     returnUs [Rec (concat new_pairs)]
 119   where
 120     do_one (binder, rhs) = wwExpr rhs   `thenUs` \ new_rhs ->
 121                            tryWW False {- recursive -} binder new_rhs
 122 \end{code}
 123
 124 @wwExpr@ basically just walks the tree, looking for appropriate
 125 annotations that can be used. Remember it is @wwBind@ that does the
 126 matching by looking for strict arguments of the correct type.
 127 @wwExpr@ is a version that just returns the ``Plain'' Tree.
 128 ???????????????? ToDo
 129
 130 \begin{code}
 131 wwExpr :: CoreExpr -> UniqSM CoreExpr
 132
 133 wwExpr e@(Type _)   = returnUs e
 134 wwExpr e@(Var _)    = returnUs e
 135
 136 wwExpr e@(Con con args)
 137  = mapUs wwExpr args    `thenUs` \ args' ->
 138    returnUs (Con con args')
 139
 140 wwExpr (Lam binder expr)
 141   = wwExpr expr                 `thenUs` \ new_expr ->
 142     returnUs (Lam binder new_expr)
 143
 144 wwExpr (App f a)
 145   = wwExpr f                    `thenUs` \ new_f ->
 146     wwExpr a                    `thenUs` \ new_a ->
 147     returnUs (App new_f new_a)
 148
 149 wwExpr (Note note expr)
 150   = wwExpr expr                 `thenUs` \ new_expr ->
 151     returnUs (Note note new_expr)
 152
 153 wwExpr (Let bind expr)
 154   = wwBind bind                 `thenUs` \ intermediate_bind ->
 155     wwExpr expr                 `thenUs` \ new_expr ->
 156     returnUs (mkLets intermediate_bind new_expr)
 157
 158 wwExpr (Case expr binder alts)
 159   = wwExpr expr                         `thenUs` \ new_expr ->
 160     mapUs ww_alt alts                   `thenUs` \ new_alts ->
 161     returnUs (Case new_expr binder new_alts)
 162   where
 163     ww_alt (con, binders, rhs)
 164       = wwExpr rhs                      `thenUs` \ new_rhs ->
 165         returnUs (con, binders, new_rhs)
 166 \end{code}
 167
 168 %************************************************************************
 169 %*                                                                      *
 170 \subsection[tryWW]{@tryWW@: attempt a worker/wrapper pair}
 171 %*                                                                      *
 172 %************************************************************************
 173
 174 @tryWW@ just accumulates arguments, converts strictness info from the
 175 front-end into the proper form, then calls @mkWwBodies@ to do
 176 the business.
 177
 178 We have to BE CAREFUL that we don't worker-wrapperize an Id that has
 179 already been w-w'd!  (You can end up with several liked-named Ids
 180 bouncing around at the same time---absolute mischief.)  So the
 181 criterion we use is: if an Id already has an unfolding (for whatever
 182 reason), then we don't w-w it.
 183
 184 The only reason this is monadised is for the unique supply.
 185
 186 \begin{code}
 187 tryWW   :: Bool                         -- True <=> a non-recursive binding
 188         -> Id                           -- The fn binder
 189         -> CoreExpr                     -- The bound rhs; its innards
 190                                         --   are already ww'd
 191         -> UniqSM [(Id, CoreExpr)]      -- either *one* or *two* pairs;
 192                                         -- if one, then no worker (only
 193                                         -- the orig "wrapper" lives on);
 194                                         -- if two, then a worker and a
 195                                         -- wrapper.
 196 tryWW non_rec fn_id rhs
 197   | (non_rec &&         -- Don't split if its non-recursive and small
 198      certainlySmallEnoughToInline (calcUnfoldingGuidance opt_UF_CreationThreshold rhs)
 199         -- No point in worker/wrappering something that is going to be
 200         -- INLINEd wholesale anyway.  If the strictness analyser is run
 201         -- twice, this test also prevents wrappers (which are INLINEd)
 202         -- from being re-done.
 203     )
 204
 205   || arity == 0         -- Don't split if it's not a function
 206   || never_inline fn_id
 207
 208   || not (do_strict_ww || do_cpr_ww || do_coerce_ww)
 209   = returnUs [ (fn_id, rhs) ]
 210
 211   | otherwise           -- Do w/w split
 212   = mkWwBodies fun_ty arity wrap_dmds one_shots cpr_info        `thenUs` \ (work_args, wrap_fn, work_fn) ->
 213     getUniqueUs                                                 `thenUs` \ work_uniq ->
 214     let
 215         work_rhs     = work_fn rhs
 216         work_demands = [getIdDemandInfo v | v <- work_args, isId v]
 217         proto_work_id            = mkWorkerId work_uniq fn_id (coreExprType work_rhs)
 218         work_id | has_strictness = proto_work_id `setIdStrictness` mkStrictnessInfo (work_demands, result_bot)
 219                 | otherwise      = proto_work_id
 220
 221         wrap_rhs = wrap_fn work_id
 222         wrap_id  = fn_id `setIdStrictness`      wrapper_strictness
 223                          `setIdWorkerInfo`      Just work_id
 224                          `setIdArity`           exactArity arity
 225                 -- Add info to the wrapper:
 226                 --      (a) we want to set its arity
 227                 --      (b) we want to pin on its revised strictness info
 228                 --      (c) we pin on its worker id
 229     in
 230     returnUs ([(work_id, work_rhs), (wrap_id, wrap_rhs)])
 231         -- Worker first, because wrapper mentions it
 232   where
 233     fun_ty = idType fn_id
 234     arity  = exprEtaExpandArity rhs
 235
 236         -- Don't split something which is marked unconditionally NOINLINE
 237     never_inline fn_id = case getInlinePragma fn_id of
 238                                 IMustNotBeINLINEd False Nothing -> True
 239                                 other                           -> False
 240
 241     strictness_info                       = getIdStrictness fn_id
 242     StrictnessInfo arg_demands result_bot = strictness_info
 243     has_strictness                        = case strictness_info of
 244                                                 StrictnessInfo _ _ -> True
 245                                                 other              -> False
 246
 247     do_strict_ww = has_strictness && worthSplitting wrap_dmds result_bot
 248
 249         -- NB: There maybe be more items in arg_demands than arity, because
 250         -- the strictness info is semantic and looks through InlineMe and Scc Notes,
 251         -- whereas arity does not
 252     demands_for_visible_args = take arity arg_demands
 253     remaining_arg_demands    = drop arity arg_demands
 254
 255     wrap_dmds | has_strictness = setUnpackStrategy demands_for_visible_args
 256               | otherwise      = take arity (repeat wwLazy)
 257
 258     wrapper_strictness | has_strictness = mkStrictnessInfo (wrap_dmds ++ remaining_arg_demands, result_bot)
 259                        | otherwise      = noStrictnessInfo
 260
 261         -------------------------------------------------------------
 262     cpr_info  = getIdCprInfo fn_id
 263     do_cpr_ww = case cpr_info of
 264                         CPRInfo _ -> True
 265                         other     -> False
 266
 267         -------------------------------------------------------------
 268     do_coerce_ww = check_for_coerce arity fun_ty
 269
 270         -------------------------------------------------------------
 271     one_shots = get_one_shots rhs
 272
 273 -- See if there's a Coerce before we run out of arity;
 274 -- if so, it's worth trying a w/w split.  Reason: we find
 275 -- functions like       f = coerce (\s -> e)
 276 --           and        g = \x -> coerce (\s -> e)
 277 -- and they may have no useful strictness or cpr info, but if we
 278 -- do the w/w thing we get rid of the coerces.
 279
 280 check_for_coerce arity ty
 281   = length arg_tys <= arity && isNewType res_ty
 282         -- Don't look further than arity args,
 283         -- but if there are arity or fewer, see if there's
 284         -- a newtype in the corner
 285   where
 286     (_, tau)          = splitForAllTys ty
 287     (arg_tys, res_ty) = splitFunTys tau
 288
 289 -- If the original function has one-shot arguments, it is important to
 290 -- make the wrapper and worker have corresponding one-shot arguments too.
 291 -- Otherwise we spuriously float stuff out of case-expression join points,
 292 -- which is very annoying.
 293 get_one_shots (Lam b e)
 294   | isId b    = isOneShotLambda b : get_one_shots e
 295   | otherwise = get_one_shots e
 296 get_one_shots (Note _ e) = get_one_shots e
 297 get_one_shots other      = noOneShotInfo
 298 \end{code}
 299
 300
 301
 302 %************************************************************************
 303 %*                                                                      *
 304 \subsection{The worker wrapper core}
 305 %*                                                                      *
 306 %************************************************************************
 307
 308 @mkWrapper@ is called when importing a function.  We have the type of
 309 the function and the name of its worker, and we want to make its body (the wrapper).
 310
 311 \begin{code}
 312 mkWrapper :: Type               -- Wrapper type
 313           -> Int                -- Arity
 314           -> [Demand]           -- Wrapper strictness info
 315           -> CprInfo            -- Wrapper cpr info
 316           -> UniqSM (Id -> CoreExpr)    -- Wrapper body, missing worker Id
 317
 318 mkWrapper fun_ty arity demands cpr_info
 319   = mkWwBodies fun_ty arity demands noOneShotInfo cpr_info      `thenUs` \ (_, wrap_fn, _) ->
 320     returnUs wrap_fn
 321
 322 noOneShotInfo = repeat False
 323 \end{code}
 324
 325