2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
5 %************************************************************************
7 Static Argument Transformation pass
9 %************************************************************************
11 May be seen as removing invariants from loops:
12 Arguments of recursive functions that do not change in recursive
13 calls are removed from the recursion, which is done locally
14 and only passes the arguments which effectively change.
17 map = /\ ab -> \f -> \xs -> case xs of
19 (a:b) -> f a : map f b
21 as map is recursively called with the same argument f (unmodified)
24 map = /\ ab -> \f -> \xs -> let map' ys = case ys of
29 Notice that for a compiler that uses lambda lifting this is
30 useless as map' will be transformed back to what map was.
32 We could possibly do the same for big lambdas, but we don't as
33 they will eventually be removed in later stages of the compiler,
34 therefore there is no penalty in keeping them.
36 We only apply the SAT when the number of static args is > 2. This
37 produces few bad cases. See
41 Here are the headline nofib results:
43 Min +0.0% -13.7% -21.4%
45 Geometric Mean +0.0% -0.2% -6.9%
47 The previous patch, to fix polymorphic floatout demand signatures, is
48 essential to make this work well!
53 module SAT ( doStaticArgs ) where
75 #include "HsVersions.h"
79 doStaticArgs :: UniqSupply -> [CoreBind] -> [CoreBind]
80 doStaticArgs us binds = snd $ mapAccumL sat_bind_threaded_us us binds
82 sat_bind_threaded_us us bind =
83 let (us1, us2) = splitUniqSupply us
84 in (us1, fst $ runSAT us2 (satBind bind emptyUniqSet))
87 -- We don't bother to SAT recursive groups since it can lead
88 -- to massive code expansion: see Andre Santos' thesis for details.
89 -- This means we only apply the actual SAT to Rec groups of one element,
90 -- but we want to recurse into the others anyway to discover other binds
91 satBind :: CoreBind -> IdSet -> SatM (CoreBind, IdSATInfo)
92 satBind (NonRec binder expr) interesting_ids = do
93 (expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
94 return (NonRec binder expr', finalizeApp expr_app sat_info_expr)
95 satBind (Rec [(binder, rhs)]) interesting_ids = do
96 let interesting_ids' = interesting_ids `addOneToUniqSet` binder
97 (rhs_binders, rhs_body) = collectBinders rhs
98 (rhs_body', sat_info_rhs_body) <- satTopLevelExpr rhs_body interesting_ids'
99 let sat_info_rhs_from_args = unitVarEnv binder (bindersToSATInfo rhs_binders)
100 sat_info_rhs' = mergeIdSATInfo sat_info_rhs_from_args sat_info_rhs_body
102 shadowing = binder `elementOfUniqSet` interesting_ids
103 sat_info_rhs'' = if shadowing
104 then sat_info_rhs' `delFromUFM` binder -- For safety
107 bind' <- saTransformMaybe binder (lookupUFM sat_info_rhs' binder)
108 rhs_binders rhs_body'
109 return (bind', sat_info_rhs'')
110 satBind (Rec pairs) interesting_ids = do
111 let (binders, rhss) = unzip pairs
112 rhss_SATed <- mapM (\e -> satTopLevelExpr e interesting_ids) rhss
113 let (rhss', sat_info_rhss') = unzip rhss_SATed
114 return (Rec (zipEqual "satBind" binders rhss'), mergeIdSATInfos sat_info_rhss')
117 data App = VarApp Id | TypeApp Type
118 data Staticness a = Static a | NotStatic
120 type IdAppInfo = (Id, SATInfo)
122 type SATInfo = [Staticness App]
123 type IdSATInfo = IdEnv SATInfo
124 emptyIdSATInfo :: IdSATInfo
125 emptyIdSATInfo = emptyUFM
128 pprIdSATInfo id_sat_info = vcat (map pprIdAndSATInfo (fmToList id_sat_info))
129 where pprIdAndSATInfo (v, sat_info) = hang (ppr v <> colon) 4 (pprSATInfo sat_info)
132 pprSATInfo :: SATInfo -> SDoc
133 pprSATInfo staticness = hcat $ map pprStaticness staticness
135 pprStaticness :: Staticness App -> SDoc
136 pprStaticness (Static (VarApp _)) = ptext (sLit "SV")
137 pprStaticness (Static (TypeApp _)) = ptext (sLit "ST")
138 pprStaticness NotStatic = ptext (sLit "NS")
141 mergeSATInfo :: SATInfo -> SATInfo -> SATInfo
142 mergeSATInfo [] _ = []
143 mergeSATInfo _ [] = []
144 mergeSATInfo (NotStatic:statics) (_:apps) = NotStatic : mergeSATInfo statics apps
145 mergeSATInfo (_:statics) (NotStatic:apps) = NotStatic : mergeSATInfo statics apps
146 mergeSATInfo ((Static (VarApp v)):statics) ((Static (VarApp v')):apps) = (if v == v' then Static (VarApp v) else NotStatic) : mergeSATInfo statics apps
147 mergeSATInfo ((Static (TypeApp t)):statics) ((Static (TypeApp t')):apps) = (if t `coreEqType` t' then Static (TypeApp t) else NotStatic) : mergeSATInfo statics apps
148 mergeSATInfo l r = pprPanic "mergeSATInfo" $ ptext (sLit "Left:") <> pprSATInfo l <> ptext (sLit ", ")
149 <> ptext (sLit "Right:") <> pprSATInfo r
151 mergeIdSATInfo :: IdSATInfo -> IdSATInfo -> IdSATInfo
152 mergeIdSATInfo = plusUFM_C mergeSATInfo
154 mergeIdSATInfos :: [IdSATInfo] -> IdSATInfo
155 mergeIdSATInfos = foldl' mergeIdSATInfo emptyIdSATInfo
157 bindersToSATInfo :: [Id] -> SATInfo
158 bindersToSATInfo vs = map (Static . binderToApp) vs
159 where binderToApp v = if isId v
161 else TypeApp $ mkTyVarTy v
163 finalizeApp :: Maybe IdAppInfo -> IdSATInfo -> IdSATInfo
164 finalizeApp Nothing id_sat_info = id_sat_info
165 finalizeApp (Just (v, sat_info')) id_sat_info =
166 let sat_info'' = case lookupUFM id_sat_info v of
168 Just sat_info -> mergeSATInfo sat_info sat_info'
169 in extendVarEnv id_sat_info v sat_info''
172 satTopLevelExpr :: CoreExpr -> IdSet -> SatM (CoreExpr, IdSATInfo)
173 satTopLevelExpr expr interesting_ids = do
174 (expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
175 return (expr', finalizeApp expr_app sat_info_expr)
177 satExpr :: CoreExpr -> IdSet -> SatM (CoreExpr, IdSATInfo, Maybe IdAppInfo)
178 satExpr var@(Var v) interesting_ids = do
179 let app_info = if v `elementOfUniqSet` interesting_ids
182 return (var, emptyIdSATInfo, app_info)
184 satExpr lit@(Lit _) _ = do
185 return (lit, emptyIdSATInfo, Nothing)
187 satExpr (Lam binders body) interesting_ids = do
188 (body', sat_info, this_app) <- satExpr body interesting_ids
189 return (Lam binders body', finalizeApp this_app sat_info, Nothing)
191 satExpr (App fn arg) interesting_ids = do
192 (fn', sat_info_fn, fn_app) <- satExpr fn interesting_ids
193 let satRemainder = boring fn' sat_info_fn
195 Nothing -> satRemainder Nothing
196 Just (fn_id, fn_app_info) ->
197 -- TODO: remove this use of append somehow (use a data structure with O(1) append but a left-to-right kind of interface)
198 let satRemainderWithStaticness arg_staticness = satRemainder $ Just (fn_id, fn_app_info ++ [arg_staticness])
200 Type t -> satRemainderWithStaticness $ Static (TypeApp t)
201 Var v -> satRemainderWithStaticness $ Static (VarApp v)
202 _ -> satRemainderWithStaticness $ NotStatic
204 boring :: CoreExpr -> IdSATInfo -> Maybe IdAppInfo -> SatM (CoreExpr, IdSATInfo, Maybe IdAppInfo)
205 boring fn' sat_info_fn app_info =
206 do (arg', sat_info_arg, arg_app) <- satExpr arg interesting_ids
207 let sat_info_arg' = finalizeApp arg_app sat_info_arg
208 sat_info = mergeIdSATInfo sat_info_fn sat_info_arg'
209 return (App fn' arg', sat_info, app_info)
211 satExpr (Case expr bndr ty alts) interesting_ids = do
212 (expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
213 let sat_info_expr' = finalizeApp expr_app sat_info_expr
215 zipped_alts' <- mapM satAlt alts
216 let (alts', sat_infos_alts) = unzip zipped_alts'
217 return (Case expr' bndr ty alts', mergeIdSATInfo sat_info_expr' (mergeIdSATInfos sat_infos_alts), Nothing)
219 satAlt (con, bndrs, expr) = do
220 (expr', sat_info_expr) <- satTopLevelExpr expr interesting_ids
221 return ((con, bndrs, expr'), sat_info_expr)
223 satExpr (Let bind body) interesting_ids = do
224 (body', sat_info_body, body_app) <- satExpr body interesting_ids
225 (bind', sat_info_bind) <- satBind bind interesting_ids
226 return (Let bind' body', mergeIdSATInfo sat_info_body sat_info_bind, body_app)
228 satExpr (Note note expr) interesting_ids = do
229 (expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
230 return (Note note expr', sat_info_expr, expr_app)
232 satExpr ty@(Type _) _ = do
233 return (ty, emptyIdSATInfo, Nothing)
235 satExpr (Cast expr coercion) interesting_ids = do
236 (expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
237 return (Cast expr' coercion, sat_info_expr, expr_app)
240 %************************************************************************
242 Static Argument Transformation Monad
244 %************************************************************************
247 type SatM result = UniqSM result
249 runSAT :: UniqSupply -> SatM a -> a
252 newUnique :: SatM Unique
253 newUnique = getUniqueUs
257 %************************************************************************
259 Static Argument Transformation Monad
261 %************************************************************************
263 To do the transformation, the game plan is to:
265 1. Create a small nonrecursive RHS that takes the
266 original arguments to the function but discards
267 the ones that are static and makes a call to the
268 SATed version with the remainder. We intend that
269 this will be inlined later, removing the overhead
271 2. Bind this nonrecursive RHS over the original body
272 WITH THE SAME UNIQUE as the original body so that
273 any recursive calls to the original now go via
276 3. Rebind the original function to a new one which contains
277 our SATed function and just makes a call to it:
278 we call the thing making this call the local body
280 Example: transform this
282 map :: forall a b. (a->b) -> [a] -> [b]
283 map = /\ab. \(f:a->b) (as:[a]) -> body[map]
285 map :: forall a b. (a->b) -> [a] -> [b]
286 map = /\ab. \(f:a->b) (as:[a]) ->
287 letrec map' :: [a] -> [b]
288 -- The "worker function
290 let map :: forall a' b'. (a -> b) -> [a] -> [b]
291 -- The "shadow function
292 map = /\a'b'. \(f':(a->b) (as:[a]).
297 Note [Shadow binding]
298 ~~~~~~~~~~~~~~~~~~~~~
299 The calls to the inner map inside body[map] should get inlined
300 by the local re-binding of 'map'. We call this the "shadow binding".
302 But we can't use the original binder 'map' unchanged, because
303 it might be exported, in which case the shadow binding won't be
304 discarded as dead code after it is inlined.
306 So we use a hack: we make a new SysLocal binder with the *same* unique
307 as binder. (Another alternative would be to reset the export flag.)
309 Note [Binder type capture]
310 ~~~~~~~~~~~~~~~~~~~~~~~~~~
311 Notice that in the inner map (the "shadow function"), the static arguments
312 are discarded -- it's as if they were underscores. Instead, mentions
313 of these arguments (notably in the types of dynamic arguments) are bound
314 by the *outer* lambdas of the main function. So we must make up fresh
315 names for the static arguments so that they do not capture variables
316 mentioned in the types of dynamic args.
318 In the map example, the shadow function must clone the static type
319 argument a,b, giving a',b', to ensure that in the \(as:[a]), the 'a'
320 is bound by the outer forall. We clone f' too for consistency, but
321 that doesn't matter either way because static Id arguments aren't
322 mentioned in the shadow binding at all.
324 If we don't we get something like this:
330 (p_a6T :: a_aiK -> GHC.Bool.Bool)
331 (f_a6V :: a_aiK -> a_aiK)
334 sat_worker_s1aU :: a_aiK -> a_aiK
337 \ (x_a6X :: a_aiK) ->
339 sat_shadow_r17 :: forall a_a3O.
340 (a_a3O -> GHC.Bool.Bool) -> (a_a3O -> a_a3O) -> a_a3O -> a_a3O
344 (p_a6T :: a_aiK -> GHC.Bool.Bool)
345 (f_a6V :: a_aiK -> a_aiK)
347 sat_worker_s1aU x_a6X } in
348 case p_a6T x_a6X of wild_X3y [ALWAYS Dead Nothing] {
349 GHC.Bool.False -> GHC.Base.until @ a_aiK p_a6T f_a6V (f_a6V x_a6X);
350 GHC.Bool.True -> x_a6X
352 sat_worker_s1aU x_a6X
354 Where sat_shadow has captured the type variables of x_a6X etc as it has a a_aiK
355 type argument. This is bad because it means the application sat_worker_s1aU x_a6X
359 saTransformMaybe :: Id -> Maybe SATInfo -> [Id] -> CoreExpr -> SatM CoreBind
360 saTransformMaybe binder maybe_arg_staticness rhs_binders rhs_body
361 | Just arg_staticness <- maybe_arg_staticness
362 , should_transform arg_staticness
363 = saTransform binder arg_staticness rhs_binders rhs_body
365 = return (Rec [(binder, mkLams rhs_binders rhs_body)])
367 should_transform staticness = n_static_args > 1 -- THIS IS THE DECISION POINT
369 n_static_args = length (filter isStaticValue staticness)
371 saTransform :: Id -> SATInfo -> [Id] -> CoreExpr -> SatM CoreBind
372 saTransform binder arg_staticness rhs_binders rhs_body
373 = do { shadow_lam_bndrs <- mapM clone binders_w_staticness
375 ; return (NonRec binder (mk_new_rhs uniq shadow_lam_bndrs)) }
377 -- Running example: foldr
378 -- foldr \alpha \beta c n xs = e, for some e
379 -- arg_staticness = [Static TypeApp, Static TypeApp, Static VarApp, Static VarApp, NonStatic]
380 -- rhs_binders = [\alpha, \beta, c, n, xs]
383 binders_w_staticness = rhs_binders `zip` (arg_staticness ++ repeat NotStatic)
384 -- Any extra args are assumed NotStatic
386 non_static_args :: [Var]
387 -- non_static_args = [xs]
388 -- rhs_binders_without_type_capture = [\alpha', \beta', c, n, xs]
389 non_static_args = [v | (v, NotStatic) <- binders_w_staticness]
391 clone (bndr, NotStatic) = return bndr
392 clone (bndr, _ ) = do { uniq <- newUnique
393 ; return (setVarUnique bndr uniq) }
395 -- new_rhs = \alpha beta c n xs ->
396 -- let sat_worker = \xs -> let sat_shadow = \alpha' beta' c n xs ->
400 mk_new_rhs uniq shadow_lam_bndrs
401 = mkLams rhs_binders $
402 Let (Rec [(rec_body_bndr, rec_body)])
405 local_body = mkVarApps (Var rec_body_bndr) non_static_args
407 rec_body = mkLams non_static_args $
408 Let (NonRec shadow_bndr shadow_rhs) rhs_body
410 -- See Note [Binder type capture]
411 shadow_rhs = mkLams shadow_lam_bndrs local_body
412 -- nonrec_rhs = \alpha' beta' c n xs -> sat_worker xs
414 rec_body_bndr = mkSysLocal (fsLit "sat_worker") uniq (exprType rec_body)
415 -- rec_body_bndr = sat_worker
417 -- See Note [Shadow binding]; make a SysLocal
418 shadow_bndr = mkSysLocal (occNameFS (getOccName binder))
420 (exprType shadow_rhs)
422 isStaticValue :: Staticness App -> Bool
423 isStaticValue (Static (VarApp _)) = True
424 isStaticValue _ = False