never executed always true always false
1
2
3 -- | Implements a selective lambda lifter, running late in the optimisation
4 -- pipeline.
5 --
6 -- If you are interested in the cost model that is employed to decide whether
7 -- to lift a binding or not, look at "GHC.Stg.Lift.Analysis".
8 -- "GHC.Stg.Lift.Monad" contains the transformation monad that hides away some
9 -- plumbing of the transformation.
10 module GHC.Stg.Lift
11 (
12 -- * Late lambda lifting in STG
13 -- $note
14 stgLiftLams
15 )
16 where
17
18 import GHC.Prelude
19
20 import GHC.Types.Basic
21 import GHC.Driver.Session
22 import GHC.Types.Id
23 import GHC.Stg.FVs ( annBindingFreeVars )
24 import GHC.Stg.Lift.Analysis
25 import GHC.Stg.Lift.Monad
26 import GHC.Stg.Syntax
27 import GHC.Utils.Outputable
28 import GHC.Types.Unique.Supply
29 import GHC.Utils.Panic
30 import GHC.Types.Var.Set
31 import Control.Monad ( when )
32 import Data.Maybe ( isNothing )
33
34 -- Note [Late lambda lifting in STG]
35 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
36 -- $note
37 -- See also the <https://gitlab.haskell.org/ghc/ghc/wikis/late-lam-lift wiki page>
38 -- and #9476.
39 --
40 -- The basic idea behind lambda lifting is to turn locally defined functions
41 -- into top-level functions. Free variables are then passed as additional
42 -- arguments at *call sites* instead of having a closure allocated for them at
43 -- *definition site*. Example:
44 --
45 -- @
46 -- let x = ...; y = ... in
47 -- let f = {x y} \a -> a + x + y in
48 -- let g = {f x} \b -> f b + x in
49 -- g 5
50 -- @
51 --
52 -- Lambda lifting @f@ would
53 --
54 -- 1. Turn @f@'s free variables into formal parameters
55 -- 2. Update @f@'s call site within @g@ to @f x y b@
56 -- 3. Update @g@'s closure: Add @y@ as an additional free variable, while
57 -- removing @f@, because @f@ no longer allocates and can be floated to
58 -- top-level.
59 -- 4. Actually float the binding of @f@ to top-level, eliminating the @let@
60 -- in the process.
61 --
62 -- This results in the following program (with free var annotations):
63 --
64 -- @
65 -- f x y a = a + x + y;
66 -- let x = ...; y = ... in
67 -- let g = {x y} \b -> f x y b + x in
68 -- g 5
69 -- @
70 --
71 -- This optimisation is all about lifting only when it is beneficial to do so.
72 -- The above seems like a worthwhile lift, judging from heap allocation:
73 -- We eliminate @f@'s closure, saving to allocate a closure with 2 words, while
74 -- not changing the size of @g@'s closure.
75 --
76 -- You can probably sense that there's some kind of cost model at play here.
77 -- And you are right! But we also employ a couple of other heuristics for the
78 -- lifting decision which are outlined in "GHC.Stg.Lift.Analysis#when".
79 --
80 -- The transformation is done in "GHC.Stg.Lift", which calls out to
81 -- 'GHC.Stg.Lift.Analysis.goodToLift' for its lifting decision. It relies on
82 -- "GHC.Stg.Lift.Monad", which abstracts some subtle STG invariants into a
83 -- monadic substrate.
84 --
85 -- Suffice to say: We trade heap allocation for stack allocation.
86 -- The additional arguments have to passed on the stack (or in registers,
87 -- depending on architecture) every time we call the function to save a single
88 -- heap allocation when entering the let binding. Nofib suggests a mean
89 -- improvement of about 1% for this pass, so it seems like a worthwhile thing to
90 -- do. Compile-times went up by 0.6%, so all in all a very modest change.
91 --
92 -- For a concrete example, look at @spectral/atom@. There's a call to 'zipWith'
93 -- that is ultimately compiled to something like this
94 -- (module desugaring/lowering to actual STG):
95 --
96 -- @
97 -- propagate dt = ...;
98 -- runExperiment ... =
99 -- let xs = ... in
100 -- let ys = ... in
101 -- let go = {dt go} \xs ys -> case (xs, ys) of
102 -- ([], []) -> []
103 -- (x:xs', y:ys') -> propagate dt x y : go xs' ys'
104 -- in go xs ys
105 -- @
106 --
107 -- This will lambda lift @go@ to top-level, speeding up the resulting program
108 -- by roughly one percent:
109 --
110 -- @
111 -- propagate dt = ...;
112 -- go dt xs ys = case (xs, ys) of
113 -- ([], []) -> []
114 -- (x:xs', y:ys') -> propagate dt x y : go dt xs' ys'
115 -- runExperiment ... =
116 -- let xs = ... in
117 -- let ys = ... in
118 -- in go dt xs ys
119 -- @
120
121
122
123 -- | Lambda lifts bindings to top-level deemed worth lifting (see 'goodToLift').
124 --
125 -- (Mostly) textbook instance of the lambda lifting transformation, selecting
126 -- which bindings to lambda lift by consulting 'goodToLift'.
127 stgLiftLams :: DynFlags -> UniqSupply -> [InStgTopBinding] -> [OutStgTopBinding]
128 stgLiftLams dflags us = runLiftM dflags us . foldr liftTopLvl (pure ())
129
130 liftTopLvl :: InStgTopBinding -> LiftM () -> LiftM ()
131 liftTopLvl (StgTopStringLit bndr lit) rest = withSubstBndr bndr $ \bndr' -> do
132 addTopStringLit bndr' lit
133 rest
134 liftTopLvl (StgTopLifted bind) rest = do
135 let is_rec = isRec $ fst $ decomposeStgBinding bind
136 when is_rec startBindingGroup
137 let bind_w_fvs = annBindingFreeVars bind
138 withLiftedBind TopLevel (tagSkeletonTopBind bind_w_fvs) NilSk $ \mb_bind' -> do
139 -- We signal lifting of a binding through returning Nothing.
140 -- Should never happen for a top-level binding, though, since we are already
141 -- at top-level.
142 case mb_bind' of
143 Nothing -> pprPanic "StgLiftLams" (text "Lifted top-level binding")
144 Just bind' -> addLiftedBinding bind'
145 when is_rec endBindingGroup
146 rest
147
148 withLiftedBind
149 :: TopLevelFlag
150 -> LlStgBinding
151 -> Skeleton
152 -> (Maybe OutStgBinding -> LiftM a)
153 -> LiftM a
154 withLiftedBind top_lvl bind scope k
155 = withLiftedBindPairs top_lvl rec pairs scope (k . fmap (mkStgBinding rec))
156 where
157 (rec, pairs) = decomposeStgBinding bind
158
159 withLiftedBindPairs
160 :: TopLevelFlag
161 -> RecFlag
162 -> [(BinderInfo, LlStgRhs)]
163 -> Skeleton
164 -> (Maybe [(Id, OutStgRhs)] -> LiftM a)
165 -> LiftM a
166 withLiftedBindPairs top rec pairs scope k = do
167 let (infos, rhss) = unzip pairs
168 let bndrs = map binderInfoBndr infos
169 expander <- liftedIdsExpander
170 dflags <- getDynFlags
171 case goodToLift dflags top rec expander pairs scope of
172 -- @abs_ids@ is the set of all variables that need to become parameters.
173 Just abs_ids -> withLiftedBndrs abs_ids bndrs $ \bndrs' -> do
174 -- Within this block, all binders in @bndrs@ will be noted as lifted, so
175 -- that the return value of @liftedIdsExpander@ in this context will also
176 -- expand the bindings in @bndrs@ to their free variables.
177 -- Now we can recurse into the RHSs and see if we can lift any further
178 -- bindings. We pass the set of expanded free variables (thus OutIds) on
179 -- to @liftRhs@ so that it can add them as parameter binders.
180 when (isRec rec) startBindingGroup
181 rhss' <- traverse (liftRhs (Just abs_ids)) rhss
182 let pairs' = zip bndrs' rhss'
183 addLiftedBinding (mkStgBinding rec pairs')
184 when (isRec rec) endBindingGroup
185 k Nothing
186 Nothing -> withSubstBndrs bndrs $ \bndrs' -> do
187 -- Don't lift the current binding, but possibly some bindings in their
188 -- RHSs.
189 rhss' <- traverse (liftRhs Nothing) rhss
190 let pairs' = zip bndrs' rhss'
191 k (Just pairs')
192
193 liftRhs
194 :: Maybe (DIdSet)
195 -- ^ @Just former_fvs@ <=> this RHS was lifted and we have to add @former_fvs@
196 -- as lambda binders, discarding all free vars.
197 -> LlStgRhs
198 -> LiftM OutStgRhs
199 liftRhs mb_former_fvs rhs@(StgRhsCon ccs con mn ts args)
200 = assertPpr (isNothing mb_former_fvs)
201 (text "Should never lift a constructor"
202 $$ pprStgRhs panicStgPprOpts rhs) $
203 StgRhsCon ccs con mn ts <$> traverse liftArgs args
204 liftRhs Nothing (StgRhsClosure _ ccs upd infos body) =
205 -- This RHS wasn't lifted.
206 withSubstBndrs (map binderInfoBndr infos) $ \bndrs' ->
207 StgRhsClosure noExtFieldSilent ccs upd bndrs' <$> liftExpr body
208 liftRhs (Just former_fvs) (StgRhsClosure _ ccs upd infos body) =
209 -- This RHS was lifted. Insert extra binders for @former_fvs@.
210 withSubstBndrs (map binderInfoBndr infos) $ \bndrs' -> do
211 let bndrs'' = dVarSetElems former_fvs ++ bndrs'
212 StgRhsClosure noExtFieldSilent ccs upd bndrs'' <$> liftExpr body
213
214 liftArgs :: InStgArg -> LiftM OutStgArg
215 liftArgs a@(StgLitArg _) = pure a
216 liftArgs (StgVarArg occ) = do
217 assertPprM (not <$> isLifted occ) (text "StgArgs should never be lifted" $$ ppr occ)
218 StgVarArg <$> substOcc occ
219
220 liftExpr :: LlStgExpr -> LiftM OutStgExpr
221 liftExpr (StgLit lit) = pure (StgLit lit)
222 liftExpr (StgTick t e) = StgTick t <$> liftExpr e
223 liftExpr (StgApp f args) = do
224 f' <- substOcc f
225 args' <- traverse liftArgs args
226 fvs' <- formerFreeVars f
227 let top_lvl_args = map StgVarArg fvs' ++ args'
228 pure (StgApp f' top_lvl_args)
229 liftExpr (StgConApp con mn args tys) = StgConApp con mn <$> traverse liftArgs args <*> pure tys
230 liftExpr (StgOpApp op args ty) = StgOpApp op <$> traverse liftArgs args <*> pure ty
231 liftExpr (StgCase scrut info ty alts) = do
232 scrut' <- liftExpr scrut
233 withSubstBndr (binderInfoBndr info) $ \bndr' -> do
234 alts' <- traverse liftAlt alts
235 pure (StgCase scrut' bndr' ty alts')
236 liftExpr (StgLet scope bind body)
237 = withLiftedBind NotTopLevel bind scope $ \mb_bind' -> do
238 body' <- liftExpr body
239 case mb_bind' of
240 Nothing -> pure body' -- withLiftedBindPairs decided to lift it and already added floats
241 Just bind' -> pure (StgLet noExtFieldSilent bind' body')
242 liftExpr (StgLetNoEscape scope bind body)
243 = withLiftedBind NotTopLevel bind scope $ \mb_bind' -> do
244 body' <- liftExpr body
245 case mb_bind' of
246 Nothing -> pprPanic "stgLiftLams" (text "Should never decide to lift LNEs")
247 Just bind' -> pure (StgLetNoEscape noExtFieldSilent bind' body')
248
249 liftAlt :: LlStgAlt -> LiftM OutStgAlt
250 liftAlt (con, infos, rhs) = withSubstBndrs (map binderInfoBndr infos) $ \bndrs' ->
251 (,,) con bndrs' <$> liftExpr rhs