Annotate CmmBranch with an optional likely target

Summary:
This allows the code generator to give hints to later code generation
steps about which branch is most likely to be taken.  Right now it
is only taken into account in one place: a special case in
CmmContFlowOpt that swapped branches over to maximise the chance of
fallthrough, which is now disabled when there is a likelihood setting.

Test Plan: validate

Reviewers: austin, simonpj, bgamari, ezyang, tibbe

Subscribers: thomie

Differential Revision: https://phabricator.haskell.org/D1273

11 files changed, 49 insertions, 36 deletions

diff --git a/compiler/cmm/CmmCommonBlockElim.hs b/compiler/cmm/CmmCommonBlockElim.hs
index 8c82fce56f..6c4742edad 100644
--- a/compiler/cmm/CmmCommonBlockElim.hs
+++ b/compiler/cmm/CmmCommonBlockElim.hs
@@ -137,7 +137,7 @@ hash_block block =
         hash_node (CmmStore e e') = hash_e e + hash_e e'
         hash_node (CmmUnsafeForeignCall t _ as) = hash_tgt t + hash_list hash_e as
         hash_node (CmmBranch _) = 23 -- NB. ignore the label
-        hash_node (CmmCondBranch p _ _) = hash_e p
+        hash_node (CmmCondBranch p _ _ _) = hash_e p
         hash_node (CmmCall e _ _ _ _ _) = hash_e e
         hash_node (CmmForeignCall t _ _ _ _ _ _) = hash_tgt t
         hash_node (CmmSwitch e _) = hash_e e
@@ -247,8 +247,8 @@ eqBlockBodyWith eqBid block block'
 
 eqLastWith :: (BlockId -> BlockId -> Bool) -> CmmNode O C -> CmmNode O C -> Bool
 eqLastWith eqBid (CmmBranch bid1) (CmmBranch bid2) = eqBid bid1 bid2
-eqLastWith eqBid (CmmCondBranch c1 t1 f1) (CmmCondBranch c2 t2 f2) =
-  c1 == c2 && eqBid t1 t2 && eqBid f1 f2
+eqLastWith eqBid (CmmCondBranch c1 t1 f1 l1) (CmmCondBranch c2 t2 f2 l2) =
+  c1 == c2 && l1 == l2 && eqBid t1 t2 && eqBid f1 f2
 eqLastWith eqBid (CmmCall t1 c1 g1 a1 r1 u1) (CmmCall t2 c2 g2 a2 r2 u2) =
   t1 == t2 && eqMaybeWith eqBid c1 c2 && a1 == a2 && r1 == r2 && u1 == u2 && g1 == g2
 eqLastWith eqBid (CmmSwitch e1 ids1) (CmmSwitch e2 ids2) =
diff --git a/compiler/cmm/CmmContFlowOpt.hs b/compiler/cmm/CmmContFlowOpt.hs
index 95c195078f..455422b47b 100644
--- a/compiler/cmm/CmmContFlowOpt.hs
+++ b/compiler/cmm/CmmContFlowOpt.hs
@@ -282,12 +282,15 @@ blockConcat splitting_procs g@CmmGraph { g_entry = entry_id }
           -- This helps the native codegen a little bit, and probably has no
           -- effect on LLVM.  It's convenient to do it here, where we have the
           -- information about predecessors.
+          --
+          -- NB., only do this if the branch does not have a
+          -- likeliness annotation.
           swapcond_last
-            | CmmCondBranch cond t f <- shortcut_last
+            | CmmCondBranch cond t f Nothing <- shortcut_last
             , numPreds f > 1
             , hasOnePredecessor t
             , Just cond' <- maybeInvertCmmExpr cond
-            = CmmCondBranch cond' f t
+            = CmmCondBranch cond' f t Nothing
 
             | otherwise
             = shortcut_last
@@ -354,21 +357,25 @@ replaceLabels env g
      lookup id = mapLookup id env `orElse` id
 
      txnode :: CmmNode e x -> CmmNode e x
-     txnode (CmmBranch bid)         = CmmBranch (lookup bid)
-     txnode (CmmCondBranch p t f)   = mkCmmCondBranch (exp p) (lookup t) (lookup f)
-     txnode (CmmSwitch e ids)       = CmmSwitch (exp e) (mapSwitchTargets lookup ids)
-     txnode (CmmCall t k rg a res r) = CmmCall (exp t) (liftM lookup k) rg a res r
-     txnode fc@CmmForeignCall{}     = fc{ args = map exp (args fc)
-                                        , succ = lookup (succ fc) }
-     txnode other                   = mapExpDeep exp other
+     txnode (CmmBranch bid) = CmmBranch (lookup bid)
+     txnode (CmmCondBranch p t f l) =
+       mkCmmCondBranch (exp p) (lookup t) (lookup f) l
+     txnode (CmmSwitch e ids) =
+       CmmSwitch (exp e) (mapSwitchTargets lookup ids)
+     txnode (CmmCall t k rg a res r) =
+       CmmCall (exp t) (liftM lookup k) rg a res r
+     txnode fc@CmmForeignCall{} =
+       fc{ args = map exp (args fc), succ = lookup (succ fc) }
+     txnode other = mapExpDeep exp other
 
      exp :: CmmExpr -> CmmExpr
      exp (CmmLit (CmmBlock bid))                = CmmLit (CmmBlock (lookup bid))
      exp (CmmStackSlot (Young id) i) = CmmStackSlot (Young (lookup id)) i
      exp e                                      = e
 
-mkCmmCondBranch :: CmmExpr -> Label -> Label -> CmmNode O C
-mkCmmCondBranch p t f = if t == f then CmmBranch t else CmmCondBranch p t f
+mkCmmCondBranch :: CmmExpr -> Label -> Label -> Maybe Bool -> CmmNode O C
+mkCmmCondBranch p t f l =
+  if t == f then CmmBranch t else CmmCondBranch p t f l
 
 -- Build a map from a block to its set of predecessors.
 predMap :: [CmmBlock] -> BlockEnv Int
diff --git a/compiler/cmm/CmmImplementSwitchPlans.hs b/compiler/cmm/CmmImplementSwitchPlans.hs
index 9fb68d8131..225c77e6d6 100644
--- a/compiler/cmm/CmmImplementSwitchPlans.hs
+++ b/compiler/cmm/CmmImplementSwitchPlans.hs
@@ -67,7 +67,7 @@ implementSwitchPlan dflags scope expr = go
         let lt | signed    = cmmSLtWord
                | otherwise = cmmULtWord
             scrut = lt dflags expr $ CmmLit $ mkWordCLit dflags i
-            lastNode = CmmCondBranch scrut bid1 bid2
+            lastNode = CmmCondBranch scrut bid1 bid2 Nothing
             lastBlock = emptyBlock `blockJoinTail` lastNode
         return (lastBlock, newBlocks1++newBlocks2)
     go (IfEqual i l ids2)
@@ -75,7 +75,7 @@ implementSwitchPlan dflags scope expr = go
         (bid2, newBlocks2) <- go' ids2
 
         let scrut = cmmNeWord dflags expr $ CmmLit $ mkWordCLit dflags i
-            lastNode = CmmCondBranch scrut bid2 l
+            lastNode = CmmCondBranch scrut bid2 l Nothing
             lastBlock = emptyBlock `blockJoinTail` lastNode
         return (lastBlock, newBlocks2)
 
diff --git a/compiler/cmm/CmmLayoutStack.hs b/compiler/cmm/CmmLayoutStack.hs
index 8439240b7e..5140aa3ae6 100644
--- a/compiler/cmm/CmmLayoutStack.hs
+++ b/compiler/cmm/CmmLayoutStack.hs
@@ -884,7 +884,7 @@ areaToSp _ _ _ _ other = other
 optStackCheck :: CmmNode O C -> CmmNode O C
 optStackCheck n = -- Note [Always false stack check]
  case n of
-   CmmCondBranch (CmmLit (CmmInt 0 _)) _true false -> CmmBranch false
+   CmmCondBranch (CmmLit (CmmInt 0 _)) _true false _ -> CmmBranch false
    other -> other
 
 
diff --git a/compiler/cmm/CmmLint.hs b/compiler/cmm/CmmLint.hs
index edce2e97bc..63a3ff5de3 100644
--- a/compiler/cmm/CmmLint.hs
+++ b/compiler/cmm/CmmLint.hs
@@ -165,7 +165,7 @@ lintCmmLast :: BlockSet -> CmmNode O C -> CmmLint ()
 lintCmmLast labels node = case node of
   CmmBranch id -> checkTarget id
 
-  CmmCondBranch e t f -> do
+  CmmCondBranch e t f _ -> do
             dflags <- getDynFlags
             mapM_ checkTarget [t,f]
             _ <- lintCmmExpr e
diff --git a/compiler/cmm/CmmNode.hs b/compiler/cmm/CmmNode.hs
index 45538d3886..40bb5a0a33 100644
--- a/compiler/cmm/CmmNode.hs
+++ b/compiler/cmm/CmmNode.hs
@@ -87,7 +87,9 @@ data CmmNode e x where
 
   CmmCondBranch :: {                 -- conditional branch
       cml_pred :: CmmExpr,
-      cml_true, cml_false :: ULabel
+      cml_true, cml_false :: ULabel,
+      cml_likely :: Maybe Bool       -- likely result of the conditional,
+                                     -- if known
   } -> CmmNode O C
 
   CmmSwitch
@@ -308,7 +310,7 @@ instance UserOfRegs LocalReg (CmmNode e x) where
     CmmAssign _ expr -> fold f z expr
     CmmStore addr rval -> fold f (fold f z addr) rval
     CmmUnsafeForeignCall t _ args -> fold f (fold f z t) args
-    CmmCondBranch expr _ _ -> fold f z expr
+    CmmCondBranch expr _ _ _ -> fold f z expr
     CmmSwitch expr _ -> fold f z expr
     CmmCall {cml_target=tgt} -> fold f z tgt
     CmmForeignCall {tgt=tgt, args=args} -> fold f (fold f z tgt) args
@@ -323,7 +325,7 @@ instance UserOfRegs GlobalReg (CmmNode e x) where
     CmmAssign _ expr -> fold f z expr
     CmmStore addr rval -> fold f (fold f z addr) rval
     CmmUnsafeForeignCall t _ args -> fold f (fold f z t) args
-    CmmCondBranch expr _ _ -> fold f z expr
+    CmmCondBranch expr _ _ _ -> fold f z expr
     CmmSwitch expr _ -> fold f z expr
     CmmCall {cml_target=tgt, cml_args_regs=args} -> fold f (fold f z args) tgt
     CmmForeignCall {tgt=tgt, args=args} -> fold f (fold f z tgt) args
@@ -463,7 +465,7 @@ mapExp f   (CmmAssign r e)                       = CmmAssign r (f e)
 mapExp f   (CmmStore addr e)                     = CmmStore (f addr) (f e)
 mapExp f   (CmmUnsafeForeignCall tgt fs as)      = CmmUnsafeForeignCall (mapForeignTarget f tgt) fs (map f as)
 mapExp _ l@(CmmBranch _)                         = l
-mapExp f   (CmmCondBranch e ti fi)               = CmmCondBranch (f e) ti fi
+mapExp f   (CmmCondBranch e ti fi l)             = CmmCondBranch (f e) ti fi l
 mapExp f   (CmmSwitch e ids)                     = CmmSwitch (f e) ids
 mapExp f   n@CmmCall {cml_target=tgt}            = n{cml_target = f tgt}
 mapExp f   (CmmForeignCall tgt fs as succ ret_args updfr intrbl) = CmmForeignCall (mapForeignTarget f tgt) fs (map f as) succ ret_args updfr intrbl
@@ -493,7 +495,7 @@ mapExpM f (CmmUnwind r e)           = CmmUnwind r `fmap` f e
 mapExpM f (CmmAssign r e)           = CmmAssign r `fmap` f e
 mapExpM f (CmmStore addr e)         = (\[addr', e'] -> CmmStore addr' e') `fmap` mapListM f [addr, e]
 mapExpM _ (CmmBranch _)             = Nothing
-mapExpM f (CmmCondBranch e ti fi)   = (\x -> CmmCondBranch x ti fi) `fmap` f e
+mapExpM f (CmmCondBranch e ti fi l) = (\x -> CmmCondBranch x ti fi l) `fmap` f e
 mapExpM f (CmmSwitch e tbl)         = (\x -> CmmSwitch x tbl)       `fmap` f e
 mapExpM f (CmmCall tgt mb_id r o i s) = (\x -> CmmCall x mb_id r o i s) `fmap` f tgt
 mapExpM f (CmmUnsafeForeignCall tgt fs as)
@@ -547,7 +549,7 @@ foldExp f (CmmAssign _ e) z                       = f e z
 foldExp f (CmmStore addr e) z                     = f addr $ f e z
 foldExp f (CmmUnsafeForeignCall t _ as) z         = foldr f (foldExpForeignTarget f t z) as
 foldExp _ (CmmBranch _) z                         = z
-foldExp f (CmmCondBranch e _ _) z                 = f e z
+foldExp f (CmmCondBranch e _ _ _) z               = f e z
 foldExp f (CmmSwitch e _) z                       = f e z
 foldExp f (CmmCall {cml_target=tgt}) z            = f tgt z
 foldExp f (CmmForeignCall {tgt=tgt, args=args}) z = foldr f (foldExpForeignTarget f tgt z) args
@@ -558,9 +560,9 @@ foldExpDeep f = foldExp (wrapRecExpf f)
 -- -----------------------------------------------------------------------------
 
 mapSuccessors :: (Label -> Label) -> CmmNode O C -> CmmNode O C
-mapSuccessors f (CmmBranch bid)        = CmmBranch (f bid)
-mapSuccessors f (CmmCondBranch p y n)  = CmmCondBranch p (f y) (f n)
-mapSuccessors f (CmmSwitch e ids)      = CmmSwitch e (mapSwitchTargets f ids)
+mapSuccessors f (CmmBranch bid)         = CmmBranch (f bid)
+mapSuccessors f (CmmCondBranch p y n l) = CmmCondBranch p (f y) (f n) l
+mapSuccessors f (CmmSwitch e ids)       = CmmSwitch e (mapSwitchTargets f ids)
 mapSuccessors _ n = n
 
 -- -----------------------------------------------------------------------------
diff --git a/compiler/cmm/CmmParse.y b/compiler/cmm/CmmParse.y
index ea0f4a5a66..dbd5d06872 100644
--- a/compiler/cmm/CmmParse.y
+++ b/compiler/cmm/CmmParse.y
@@ -1268,7 +1268,7 @@ cmmRawIf cond then_id = do
 -- branching to true_id if so, and falling through otherwise.
 emitCond (BoolTest e) then_id = do
   else_id <- newBlockId
-  emit (mkCbranch e then_id else_id)
+  emit (mkCbranch e then_id else_id Nothing)
   emitLabel else_id
 emitCond (BoolNot (BoolTest (CmmMachOp op args))) then_id
   | Just op' <- maybeInvertComparison op
diff --git a/compiler/cmm/CmmProcPoint.hs b/compiler/cmm/CmmProcPoint.hs
index a31048206b..0e772c41d0 100644
--- a/compiler/cmm/CmmProcPoint.hs
+++ b/compiler/cmm/CmmProcPoint.hs
@@ -295,7 +295,7 @@ splitAtProcPoints dflags entry_label callPPs procPoints procMap
                   add_if_branch_to_pp block rst =
                     case lastNode block of
                       CmmBranch id          -> add_if_pp id rst
-                      CmmCondBranch _ ti fi -> add_if_pp ti (add_if_pp fi rst)
+                      CmmCondBranch _ ti fi _ -> add_if_pp ti (add_if_pp fi rst)
                       CmmSwitch _ ids       -> foldr add_if_pp rst $ switchTargetsToList ids
                       _                     -> rst
 
@@ -382,7 +382,7 @@ replaceBranches env cmmg
 
     last :: CmmNode O C -> CmmNode O C
     last (CmmBranch id)          = CmmBranch (lookup id)
-    last (CmmCondBranch e ti fi) = CmmCondBranch e (lookup ti) (lookup fi)
+    last (CmmCondBranch e ti fi l) = CmmCondBranch e (lookup ti) (lookup fi) l
     last (CmmSwitch e ids)       = CmmSwitch e (mapSwitchTargets lookup ids)
     last l@(CmmCall {})          = l { cml_cont = Nothing }
             -- NB. remove the continuation of a CmmCall, since this
diff --git a/compiler/cmm/MkGraph.hs b/compiler/cmm/MkGraph.hs
index d2aa4aa057..657585e75a 100644
--- a/compiler/cmm/MkGraph.hs
+++ b/compiler/cmm/MkGraph.hs
@@ -35,7 +35,7 @@ import OrdList
 import Control.Monad
 import Data.List
 import Data.Maybe
-import Prelude (($),Int,Eq(..)) -- avoid importing (<*>)
+import Prelude (($),Int,Bool,Eq(..)) -- avoid importing (<*>)
 
 #include "HsVersions.h"
 
@@ -221,8 +221,9 @@ mkJumpExtra dflags conv e actuals updfr_off extra_stack =
   lastWithArgsAndExtraStack dflags Jump Old conv actuals updfr_off extra_stack $
     toCall e Nothing updfr_off 0
 
-mkCbranch       :: CmmExpr -> BlockId -> BlockId -> CmmAGraph
-mkCbranch pred ifso ifnot = mkLast (CmmCondBranch pred ifso ifnot)
+mkCbranch       :: CmmExpr -> BlockId -> BlockId -> Maybe Bool -> CmmAGraph
+mkCbranch pred ifso ifnot likely =
+  mkLast (CmmCondBranch pred ifso ifnot likely)
 
 mkSwitch        :: CmmExpr -> SwitchTargets -> CmmAGraph
 mkSwitch e tbl   = mkLast $ CmmSwitch e tbl
diff --git a/compiler/cmm/PprC.hs b/compiler/cmm/PprC.hs
index ee0680d3e9..c96b7076bf 100644
--- a/compiler/cmm/PprC.hs
+++ b/compiler/cmm/PprC.hs
@@ -249,7 +249,7 @@ pprStmt stmt =
           = pprCall fn cconv hresults hargs
 
     CmmBranch ident          -> pprBranch ident
-    CmmCondBranch expr yes no -> pprCondBranch expr yes no
+    CmmCondBranch expr yes no _ -> pprCondBranch expr yes no
     CmmCall { cml_target = expr } -> mkJMP_ (pprExpr expr) <> semi
     CmmSwitch arg ids        -> sdocWithDynFlags $ \dflags ->
                                 pprSwitch dflags arg ids
@@ -1042,7 +1042,7 @@ te_Stmt (CmmUnsafeForeignCall target rs es)
   = do  te_Target target
         mapM_ te_temp rs
         mapM_ te_Expr es
-te_Stmt (CmmCondBranch e _ _)   = te_Expr e
+te_Stmt (CmmCondBranch e _ _ _) = te_Expr e
 te_Stmt (CmmSwitch e _)         = te_Expr e
 te_Stmt (CmmCall { cml_target = e }) = te_Expr e
 te_Stmt _                       = return ()
diff --git a/compiler/cmm/PprCmm.hs b/compiler/cmm/PprCmm.hs
index d5999f53fa..5caea90db4 100644
--- a/compiler/cmm/PprCmm.hs
+++ b/compiler/cmm/PprCmm.hs
@@ -220,9 +220,12 @@ pprNode node = pp_node <+> pp_debug
       CmmBranch ident -> ptext (sLit "goto") <+> ppr ident <> semi
 
       -- if (expr) goto t; else goto f;
-      CmmCondBranch expr t f ->
+      CmmCondBranch expr t f l ->
           hsep [ ptext (sLit "if")
                , parens(ppr expr)
+               , case l of
+                   Nothing -> empty
+                   Just b -> parens (ptext (sLit "likely:") <+> ppr b)
                , ptext (sLit "goto")
                , ppr t <> semi
                , ptext (sLit "else goto")