1{-
2(c) The GRASP/AQUA Project, Glasgow University, 1992-1998
3
4
5************************************************************************
6
7               Static Argument Transformation pass
8
9************************************************************************
10
11May be seen as removing invariants from loops:
12Arguments of recursive functions that do not change in recursive
13calls are removed from the recursion, which is done locally
14and only passes the arguments which effectively change.
15
16Example:
17map = /\ ab -> \f -> \xs -> case xs of
18                 []       -> []
19                 (a:b) -> f a : map f b
20
21as map is recursively called with the same argument f (unmodified)
22we transform it to
23
24map = /\ ab -> \f -> \xs -> let map' ys = case ys of
25                       []     -> []
26                       (a:b) -> f a : map' b
27                in map' xs
28
29Notice that for a compiler that uses lambda lifting this is
30useless as map' will be transformed back to what map was.
31
32We could possibly do the same for big lambdas, but we don't as
33they will eventually be removed in later stages of the compiler,
34therefore there is no penalty in keeping them.
35
36We only apply the SAT when the number of static args is > 2. This
37produces few bad cases.  See
38                should_transform
39in saTransform.
40
41Here are the headline nofib results:
42                  Size    Allocs   Runtime
43Min             +0.0%    -13.7%    -21.4%
44Max             +0.1%     +0.0%     +5.4%
45Geometric Mean  +0.0%     -0.2%     -6.9%
46
47The previous patch, to fix polymorphic floatout demand signatures, is
48essential to make this work well!
49-}
50
51{-# LANGUAGE CPP #-}
52module SAT ( doStaticArgs ) where
53
54import GhcPrelude
55
56import Var
57import CoreSyn
58import CoreUtils
59import Type
60import Coercion
61import Id
62import Name
63import VarEnv
64import UniqSupply
65import Util
66import UniqFM
67import VarSet
68import Unique
69import UniqSet
70import Outputable
71
72import Data.List (mapAccumL)
73import FastString
74
75#include "HsVersions.h"
76
77doStaticArgs :: UniqSupply -> CoreProgram -> CoreProgram
78doStaticArgs 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
88satBind :: CoreBind -> IdSet -> SatM (CoreBind, IdSATInfo)
89satBind (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)
92satBind (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'')
107satBind (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
113data App = VarApp Id | TypeApp Type | CoApp Coercion
114data Staticness a = Static a | NotStatic
115
116type IdAppInfo = (Id, SATInfo)
117
118type SATInfo = [Staticness App]
119type IdSATInfo = IdEnv SATInfo
120emptyIdSATInfo :: IdSATInfo
121emptyIdSATInfo = emptyUFM
122
123{-
124pprIdSATInfo 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
128pprSATInfo :: SATInfo -> SDoc
129pprSATInfo staticness = hcat $ map pprStaticness staticness
130
131pprStaticness :: Staticness App -> SDoc
132pprStaticness (Static (VarApp _))  = text "SV"
133pprStaticness (Static (TypeApp _)) = text "ST"
134pprStaticness (Static (CoApp _))   = text "SC"
135pprStaticness NotStatic            = text "NS"
136
137
138mergeSATInfo :: SATInfo -> SATInfo -> SATInfo
139mergeSATInfo 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
158mergeIdSATInfo :: IdSATInfo -> IdSATInfo -> IdSATInfo
159mergeIdSATInfo = plusUFM_C mergeSATInfo
160
161mergeIdSATInfos :: [IdSATInfo] -> IdSATInfo
162mergeIdSATInfos = foldl' mergeIdSATInfo emptyIdSATInfo
163
164bindersToSATInfo :: [Id] -> SATInfo
165bindersToSATInfo 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
170finalizeApp :: Maybe IdAppInfo -> IdSATInfo -> IdSATInfo
171finalizeApp Nothing id_sat_info = id_sat_info
172finalizeApp (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
178satTopLevelExpr :: CoreExpr -> IdSet -> SatM (CoreExpr, IdSATInfo)
179satTopLevelExpr 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
183satExpr :: CoreExpr -> IdSet -> SatM (CoreExpr, IdSATInfo, Maybe IdAppInfo)
184satExpr 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
190satExpr lit@(Lit _) _ = do
191    return (lit, emptyIdSATInfo, Nothing)
192
193satExpr (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
197satExpr (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
218satExpr (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 (con, bndrs, expr) = do
227        (expr', sat_info_expr) <- satTopLevelExpr expr interesting_ids
228        return ((con, bndrs, expr'), sat_info_expr)
229
230satExpr (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
235satExpr (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
239satExpr ty@(Type _) _ = do
240    return (ty, emptyIdSATInfo, Nothing)
241
242satExpr co@(Coercion _) _ = do
243    return (co, emptyIdSATInfo, Nothing)
244
245satExpr (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
257type SatM result = UniqSM result
258
259runSAT :: UniqSupply -> SatM a -> a
260runSAT = initUs_
261
262newUnique :: SatM Unique
263newUnique = getUniqueM
264
265{-
266************************************************************************
267
268                Static Argument Transformation Monad
269
270************************************************************************
271
272To do the transformation, the game plan is to:
273
2741. 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
2802. 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
2853. 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
289Example: transform this
290
291    map :: forall a b. (a->b) -> [a] -> [b]
292    map = /\ab. \(f:a->b) (as:[a]) -> body[map]
293to
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
306Note [Shadow binding]
307~~~~~~~~~~~~~~~~~~~~~
308The calls to the inner map inside body[map] should get inlined
309by the local re-binding of 'map'.  We call this the "shadow binding".
310
311But we can't use the original binder 'map' unchanged, because
312it might be exported, in which case the shadow binding won't be
313discarded as dead code after it is inlined.
314
315So we use a hack: we make a new SysLocal binder with the *same* unique
316as binder.  (Another alternative would be to reset the export flag.)
317
318Note [Binder type capture]
319~~~~~~~~~~~~~~~~~~~~~~~~~~
320Notice that in the inner map (the "shadow function"), the static arguments
321are discarded -- it's as if they were underscores.  Instead, mentions
322of these arguments (notably in the types of dynamic arguments) are bound
323by the *outer* lambdas of the main function.  So we must make up fresh
324names for the static arguments so that they do not capture variables
325mentioned in the types of dynamic args.
326
327In the map example, the shadow function must clone the static type
328argument a,b, giving a',b', to ensure that in the \(as:[a]), the 'a'
329is bound by the outer forall.  We clone f' too for consistency, but
330that doesn't matter either way because static Id arguments aren't
331mentioned in the shadow binding at all.
332
333If we don't we get something like this:
334
335[Exported]
336[Arity 3]
337GHC.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
363Where sat_shadow has captured the type variables of x_a6X etc as it has a a_aiK
364type argument. This is bad because it means the application sat_worker_s1aU x_a6X
365is not well typed.
366-}
367
368saTransformMaybe :: Id -> Maybe SATInfo -> [Id] -> CoreExpr -> SatM CoreBind
369saTransformMaybe 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
380saTransform :: Id -> SATInfo -> [Id] -> CoreExpr -> SatM CoreBind
381saTransform 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 (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                                   (exprType shadow_rhs)
430
431isStaticValue :: Staticness App -> Bool
432isStaticValue (Static (VarApp _)) = True
433isStaticValue _                   = False
434