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