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