02b3891b7be8e1739388aa3d20b0bad03aee306b
[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, nameUnique)
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, _) | traceLookup dataCon = undefined
139 lookupTM (dataCon, args) = un_cam >.> lkDNamed dataCon >=> lookupTM args
140 alterTM (dataCon, args) f m =
141 m { un_cam = un_cam m |> xtDNamed dataCon |>> alterTM args f }
142 foldTM k = un_cam >.> foldTM (foldTM k)
143 mapTM f = un_cam >.> mapTM (mapTM f) >.> CAM
144
145 traceLookup :: LaxDataCon -> Bool
146 traceLookup _ = False
147 {-
148 traceLookup l@(Lax dc) = pprTrace "lookupTM" (ppr dc <> (if getKey u < 0 then text " -" else text " ") <> ppr u') False
149 where u = nameUnique . getName $ l
150 u' = mkUniqueGrimily (abs(getKey u))
151 -}
152 {-# NOINLINE traceLookup #-}
153
154 -----------------
155 -- The CSE Env --
156 -----------------
157
158 -- | The CSE environment. See note [CseEnv Example]
159 data CseEnv = CseEnv
160 { ce_conAppMap :: ConAppMap OutId
161 -- ^ The main component of the environment is the trie that maps
162 -- data constructor applications (with their `OutId` arguments)
163 -- to an in-scope name that can be used instead.
164 -- This name is always either a let-bound variable or a case binder.
165 , ce_subst :: IdEnv OutId
166 -- ^ This substitution is applied to the code as we traverse it.
167 -- Entries have one of two reasons:
168 --
169 -- * The input might have shadowing (see Note [Shadowing]), so we have
170 -- to rename some binders as we traverse the tree.
171 -- * If we remove `let x = Con z` because `let y = Con z` is in scope,
172 -- we note this here as x ↦ y.
173 , ce_bndrMap :: IdEnv OutId
174 -- ^ If we come across a case expression case x as b of … with a trivial
175 -- binder, we add b ↦ x to this.
176 -- This map is *only* used when looking something up in the ce_conAppMap.
177 -- See Note [Trivial case scrutinee]
178 , ce_in_scope :: InScopeSet
179 -- ^ The third component is an in-scope set, to rename away any
180 -- shadowing binders
181 }
182
183 {-|
184 Note [CseEnv Example]
185 ~~~~~~~~~~~~~~~~~~~~~
186 The following tables shows how the CseEnvironment changes as code is traversed,
187 as well as the changes to that code.
188
189 InExpr OutExpr
190 conAppMap subst in_scope
191 ───────────────────────────────────────────────────────────
192 -- empty {} {}
193 case … as a of {Con x y -> case … as a of {Con x y ->
194 -- Con x y ↦ a {} {a,x,y}
195 let b = Con x y (removed)
196 -- Con x y ↦ a b↦a {a,x,y,b}
197 let c = Bar a let c = Bar a
198 -- Con x y ↦ a, Bar a ↦ c b↦a {a,x,y,b,c}
199 let c = some expression let c' = some expression
200 -- Con x y ↦ a, Bar a ↦ c b↦a, c↦c', {a,x,y,b,c,c'}
201 let d = Bar b (removed)
202 -- Con x y ↦ a, Bar a ↦ c b↦a, c↦c', d↦c {a,x,y,b,c,c',d}
203 (a, b, c d) (a, a, c' c)
204 -}
205
206 initEnv :: InScopeSet -> CseEnv
207 initEnv in_scope = CseEnv
208 { ce_conAppMap = emptyTM
209 , ce_subst = emptyVarEnv
210 , ce_bndrMap = emptyVarEnv
211 , ce_in_scope = in_scope
212 }
213
214 envLookup :: LaxDataCon -> [OutStgArg] -> CseEnv -> Maybe OutId
215 envLookup dataCon args env = lookupTM (dataCon, args') (ce_conAppMap env)
216 where args' = map go args -- See Note [Trivial case scrutinee]
217 go (StgVarArg v ) = StgVarArg (fromMaybe v $ lookupVarEnv (ce_bndrMap env) v)
218 go (StgLitArg lit) = StgLitArg lit
219
220 addDataCon :: OutId -> LaxDataCon -> [OutStgArg] -> CseEnv -> CseEnv
221 -- do not bother with nullary data constructors, they are static anyways
222 addDataCon bndr dataCon [] env = env { ce_conAppMap = new_env }
223 where new_env = alterTM (dataCon, []) (\case Nothing -> pure bndr; p -> p) (ce_conAppMap env)
224 --addDataCon _ _ [] env = env
225 addDataCon bndr dataCon args env = env { ce_conAppMap = new_env }
226 where
227 new_env = insertTM (dataCon, args) bndr (ce_conAppMap env)
228
229 forgetCse :: CseEnv -> CseEnv
230 forgetCse env = env { ce_conAppMap = emptyTM }
231 -- See note [Free variables of an StgClosure]
232
233 addSubst :: OutId -> OutId -> CseEnv -> CseEnv
234 addSubst from to env
235 = env { ce_subst = extendVarEnv (ce_subst env) from to }
236
237 addTrivCaseBndr :: OutId -> OutId -> CseEnv -> CseEnv
238 addTrivCaseBndr from to env
239 = env { ce_bndrMap = extendVarEnv (ce_bndrMap env) from to }
240
241 substArgs :: CseEnv -> [InStgArg] -> [OutStgArg]
242 substArgs env = map (substArg env)
243
244 substArg :: CseEnv -> InStgArg -> OutStgArg
245 substArg env (StgVarArg from) = StgVarArg (substVar env from)
246 substArg _ (StgLitArg lit) = StgLitArg lit
247
248 substVars :: CseEnv -> [InId] -> [OutId]
249 substVars env = map (substVar env)
250
251 substVar :: CseEnv -> InId -> OutId
252 substVar env id = fromMaybe id $ lookupVarEnv (ce_subst env) id
253
254 -- Functions to enter binders
255
256 -- This is much simpler than the equivalent code in CoreSubst:
257 -- * We do not substitute type variables, and
258 -- * There is nothing relevant in IdInfo at this stage
259 -- that needs substitutions.
260 -- Therefore, no special treatment for a recursive group is required.
261
262 substBndr :: CseEnv -> InId -> (CseEnv, OutId)
263 substBndr env old_id
264 = (new_env, new_id)
265 where
266 new_id = uniqAway (ce_in_scope env) old_id
267 no_change = new_id == old_id
268 env' = env { ce_in_scope = ce_in_scope env `extendInScopeSet` new_id }
269 new_env | no_change = env' { ce_subst = extendVarEnv (ce_subst env) old_id new_id }
270 | otherwise = env'
271
272 substBndrs :: CseEnv -> [InVar] -> (CseEnv, [OutVar])
273 substBndrs env bndrs = mapAccumL substBndr env bndrs
274
275 substPairs :: CseEnv -> [(InVar, a)] -> (CseEnv, [(OutVar, a)])
276 substPairs env bndrs = mapAccumL go env bndrs
277 where go env (id, x) = let (env', id') = substBndr env id
278 in (env', (id', x))
279
280 -- Main entry point
281
282 stgCse :: [InStgTopBinding] -> [OutStgTopBinding]
283 stgCse binds = snd $ mapAccumL stgCseTopLvl emptyInScopeSet binds
284
285 -- Top level bindings.
286 --
287 -- We do not CSE these, as top-level closures are allocated statically anyways.
288 -- Also, they might be exported.
289 -- But we still have to collect the set of in-scope variables, otherwise
290 -- uniqAway might shadow a top-level closure.
291
292 stgCseTopLvl :: InScopeSet -> InStgTopBinding -> (InScopeSet, OutStgTopBinding)
293 stgCseTopLvl in_scope t@(StgTopStringLit _ _) = (in_scope, t)
294 stgCseTopLvl in_scope (StgTopLifted (StgNonRec bndr rhs))
295 = (in_scope'
296 , StgTopLifted (StgNonRec bndr (stgCseTopLvlRhs in_scope rhs)))
297 where in_scope' = in_scope `extendInScopeSet` bndr
298
299 stgCseTopLvl in_scope (StgTopLifted (StgRec eqs))
300 = ( in_scope'
301 , StgTopLifted (StgRec [ (bndr, stgCseTopLvlRhs in_scope' rhs) | (bndr, rhs) <- eqs ]))
302 where in_scope' = in_scope `extendInScopeSetList` [ bndr | (bndr, _) <- eqs ]
303
304 stgCseTopLvlRhs :: InScopeSet -> InStgRhs -> OutStgRhs
305 stgCseTopLvlRhs in_scope (StgRhsClosure ccs info occs upd args body)
306 = let body' = stgCseExpr (initEnv in_scope) body
307 in StgRhsClosure ccs info occs upd args body'
308 stgCseTopLvlRhs _ (StgRhsCon ccs dataCon args)
309 = StgRhsCon ccs dataCon args
310
311 ------------------------------
312 -- The actual AST traversal --
313 ------------------------------
314
315 -- Trivial cases
316 stgCseExpr :: CseEnv -> InStgExpr -> OutStgExpr
317 stgCseExpr env (StgApp fun args)
318 = StgApp fun' args'
319 where fun' = substVar env fun
320 args' = substArgs env args
321 stgCseExpr _ (StgLit lit)
322 = StgLit lit
323 stgCseExpr env (StgOpApp op args tys)
324 = StgOpApp op args' tys
325 where args' = substArgs env args
326 stgCseExpr _ (StgLam _ _)
327 = pprPanic "stgCseExp" (text "StgLam")
328 stgCseExpr env (StgTick tick body)
329 = let body' = stgCseExpr env body
330 in StgTick tick body'
331 stgCseExpr env (StgCase scrut bndr ty alts)
332 = mkStgCase scrut' bndr' ty alts'
333 where
334 scrut' = stgCseExpr env scrut
335 (env1, bndr') = substBndr env bndr
336 env2 | StgApp trivial_scrut [] <- scrut' = addTrivCaseBndr bndr trivial_scrut env1
337 -- See Note [Trivial case scrutinee]
338 | otherwise = env1
339 alts' = map (stgCseAlt env2 bndr') alts
340
341
342 -- A constructor application.
343 -- To be removed by a variable use when found in the CSE environment
344 stgCseExpr env orig@(StgConApp dataCon args tys)
345 | Just bndr' <- envLookup dc args' env
346 = (if getKey u < 0 then pprTrace "stgCseExpr" (ppr dataCon <+> text (show $ length (dataConOrigArgTys dataCon)) <+> (text . show $ tyConFamilySize (dataConTyCon dataCon))) else id) $ StgApp bndr' []
347 | otherwise
348 = StgConApp dataCon args' tys
349 where args' = substArgs env args
350 dc = Lax dataCon
351 u = getUnique (getName dc)
352
353 -- Let bindings
354 -- The binding might be removed due to CSE (we do not want trivial bindings on
355 -- the STG level), so use the smart constructor `mkStgLet` to remove the binding
356 -- if empty.
357 stgCseExpr env (StgLet binds body)
358 = let (binds', env') = stgCseBind env binds
359 body' = stgCseExpr env' body
360 in mkStgLet StgLet binds' body'
361 stgCseExpr env (StgLetNoEscape binds body)
362 = let (binds', env') = stgCseBind env binds
363 body' = stgCseExpr env' body
364 in mkStgLet StgLetNoEscape binds' body'
365
366 -- Case alternatives
367 -- Extend the CSE environment
368 stgCseAlt :: CseEnv -> OutId -> InStgAlt -> OutStgAlt
369 stgCseAlt env case_bndr (DataAlt dataCon, args, rhs)
370 = let (env1, args') = substBndrs env args
371 env2 = addDataCon case_bndr (Lax dataCon) (map StgVarArg args') env1
372 -- see note [Case 2: CSEing case binders]
373 rhs' = stgCseExpr env2 rhs
374 in (DataAlt dataCon, args', rhs')
375 stgCseAlt env _ (altCon, args, rhs)
376 = let (env1, args') = substBndrs env args
377 rhs' = stgCseExpr env1 rhs
378 in (altCon, args', rhs')
379
380 -- Bindings
381 stgCseBind :: CseEnv -> InStgBinding -> (Maybe OutStgBinding, CseEnv)
382 stgCseBind env (StgNonRec b e)
383 = let (env1, b') = substBndr env b
384 in case stgCseRhs env1 b' e of
385 (Nothing, env2) -> (Nothing, env2)
386 (Just (b2,e'), env2) -> (Just (StgNonRec b2 e'), env2)
387 stgCseBind env (StgRec pairs)
388 = let (env1, pairs1) = substPairs env pairs
389 in case stgCsePairs env1 pairs1 of
390 ([], env2) -> (Nothing, env2)
391 (pairs2, env2) -> (Just (StgRec pairs2), env2)
392
393 stgCsePairs :: CseEnv -> [(OutId, InStgRhs)] -> ([(OutId, OutStgRhs)], CseEnv)
394 stgCsePairs env [] = ([], env)
395 stgCsePairs env0 ((b,e):pairs)
396 = let (pairMB, env1) = stgCseRhs env0 b e
397 (pairs', env2) = stgCsePairs env1 pairs
398 in (pairMB `mbCons` pairs', env2)
399 where
400 mbCons = maybe id (:)
401
402 -- The RHS of a binding.
403 -- If it is a constructor application, either short-cut it or extend the environment
404 stgCseRhs :: CseEnv -> OutId -> InStgRhs -> (Maybe (OutId, OutStgRhs), CseEnv)
405 stgCseRhs env bndr (StgRhsCon ccs dataCon args)
406 | Just other_bndr <- envLookup (Lax dataCon) args' env
407 = let env' = addSubst bndr other_bndr env
408 in (Nothing, env')
409 | otherwise
410 = let env' = addDataCon bndr (Lax dataCon) args' env
411 -- see note [Case 1: CSEing allocated closures]
412 pair = (bndr, StgRhsCon ccs dataCon args')
413 in (Just pair, env')
414 where args' = substArgs env args
415 stgCseRhs env bndr (StgRhsClosure ccs info occs upd args body)
416 = let (env1, args') = substBndrs env args
417 env2 = forgetCse env1 -- See note [Free variables of an StgClosure]
418 body' = stgCseExpr env2 body
419 in (Just (substVar env bndr, StgRhsClosure ccs info occs' upd args' body'), env)
420 where occs' = substVars env occs
421
422
423 mkStgCase :: StgExpr -> OutId -> AltType -> [StgAlt] -> StgExpr
424 mkStgCase scrut bndr ty alts | all isBndr alts = scrut
425 | otherwise = StgCase scrut bndr ty alts
426
427 where
428 -- see Note [All alternatives are the binder]
429 isBndr (_, _, StgApp f []) = f == bndr
430 isBndr _ = False
431
432
433 -- Utilities
434
435 -- | This function short-cuts let-bindings that are now obsolete
436 mkStgLet :: (a -> b -> b) -> Maybe a -> b -> b
437 mkStgLet _ Nothing body = body
438 mkStgLet stgLet (Just binds) body = stgLet binds body
439
440
441 {-
442 Note [All alternatives are the binder]
443 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
444
445 When all alternatives simply refer to the case binder, then we do not have
446 to bother with the case expression at all (#13588). CoreSTG does this as well,
447 but sometimes, types get into the way:
448
449 newtype T = MkT Int
450 f :: (Int, Int) -> (T, Int)
451 f (x, y) = (MkT x, y)
452
453 Core cannot just turn this into
454
455 f p = p
456
457 as this would not be well-typed. But to STG, where MkT is no longer in the way,
458 we can.
459
460 Note [Trivial case scrutinee]
461 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
462 We want to be able to handle nested reconstruction of constructors as in
463
464 nested :: Either Int (Either Int a) -> Either Bool (Either Bool a)
465 nested (Right (Right v)) = Right (Right v)
466 nested _ = Left True
467
468 So if we come across
469
470 case x of r1
471 Right a -> case a of r2
472 Right b -> let v = Right b
473 in Right v
474
475 we first replace v with r2. Next we want to replace Right r2 with r1. But the
476 ce_conAppMap contains Right a!
477
478 Therefore, we add r1 ↦ x to ce_bndrMap when analysing the outer case, and use
479 this substitution before looking Right r2 up in ce_conAppMap, and everything
480 works out.
481
482 Note [Free variables of an StgClosure]
483 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
484 StgClosures (function and thunks) have an explicit list of free variables:
485
486 foo [x] =
487 let not_a_free_var = Left [x]
488 let a_free_var = Right [x]
489 let closure = \[x a_free_var] -> \[y] -> bar y (Left [x]) a_free_var
490 in closure
491
492 If we were to CSE `Left [x]` in the body of `closure` with `not_a_free_var`,
493 then the list of free variables would be wrong, so for now, we do not CSE
494 across such a closure, simply because I (Joachim) was not sure about possible
495 knock-on effects. If deemed safe and worth the slight code complication of
496 re-calculating this list during or after this pass, this can surely be done.
497 -}