Commit 0f58d348 authored by Joachim Breitner's avatar Joachim Breitner

Demand Analyzer: Do not set OneShot information (second try)

as suggested in ticket:11770#comment:1. This code was buggy
(#11770), and the occurrence analyzer does the same job anyways.

This also elaborates the notes in the occurrence analyzer accordingly.

Previously, the worker/wrapper code would go through lengths to transfer
the oneShot annotations from the original function to both the worker
and the wrapper. We now simply transfer the demand on the worker, and
let the subsequent occurrence analyzer push this onto the lambda
binders.

This also requires the occurrence analyzer to do this more reliably.
Previously, it would not hand out OneShot annotatoins to things that
would not `certainly_inline` (and it might not have mattered, as the
Demand Analysis might have handed out the annotations). Now we hand out
one-shot annotations unconditionally.

Differential Revision: https://phabricator.haskell.org/D2085
parent 5b986a4d
...@@ -46,9 +46,9 @@ module Demand ( ...@@ -46,9 +46,9 @@ module Demand (
deferAfterIO, deferAfterIO,
postProcessUnsat, postProcessDmdType, postProcessUnsat, postProcessDmdType,
splitProdDmd_maybe, peelCallDmd, mkCallDmd, splitProdDmd_maybe, peelCallDmd, mkCallDmd, mkWorkerDemand,
dmdTransformSig, dmdTransformDataConSig, dmdTransformDictSelSig, dmdTransformSig, dmdTransformDataConSig, dmdTransformDictSelSig,
argOneShots, argsOneShots, argOneShots, argsOneShots, saturatedByOneShots,
trimToType, TypeShape(..), trimToType, TypeShape(..),
useCount, isUsedOnce, reuseEnv, useCount, isUsedOnce, reuseEnv,
...@@ -668,6 +668,12 @@ mkCallDmd :: CleanDemand -> CleanDemand ...@@ -668,6 +668,12 @@ mkCallDmd :: CleanDemand -> CleanDemand
mkCallDmd (JD {sd = d, ud = u}) mkCallDmd (JD {sd = d, ud = u})
= JD { sd = mkSCall d, ud = mkUCall One u } = JD { sd = mkSCall d, ud = mkUCall One u }
-- See Note [Demand on the worker] in WorkWrap
mkWorkerDemand :: Int -> Demand
mkWorkerDemand n = JD { sd = Lazy, ud = Use One (go n) }
where go 0 = Used
go n = mkUCall One $ go (n-1)
cleanEvalDmd :: CleanDemand cleanEvalDmd :: CleanDemand
cleanEvalDmd = JD { sd = HeadStr, ud = Used } cleanEvalDmd = JD { sd = HeadStr, ud = Used }
...@@ -1776,6 +1782,20 @@ argsOneShots (StrictSig (DmdType _ arg_ds _)) n_val_args ...@@ -1776,6 +1782,20 @@ argsOneShots (StrictSig (DmdType _ arg_ds _)) n_val_args
cons [] [] = [] cons [] [] = []
cons a as = a:as cons a as = a:as
-- saturatedByOneShots n C1(C1(...)) = True,
-- <=>
-- there are at least n nested C1(..) calls
-- See Note [Demand on the worker] in WorkWrap
saturatedByOneShots :: Int -> Demand -> Bool
saturatedByOneShots n (JD { ud = usg })
= case usg of
Use _ arg_usg -> go n arg_usg
_ -> False
where
go 0 _ = True
go n (UCall One u) = go (n-1) u
go _ _ = False
argOneShots :: OneShotInfo -- OneShotLam or ProbOneShot, argOneShots :: OneShotInfo -- OneShotLam or ProbOneShot,
-> Demand -- depending on saturation -> Demand -- depending on saturation
-> [OneShotInfo] -> [OneShotInfo]
......
...@@ -1104,12 +1104,13 @@ occAnalNonRecRhs :: OccEnv ...@@ -1104,12 +1104,13 @@ occAnalNonRecRhs :: OccEnv
occAnalNonRecRhs env bndr rhs occAnalNonRecRhs env bndr rhs
= occAnal rhs_env rhs = occAnal rhs_env rhs
where where
-- See Note [Cascading inlines]
env1 | certainly_inline = env
| otherwise = rhsCtxt env
-- See Note [Use one-shot info] -- See Note [Use one-shot info]
env1 = env { occ_one_shots = argOneShots OneShotLam dmd } rhs_env = env1 { occ_one_shots = argOneShots OneShotLam dmd }
-- See Note [Cascading inlines]
rhs_env | certainly_inline = env1
| otherwise = rhsCtxt env1
certainly_inline -- See Note [Cascading inlines] certainly_inline -- See Note [Cascading inlines]
= case idOccInfo bndr of = case idOccInfo bndr of
...@@ -1395,19 +1396,29 @@ markManyIf False uds = uds ...@@ -1395,19 +1396,29 @@ markManyIf False uds = uds
{- {-
Note [Use one-shot information] Note [Use one-shot information]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The occurrrence analyser propagates one-shot-lambda information in two situation The occurrrence analyser propagates one-shot-lambda information in two
* Applications: eg build (\cn -> blah) situations:
* Applications: eg build (\c n -> blah)
Propagate one-shot info from the strictness signature of 'build' to Propagate one-shot info from the strictness signature of 'build' to
the \cn the \c n.
This strictness signature can come from a module interface, in the case of
an imported function, or from a previous run of the demand analyser.
* Let-bindings: eg let f = \c. let ... in \n -> blah * Let-bindings: eg let f = \c. let ... in \n -> blah
in (build f, build f) in (build f, build f)
Propagate one-shot info from the demanand-info on 'f' to the Propagate one-shot info from the demanand-info on 'f' to the
lambdas in its RHS (which may not be syntactically at the top) lambdas in its RHS (which may not be syntactically at the top)
Some of this is done by the demand analyser, but this way it happens This information must have come from a previous run of the demanand
much earlier, taking advantage of the strictness signature of analyser.
imported functions.
Previously, the demand analyser would *also* set the one-shot information, but
that code was buggy (see #11770), so doing it only in on place, namely here, is
saner.
Note [Binders in case alternatives] Note [Binders in case alternatives]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
...@@ -1534,7 +1545,7 @@ oneShotGroup :: OccEnv -> [CoreBndr] ...@@ -1534,7 +1545,7 @@ oneShotGroup :: OccEnv -> [CoreBndr]
-> ( OccEnv -> ( OccEnv
, [CoreBndr] ) , [CoreBndr] )
-- The result binders have one-shot-ness set that they might not have had originally. -- The result binders have one-shot-ness set that they might not have had originally.
-- This happens in (build (\cn -> e)). Here the occurrence analyser -- This happens in (build (\c n -> e)). Here the occurrence analyser
-- linearity context knows that c,n are one-shot, and it records that fact in -- linearity context knows that c,n are one-shot, and it records that fact in
-- the binder. This is useful to guide subsequent float-in/float-out tranformations -- the binder. This is useful to guide subsequent float-in/float-out tranformations
...@@ -1555,8 +1566,13 @@ oneShotGroup env@(OccEnv { occ_one_shots = ctxt }) bndrs ...@@ -1555,8 +1566,13 @@ oneShotGroup env@(OccEnv { occ_one_shots = ctxt }) bndrs
= case ctxt of = case ctxt of
[] -> go [] bndrs (bndr : rev_bndrs) [] -> go [] bndrs (bndr : rev_bndrs)
(one_shot : ctxt) -> go ctxt bndrs (bndr': rev_bndrs) (one_shot : ctxt) -> go ctxt bndrs (bndr': rev_bndrs)
where where
bndr' = updOneShotInfo bndr one_shot bndr' = updOneShotInfo bndr one_shot
-- Use updOneShotInfo, not setOneShotInfo, as pre-existing
-- one-shot info might be better than what we can infer, e.g.
-- due to explicit use of the magic 'oneShot' function.
-- See Note [The oneShot function]
| otherwise | otherwise
= go ctxt bndrs (bndr:rev_bndrs) = go ctxt bndrs (bndr:rev_bndrs)
......
...@@ -1651,7 +1651,7 @@ spec_one env fn arg_bndrs body (call_pat@(qvars, pats), rule_number) ...@@ -1651,7 +1651,7 @@ spec_one env fn arg_bndrs body (call_pat@(qvars, pats), rule_number)
-- Conditionally use result of new worker-wrapper transform -- Conditionally use result of new worker-wrapper transform
(spec_lam_args, spec_call_args) = mkWorkerArgs (sc_dflags env) qvars NoOneShotInfo body_ty (spec_lam_args, spec_call_args) = mkWorkerArgs (sc_dflags env) qvars body_ty
-- Usual w/w hack to avoid generating -- Usual w/w hack to avoid generating
-- a spec_rhs of unlifted type and no args -- a spec_rhs of unlifted type and no args
......
...@@ -200,13 +200,9 @@ dmdAnal' env dmd (Lam var body) ...@@ -200,13 +200,9 @@ dmdAnal' env dmd (Lam var body)
= let (body_dmd, defer_and_use) = peelCallDmd dmd = let (body_dmd, defer_and_use) = peelCallDmd dmd
-- body_dmd: a demand to analyze the body -- body_dmd: a demand to analyze the body
one_shot = useCount (getUseDmd defer_and_use)
-- one_shot: one-shotness of the lambda
-- hence, cardinality of its free vars
env' = extendSigsWithLam env var env' = extendSigsWithLam env var
(body_ty, body') = dmdAnal env' body_dmd body (body_ty, body') = dmdAnal env' body_dmd body
(lam_ty, var') = annotateLamIdBndr env notArgOfDfun body_ty one_shot var (lam_ty, var') = annotateLamIdBndr env notArgOfDfun body_ty var
in in
(postProcessUnsat defer_and_use lam_ty, Lam var' body') (postProcessUnsat defer_and_use lam_ty, Lam var' body')
...@@ -260,17 +256,13 @@ dmdAnal' env dmd (Case scrut case_bndr ty alts) ...@@ -260,17 +256,13 @@ dmdAnal' env dmd (Case scrut case_bndr ty alts)
(res_ty, Case scrut' case_bndr' ty alts') (res_ty, Case scrut' case_bndr' ty alts')
dmdAnal' env dmd (Let (NonRec id rhs) body) dmdAnal' env dmd (Let (NonRec id rhs) body)
= (body_ty2, Let (NonRec id2 annotated_rhs) body') = (body_ty2, Let (NonRec id2 rhs') body')
where where
(sig, lazy_fv, id1, rhs') = dmdAnalRhs NotTopLevel Nothing env id rhs (sig, lazy_fv, id1, rhs') = dmdAnalRhs NotTopLevel Nothing env id rhs
(body_ty, body') = dmdAnal (extendAnalEnv NotTopLevel env id sig) dmd body (body_ty, body') = dmdAnal (extendAnalEnv NotTopLevel env id sig) dmd body
(body_ty1, id2) = annotateBndr env body_ty id1 (body_ty1, id2) = annotateBndr env body_ty id1
body_ty2 = addLazyFVs body_ty1 lazy_fv body_ty2 = addLazyFVs body_ty1 lazy_fv
-- Annotate top-level lambdas at RHS basing on the aggregated demand info
-- See Note [Annotating lambdas at right-hand side]
annotated_rhs = annLamWithShotness (idDemandInfo id2) rhs'
-- If the actual demand is better than the vanilla call -- If the actual demand is better than the vanilla call
-- demand, you might think that we might do better to re-analyse -- demand, you might think that we might do better to re-analyse
-- the RHS with the stronger demand. -- the RHS with the stronger demand.
...@@ -307,25 +299,6 @@ io_hack_reqd scrut con bndrs ...@@ -307,25 +299,6 @@ io_hack_reqd scrut con bndrs
| otherwise | otherwise
= False = False
annLamWithShotness :: Demand -> CoreExpr -> CoreExpr
annLamWithShotness d e
| Just u <- cleanUseDmd_maybe d
= go u e
| otherwise = e
where
go u e
| Just (c, u') <- peelUseCall u
, Lam bndr body <- e
= if isTyVar bndr
then Lam bndr (go u body)
else Lam (setOneShotness c bndr) (go u' body)
| otherwise
= e
setOneShotness :: Count -> Id -> Id
setOneShotness One bndr = setOneShotLambda bndr
setOneShotness Many bndr = bndr
dmdAnalAlt :: AnalEnv -> CleanDemand -> Id -> Alt Var -> (DmdType, Alt Var) dmdAnalAlt :: AnalEnv -> CleanDemand -> Id -> Alt Var -> (DmdType, Alt Var)
dmdAnalAlt env dmd case_bndr (con,bndrs,rhs) dmdAnalAlt env dmd case_bndr (con,bndrs,rhs)
| null bndrs -- Literals, DEFAULT, and nullary constructors | null bndrs -- Literals, DEFAULT, and nullary constructors
...@@ -432,23 +405,6 @@ free variable |y|. Conversely, if the demand on |h| is unleashed right ...@@ -432,23 +405,6 @@ free variable |y|. Conversely, if the demand on |h| is unleashed right
on the spot, we will get the desired result, namely, that |f| is on the spot, we will get the desired result, namely, that |f| is
strict in |y|. strict in |y|.
Note [Annotating lambdas at right-hand side]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Let us take a look at the following example:
g f = let x = 100
h = \y -> f x y
in h 5
One can see that |h| is called just once, therefore the RHS of h can
be annotated as a one-shot lambda. This is done by the function
annLamWithShotness *a posteriori*, i.e., basing on the aggregated
usage demand on |h| from the body of |let|-expression, which is C1(U)
in this case.
In other words, for locally-bound lambdas we can infer
one-shotness.
************************************************************************ ************************************************************************
* * * *
...@@ -749,23 +705,22 @@ annotateLamBndrs :: AnalEnv -> DFunFlag -> DmdType -> [Var] -> (DmdType, [Var]) ...@@ -749,23 +705,22 @@ annotateLamBndrs :: AnalEnv -> DFunFlag -> DmdType -> [Var] -> (DmdType, [Var])
annotateLamBndrs env args_of_dfun ty bndrs = mapAccumR annotate ty bndrs annotateLamBndrs env args_of_dfun ty bndrs = mapAccumR annotate ty bndrs
where where
annotate dmd_ty bndr annotate dmd_ty bndr
| isId bndr = annotateLamIdBndr env args_of_dfun dmd_ty Many bndr | isId bndr = annotateLamIdBndr env args_of_dfun dmd_ty bndr
| otherwise = (dmd_ty, bndr) | otherwise = (dmd_ty, bndr)
annotateLamIdBndr :: AnalEnv annotateLamIdBndr :: AnalEnv
-> DFunFlag -- is this lambda at the top of the RHS of a dfun? -> DFunFlag -- is this lambda at the top of the RHS of a dfun?
-> DmdType -- Demand type of body -> DmdType -- Demand type of body
-> Count -- One-shot-ness of the lambda
-> Id -- Lambda binder -> Id -- Lambda binder
-> (DmdType, -- Demand type of lambda -> (DmdType, -- Demand type of lambda
Id) -- and binder annotated with demand Id) -- and binder annotated with demand
annotateLamIdBndr env arg_of_dfun dmd_ty one_shot id annotateLamIdBndr env arg_of_dfun dmd_ty id
-- For lambdas we add the demand to the argument demands -- For lambdas we add the demand to the argument demands
-- Only called for Ids -- Only called for Ids
= ASSERT( isId id ) = ASSERT( isId id )
-- pprTrace "annLamBndr" (vcat [ppr id, ppr _dmd_ty]) $ -- pprTrace "annLamBndr" (vcat [ppr id, ppr _dmd_ty]) $
(final_ty, setOneShotness one_shot (setIdDemandInfo id dmd)) (final_ty, setIdDemandInfo id dmd)
where where
-- Watch out! See note [Lambda-bound unfoldings] -- Watch out! See note [Lambda-bound unfoldings]
final_ty = case maybeUnfoldingTemplate (idUnfolding id) of final_ty = case maybeUnfoldingTemplate (idUnfolding id) of
......
...@@ -10,7 +10,6 @@ module WorkWrap ( wwTopBinds ) where ...@@ -10,7 +10,6 @@ module WorkWrap ( wwTopBinds ) where
import CoreSyn import CoreSyn
import CoreUnfold ( certainlyWillInline, mkWwInlineRule, mkWorkerUnfolding ) import CoreUnfold ( certainlyWillInline, mkWwInlineRule, mkWorkerUnfolding )
import CoreUtils ( exprType, exprIsHNF ) import CoreUtils ( exprType, exprIsHNF )
import CoreArity ( exprArity )
import Var import Var
import Id import Id
import IdInfo import IdInfo
...@@ -330,7 +329,7 @@ splitFun :: DynFlags -> FamInstEnvs -> Id -> IdInfo -> [Demand] -> DmdResult -> ...@@ -330,7 +329,7 @@ splitFun :: DynFlags -> FamInstEnvs -> Id -> IdInfo -> [Demand] -> DmdResult ->
splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs splitFun dflags fam_envs 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 = WARN( not (wrap_dmds `lengthIs` arity), ppr fn_id <+> (ppr arity $$ ppr wrap_dmds $$ ppr res_info) ) do
-- The arity should match the signature -- The arity should match the signature
stuff <- mkWwBodies dflags fam_envs fun_ty wrap_dmds res_info one_shots stuff <- mkWwBodies dflags fam_envs fun_ty wrap_dmds res_info
case stuff of case stuff of
Just (work_demands, wrap_fn, work_fn) -> do Just (work_demands, wrap_fn, work_fn) -> do
work_uniq <- getUniqueM work_uniq <- getUniqueM
...@@ -360,8 +359,18 @@ splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs ...@@ -360,8 +359,18 @@ splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs
-- Even though we may not be at top level, -- Even though we may not be at top level,
-- it's ok to give it an empty DmdEnv -- it's ok to give it an empty DmdEnv
`setIdArity` exprArity work_rhs `setIdDemandInfo` worker_demand
`setIdArity` work_arity
-- Set the arity so that the Core Lint check that the -- Set the arity so that the Core Lint check that the
work_arity = length work_demands
-- See Note [Demand on the Worker]
single_call = saturatedByOneShots arity (demandInfo fn_info)
worker_demand | single_call = mkWorkerDemand work_arity
| otherwise = topDmd
-- arity is consistent with the demand type goes through -- arity is consistent with the demand type goes through
wrap_act = ActiveAfter "0" 0 wrap_act = ActiveAfter "0" 0
...@@ -380,6 +389,8 @@ splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs ...@@ -380,6 +389,8 @@ splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs
-- Zap any loop-breaker-ness, to avoid bleating from Lint -- Zap any loop-breaker-ness, to avoid bleating from Lint
-- about a loop breaker with an INLINE rule -- about a loop breaker with an INLINE rule
return $ [(work_id, work_rhs), (wrap_id, wrap_rhs)] return $ [(work_id, work_rhs), (wrap_id, wrap_rhs)]
-- Worker first, because wrapper mentions it -- Worker first, because wrapper mentions it
...@@ -396,20 +407,39 @@ splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs ...@@ -396,20 +407,39 @@ splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs
Just _ -> topRes -- Cpr stuff done by wrapper; kill it here Just _ -> topRes -- Cpr stuff done by wrapper; kill it here
Nothing -> res_info -- Preserve exception/divergence Nothing -> res_info -- Preserve exception/divergence
one_shots = get_one_shots rhs
-- If the original function has one-shot arguments, it is important to
-- make the wrapper and worker have corresponding one-shot arguments too.
-- Otherwise we spuriously float stuff out of case-expression join points,
-- which is very annoying.
get_one_shots :: Expr Var -> [OneShotInfo]
get_one_shots (Lam b e)
| isId b = idOneShotInfo b : get_one_shots e
| otherwise = get_one_shots e
get_one_shots (Tick _ e) = get_one_shots e
get_one_shots _ = []
{- {-
Note [Demand on the worker]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
If the original function is called once, according to its demand info, then
so is the worker. This is important so that the occurrence analyser can
attach OneShot annotations to the worker’s lambda binders.
Example:
-- Original function
f [Demand=<L,1*C1(U)>] :: (a,a) -> a
f = \p -> ...
-- Wrapper
f [Demand=<L,1*C1(U)>] :: a -> a -> a
f = \p -> case p of (a,b) -> $wf a b
-- Worker
$wf [Demand=<L,1*C1(C1(U))>] :: Int -> Int
$wf = \a b -> ...
We need to check whether the original function is called once, with
sufficiently many arguments. This is done using saturatedByOneShots, which
takes the arity of the original function (resp. the wrapper) and the demand on
the original function.
The demand on the worker is then calculated using mkWorkerDemand, and always of
the form [Demand=<L,1*(C1(...(C1(U))))>]
Note [Do not split void functions] Note [Do not split void functions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider this rather common form of binding: Consider this rather common form of binding:
......
...@@ -25,7 +25,7 @@ import TysWiredIn ( tupleDataCon ) ...@@ -25,7 +25,7 @@ import TysWiredIn ( tupleDataCon )
import Type import Type
import Coercion import Coercion
import FamInstEnv import FamInstEnv
import BasicTypes ( Boxity(..), OneShotInfo(..), worstOneShot ) import BasicTypes ( Boxity(..) )
import Literal ( absentLiteralOf ) import Literal ( absentLiteralOf )
import TyCon import TyCon
import UniqSupply import UniqSupply
...@@ -111,7 +111,6 @@ mkWwBodies :: DynFlags ...@@ -111,7 +111,6 @@ mkWwBodies :: DynFlags
-> Type -- Type of original function -> Type -- Type of original function
-> [Demand] -- Strictness of original function -> [Demand] -- Strictness of original function
-> DmdResult -- Info about function result -> DmdResult -- Info about function result
-> [OneShotInfo] -- One-shot-ness of the function, value args only
-> UniqSM (Maybe ([Demand], -- Demands for worker (value) args -> UniqSM (Maybe ([Demand], -- Demands for worker (value) args
Id -> CoreExpr, -- Wrapper body, lacking only the worker Id Id -> CoreExpr, -- Wrapper body, lacking only the worker Id
CoreExpr -> CoreExpr)) -- Worker body, lacking the original function rhs CoreExpr -> CoreExpr)) -- Worker body, lacking the original function rhs
...@@ -127,22 +126,20 @@ mkWwBodies :: DynFlags ...@@ -127,22 +126,20 @@ mkWwBodies :: DynFlags
-- let x = (a,b) in -- let x = (a,b) in
-- E -- E
mkWwBodies dflags fam_envs fun_ty demands res_info one_shots mkWwBodies dflags fam_envs fun_ty demands res_info
= do { let arg_info = demands `zip` (one_shots ++ repeat NoOneShotInfo) = do { (wrap_args, wrap_fn_args, work_fn_args, res_ty) <- mkWWargs emptyTCvSubst fun_ty demands
all_one_shots = foldr (worstOneShot . snd) OneShotLam arg_info
; (wrap_args, wrap_fn_args, work_fn_args, res_ty) <- mkWWargs emptyTCvSubst fun_ty arg_info
; (useful1, work_args, wrap_fn_str, work_fn_str) <- mkWWstr dflags fam_envs wrap_args ; (useful1, work_args, wrap_fn_str, work_fn_str) <- mkWWstr dflags fam_envs wrap_args
-- Do CPR w/w. See Note [Always do CPR w/w] -- Do CPR w/w. See Note [Always do CPR w/w]
; (useful2, wrap_fn_cpr, work_fn_cpr, cpr_res_ty) ; (useful2, wrap_fn_cpr, work_fn_cpr, cpr_res_ty)
<- mkWWcpr (gopt Opt_CprAnal dflags) fam_envs res_ty res_info <- mkWWcpr (gopt Opt_CprAnal dflags) fam_envs res_ty res_info
; let (work_lam_args, work_call_args) = mkWorkerArgs dflags work_args all_one_shots cpr_res_ty ; let (work_lam_args, work_call_args) = mkWorkerArgs dflags work_args cpr_res_ty
worker_args_dmds = [idDemandInfo v | v <- work_call_args, isId v] worker_args_dmds = [idDemandInfo v | v <- work_call_args, isId v]
wrapper_body = wrap_fn_args . wrap_fn_cpr . wrap_fn_str . applyToVars work_call_args . Var wrapper_body = wrap_fn_args . wrap_fn_cpr . wrap_fn_str . applyToVars work_call_args . Var
worker_body = mkLams work_lam_args. work_fn_str . work_fn_cpr . work_fn_args worker_body = mkLams work_lam_args. work_fn_str . work_fn_cpr . work_fn_args
; if useful1 && not (only_one_void_argument) || useful2 ; if useful1 && not only_one_void_argument || useful2
then return (Just (worker_args_dmds, wrapper_body, worker_body)) then return (Just (worker_args_dmds, wrapper_body, worker_body))
else return Nothing else return Nothing
} }
...@@ -196,24 +193,20 @@ We use the state-token type which generates no code. ...@@ -196,24 +193,20 @@ We use the state-token type which generates no code.
-} -}
mkWorkerArgs :: DynFlags -> [Var] mkWorkerArgs :: DynFlags -> [Var]
-> OneShotInfo -- Whether all arguments are one-shot
-> Type -- Type of body -> Type -- Type of body
-> ([Var], -- Lambda bound args -> ([Var], -- Lambda bound args
[Var]) -- Args at call site [Var]) -- Args at call site
mkWorkerArgs dflags args all_one_shot res_ty mkWorkerArgs dflags args res_ty
| any isId args || not needsAValueLambda | any isId args || not needsAValueLambda
= (args, args) = (args, args)
| otherwise | otherwise
= (args ++ [newArg], args ++ [voidPrimId]) = (args ++ [voidArgId], args ++ [voidPrimId])
where where
needsAValueLambda = needsAValueLambda =
isUnliftedType res_ty isUnliftedType res_ty
|| not (gopt Opt_FunToThunk dflags) || not (gopt Opt_FunToThunk dflags)
-- see Note [Protecting the last value argument] -- see Note [Protecting the last value argument]
-- see Note [All One-Shot Arguments of a Worker]
newArg = setIdOneShotInfo voidArgId all_one_shot
{- {-
Note [Protecting the last value argument] Note [Protecting the last value argument]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
...@@ -228,29 +221,6 @@ create a space leak. 2) It can prevent inlining *under a lambda*. If w/w ...@@ -228,29 +221,6 @@ create a space leak. 2) It can prevent inlining *under a lambda*. If w/w
removes the last argument from a function f, then f now looks like a thunk, and removes the last argument from a function f, then f now looks like a thunk, and
so f can't be inlined *under a lambda*. so f can't be inlined *under a lambda*.
Note [All One-Shot Arguments of a Worker]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Sometimes, derived join-points are just lambda-lifted thunks, whose
only argument is of the unit type and is never used. This might
interfere with the absence analysis, basing on which results these
never-used arguments are eliminated in the worker. The additional
argument `all_one_shot` of `mkWorkerArgs` is to prevent this.
Example. Suppose we have
foo = \p(one-shot) q(one-shot). y + 3
Then we drop the unused args to give
foo = \pq. $wfoo void#
$wfoo = \void(one-shot). y + 3
But suppse foo didn't have all one-shot args:
foo = \p(not-one-shot) q(one-shot). expensive y + 3
Then we drop the unused args to give
foo = \pq. $wfoo void#
$wfoo = \void(not-one-shot). y + 3
If we made the void-arg one-shot we might inline an expensive
computation for y, which would be terrible!
************************************************************************ ************************************************************************
* * * *
...@@ -292,23 +262,23 @@ the \x to get what we want. ...@@ -292,23 +262,23 @@ the \x to get what we want.
mkWWargs :: TCvSubst -- Freshening substitution to apply to the type mkWWargs :: TCvSubst -- Freshening substitution to apply to the type
-- See Note [Freshen type variables] -- See Note [Freshen type variables]
-> Type -- The type of the function -> Type -- The type of the function
-> [(Demand,OneShotInfo)] -- Demands and one-shot info for value arguments -> [Demand] -- Demands and one-shot info for value arguments
-> UniqSM ([Var], -- Wrapper args -> UniqSM ([Var], -- Wrapper args
CoreExpr -> CoreExpr, -- Wrapper fn CoreExpr -> CoreExpr, -- Wrapper fn
CoreExpr -> CoreExpr, -- Worker fn CoreExpr -> CoreExpr, -- Worker fn
Type) -- Type of wrapper body Type) -- Type of wrapper body
mkWWargs subst fun_ty arg_info mkWWargs subst fun_ty demands
| null arg_info | null demands
= return ([], id, id, substTy subst fun_ty) = return ([], id, id, substTy subst fun_ty)
| ((dmd,one_shot):arg_info') <- arg_info | (dmd:demands') <- demands
, Just (arg_ty, fun_ty') <- splitFunTy_maybe fun_ty , Just (arg_ty, fun_ty') <- splitFunTy_maybe fun_ty
= do { uniq <- getUniqueM = do { uniq <- getUniqueM
; let arg_ty' = substTy subst arg_ty ; let arg_ty' = substTy subst arg_ty
id = mk_wrap_arg uniq arg_ty' dmd one_shot id = mk_wrap_arg uniq arg_ty' dmd
; (wrap_args, wrap_fn_args, work_fn_args, res_ty) ; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
<- mkWWargs subst fun_ty' arg_info' <- mkWWargs subst fun_ty' demands'
; return (id : wrap_args, ; return (id : wrap_args,
Lam id . wrap_fn_args, Lam id . wrap_fn_args,
work_fn_args . (`App` varToCoreExpr id), work_fn_args . (`App` varToCoreExpr id),
...@@ -319,7 +289,7 @@ mkWWargs subst fun_ty arg_info ...@@ -319,7 +289,7 @@ mkWWargs subst fun_ty arg_info
-- This substTyVarBndr clones the type variable when necy -- This substTyVarBndr clones the type variable when necy
-- See Note [Freshen type variables] -- See Note [Freshen type variables]
; (wrap_args, wrap_fn_args, work_fn_args, res_ty) ; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
<- mkWWargs subst' fun_ty' arg_info