Statically estimate execution frequency of CoreAlts (#20378)wip/exec-freq

This patch implements #20378. See `Note [Estimating CoreAlt frequencies]` in the
new module GHC.Core.Opt.ExecFreq for details.

These were the changes:

1. Introduce `newtype Freq = Freq Float` as a type that captures relative
   execution frequency and use it as an additional field in `CoreAlt`.
   The default when we have no estimate available is `NoFreq`, e.g., NaN.
   Otherwise, all `Freq`s of a `Case` should add up to 1.
   Then fix up a whole bunch of use sites.
2. Introduce a new enum `Comparison` for the different kinds of comparison
   operators (LessThan, GreaterOrEqual, NotEqual, ...). Then make `Compare`
   primops also declare what kind of `Comparison` they do. Then introduce a
   function `isComparisonApp_maybe` in GHC.Core.Utils that we can use for our
   estimates (see below).
3. Write a static analysis pass `estimateAltFreqs`, that annotates `CoreAlt`s
   with their relative execution `Freq`. These `Freq`s are determined by
   combining the estimates of different branch heuristics, one of which uses
   the new `isComparisonApp_maybe`.

The main function `estimateAltFreqs` is currently dead, but that is bound to
change in follow-up MRs.

Fixes #20378.

1 files changed, 43 insertions, 31 deletions

diff --git a/compiler/GHC/Core/Utils.hs b/compiler/GHC/Core/Utils.hs
index b2af755e78..68378c6e2e 100644
--- a/compiler/GHC/Core/Utils.hs
+++ b/compiler/GHC/Core/Utils.hs
@@ -33,6 +33,7 @@ module GHC.Core.Utils (
         exprIsTickedString, exprIsTickedString_maybe,
         exprIsTopLevelBindable,
         altsAreExhaustive,
+        isComparisonApp_maybe,
 
         -- * Equality
         cheapEqExpr, cheapEqExpr', eqExpr,
@@ -93,7 +94,7 @@ import GHC.Types.Tickish
 import GHC.Types.Id
 import GHC.Types.Id.Info
 import GHC.Types.Unique
-import GHC.Types.Basic     ( Arity, FullArgCount )
+import GHC.Types.Basic
 import GHC.Types.Unique.Set
 
 import GHC.Data.FastString
@@ -146,7 +147,7 @@ exprType other = pprPanic "exprType" (pprCoreExpr other)
 
 coreAltType :: CoreAlt -> Type
 -- ^ Returns the type of the alternatives right hand side
-coreAltType alt@(Alt _ bs rhs)
+coreAltType alt@(Alt _ _ bs rhs)
   = case occCheckExpand bs rhs_ty of
       -- Note [Existential variables and silly type synonyms]
       Just ty -> ty
@@ -499,7 +500,7 @@ stripTicksE p expr = go expr
         go_bs (NonRec b e)  = NonRec b (go e)
         go_bs (Rec bs)      = Rec (map go_b bs)
         go_b (b, e)         = (b, go e)
-        go_a (Alt c bs e)   = Alt c bs (go e)
+        go_a (Alt c f bs e) = Alt c f bs (go e)
 
 stripTicksT :: (CoreTickish -> Bool) -> Expr b -> [CoreTickish]
 stripTicksT p expr = fromOL $ go expr
@@ -515,7 +516,7 @@ stripTicksT p expr = fromOL $ go expr
         go_bs (NonRec _ e)  = go e
         go_bs (Rec bs)      = concatOL (map go_b bs)
         go_b (_, e)         = go e
-        go_a (Alt _ _ e)    = go e
+        go_a (Alt _ _ _ e)  = go e
 
 {-
 ************************************************************************
@@ -575,7 +576,7 @@ mkAltExpr DEFAULT _ _ = panic "mkAltExpr DEFAULT"
 mkDefaultCase :: CoreExpr -> Id -> CoreExpr -> CoreExpr
 -- Make (case x of y { DEFAULT -> e }
 mkDefaultCase scrut case_bndr body
-  = Case scrut case_bndr (exprType body) [Alt DEFAULT [] body]
+  = Case scrut case_bndr (exprType body) [Alt DEFAULT NoFreq [] body]
 
 mkSingleAltCase :: CoreExpr -> Id -> AltCon -> [Var] -> CoreExpr -> CoreExpr
 -- Use this function if possible, when building a case,
@@ -583,7 +584,7 @@ mkSingleAltCase :: CoreExpr -> Id -> AltCon -> [Var] -> CoreExpr -> CoreExpr
 -- doesn't mention variables bound by the case
 -- See Note [Care with the type of a case expression]
 mkSingleAltCase scrut case_bndr con bndrs body
-  = Case scrut case_bndr case_ty [Alt con bndrs body]
+  = Case scrut case_bndr case_ty [Alt con NoFreq bndrs body]
   where
     body_ty = exprType body
 
@@ -627,16 +628,16 @@ This makes it easy to find, though it makes matching marginally harder.
 
 -- | Extract the default case alternative
 findDefault :: [Alt b] -> ([Alt b], Maybe (Expr b))
-findDefault (Alt DEFAULT args rhs : alts) = assert (null args) (alts, Just rhs)
-findDefault alts                          =                    (alts, Nothing)
+findDefault (Alt DEFAULT _ args rhs : alts) = assert (null args) (alts, Just rhs)
+findDefault alts                            =                    (alts, Nothing)
 
 addDefault :: [Alt b] -> Maybe (Expr b) -> [Alt b]
 addDefault alts Nothing    = alts
-addDefault alts (Just rhs) = Alt DEFAULT [] rhs : alts
+addDefault alts (Just rhs) = Alt DEFAULT NoFreq [] rhs : alts
 
 isDefaultAlt :: Alt b -> Bool
-isDefaultAlt (Alt DEFAULT _ _) = True
-isDefaultAlt _                 = False
+isDefaultAlt (Alt DEFAULT _ _ _) = True
+isDefaultAlt _                   = False
 
 -- | Find the case alternative corresponding to a particular
 -- constructor: panics if no such constructor exists
@@ -645,11 +646,11 @@ findAlt :: AltCon -> [Alt b] -> Maybe (Alt b)
     -- See Note [Unreachable code]
 findAlt con alts
   = case alts of
-        (deflt@(Alt DEFAULT _ _):alts) -> go alts (Just deflt)
-        _                              -> go alts Nothing
+        (deflt@(Alt DEFAULT _ _ _):alts) -> go alts (Just deflt)
+        _                                -> go alts Nothing
   where
     go []                     deflt = deflt
-    go (alt@(Alt con1 _ _) : alts) deflt
+    go (alt@(Alt con1 _ _ _) : alts) deflt
       = case con `cmpAltCon` con1 of
           LT -> deflt   -- Missed it already; the alts are in increasing order
           EQ -> Just alt
@@ -736,7 +737,7 @@ filterAlts _tycon inst_tys imposs_cons alts
   = (imposs_deflt_cons, addDefault trimmed_alts maybe_deflt)
   where
     (alts_wo_default, maybe_deflt) = findDefault alts
-    alt_cons = [con | Alt con _ _ <- alts_wo_default]
+    alt_cons = [con | Alt con _ _ _ <- alts_wo_default]
 
     trimmed_alts = filterOut (impossible_alt inst_tys) alts_wo_default
 
@@ -748,9 +749,9 @@ filterAlts _tycon inst_tys imposs_cons alts
          --   OR by a non-DEFAULT branch in this case expression.
 
     impossible_alt :: [Type] -> Alt b -> Bool
-    impossible_alt _ (Alt con _ _) | con `Set.member` imposs_cons_set = True
-    impossible_alt inst_tys (Alt (DataAlt con) _ _) = dataConCannotMatch inst_tys con
-    impossible_alt _  _                             = False
+    impossible_alt _ (Alt con _ _ _) | con `Set.member` imposs_cons_set = True
+    impossible_alt inst_tys (Alt (DataAlt con) _ _ _) = dataConCannotMatch inst_tys con
+    impossible_alt _  _                               = False
 
 -- | Refine the default alternative to a 'DataAlt', if there is a unique way to do so.
 -- See Note [Refine DEFAULT case alternatives]
@@ -762,7 +763,7 @@ refineDefaultAlt :: [Unique]          -- ^ Uniques for constructing new binders
                  -> [CoreAlt]
                  -> (Bool, [CoreAlt]) -- ^ 'True', if a default alt was replaced with a 'DataAlt'
 refineDefaultAlt us mult tycon tys imposs_deflt_cons all_alts
-  | Alt DEFAULT _ rhs : rest_alts <- all_alts
+  | Alt DEFAULT _ _ rhs : rest_alts <- all_alts
   , isAlgTyCon tycon            -- It's a data type, tuple, or unboxed tuples.
   , not (isNewTyCon tycon)      -- We can have a newtype, if we are just doing an eval:
                                 --      case x of { DEFAULT -> e }
@@ -779,7 +780,7 @@ refineDefaultAlt us mult tycon tys imposs_deflt_cons all_alts
        []    -> (False, rest_alts)
 
        -- It matches exactly one constructor, so fill it in:
-       [con] -> (True, mergeAlts rest_alts [Alt (DataAlt con) (ex_tvs ++ arg_ids) rhs])
+       [con] -> (True, mergeAlts rest_alts [Alt (DataAlt con) NoFreq (ex_tvs ++ arg_ids) rhs])
                        -- We need the mergeAlts to keep the alternatives in the right order
              where
                 (ex_tvs, arg_ids) = dataConRepInstPat us mult con tys
@@ -962,25 +963,26 @@ combineIdenticalAlts :: [AltCon]    -- Constructors that cannot match DEFAULT
                          [CoreAlt]) -- New alternatives
 -- See Note [Combine identical alternatives]
 -- True <=> we did some combining, result is a single DEFAULT alternative
-combineIdenticalAlts imposs_deflt_cons (Alt con1 bndrs1 rhs1 : rest_alts)
+combineIdenticalAlts imposs_deflt_cons (Alt con1 freq1 bndrs1 rhs1 : rest_alts)
   | all isDeadBinder bndrs1    -- Remember the default
   , not (null elim_rest) -- alternative comes first
   = (True, imposs_deflt_cons', deflt_alt : filtered_rest)
   where
     (elim_rest, filtered_rest) = partition identical_to_alt1 rest_alts
-    deflt_alt = Alt DEFAULT [] (mkTicks (concat tickss) rhs1)
+    elim_freqs = sum [ freq | Alt _ freq _ _ <- elim_rest ]
+    deflt_alt = Alt DEFAULT (elim_freqs + freq1) [] (mkTicks (concat tickss) rhs1)
 
      -- See Note [Care with impossible-constructors when combining alternatives]
     imposs_deflt_cons' = imposs_deflt_cons `minusList` elim_cons
-    elim_cons = elim_con1 ++ map (\(Alt con _ _) -> con) elim_rest
+    elim_cons = elim_con1 ++ map (\(Alt con _ _ _) -> con) elim_rest
     elim_con1 = case con1 of     -- Don't forget con1!
                   DEFAULT -> []  -- See Note [
                   _       -> [con1]
 
     cheapEqTicked e1 e2 = cheapEqExpr' tickishFloatable e1 e2
-    identical_to_alt1 (Alt _con bndrs rhs)
+    identical_to_alt1 (Alt _con _freq bndrs rhs)
       = all isDeadBinder bndrs && rhs `cheapEqTicked` rhs1
-    tickss = map (\(Alt _ _ rhs) -> stripTicksT tickishFloatable rhs) elim_rest
+    tickss = map (\(Alt _ _ _ rhs) -> stripTicksT tickishFloatable rhs) elim_rest
 
 combineIdenticalAlts imposs_cons alts
   = (False, imposs_cons, alts)
@@ -991,7 +993,7 @@ scaleAltsBy :: Mult -> [CoreAlt] -> [CoreAlt]
 scaleAltsBy w alts = map scaleAlt alts
   where
     scaleAlt :: CoreAlt -> CoreAlt
-    scaleAlt (Alt con bndrs rhs) = Alt con (map scaleBndr bndrs) rhs
+    scaleAlt (Alt con freq bndrs rhs) = Alt con freq (map scaleBndr bndrs) rhs
 
     scaleBndr :: CoreBndr -> CoreBndr
     scaleBndr b = scaleVarBy w b
@@ -1332,7 +1334,7 @@ exprIsCheapX ok_app e
     go _ (Coercion {})                = True
     go n (Cast e _)                   = go n e
     go n (Case scrut _ _ alts)        = ok scrut &&
-                                        and [ go n rhs | Alt _ _ rhs <- alts ]
+                                        and [ go n rhs | Alt _ _ _ rhs <- alts ]
     go n (Tick t e) | tickishCounts t = False
                     | otherwise       = go n e
     go n (Lam x e)  | isRuntimeVar x  = n==0 || go (n-1) e
@@ -1617,7 +1619,7 @@ expr_ok primop_ok (Case scrut bndr _ alts)
   =  -- See Note [exprOkForSpeculation: case expressions]
      expr_ok primop_ok scrut
   && isUnliftedType (idType bndr)
-  && all (\(Alt _ _ rhs) -> expr_ok primop_ok rhs) alts
+  && all (\(Alt _ _ _ rhs) -> expr_ok primop_ok rhs) alts
   && altsAreExhaustive alts
 
 expr_ok primop_ok other_expr
@@ -1707,7 +1709,7 @@ altsAreExhaustive :: [Alt b] -> Bool
 -- False <=> they may or may not be
 altsAreExhaustive []
   = False    -- Should not happen
-altsAreExhaustive (Alt con1 _ _ : alts)
+altsAreExhaustive (Alt con1 _ _ _ : alts)
   = case con1 of
       DEFAULT   -> True
       LitAlt {} -> False
@@ -2022,6 +2024,16 @@ exprIsTickedString_maybe (Tick t e)
   | otherwise = exprIsTickedString_maybe e
 exprIsTickedString_maybe _ = Nothing
 
+-- | Is the expression an application to a primitive comparison operator
+-- ('primOpIsComparison_maybe')? If so, return the kind of 'Comparison'
+-- and the two argument expressions.
+isComparisonApp_maybe :: CoreExpr -> Maybe (Comparison, CoreExpr, CoreExpr)
+isComparisonApp_maybe e = do
+  App (App (Var f) a1) a2 <- pure e
+  op <- isPrimOpId_maybe f
+  cmp <- primOpIsComparison_maybe op
+  pure (cmp, a1, a2)
+
 {-
 ************************************************************************
 *                                                                      *
@@ -2201,7 +2213,7 @@ eqExpr in_scope e1 e2
     go _ _ _ = False
 
     -----------
-    go_alt env (Alt c1 bs1 e1) (Alt c2 bs2 e2)
+    go_alt env (Alt c1 _f1 bs1 e1) (Alt c2 _f2 bs2 e2)
       = c1 == c2 && go (rnBndrs2 env bs1 bs2) e1 e2
 
 eqTickish :: RnEnv2 -> CoreTickish -> CoreTickish -> Bool
@@ -2246,7 +2258,7 @@ diffExpr top env (Case e1 b1 t1 a1) (Case e2 b2 t2 a2)
     -- See Note [Empty case alternatives] in GHC.Data.TrieMap
   = diffExpr top env e1 e2 ++ concat (zipWith diffAlt a1 a2)
   where env' = rnBndr2 env b1 b2
-        diffAlt (Alt c1 bs1 e1) (Alt c2 bs2 e2)
+        diffAlt (Alt c1 _f1 bs1 e1) (Alt c2 _f2 bs2 e2)
           | c1 /= c2  = [text "alt-cons " <> ppr c1 <> text " /= " <> ppr c2]
           | otherwise = diffExpr top (rnBndrs2 env' bs1 bs2) e1 e2
 diffExpr _  _ e1 e2