%
% (c) The GRASP/AQUA Project, Glasgow University, 1993-1998
%
\section[WorkWrap]{Worker/wrapper-generating back-end of strictness analyser}
\begin{code}
module WorkWrap ( wwTopBinds, mkWrapper ) where
import CoreSyn
import CoreUnfold ( certainlyWillInline, mkInlineUnfolding, mkWwInlineRule )
import CoreUtils ( exprType, exprIsHNF )
import CoreArity ( exprArity )
import Var
import Id
import Type ( Type )
import IdInfo
import Demand
import UniqSupply
import BasicTypes
import VarEnv ( isEmptyVarEnv )
import Maybes ( orElse )
import WwLib
import Util ( lengthIs, notNull )
import Outputable
import MonadUtils
#include "HsVersions.h"
\end{code}
We take Core bindings whose binders have:
\begin{enumerate}
\item Strictness attached (by the front-end of the strictness
analyser), and / or
\item Constructed Product Result information attached by the CPR
analysis pass.
\end{enumerate}
and we return some ``plain'' bindings which have been
worker/wrapper-ified, meaning:
\begin{enumerate}
\item Functions have been split into workers and wrappers where
appropriate. If a function has both strictness and CPR properties
then only one worker/wrapper doing both transformations is produced;
\item Binders' @IdInfos@ have been updated to reflect the existence of
these workers/wrappers (this is where we get STRICTNESS and CPR pragma
info for exported values).
\end{enumerate}
\begin{code}
wwTopBinds :: UniqSupply -> [CoreBind] -> [CoreBind]
wwTopBinds us top_binds
= initUs_ us $ do
top_binds' <- mapM wwBind top_binds
return (concat top_binds')
\end{code}
%************************************************************************
%* *
\subsection[wwBind-wwExpr]{@wwBind@ and @wwExpr@}
%* *
%************************************************************************
@wwBind@ works on a binding, trying each \tr{(binder, expr)} pair in
turn. Non-recursive case first, then recursive...
\begin{code}
wwBind :: CoreBind
-> UniqSM [CoreBind]
wwBind (NonRec binder rhs) = do
new_rhs <- wwExpr rhs
new_pairs <- tryWW NonRecursive binder new_rhs
return [NonRec b e | (b,e) <- new_pairs]
wwBind (Rec pairs)
= return . Rec <$> concatMapM do_one pairs
where
do_one (binder, rhs) = do new_rhs <- wwExpr rhs
tryWW Recursive binder new_rhs
\end{code}
@wwExpr@ basically just walks the tree, looking for appropriate
annotations that can be used. Remember it is @wwBind@ that does the
matching by looking for strict arguments of the correct type.
@wwExpr@ is a version that just returns the ``Plain'' Tree.
\begin{code}
wwExpr :: CoreExpr -> UniqSM CoreExpr
wwExpr e@(Type {}) = return e
wwExpr e@(Lit {}) = return e
wwExpr e@(Var {}) = return e
wwExpr (Lam binder expr)
= Lam binder <$> wwExpr expr
wwExpr (App f a)
= App <$> wwExpr f <*> wwExpr a
wwExpr (Note note expr)
= Note note <$> wwExpr expr
wwExpr (Cast expr co) = do
new_expr <- wwExpr expr
return (Cast new_expr co)
wwExpr (Let bind expr)
= mkLets <$> wwBind bind <*> wwExpr expr
wwExpr (Case expr binder ty alts) = do
new_expr <- wwExpr expr
new_alts <- mapM ww_alt alts
return (Case new_expr binder ty new_alts)
where
ww_alt (con, binders, rhs) = do
new_rhs <- wwExpr rhs
return (con, binders, new_rhs)
\end{code}
%************************************************************************
%* *
\subsection[tryWW]{@tryWW@: attempt a worker/wrapper pair}
%* *
%************************************************************************
@tryWW@ just accumulates arguments, converts strictness info from the
front-end into the proper form, then calls @mkWwBodies@ to do
the business.
The only reason this is monadised is for the unique supply.
Note [Don't w/w INLINE things]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
It's very important to refrain from w/w-ing an INLINE function (ie one
with an InlineRule) because the wrapper will then overwrite the
InlineRule unfolding.
Furthermore, if the programmer has marked something as INLINE,
we may lose by w/w'ing it.
If the strictness analyser is run twice, this test also prevents
wrappers (which are INLINEd) from being re-done. (You can end up with
several liked-named Ids bouncing around at the same time---absolute
mischief.)
Notice that we refrain from w/w'ing an INLINE function even if it is
in a recursive group. It might not be the loop breaker. (We could
test for loop-breaker-hood, but I'm not sure that ever matters.)
Note [Don't w/w INLINABLE things]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
If we have
{-# INLINABLE f #-}
f x y = ....
then in principle we might get a more efficient loop by w/w'ing f.
But that would make a new unfolding which would overwrite the old
one. So we leave INLINABLE things alone too.
This is a slight infelicity really, because it means that adding
an INLINABLE pragma could make a program a bit less efficient,
because you lose the worker/wrapper stuff. But I don't see a way
to avoid that.
Note [Don't w/w inline small non-loop-breker things]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
In general, we refrain from w/w-ing *small* functions, which are not
loop breakers, because they'll inline anyway. But we must take care:
it may look small now, but get to be big later after other inlining
has happened. So we take the precaution of adding an INLINE pragma to
any such functions.
I made this change when I observed a big function at the end of
compilation with a useful strictness signature but no w-w. When
I measured it on nofib, it didn't make much difference; just a few
percent improved allocation on one benchmark (bspt/Euclid.space).
But nothing got worse.
Note [Wrapper activation]
~~~~~~~~~~~~~~~~~~~~~~~~~
When should the wrapper inlining be active? It must not be active
earlier than the current Activation of the Id (eg it might have a
NOINLINE pragma). But in fact strictness analysis happens fairly
late in the pipeline, and we want to prioritise specialisations over
strictness. Eg if we have
module Foo where
f :: Num a => a -> Int -> a
f n 0 = n -- Strict in the Int, hence wrapper
f n x = f (n+n) (x-1)
g :: Int -> Int
g x = f x x -- Provokes a specialisation for f
module Bsr where
import Foo
h :: Int -> Int
h x = f 3 x
Then we want the specialisation for 'f' to kick in before the wrapper does.
Now in fact the 'gentle' simplification pass encourages this, by
having rules on, but inlinings off. But that's kind of lucky. It seems
more robust to give the wrapper an Activation of (ActiveAfter 0),
so that it becomes active in an importing module at the same time that
it appears in the first place in the defining module.
\begin{code}
tryWW :: RecFlag
-> Id
-> CoreExpr
-> UniqSM [(Id, CoreExpr)]
tryWW is_rec fn_id rhs
| isNeverActive inline_act
= return [ (fn_id, rhs) ]
| is_thunk && worthSplittingThunk maybe_fn_dmd res_info
= ASSERT2( isNonRec is_rec, ppr new_fn_id )
checkSize new_fn_id rhs $
splitThunk new_fn_id rhs
| is_fun && worthSplittingFun wrap_dmds res_info
= checkSize new_fn_id rhs $
splitFun new_fn_id fn_info wrap_dmds res_info rhs
| otherwise
= return [ (new_fn_id, rhs) ]
where
fn_info = idInfo fn_id
maybe_fn_dmd = demandInfo fn_info
inline_act = inlinePragmaActivation (inlinePragInfo fn_info)
strict_sig = strictnessInfo fn_info `orElse` topSig
StrictSig (DmdType env wrap_dmds res_info) = strict_sig
new_fn_id | isEmptyVarEnv env = fn_id
| otherwise = fn_id `setIdStrictness`
StrictSig (mkTopDmdType wrap_dmds res_info)
is_fun = notNull wrap_dmds
is_thunk = not is_fun && not (exprIsHNF rhs)
checkSize :: Id -> CoreExpr
-> UniqSM [(Id,CoreExpr)] -> UniqSM [(Id,CoreExpr)]
checkSize fn_id rhs thing_inside
| isStableUnfolding (realIdUnfolding fn_id)
= return [ (fn_id, rhs) ]
| certainlyWillInline (idUnfolding fn_id)
= return [ (fn_id `setIdUnfolding` inline_rule, rhs) ]
| otherwise = thing_inside
where
inline_rule = mkInlineUnfolding Nothing rhs
splitFun :: Id -> IdInfo -> [Demand] -> DmdResult -> Expr Var
-> UniqSM [(Id, CoreExpr)]
splitFun fn_id fn_info wrap_dmds res_info rhs
= WARN( not (wrap_dmds `lengthIs` arity), ppr fn_id <+> (ppr arity $$ ppr wrap_dmds $$ ppr res_info) )
(do {
(work_demands, wrap_fn, work_fn) <- mkWwBodies fun_ty wrap_dmds res_info one_shots
; work_uniq <- getUniqueM
; let
work_rhs = work_fn rhs
work_id = mkWorkerId work_uniq fn_id (exprType work_rhs)
`setIdOccInfo` occInfo fn_info
`setInlineActivation` (inlinePragmaActivation inl_prag)
`setIdStrictness` StrictSig (mkTopDmdType work_demands work_res_info)
`setIdArity` (exprArity work_rhs)
wrap_rhs = wrap_fn work_id
wrap_prag = InlinePragma { inl_inline = Inline
, inl_sat = Nothing
, inl_act = ActiveAfter 0
, inl_rule = rule_match_info }
wrap_id = fn_id `setIdUnfolding` mkWwInlineRule work_id wrap_rhs arity
`setInlinePragma` wrap_prag
`setIdOccInfo` NoOccInfo
; return ([(work_id, work_rhs), (wrap_id, wrap_rhs)]) })
where
fun_ty = idType fn_id
inl_prag = inlinePragInfo fn_info
rule_match_info = inlinePragmaRuleMatchInfo inl_prag
arity = arityInfo fn_info
work_res_info | isBotRes res_info = BotRes
| otherwise = TopRes
one_shots = get_one_shots rhs
get_one_shots :: Expr Var -> [Bool]
get_one_shots (Lam b e)
| isId b = isOneShotLambda b : get_one_shots e
| otherwise = get_one_shots e
get_one_shots (Note _ e) = get_one_shots e
get_one_shots _ = noOneShotInfo
\end{code}
Note [Thunk splitting]
~~~~~~~~~~~~~~~~~~~~~~
Suppose x is used strictly (never mind whether it has the CPR
property).
let
x* = x-rhs
in body
splitThunk transforms like this:
let
x* = case x-rhs of { I# a -> I# a }
in body
Now simplifier will transform to
case x-rhs of
I# a -> let x* = I# a
in body
which is what we want. Now suppose x-rhs is itself a case:
x-rhs = case e of { T -> I# a; F -> I# b }
The join point will abstract over a, rather than over (which is
what would have happened before) which is fine.
Notice that x certainly has the CPR property now!
In fact, splitThunk uses the function argument w/w splitting
function, so that if x's demand is deeper (say U(U(L,L),L))
then the splitting will go deeper too.
\begin{code}
splitThunk :: Var -> Expr Var -> UniqSM [(Var, Expr Var)]
splitThunk fn_id rhs = do
(_, wrap_fn, work_fn) <- mkWWstr [fn_id]
return [ (fn_id, Let (NonRec fn_id rhs) (wrap_fn (work_fn (Var fn_id)))) ]
\end{code}
%************************************************************************
%* *
\subsection{Functions over Demands}
%* *
%************************************************************************
\begin{code}
worthSplittingFun :: [Demand] -> DmdResult -> Bool
worthSplittingFun ds res
= any worth_it ds || returnsCPR res
where
worth_it Abs = True
worth_it (Eval (Prod _)) = True
worth_it _ = False
worthSplittingThunk :: Maybe Demand
-> DmdResult
-> Bool
worthSplittingThunk maybe_dmd res
= worth_it maybe_dmd || returnsCPR res
where
worth_it (Just (Eval (Prod ds))) = not (all isAbsent ds)
worth_it _ = False
\end{code}
Note [Worker-wrapper for bottoming functions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We used not to split if the result is bottom.
[Justification: there's no efficiency to be gained.]
But it's sometimes bad not to make a wrapper. Consider
fw = \x# -> let x = I# x# in case e of
p1 -> error_fn x
p2 -> error_fn x
p3 -> the real stuff
The re-boxing code won't go away unless error_fn gets a wrapper too.
[We don't do reboxing now, but in general it's better to pass an
unboxed thing to f, and have it reboxed in the error cases....]
%************************************************************************
%* *
\subsection{The worker wrapper core}
%* *
%************************************************************************
@mkWrapper@ is called when importing a function. We have the type of
the function and the name of its worker, and we want to make its body (the wrapper).
\begin{code}
mkWrapper :: Type
-> StrictSig
-> UniqSM (Id -> CoreExpr)
mkWrapper fun_ty (StrictSig (DmdType _ demands res_info)) = do
(_, wrap_fn, _) <- mkWwBodies fun_ty demands res_info noOneShotInfo
return wrap_fn
noOneShotInfo :: [Bool]
noOneShotInfo = repeat False
\end{code}