[Git][ghc/ghc][wip/T22274] Identify exit cases in OccurAnal

Sebastian Graf (@sgraf812) gitlab at gitlab.haskell.org
Mon Nov 7 13:28:59 UTC 2022



Sebastian Graf pushed to branch wip/T22274 at Glasgow Haskell Compiler / GHC


Commits:
89c11517 by Sebastian Graf at 2022-11-07T14:28:26+01:00
Identify exit cases in OccurAnal

Also had to mark a few key WordArray functions as INLINE so that they don't allocate a closure for the continuation.

- - - - -


3 changed files:

- compiler/GHC/Core/Opt/OccurAnal.hs
- compiler/GHC/Types/Unique/SlimSet.hs
- libraries/ghc-bignum/src/GHC/Num/WordArray.hs


Changes:

=====================================
compiler/GHC/Core/Opt/OccurAnal.hs
=====================================
@@ -34,7 +34,7 @@ import GHC.Core.Predicate   ( isDictId )
 import GHC.Core.Type
 import GHC.Core.TyCo.FVs    ( tyCoVarsOfMCo )
 
-import GHC.Data.Maybe( isJust, orElse )
+import GHC.Data.Maybe( isJust, orElse, mapMaybe )
 import GHC.Data.Graph.Directed ( SCC(..), Node(..)
                                , stronglyConnCompFromEdgedVerticesUniq
                                , stronglyConnCompFromEdgedVerticesUniqR )
@@ -48,6 +48,7 @@ import GHC.Types.Tickish
 import GHC.Types.Var.Set
 import GHC.Types.Var.Env
 import GHC.Types.Var
+import GHC.Types.Unique.SlimSet
 import GHC.Types.Demand ( argOneShots, argsOneShots )
 
 import GHC.Utils.Outputable
@@ -59,6 +60,8 @@ import GHC.Utils.Trace
 import GHC.Builtin.Names( runRWKey )
 import GHC.Unit.Module( Module )
 
+import Data.IntMap.Strict (IntMap)
+import qualified Data.IntMap.Strict as IntMap
 import Data.List (mapAccumL, mapAccumR)
 import Data.List.NonEmpty (NonEmpty (..), nonEmpty)
 import qualified Data.List.NonEmpty as NE
@@ -831,7 +834,7 @@ occAnalRecBind !env lvl imp_rule_edges pairs body_usage
 
     bndrs    = map fst pairs
     bndr_set = mkVarSet bndrs
-    rhs_env  = env `addInScope` bndrs
+    rhs_env  = env `addInScope` bndrs `enterLoop` bndrs
 
 
 -----------------------------
