travis compilation
[ghc.git] / compiler / simplStg / StgCse.hs
1 {-# LANGUAGE TypeFamilies, LambdaCase #-}
2 {-# OPTIONS -Wno-error=unused-imports -Wno-error=unused-top-binds #-}
3
4 {-|
5 Note [CSE for Stg]
6 ~~~~~~~~~~~~~~~~~~
7 This module implements a simple common subexpression elimination pass for STG.
8 This is useful because there are expressions that we want to common up (because
9 they are operationally equivalent), but that we cannot common up in Core, because
10 their types differ.
11 This was originally reported as #9291.
12
13 There are two types of common code occurrences that we aim for, see
14 note [Case 1: CSEing allocated closures] and
15 note [Case 2: CSEing case binders] below.
16
17 TODOs:
18 - vanilla for unpacked tuples?
19 - rerun occurrence analysis
20 - dumping of STG misses binder
21 - does not look up in scope to find low-hanging fruit
22 - can we dedup info tables for representationally equal data constructors?
23
24
25 Note [Case 1: CSEing allocated closures]
26 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
27 The first kind of CSE opportunity we aim for is generated by this Haskell code:
28
29 bar :: a -> (Either Int a, Either Bool a)
30 bar x = (Right x, Right x)
31
32 which produces this Core:
33
34 bar :: forall a. a -> (Either Int a, Either Bool a)
35 bar @a x = (Right @Int @a x, Right @Bool @a x)
36
37 where the two components of the tuple are different terms, and cannot be
38 commoned up (easily). On the STG level we have
39
40 bar [x] = let c1 = Right [x]
41 c2 = Right [x]
42 in (c1,c2)
43
44 and now it is obvious that we can write
45
46 bar [x] = let c1 = Right [x]
47 in (c1,c1)
48
49 instead.
50
51
52 Note [Case 2: CSEing case binders]
53 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
54 The second kind of CSE opportunity we aim for is more interesting, and
55 came up in #9291 and #5344: The Haskell code
56
57 foo :: Either Int a -> Either Bool a
58 foo (Right x) = Right x
59 foo _ = Left False
60
61 produces this Core
62
63 foo :: forall a. Either Int a -> Either Bool a
64 foo @a e = case e of b { Left n -> …
65 , Right x -> Right @Bool @a x }
66
67 where we cannot CSE `Right @Bool @a x` with the case binder `b` as they have
68 different types. But in STG we have
69
70 foo [e] = case e of b { Left [n] -> …
71 , Right [x] -> Right [x] }
72
73 and nothing stops us from transforming that to
74
75 foo [e] = case e of b { Left [n] -> …
76 , Right [x] -> b}
77
78 -}
79 module StgCse (stgCse) where
80
81 import GhcPrelude
82
83 import DataCon
84 import Id
85 import StgSyn
86 import Outputable
87 import VarEnv
88 import CoreSyn (AltCon(..))
89 import Data.List (mapAccumL)
90 import Data.Maybe (fromMaybe)
91 import TrieMap
92 import NameEnv
93 import Control.Monad( (>=>) )
94 import Name (NamedThing (..), mkFCallName)
95 import Unique (mkUniqueGrimily, getKey, getUnique)
96 import TyCon (tyConFamilySize)
97
98 --------------
99 -- The Trie --
100 --------------
101
102 -- A lookup trie for data constructor applications, i.e.
103 -- keys of type `(LaxDataCon, [StgArg])`, following the patterns in TrieMap.
104
105 data StgArgMap a = SAM
106 { sam_var :: DVarEnv a
107 , sam_lit :: LiteralMap a
108 }
109
110 instance TrieMap StgArgMap where
111 type Key StgArgMap = StgArg
112 emptyTM = SAM { sam_var = emptyTM
113 , sam_lit = emptyTM }
114 lookupTM (StgVarArg var) = sam_var >.> lkDFreeVar var
115 lookupTM (StgLitArg lit) = sam_lit >.> lookupTM lit
116 alterTM (StgVarArg var) f m = m { sam_var = sam_var m |> xtDFreeVar var f }
117 alterTM (StgLitArg lit) f m = m { sam_lit = sam_lit m |> alterTM lit f }
118 foldTM k m = foldTM k (sam_var m) . foldTM k (sam_lit m)
119 mapTM f (SAM {sam_var = varm, sam_lit = litm}) =
120 SAM { sam_var = mapTM f varm, sam_lit = mapTM f litm }
121
122 newtype ConAppMap a = CAM { un_cam :: DNameEnv (ListMap StgArgMap a) }
123
124 newtype LaxDataCon = Lax DataCon
125
126 instance NamedThing LaxDataCon where
127 getName (Lax dc) | long && isVanillaDataCon dc && not hasStrict && not unpacked = mkFCallName uniq "" -- FIXME: is there a better way?
128 where uniq = mkUniqueGrimily . negate $ dataConTag dc * 1048576 + length (dataConOrigArgTys dc) -- FIXME
129 hasStrict = any (\case HsLazy -> False; _ -> True) (dataConImplBangs dc)
130 unpacked = isUnboxedTupleCon dc || isUnboxedSumCon dc
131 long = dataConTag dc < 7 && (null $ dataConOrigArgTys dc) -- True -- length (dataConOrigArgTys dc) > 0
132 getName (Lax dc) = getName dc
133
134
135 instance TrieMap ConAppMap where
136 type Key ConAppMap = (LaxDataCon, [StgArg])
137 emptyTM = CAM emptyTM
138 lookupTM (dataCon, args) = un_cam >.> lkDNamed dataCon >=> lookupTM args
139 alterTM (dataCon, args) f m =
140 m { un_cam = un_cam m |> xtDNamed dataCon |>> alterTM args f }
141 foldTM k = un_cam >.> foldTM (foldTM k)
142 mapTM f = un_cam >.> mapTM (mapTM f) >.> CAM
143
144
145 -----------------
146 -- The CSE Env --
147 -----------------
148
149 -- | The CSE environment. See note [CseEnv Example]
150 data CseEnv = CseEnv
151 { ce_conAppMap :: ConAppMap OutId
152 -- ^ The main component of the environment is the trie that maps
153 -- data constructor applications (with their `OutId` arguments)
154 -- to an in-scope name that can be used instead.
155 -- This name is always either a let-bound variable or a case binder.
156 , ce_subst :: IdEnv OutId
157 -- ^ This substitution is applied to the code as we traverse it.
158 -- Entries have one of two reasons:
159 --
160 -- * The input might have shadowing (see Note [Shadowing]), so we have
161 -- to rename some binders as we traverse the tree.
162 -- * If we remove `let x = Con z` because `let y = Con z` is in scope,
163 -- we note this here as x ↦ y.
164 , ce_bndrMap :: IdEnv OutId
165 -- ^ If we come across a case expression case x as b of … with a trivial
166 -- binder, we add b ↦ x to this.
167 -- This map is *only* used when looking something up in the ce_conAppMap.
168 -- See Note [Trivial case scrutinee]
169 , ce_in_scope :: InScopeSet
170 -- ^ The third component is an in-scope set, to rename away any
171 -- shadowing binders
172 }
173
174 {-|
175 Note [CseEnv Example]
176 ~~~~~~~~~~~~~~~~~~~~~
177 The following tables shows how the CseEnvironment changes as code is traversed,
178 as well as the changes to that code.
179
180 InExpr OutExpr
181 conAppMap subst in_scope
182 ───────────────────────────────────────────────────────────
183 -- empty {} {}
184 case … as a of {Con x y -> case … as a of {Con x y ->
185 -- Con x y ↦ a {} {a,x,y}
186 let b = Con x y (removed)
187 -- Con x y ↦ a b↦a {a,x,y,b}
188 let c = Bar a let c = Bar a
189 -- Con x y ↦ a, Bar a ↦ c b↦a {a,x,y,b,c}
190 let c = some expression let c' = some expression
191 -- Con x y ↦ a, Bar a ↦ c b↦a, c↦c', {a,x,y,b,c,c'}
192 let d = Bar b (removed)
193 -- Con x y ↦ a, Bar a ↦ c b↦a, c↦c', d↦c {a,x,y,b,c,c',d}
194 (a, b, c d) (a, a, c' c)
195 -}
196
197 initEnv :: InScopeSet -> CseEnv
198 initEnv in_scope = CseEnv
199 { ce_conAppMap = emptyTM
200 , ce_subst = emptyVarEnv
201 , ce_bndrMap = emptyVarEnv
202 , ce_in_scope = in_scope
203 }
204
205 envLookup :: LaxDataCon -> [OutStgArg] -> CseEnv -> Maybe OutId
206 envLookup dataCon args env = lookupTM (dataCon, args') (ce_conAppMap env)
207 where args' = map go args -- See Note [Trivial case scrutinee]
208 go (StgVarArg v ) = StgVarArg (fromMaybe v $ lookupVarEnv (ce_bndrMap env) v)
209 go (StgLitArg lit) = StgLitArg lit
210
211 addDataCon :: OutId -> LaxDataCon -> [OutStgArg] -> CseEnv -> CseEnv
212 -- do not bother with nullary data constructors, they are static anyways
213 addDataCon bndr dataCon [] env = env { ce_conAppMap = new_env }
214 where new_env = alterTM (dataCon, []) (\case Nothing -> pure bndr; p -> p) (ce_conAppMap env)
215 --addDataCon _ _ [] env = env
216 addDataCon bndr dataCon args env = env { ce_conAppMap = new_env }
217 where
218 new_env = insertTM (dataCon, args) bndr (ce_conAppMap env)
219
220 forgetCse :: CseEnv -> CseEnv
221 forgetCse env = env { ce_conAppMap = emptyTM }
222 -- See note [Free variables of an StgClosure]
223
224 addSubst :: OutId -> OutId -> CseEnv -> CseEnv
225 addSubst from to env
226 = env { ce_subst = extendVarEnv (ce_subst env) from to }
227
228 addTrivCaseBndr :: OutId -> OutId -> CseEnv -> CseEnv
229 addTrivCaseBndr from to env
230 = env { ce_bndrMap = extendVarEnv (ce_bndrMap env) from to }
231
232 substArgs :: CseEnv -> [InStgArg] -> [OutStgArg]
233 substArgs env = map (substArg env)
234
235 substArg :: CseEnv -> InStgArg -> OutStgArg
236 substArg env (StgVarArg from) = StgVarArg (substVar env from)
237 substArg _ (StgLitArg lit) = StgLitArg lit
238
239 substVars :: CseEnv -> [InId] -> [OutId]
240 substVars env = map (substVar env)
241
242 substVar :: CseEnv -> InId -> OutId
243 substVar env id = fromMaybe id $ lookupVarEnv (ce_subst env) id
244
245 -- Functions to enter binders
246
247 -- This is much simpler than the equivalent code in CoreSubst:
248 -- * We do not substitute type variables, and
249 -- * There is nothing relevant in IdInfo at this stage
250 -- that needs substitutions.
251 -- Therefore, no special treatment for a recursive group is required.
252
253 substBndr :: CseEnv -> InId -> (CseEnv, OutId)
254 substBndr env old_id
255 = (new_env, new_id)
256 where
257 new_id = uniqAway (ce_in_scope env) old_id
258 no_change = new_id == old_id
259 env' = env { ce_in_scope = ce_in_scope env `extendInScopeSet` new_id }
260 new_env | no_change = env' { ce_subst = extendVarEnv (ce_subst env) old_id new_id }
261 | otherwise = env'
262
263 substBndrs :: CseEnv -> [InVar] -> (CseEnv, [OutVar])
264 substBndrs env bndrs = mapAccumL substBndr env bndrs
265
266 substPairs :: CseEnv -> [(InVar, a)] -> (CseEnv, [(OutVar, a)])
267 substPairs env bndrs = mapAccumL go env bndrs
268 where go env (id, x) = let (env', id') = substBndr env id
269 in (env', (id', x))
270
271 -- Main entry point
272
273 stgCse :: [InStgTopBinding] -> [OutStgTopBinding]
274 stgCse binds = snd $ mapAccumL stgCseTopLvl emptyInScopeSet binds
275
276 -- Top level bindings.
277 --
278 -- We do not CSE these, as top-level closures are allocated statically anyways.
279 -- Also, they might be exported.
280 -- But we still have to collect the set of in-scope variables, otherwise
281 -- uniqAway might shadow a top-level closure.
282
283 stgCseTopLvl :: InScopeSet -> InStgTopBinding -> (InScopeSet, OutStgTopBinding)
284 stgCseTopLvl in_scope t@(StgTopStringLit _ _) = (in_scope, t)
285 stgCseTopLvl in_scope (StgTopLifted (StgNonRec bndr rhs))
286 = (in_scope'
287 , StgTopLifted (StgNonRec bndr (stgCseTopLvlRhs in_scope rhs)))
288 where in_scope' = in_scope `extendInScopeSet` bndr
289
290 stgCseTopLvl in_scope (StgTopLifted (StgRec eqs))
291 = ( in_scope'
292 , StgTopLifted (StgRec [ (bndr, stgCseTopLvlRhs in_scope' rhs) | (bndr, rhs) <- eqs ]))
293 where in_scope' = in_scope `extendInScopeSetList` [ bndr | (bndr, _) <- eqs ]
294
295 stgCseTopLvlRhs :: InScopeSet -> InStgRhs -> OutStgRhs
296 stgCseTopLvlRhs in_scope (StgRhsClosure ccs info occs upd args body)
297 = let body' = stgCseExpr (initEnv in_scope) body
298 in StgRhsClosure ccs info occs upd args body'
299 stgCseTopLvlRhs _ (StgRhsCon ccs dataCon args)
300 = StgRhsCon ccs dataCon args
301
302 ------------------------------
303 -- The actual AST traversal --
304 ------------------------------
305
306 -- Trivial cases
307 stgCseExpr :: CseEnv -> InStgExpr -> OutStgExpr
308 stgCseExpr env (StgApp fun args)
309 = StgApp fun' args'
310 where fun' = substVar env fun
311 args' = substArgs env args
312 stgCseExpr _ (StgLit lit)
313 = StgLit lit
314 stgCseExpr env (StgOpApp op args tys)
315 = StgOpApp op args' tys
316 where args' = substArgs env args
317 stgCseExpr _ (StgLam _ _)
318 = pprPanic "stgCseExp" (text "StgLam")
319 stgCseExpr env (StgTick tick body)
320 = let body' = stgCseExpr env body
321 in StgTick tick body'
322 stgCseExpr env (StgCase scrut bndr ty alts)
323 = mkStgCase scrut' bndr' ty alts'
324 where
325 scrut' = stgCseExpr env scrut
326 (env1, bndr') = substBndr env bndr
327 env2 | StgApp trivial_scrut [] <- scrut' = addTrivCaseBndr bndr trivial_scrut env1
328 -- See Note [Trivial case scrutinee]
329 | otherwise = env1
330 alts' = map (stgCseAlt env2 bndr') alts
331
332
333 -- A constructor application.
334 -- To be removed by a variable use when found in the CSE environment
335 stgCseExpr env (StgConApp dataCon args tys)
336 | Just bndr' <- envLookup dc args' env
337 = (if getKey u < 0 then pprTrace "stgCseExpr" (ppr dataCon <+> text (show $ length (dataConOrigArgTys dataCon)) <+> (text . show $ tyConFamilySize (dataConTyCon dataCon))) else id) $ StgApp bndr' []
338 | otherwise
339 = StgConApp dataCon args' tys
340 where args' = substArgs env args
341 dc = Lax dataCon
342 u = getUnique (getName dc)
343
344 -- Let bindings
345 -- The binding might be removed due to CSE (we do not want trivial bindings on
346 -- the STG level), so use the smart constructor `mkStgLet` to remove the binding
347 -- if empty.
348 stgCseExpr env (StgLet binds body)
349 = let (binds', env') = stgCseBind env binds
350 body' = stgCseExpr env' body
351 in mkStgLet StgLet binds' body'
352 stgCseExpr env (StgLetNoEscape binds body)
353 = let (binds', env') = stgCseBind env binds
354 body' = stgCseExpr env' body
355 in mkStgLet StgLetNoEscape binds' body'
356
357 -- Case alternatives
358 -- Extend the CSE environment
359 stgCseAlt :: CseEnv -> OutId -> InStgAlt -> OutStgAlt
360 stgCseAlt env case_bndr (DataAlt dataCon, args, rhs)
361 = let (env1, args') = substBndrs env args
362 env2 = addDataCon case_bndr (Lax dataCon) (map StgVarArg args') env1
363 -- see note [Case 2: CSEing case binders]
364 rhs' = stgCseExpr env2 rhs
365 in (DataAlt dataCon, args', rhs')
366 stgCseAlt env _ (altCon, args, rhs)
367 = let (env1, args') = substBndrs env args
368 rhs' = stgCseExpr env1 rhs
369 in (altCon, args', rhs')
370
371 -- Bindings
372 stgCseBind :: CseEnv -> InStgBinding -> (Maybe OutStgBinding, CseEnv)
373 stgCseBind env (StgNonRec b e)
374 = let (env1, b') = substBndr env b
375 in case stgCseRhs env1 b' e of
376 (Nothing, env2) -> (Nothing, env2)
377 (Just (b2,e'), env2) -> (Just (StgNonRec b2 e'), env2)
378 stgCseBind env (StgRec pairs)
379 = let (env1, pairs1) = substPairs env pairs
380 in case stgCsePairs env1 pairs1 of
381 ([], env2) -> (Nothing, env2)
382 (pairs2, env2) -> (Just (StgRec pairs2), env2)
383
384 stgCsePairs :: CseEnv -> [(OutId, InStgRhs)] -> ([(OutId, OutStgRhs)], CseEnv)
385 stgCsePairs env [] = ([], env)
386 stgCsePairs env0 ((b,e):pairs)
387 = let (pairMB, env1) = stgCseRhs env0 b e
388 (pairs', env2) = stgCsePairs env1 pairs
389 in (pairMB `mbCons` pairs', env2)
390 where
391 mbCons = maybe id (:)
392
393 -- The RHS of a binding.
394 -- If it is a constructor application, either short-cut it or extend the environment
395 stgCseRhs :: CseEnv -> OutId -> InStgRhs -> (Maybe (OutId, OutStgRhs), CseEnv)
396 stgCseRhs env bndr (StgRhsCon ccs dataCon args)
397 | Just other_bndr <- envLookup (Lax dataCon) args' env
398 = let env' = addSubst bndr other_bndr env
399 in (Nothing, env')
400 | otherwise
401 = let env' = addDataCon bndr (Lax dataCon) args' env
402 -- see note [Case 1: CSEing allocated closures]
403 pair = (bndr, StgRhsCon ccs dataCon args')
404 in (Just pair, env')
405 where args' = substArgs env args
406 stgCseRhs env bndr (StgRhsClosure ccs info occs upd args body)
407 = let (env1, args') = substBndrs env args
408 env2 = forgetCse env1 -- See note [Free variables of an StgClosure]
409 body' = stgCseExpr env2 body
410 in (Just (substVar env bndr, StgRhsClosure ccs info occs' upd args' body'), env)
411 where occs' = substVars env occs
412
413
414 mkStgCase :: StgExpr -> OutId -> AltType -> [StgAlt] -> StgExpr
415 mkStgCase scrut bndr ty alts | all isBndr alts = scrut
416 | otherwise = StgCase scrut bndr ty alts
417
418 where
419 -- see Note [All alternatives are the binder]
420 isBndr (_, _, StgApp f []) = f == bndr
421 isBndr _ = False
422
423
424 -- Utilities
425
426 -- | This function short-cuts let-bindings that are now obsolete
427 mkStgLet :: (a -> b -> b) -> Maybe a -> b -> b
428 mkStgLet _ Nothing body = body
429 mkStgLet stgLet (Just binds) body = stgLet binds body
430
431
432 {-
433 Note [All alternatives are the binder]
434 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
435
436 When all alternatives simply refer to the case binder, then we do not have
437 to bother with the case expression at all (#13588). CoreSTG does this as well,
438 but sometimes, types get into the way:
439
440 newtype T = MkT Int
441 f :: (Int, Int) -> (T, Int)
442 f (x, y) = (MkT x, y)
443
444 Core cannot just turn this into
445
446 f p = p
447
448 as this would not be well-typed. But to STG, where MkT is no longer in the way,
449 we can.
450
451 Note [Trivial case scrutinee]
452 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
453 We want to be able to handle nested reconstruction of constructors as in
454
455 nested :: Either Int (Either Int a) -> Either Bool (Either Bool a)
456 nested (Right (Right v)) = Right (Right v)
457 nested _ = Left True
458
459 So if we come across
460
461 case x of r1
462 Right a -> case a of r2
463 Right b -> let v = Right b
464 in Right v
465
466 we first replace v with r2. Next we want to replace Right r2 with r1. But the
467 ce_conAppMap contains Right a!
468
469 Therefore, we add r1 ↦ x to ce_bndrMap when analysing the outer case, and use
470 this substitution before looking Right r2 up in ce_conAppMap, and everything
471 works out.
472
473 Note [Free variables of an StgClosure]
474 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
475 StgClosures (function and thunks) have an explicit list of free variables:
476
477 foo [x] =
478 let not_a_free_var = Left [x]
479 let a_free_var = Right [x]
480 let closure = \[x a_free_var] -> \[y] -> bar y (Left [x]) a_free_var
481 in closure
482
483 If we were to CSE `Left [x]` in the body of `closure` with `not_a_free_var`,
484 then the list of free variables would be wrong, so for now, we do not CSE
485 across such a closure, simply because I (Joachim) was not sure about possible
486 knock-on effects. If deemed safe and worth the slight code complication of
487 re-calculating this list during or after this pass, this can surely be done.
488 -}