never executed always true always false
1
2
3 {-
4 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
5
6
7 ************************************************************************
8
9 Static Argument Transformation pass
10
11 ************************************************************************
12
13 May be seen as removing invariants from loops:
14 Arguments of recursive functions that do not change in recursive
15 calls are removed from the recursion, which is done locally
16 and only passes the arguments which effectively change.
17
18 Example:
19 map = /\ ab -> \f -> \xs -> case xs of
20 [] -> []
21 (a:b) -> f a : map f b
22
23 as map is recursively called with the same argument f (unmodified)
24 we transform it to
25
26 map = /\ ab -> \f -> \xs -> let map' ys = case ys of
27 [] -> []
28 (a:b) -> f a : map' b
29 in map' xs
30
31 Notice that for a compiler that uses lambda lifting this is
32 useless as map' will be transformed back to what map was.
33
34 We could possibly do the same for big lambdas, but we don't as
35 they will eventually be removed in later stages of the compiler,
36 therefore there is no penalty in keeping them.
37
38 We only apply the SAT when the number of static args is > 2. This
39 produces few bad cases. See
40 should_transform
41 in saTransform.
42
43 Here are the headline nofib results:
44 Size Allocs Runtime
45 Min +0.0% -13.7% -21.4%
46 Max +0.1% +0.0% +5.4%
47 Geometric Mean +0.0% -0.2% -6.9%
48
49 The previous patch, to fix polymorphic floatout demand signatures, is
50 essential to make this work well!
51 -}
52
53 module GHC.Core.Opt.StaticArgs ( doStaticArgs ) where
54
55 import GHC.Prelude
56
57 import GHC.Types.Var
58 import GHC.Core
59 import GHC.Core.Utils
60 import GHC.Core.Type
61 import GHC.Core.Coercion
62 import GHC.Types.Id
63 import GHC.Types.Name
64 import GHC.Types.Var.Env
65 import GHC.Types.Unique.Supply
66 import GHC.Utils.Misc
67 import GHC.Types.Unique.FM
68 import GHC.Types.Var.Set
69 import GHC.Types.Unique
70 import GHC.Types.Unique.Set
71 import GHC.Utils.Outputable
72 import GHC.Utils.Panic
73
74 import Data.List (mapAccumL)
75 import GHC.Data.FastString
76
77 doStaticArgs :: UniqSupply -> CoreProgram -> CoreProgram
78 doStaticArgs us binds = snd $ mapAccumL sat_bind_threaded_us us binds
79 where
80 sat_bind_threaded_us us bind =
81 let (us1, us2) = splitUniqSupply us
82 in (us1, fst $ runSAT us2 (satBind bind emptyUniqSet))
83
84 -- We don't bother to SAT recursive groups since it can lead
85 -- to massive code expansion: see Andre Santos' thesis for details.
86 -- This means we only apply the actual SAT to Rec groups of one element,
87 -- but we want to recurse into the others anyway to discover other binds
88 satBind :: CoreBind -> IdSet -> SatM (CoreBind, IdSATInfo)
89 satBind (NonRec binder expr) interesting_ids = do
90 (expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
91 return (NonRec binder expr', finalizeApp expr_app sat_info_expr)
92 satBind (Rec [(binder, rhs)]) interesting_ids = do
93 let interesting_ids' = interesting_ids `addOneToUniqSet` binder
94 (rhs_binders, rhs_body) = collectBinders rhs
95 (rhs_body', sat_info_rhs_body) <- satTopLevelExpr rhs_body interesting_ids'
96 let sat_info_rhs_from_args = unitVarEnv binder (bindersToSATInfo rhs_binders)
97 sat_info_rhs' = mergeIdSATInfo sat_info_rhs_from_args sat_info_rhs_body
98
99 shadowing = binder `elementOfUniqSet` interesting_ids
100 sat_info_rhs'' = if shadowing
101 then sat_info_rhs' `delFromUFM` binder -- For safety
102 else sat_info_rhs'
103
104 bind' <- saTransformMaybe binder (lookupUFM sat_info_rhs' binder)
105 rhs_binders rhs_body'
106 return (bind', sat_info_rhs'')
107 satBind (Rec pairs) interesting_ids = do
108 let (binders, rhss) = unzip pairs
109 rhss_SATed <- mapM (\e -> satTopLevelExpr e interesting_ids) rhss
110 let (rhss', sat_info_rhss') = unzip rhss_SATed
111 return (Rec (zipEqual "satBind" binders rhss'), mergeIdSATInfos sat_info_rhss')
112
113 data App = VarApp Id | TypeApp Type | CoApp Coercion
114 data Staticness a = Static a | NotStatic
115
116 type IdAppInfo = (Id, SATInfo)
117
118 type SATInfo = [Staticness App]
119 type IdSATInfo = IdEnv SATInfo
120 emptyIdSATInfo :: IdSATInfo
121 emptyIdSATInfo = emptyUFM
122
123 {-
124 pprIdSATInfo id_sat_info = vcat (map pprIdAndSATInfo (Map.toList id_sat_info))
125 where pprIdAndSATInfo (v, sat_info) = hang (ppr v <> colon) 4 (pprSATInfo sat_info)
126 -}
127
128 pprSATInfo :: SATInfo -> SDoc
129 pprSATInfo staticness = hcat $ map pprStaticness staticness
130
131 pprStaticness :: Staticness App -> SDoc
132 pprStaticness (Static (VarApp _)) = text "SV"
133 pprStaticness (Static (TypeApp _)) = text "ST"
134 pprStaticness (Static (CoApp _)) = text "SC"
135 pprStaticness NotStatic = text "NS"
136
137
138 mergeSATInfo :: SATInfo -> SATInfo -> SATInfo
139 mergeSATInfo l r = zipWith mergeSA l r
140 where
141 mergeSA NotStatic _ = NotStatic
142 mergeSA _ NotStatic = NotStatic
143 mergeSA (Static (VarApp v)) (Static (VarApp v'))
144 | v == v' = Static (VarApp v)
145 | otherwise = NotStatic
146 mergeSA (Static (TypeApp t)) (Static (TypeApp t'))
147 | t `eqType` t' = Static (TypeApp t)
148 | otherwise = NotStatic
149 mergeSA (Static (CoApp c)) (Static (CoApp c'))
150 | c `eqCoercion` c' = Static (CoApp c)
151 | otherwise = NotStatic
152 mergeSA _ _ = pprPanic "mergeSATInfo" $
153 text "Left:"
154 <> pprSATInfo l <> text ", "
155 <> text "Right:"
156 <> pprSATInfo r
157
158 mergeIdSATInfo :: IdSATInfo -> IdSATInfo -> IdSATInfo
159 mergeIdSATInfo = plusUFM_C mergeSATInfo
160
161 mergeIdSATInfos :: [IdSATInfo] -> IdSATInfo
162 mergeIdSATInfos = foldl' mergeIdSATInfo emptyIdSATInfo
163
164 bindersToSATInfo :: [Id] -> SATInfo
165 bindersToSATInfo vs = map (Static . binderToApp) vs
166 where binderToApp v | isId v = VarApp v
167 | isTyVar v = TypeApp $ mkTyVarTy v
168 | otherwise = CoApp $ mkCoVarCo v
169
170 finalizeApp :: Maybe IdAppInfo -> IdSATInfo -> IdSATInfo
171 finalizeApp Nothing id_sat_info = id_sat_info
172 finalizeApp (Just (v, sat_info')) id_sat_info =
173 let sat_info'' = case lookupUFM id_sat_info v of
174 Nothing -> sat_info'
175 Just sat_info -> mergeSATInfo sat_info sat_info'
176 in extendVarEnv id_sat_info v sat_info''
177
178 satTopLevelExpr :: CoreExpr -> IdSet -> SatM (CoreExpr, IdSATInfo)
179 satTopLevelExpr expr interesting_ids = do
180 (expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
181 return (expr', finalizeApp expr_app sat_info_expr)
182
183 satExpr :: CoreExpr -> IdSet -> SatM (CoreExpr, IdSATInfo, Maybe IdAppInfo)
184 satExpr var@(Var v) interesting_ids = do
185 let app_info = if v `elementOfUniqSet` interesting_ids
186 then Just (v, [])
187 else Nothing
188 return (var, emptyIdSATInfo, app_info)
189
190 satExpr lit@(Lit _) _ =
191 return (lit, emptyIdSATInfo, Nothing)
192
193 satExpr (Lam binders body) interesting_ids = do
194 (body', sat_info, this_app) <- satExpr body interesting_ids
195 return (Lam binders body', finalizeApp this_app sat_info, Nothing)
196
197 satExpr (App fn arg) interesting_ids = do
198 (fn', sat_info_fn, fn_app) <- satExpr fn interesting_ids
199 let satRemainder = boring fn' sat_info_fn
200 case fn_app of
201 Nothing -> satRemainder Nothing
202 Just (fn_id, fn_app_info) ->
203 -- TODO: remove this use of append somehow (use a data structure with O(1) append but a left-to-right kind of interface)
204 let satRemainderWithStaticness arg_staticness = satRemainder $ Just (fn_id, fn_app_info ++ [arg_staticness])
205 in case arg of
206 Type t -> satRemainderWithStaticness $ Static (TypeApp t)
207 Coercion c -> satRemainderWithStaticness $ Static (CoApp c)
208 Var v -> satRemainderWithStaticness $ Static (VarApp v)
209 _ -> satRemainderWithStaticness $ NotStatic
210 where
211 boring :: CoreExpr -> IdSATInfo -> Maybe IdAppInfo -> SatM (CoreExpr, IdSATInfo, Maybe IdAppInfo)
212 boring fn' sat_info_fn app_info =
213 do (arg', sat_info_arg, arg_app) <- satExpr arg interesting_ids
214 let sat_info_arg' = finalizeApp arg_app sat_info_arg
215 sat_info = mergeIdSATInfo sat_info_fn sat_info_arg'
216 return (App fn' arg', sat_info, app_info)
217
218 satExpr (Case expr bndr ty alts) interesting_ids = do
219 (expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
220 let sat_info_expr' = finalizeApp expr_app sat_info_expr
221
222 zipped_alts' <- mapM satAlt alts
223 let (alts', sat_infos_alts) = unzip zipped_alts'
224 return (Case expr' bndr ty alts', mergeIdSATInfo sat_info_expr' (mergeIdSATInfos sat_infos_alts), Nothing)
225 where
226 satAlt (Alt con bndrs expr) = do
227 (expr', sat_info_expr) <- satTopLevelExpr expr interesting_ids
228 return (Alt con bndrs expr', sat_info_expr)
229
230 satExpr (Let bind body) interesting_ids = do
231 (body', sat_info_body, body_app) <- satExpr body interesting_ids
232 (bind', sat_info_bind) <- satBind bind interesting_ids
233 return (Let bind' body', mergeIdSATInfo sat_info_body sat_info_bind, body_app)
234
235 satExpr (Tick tickish expr) interesting_ids = do
236 (expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
237 return (Tick tickish expr', sat_info_expr, expr_app)
238
239 satExpr ty@(Type _) _ =
240 return (ty, emptyIdSATInfo, Nothing)
241
242 satExpr co@(Coercion _) _ =
243 return (co, emptyIdSATInfo, Nothing)
244
245 satExpr (Cast expr coercion) interesting_ids = do
246 (expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
247 return (Cast expr' coercion, sat_info_expr, expr_app)
248
249 {-
250 ************************************************************************
251
252 Static Argument Transformation Monad
253
254 ************************************************************************
255 -}
256
257 type SatM result = UniqSM result
258
259 runSAT :: UniqSupply -> SatM a -> a
260 runSAT = initUs_
261
262 newUnique :: SatM Unique
263 newUnique = getUniqueM
264
265 {-
266 ************************************************************************
267
268 Static Argument Transformation Monad
269
270 ************************************************************************
271
272 To do the transformation, the game plan is to:
273
274 1. Create a small nonrecursive RHS that takes the
275 original arguments to the function but discards
276 the ones that are static and makes a call to the
277 SATed version with the remainder. We intend that
278 this will be inlined later, removing the overhead
279
280 2. Bind this nonrecursive RHS over the original body
281 WITH THE SAME UNIQUE as the original body so that
282 any recursive calls to the original now go via
283 the small wrapper
284
285 3. Rebind the original function to a new one which contains
286 our SATed function and just makes a call to it:
287 we call the thing making this call the local body
288
289 Example: transform this
290
291 map :: forall a b. (a->b) -> [a] -> [b]
292 map = /\ab. \(f:a->b) (as:[a]) -> body[map]
293 to
294 map :: forall a b. (a->b) -> [a] -> [b]
295 map = /\ab. \(f:a->b) (as:[a]) ->
296 letrec map' :: [a] -> [b]
297 -- The "worker function
298 map' = \(as:[a]) ->
299 let map :: forall a' b'. (a -> b) -> [a] -> [b]
300 -- The "shadow function
301 map = /\a'b'. \(f':(a->b) (as:[a]).
302 map' as
303 in body[map]
304 in map' as
305
306 Note [Shadow binding]
307 ~~~~~~~~~~~~~~~~~~~~~
308 The calls to the inner map inside body[map] should get inlined
309 by the local re-binding of 'map'. We call this the "shadow binding".
310
311 But we can't use the original binder 'map' unchanged, because
312 it might be exported, in which case the shadow binding won't be
313 discarded as dead code after it is inlined.
314
315 So we use a hack: we make a new SysLocal binder with the *same* unique
316 as binder. (Another alternative would be to reset the export flag.)
317
318 Note [Binder type capture]
319 ~~~~~~~~~~~~~~~~~~~~~~~~~~
320 Notice that in the inner map (the "shadow function"), the static arguments
321 are discarded -- it's as if they were underscores. Instead, mentions
322 of these arguments (notably in the types of dynamic arguments) are bound
323 by the *outer* lambdas of the main function. So we must make up fresh
324 names for the static arguments so that they do not capture variables
325 mentioned in the types of dynamic args.
326
327 In the map example, the shadow function must clone the static type
328 argument a,b, giving a',b', to ensure that in the \(as:[a]), the 'a'
329 is bound by the outer forall. We clone f' too for consistency, but
330 that doesn't matter either way because static Id arguments aren't
331 mentioned in the shadow binding at all.
332
333 If we don't we get something like this:
334
335 [Exported]
336 [Arity 3]
337 GHC.Base.until =
338 \ (@ a_aiK)
339 (p_a6T :: a_aiK -> GHC.Types.Bool)
340 (f_a6V :: a_aiK -> a_aiK)
341 (x_a6X :: a_aiK) ->
342 letrec {
343 sat_worker_s1aU :: a_aiK -> a_aiK
344 []
345 sat_worker_s1aU =
346 \ (x_a6X :: a_aiK) ->
347 let {
348 sat_shadow_r17 :: forall a_a3O.
349 (a_a3O -> GHC.Types.Bool) -> (a_a3O -> a_a3O) -> a_a3O -> a_a3O
350 []
351 sat_shadow_r17 =
352 \ (@ a_aiK)
353 (p_a6T :: a_aiK -> GHC.Types.Bool)
354 (f_a6V :: a_aiK -> a_aiK)
355 (x_a6X :: a_aiK) ->
356 sat_worker_s1aU x_a6X } in
357 case p_a6T x_a6X of wild_X3y [ALWAYS Dead Nothing] {
358 GHC.Types.False -> GHC.Base.until @ a_aiK p_a6T f_a6V (f_a6V x_a6X);
359 GHC.Types.True -> x_a6X
360 }; } in
361 sat_worker_s1aU x_a6X
362
363 Where sat_shadow has captured the type variables of x_a6X etc as it has a a_aiK
364 type argument. This is bad because it means the application sat_worker_s1aU x_a6X
365 is not well typed.
366 -}
367
368 saTransformMaybe :: Id -> Maybe SATInfo -> [Id] -> CoreExpr -> SatM CoreBind
369 saTransformMaybe binder maybe_arg_staticness rhs_binders rhs_body
370 | Just arg_staticness <- maybe_arg_staticness
371 , should_transform arg_staticness
372 = saTransform binder arg_staticness rhs_binders rhs_body
373 | otherwise
374 = return (Rec [(binder, mkLams rhs_binders rhs_body)])
375 where
376 should_transform staticness = n_static_args > 1 -- THIS IS THE DECISION POINT
377 where
378 n_static_args = count isStaticValue staticness
379
380 saTransform :: Id -> SATInfo -> [Id] -> CoreExpr -> SatM CoreBind
381 saTransform binder arg_staticness rhs_binders rhs_body
382 = do { shadow_lam_bndrs <- mapM clone binders_w_staticness
383 ; uniq <- newUnique
384 ; return (NonRec binder (mk_new_rhs uniq shadow_lam_bndrs)) }
385 where
386 -- Running example: foldr
387 -- foldr \alpha \beta c n xs = e, for some e
388 -- arg_staticness = [Static TypeApp, Static TypeApp, Static VarApp, Static VarApp, NonStatic]
389 -- rhs_binders = [\alpha, \beta, c, n, xs]
390 -- rhs_body = e
391
392 binders_w_staticness = rhs_binders `zip` (arg_staticness ++ repeat NotStatic)
393 -- Any extra args are assumed NotStatic
394
395 non_static_args :: [Var]
396 -- non_static_args = [xs]
397 -- rhs_binders_without_type_capture = [\alpha', \beta', c, n, xs]
398 non_static_args = [v | (v, NotStatic) <- binders_w_staticness]
399
400 clone (bndr, NotStatic) = return bndr
401 clone (bndr, _ ) = do { uniq <- newUnique
402 ; return (setVarUnique bndr uniq) }
403
404 -- new_rhs = \alpha beta c n xs ->
405 -- let sat_worker = \xs -> let sat_shadow = \alpha' beta' c n xs ->
406 -- sat_worker xs
407 -- in e
408 -- in sat_worker xs
409 mk_new_rhs uniq shadow_lam_bndrs
410 = mkLams rhs_binders $
411 Let (Rec [(rec_body_bndr, rec_body)])
412 local_body
413 where
414 local_body = mkVarApps (Var rec_body_bndr) non_static_args
415
416 rec_body = mkLams non_static_args $
417 Let (NonRec shadow_bndr shadow_rhs) rhs_body
418
419 -- See Note [Binder type capture]
420 shadow_rhs = mkLams shadow_lam_bndrs local_body
421 -- nonrec_rhs = \alpha' beta' c n xs -> sat_worker xs
422
423 rec_body_bndr = mkSysLocal (fsLit "sat_worker") uniq Many (exprType rec_body)
424 -- rec_body_bndr = sat_worker
425
426 -- See Note [Shadow binding]; make a SysLocal
427 shadow_bndr = mkSysLocal (occNameFS (getOccName binder))
428 (idUnique binder)
429 Many
430 (exprType shadow_rhs)
431
432 isStaticValue :: Staticness App -> Bool
433 isStaticValue (Static (VarApp _)) = True
434 isStaticValue _ = False