@@ -2299,7 +2302,7 @@ occAnalApp env (Var fun_id, args, ticks)
     !(fun', fun_id')  = lookupBndrSwap env fun_id
     !(WithUsageDetails args_uds app') = occAnalArgs env fun' args one_shots
 
-    fun_uds = mkOneOcc fun_id' int_cxt n_args
+    fun_uds = mkOneOcc env fun_id' int_cxt n_args
        -- NB: fun_uds is computed for fun_id', not fun_id
        -- See (BS1) in Note [The binder-swap substitution]
 
@@ -2459,6 +2462,8 @@ scrutinised y).
 data OccEnv
   = OccEnv { occ_encl       :: !OccEncl      -- Enclosing context information
            , occ_one_shots  :: !OneShots     -- See Note [OneShots]
+           , occ_cur_lvl    :: !LoopLevel         -- ^ Current loop level
+           , occ_loops      :: !(IdEnv LoopLevel) -- ^ The loop levels of enclosing letrec binders
            , occ_unf_act    :: Id -> Bool          -- Which Id unfoldings are active
            , occ_rule_act   :: Activation -> Bool  -- Which rules are active
              -- See Note [Finding rule RHS free vars]
@@ -2507,6 +2512,8 @@ initOccEnv :: OccEnv
 initOccEnv
   = OccEnv { occ_encl      = OccVanilla
            , occ_one_shots = []
+           , occ_cur_lvl   = 0
+           , occ_loops     = emptyVarEnv
 
                  -- To be conservative, we say that all
                  -- inlines and rules are active
@@ -2546,6 +2553,11 @@ isRhsEnv (OccEnv { occ_encl = cxt }) = case cxt of
                                           OccRhs -> True
                                           _      -> False
 
+lookupLoopLevel :: OccEnv -> Id -> Int
+lookupLoopLevel (OccEnv { occ_loops = loops }) id
+  | Just lvl <- lookupVarEnv loops id = lvl
+  | otherwise                         = 0
+
 addOneInScope :: OccEnv -> CoreBndr -> OccEnv
 addOneInScope env@(OccEnv { occ_bs_env = swap_env, occ_bs_rng = rng_vars }) bndr
   | bndr `elemVarSet` rng_vars = env { occ_bs_env = emptyVarEnv, occ_bs_rng = emptyVarSet }
@@ -2559,6 +2571,24 @@ addInScope env@(OccEnv { occ_bs_env = swap_env, occ_bs_rng = rng_vars }) bndrs
   | any (`elemVarSet` rng_vars) bndrs = env { occ_bs_env = emptyVarEnv, occ_bs_rng = emptyVarSet }
   | otherwise                         = env { occ_bs_env = swap_env `delVarEnvList` bndrs }
 
+enterLoop :: OccEnv -> [Var] -> OccEnv
+enterLoop env vs
+  = env { occ_cur_lvl = new_lvl
+        , occ_loops   = extendVarEnvList (occ_loops env) [(v,new_lvl) | v<-vs] }
+  where
+    new_lvl = occ_cur_lvl env + 1
+
+leaveLoop :: UsageDetails -> UsageDetails
+leaveLoop ud = alterLoopLevelInfo ud f
+  where
+    f lli
+      | Just ((lvl, new_max_occs), inv') <- IntMap.maxViewWithKey (lli_inv lli)
+      = lli { lli_max = lvl
+            , lli_inv = inv'
+            , lli_non_max_lvls = nonDetFoldUniqSlimSet (\u nml -> delFromUFM_Directly nml u) (lli_non_max_lvls lli) new_max_occs
+            }
+      | otherwise
+      = emptyLoopLevelInfo
 
 --------------------
 transClosureFV :: VarEnv VarSet -> VarEnv VarSet
@@ -2976,14 +3006,42 @@ info then simply means setting the corresponding zapped set to the whole
 'OccInfoEnv', a fast O(1) operation.
 -}
 
-type OccInfoEnv = IdEnv OccInfo -- A finite map from ids to their usage
-                -- INVARIANT: never IAmDead
-                -- (Deadness is signalled by not being in the map at all)
+type LoopLevel = Int
+
+-- | Level 0 is the loop level we never exit. Every letrec binder will have loop
+-- level at least 1.
+notLooping :: LoopLevel
+notLooping = 0
+
+type LoopLevelMap = IntMap
+
+type OccInfoEnv = IdEnv OccInfo
+  -- ^ A finite map from ids to their usage.
+  -- INVARIANT: The OccInfo is never IAmDead
+  -- (Deadness is signalled by not being in the map at all)
 
 type ZappedSet = OccInfoEnv -- Values are ignored
 
+-- | Represents an efficient bidirectional mapping between occuring 'Id's
+-- and the maximum 'LoopLevel' of the recursive binders with which they
+-- co-occur.
+data LoopLevelInfo
+  = LLI { lli_max :: !LoopLevel
+        -- ^ Maximum loop level of a rec binder occuring in the expression
+        , lli_non_max_lvls :: !(IdEnv Int)
+        -- ^ Binders that (are not dead, and) do not occur at loop level
+        -- 'lli_max' will have their loop-level stated here.
+        , lli_inv :: !(LoopLevelMap VarSlimSet)
+        -- ^ Inverse mapping of 'lli_non_max_lvls'.
+        -- If a binder has max loop level l, it will be regarded as "used on an
+        -- exit path" wrt. the loop with level l.
+        -- INVARIANT: The sets for different levels are disjoint
+        }
+
+
 data UsageDetails
   = UD { ud_env       :: !OccInfoEnv
+       , ud_loop_info :: !LoopLevelInfo
        , ud_z_many    :: !ZappedSet   -- apply 'markMany' to these
        , ud_z_in_lam  :: !ZappedSet   -- apply 'markInsideLam' to these
        , ud_z_no_tail :: !ZappedSet } -- apply 'markNonTail' to these
@@ -2997,13 +3055,40 @@ instance Outputable UsageDetails where
 
 andUDs, orUDs
         :: UsageDetails -> UsageDetails -> UsageDetails
-andUDs = combineUsageDetailsWith addOccInfo
-orUDs  = combineUsageDetailsWith orOccInfo
+andUDs = combineUsageDetailsWith addOccInfo andLoopLevelInfo
+orUDs  = combineUsageDetailsWith orOccInfo orLoopLevelInfo
+
+andLoopLevelInfo :: LoopLevel -> OccInfoEnv -> LoopLevelInfo -> LoopLevelInfo
+andLoopLevelInfo lvl _occs lli
+  | lvl >= lli_max lli = LLI { lli_max = lvl, lli_non_max_lvls = emptyVarEnv, lli_inv = IntMap.empty }
+  | otherwise          = LLI { lli_max = lli_max lli
+                             , lli_non_max_lvls = non_max_lvls'
+                             , lli_inv = inv'
+                             }
+  where
+    (lower, mb_exact, higher) = IntMap.splitLookup lvl (lli_inv lli)
+    raised_vars = IntMap.foldr unionUniqSlimSet (mb_exact `orElse` emptyUniqSlimSet) lower
+    inv' = IntMap.insert lvl raised_vars higher
+    non_max_lvls' = nonDetFoldUniqSlimSet (\u lvls -> addToUFM_Directly lvls u lvl) (lli_non_max_lvls lli) raised_vars
+
+orLoopLevelInfo :: LoopLevel -> OccInfoEnv -> LoopLevelInfo -> LoopLevelInfo
+orLoopLevelInfo lvl occs lli
+  | lvl <= lli_max lli = lli
+  | otherwise          = LLI { lli_max = lvl
+                             , lli_non_max_lvls = non_max_lvls'
+                             , lli_inv = inv'
+                             }
+  where
+    old_max_occs = occs `minusVarEnv` lli_non_max_lvls lli
+    inv' = IntMap.insert lvl (ufmDom old_max_occs) (lli_inv lli)
+    non_max_lvls' = mapVarEnv (const lvl) old_max_occs `plusVarEnv` lli_non_max_lvls lli -- NB: plusVarEnv is right-biased
 
-mkOneOcc :: Id -> InterestingCxt -> JoinArity -> UsageDetails
-mkOneOcc id int_cxt arity
+mkOneOcc :: OccEnv -> Id -> InterestingCxt -> JoinArity -> UsageDetails
+mkOneOcc env id int_cxt arity
   | isLocalId id
-  = emptyDetails { ud_env = unitVarEnv id occ_info }
+  , let !lvl = lookupLoopLevel env id
+  = emptyDetails { ud_env = unitVarEnv id occ_info
+                 , ud_loop_info = emptyLoopLevelInfo { lli_max = lvl } }
   | otherwise
   = emptyDetails
   where
@@ -3038,14 +3123,36 @@ addLamCoVarOccs uds bndrs
 
 delDetails :: UsageDetails -> Id -> UsageDetails
 delDetails ud bndr
-  = ud `alterUsageDetails` (`delVarEnv` bndr)
+  = ud `alterUsageDetails` (`delVarEnv` bndr) `alterLoopLevelInfo` (`delLoopLevel` bndr)
 
 delDetailsList :: UsageDetails -> [Id] -> UsageDetails
 delDetailsList ud bndrs
-  = ud `alterUsageDetails` (`delVarEnvList` bndrs)
+  = ud `alterUsageDetails` (`delVarEnvList` bndrs) `alterLoopLevelInfo` (`delLoopLevelList` bndrs)
+
+delLoopLevel :: LoopLevelInfo -> Id -> LoopLevelInfo
+delLoopLevel lli@(LLI { lli_non_max_lvls = nml, lli_inv = inv }) id
+  | Just lvl <- lookupVarEnv (lli_non_max_lvls lli) id
+  = lli { lli_non_max_lvls = delVarEnv nml id
+        , lli_inv = IntMap.adjust (`delUniqSlimSet` id) lvl inv }
+  | otherwise
+  = lli
+
+delLoopLevelList :: LoopLevelInfo -> [Id] -> LoopLevelInfo
+delLoopLevelList lli@(LLI { lli_non_max_lvls = nml, lli_inv = inv }) ids
+  = lli { lli_non_max_lvls = delVarEnvList nml ids
+        , lli_inv = foldr (IntMap.adjust (`minusUniqSlimSet` ids_set)) inv lvls }
+  where
+    ids_set = mkUniqSlimSet ids
+    lvls = mapMaybe (lookupVarEnv (lli_non_max_lvls lli)) ids
+
+emptyLoopLevelInfo :: LoopLevelInfo
+emptyLoopLevelInfo = LLI { lli_max = notLooping
+                         , lli_non_max_lvls = emptyVarEnv
+                         , lli_inv = IntMap.empty }
 
 emptyDetails :: UsageDetails
 emptyDetails = UD { ud_env       = emptyVarEnv
+                  , ud_loop_info = emptyLoopLevelInfo
                   , ud_z_many    = emptyVarEnv
                   , ud_z_in_lam  = emptyVarEnv
                   , ud_z_no_tail = emptyVarEnv }
@@ -3059,7 +3166,7 @@ markAllMany          ud = ud { ud_z_many    = ud_env ud }
 markAllInsideLam     ud = ud { ud_z_in_lam  = ud_env ud }
 markAllNonTail ud = ud { ud_z_no_tail = ud_env ud }
 
-markAllInsideLamIf, markAllNonTailIf :: Bool -> UsageDetails -> UsageDetails
+markAllInsideLamIf, markMaxLvlInsideLamIf, markAllNonTailIf :: Bool -> UsageDetails -> UsageDetails
 
 markAllInsideLamIf  True  ud = markAllInsideLam ud
 markAllInsideLamIf  False ud = ud
@@ -3067,6 +3174,10 @@ markAllInsideLamIf  False ud = ud
 markAllNonTailIf True  ud = markAllNonTail ud
 markAllNonTailIf False ud = ud
 
+markMaxLvlInsideLamIf False ud = ud
+markMaxLvlInsideLamIf True  ud@(UD{ud_loop_info=LLI{lli_non_max_lvls=nml}}) =
+  ud { ud_z_in_lam = ud_env ud `minusVarEnv` nml }
+
 
 markAllManyNonTail = markAllMany . markAllNonTail -- effectively sets to noOccInfo
 
@@ -3090,16 +3201,32 @@ restrictFreeVars bndrs fvs = restrictUniqSetToUFM bndrs fvs
 -- Auxiliary functions for UsageDetails implementation
 
 combineUsageDetailsWith :: (OccInfo -> OccInfo -> OccInfo)
+                        -> (LoopLevel -> OccInfoEnv -> LoopLevelInfo -> LoopLevelInfo)
                         -> UsageDetails -> UsageDetails -> UsageDetails
-combineUsageDetailsWith plus_occ_info ud1 ud2
+combineUsageDetailsWith plus_occ_info bump_loop_info ud1 ud2
   | isEmptyDetails ud1 = ud2
   | isEmptyDetails ud2 = ud1
   | otherwise
   = UD { ud_env       = plusVarEnv_C plus_occ_info (ud_env ud1) (ud_env ud2)
+       , ud_loop_info = combineLoopLevelInfoWith bump_loop_info (ud_env ud1) (ud_loop_info ud1) (ud_env ud2) (ud_loop_info ud2)
        , ud_z_many    = plusVarEnv (ud_z_many    ud1) (ud_z_many    ud2)
        , ud_z_in_lam  = plusVarEnv (ud_z_in_lam  ud1) (ud_z_in_lam  ud2)
        , ud_z_no_tail = plusVarEnv (ud_z_no_tail ud1) (ud_z_no_tail ud2) }
 
+combineLoopLevelInfoWith :: (LoopLevel -> OccInfoEnv -> LoopLevelInfo -> LoopLevelInfo)
+                         -> OccInfoEnv -> LoopLevelInfo
+                         -> OccInfoEnv -> LoopLevelInfo
+                         -> LoopLevelInfo
+combineLoopLevelInfoWith bump_loop_info u1 lli1 u2 lli2
+  = assert (lli_max lli1' == lli_max lli2') $
+    LLI { lli_max = lli_max lli1 `max` lli_max lli2
+        , lli_non_max_lvls = plusVarEnv_C max (lli_non_max_lvls lli1') (lli_non_max_lvls lli2')
+        , lli_inv = IntMap.unionWith unionUniqSlimSet (lli_inv lli1') (lli_inv lli2')
+        }
+  where
+    lli1' = bump_loop_info (lli_max lli2) u1 lli1
+    lli2' = bump_loop_info (lli_max lli1) u2 lli2
+
 doZapping :: UsageDetails -> Var -> OccInfo -> OccInfo
 doZapping ud var occ
   = doZappingByUnique ud (varUnique var) occ
@@ -3119,14 +3246,19 @@ doZappingByUnique (UD { ud_z_many = many
 
 alterUsageDetails :: UsageDetails -> (OccInfoEnv -> OccInfoEnv) -> UsageDetails
 alterUsageDetails !ud f
-  = UD { ud_env       = f (ud_env       ud)
+  = ud { ud_env       = f (ud_env       ud)
        , ud_z_many    = f (ud_z_many    ud)
        , ud_z_in_lam  = f (ud_z_in_lam  ud)
        , ud_z_no_tail = f (ud_z_no_tail ud) }
 
+alterLoopLevelInfo :: UsageDetails -> (LoopLevelInfo -> LoopLevelInfo) -> UsageDetails
+alterLoopLevelInfo !ud f
+  = ud { ud_loop_info = f (ud_loop_info ud) }
+
 flattenUsageDetails :: UsageDetails -> UsageDetails
 flattenUsageDetails ud@(UD { ud_env = env })
   = UD { ud_env       = mapUFM_Directly (doZappingByUnique ud) env
+       , ud_loop_info = ud_loop_info ud
        , ud_z_many    = emptyVarEnv
        , ud_z_in_lam  = emptyVarEnv
        , ud_z_no_tail = emptyVarEnv }
@@ -3139,7 +3271,7 @@ adjustRhsUsage :: Maybe JoinArity
                -> UsageDetails
 adjustRhsUsage mb_join_arity rhs usage
   = -- c.f. occAnal (Lam {})
-    markAllInsideLamIf (not one_shot) $
+    markMaxLvlInsideLamIf (not one_shot) $
     markAllNonTailIf (not exact_join) $
     usage
   where
@@ -3249,7 +3381,7 @@ tagRecBinders lvl body_uds details_s
                  | bndr <- bndrs ]
 
      -- 5. Drop the binders from the adjusted details and return
-     usage'    = adj_uds `delDetailsList` bndrs
+     usage'    = leaveLoop $ adj_uds `delDetailsList` bndrs
    in
    (usage', bndrs')
 


=====================================
compiler/GHC/Types/Unique/SlimSet.hs
=====================================
@@ -11,7 +11,8 @@ module GHC.Types.Unique.SlimSet (
     minusUniqSlimSet, unionUniqSlimSet, unionUniqSlimSets,
     ufmDom,
     -- * Querying
-    isEmptyUniqSlimSet, sizeUniqSlimSet, elemUniqSlimSet
+    isEmptyUniqSlimSet, sizeUniqSlimSet, elemUniqSlimSet,
+    nonDetEltsUniqSlimSet, nonDetFoldUniqSlimSet
   ) where
 
 import GHC.Prelude
@@ -76,6 +77,12 @@ unionUniqSlimSet (UniqSlimSet set1) (UniqSlimSet set2) = UniqSlimSet (set1 `S.un
 unionUniqSlimSets :: [UniqSlimSet a] -> UniqSlimSet a
 unionUniqSlimSets = foldl' (flip unionUniqSlimSet) emptyUniqSlimSet
 
+nonDetEltsUniqSlimSet :: UniqSlimSet a -> [Unique]
+nonDetEltsUniqSlimSet (UniqSlimSet s) = map mkUniqueGrimily (S.elems s)
+
+nonDetFoldUniqSlimSet :: (Unique -> acc -> acc) -> acc -> UniqSlimSet a -> acc
+nonDetFoldUniqSlimSet f acc (UniqSlimSet s) = S.foldr (f . mkUniqueGrimily) acc s
+
 instance Outputable (UniqSlimSet a) where
     ppr (UniqSlimSet s) = braces $
         hcat $ punctuate comma [ ppr (getUnique i) | i <- S.toList s]


=====================================
libraries/ghc-bignum/src/GHC/Num/WordArray.hs
=====================================
@@ -51,6 +51,7 @@ withNewWordArray# sz act = case runRW# io of (# _, a #) -> a
          case act mwa s          of { s ->
          unsafeFreezeByteArray# mwa s
          }}
+{-# INLINE withNewWordArray# #-}
 
 -- | Create two new WordArray# of the given sizes (*in Word#*) and apply the
 -- action to them before returning them frozen
@@ -86,6 +87,7 @@ withNewWordArrayTrimmed#
 withNewWordArrayTrimmed# sz act = withNewWordArray# sz \mwa s ->
    case act mwa s of
       s' -> mwaTrimZeroes# mwa s'
+{-# INLINE withNewWordArrayTrimmed# #-}
 
 -- | Create two new WordArray# of the given sizes (*in Word#*), apply the action
 -- to them, trim their most significant zeroes, then return them frozen
@@ -101,6 +103,7 @@ withNewWordArray2Trimmed# sz1 sz2 act = withNewWordArray2# sz1 sz2 \mwa1 mwa2 s
    case act mwa1 mwa2 s of
       s' -> case mwaTrimZeroes# mwa1 s' of
          s'' -> mwaTrimZeroes# mwa2 s''
+{-# INLINE withNewWordArray2Trimmed# #-}
 
 -- | Create a new WordArray# of the given size (*in Word#*), apply the action to
 -- it. If the action returns true#, trim its most significant zeroes, then
@@ -118,6 +121,7 @@ withNewWordArrayTrimmedMaybe# sz act = case runRW# io of (# _, a #) -> a
                (# s, _  #) -> case mwaTrimZeroes# mwa s of
                   s -> case unsafeFreezeByteArray# mwa s of
                      (# s, ba #) -> (# s, (# | ba #) #)
+{-# INLINE withNewWordArrayTrimmedMaybe# #-}
 
 -- | Create a WordArray# from two Word#
 --
@@ -296,6 +300,7 @@ mwaInitArrayBinOp mwa wa wb op s = go 0# s
             case indexWordArray# wa i `op` indexWordArray# wb i of
                v -> case mwaWrite# mwa i v s' of
                   s'' -> go (i +# 1#) s''
+{-# INLINE mwaInitArrayBinOp #-}
 
 -- | Write an element of the MutableWordArray
 mwaWrite# :: MutableWordArray# s -> Int# -> Word# -> State# s -> State# s



View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/89c1151745e514c3f200ce322f442f23f2c4ab62

-- 
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/89c1151745e514c3f200ce322f442f23f2c4ab62
You're receiving this email because of your account on gitlab.haskell.org.


-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.haskell.org/pipermail/ghc-commits/attachments/20221107/3451dc41/attachment-0001.html>


More information about the ghc-commits mailing list