WIP: Add likelyhood to alternatives from stg onwardswip/D4327

Summary:
Adds a Freq value to Stg/Cmm cases/switches/conditionals.

Currently only generates these values by checking alternatives for
bottom expressions.

They are passed along to the backend where they affect conditional generation
slightly.

As it stands runtime improvements seem to be less than expected. This might only be worth merging once we have more branch weights available.

Reviewers: hvr, goldfire, bgamari, simonmar, simonpj, erikd

Reviewed By: simonpj

Subscribers: simonpj, rwbarton, thomie, carter

GHC Trac Issues: #14672

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

6 files changed, 124 insertions, 80 deletions

diff --git a/compiler/codeGen/StgCmmBind.hs b/compiler/codeGen/StgCmmBind.hs
index cf602ef0b8..9b3c166d1e 100644
--- a/compiler/codeGen/StgCmmBind.hs
+++ b/compiler/codeGen/StgCmmBind.hs
@@ -272,7 +272,7 @@ mkRhsClosure    dflags bndr _cc _bi
   , StgCase (StgApp scrutinee [{-no args-}])
          _   -- ignore bndr
          (AlgAlt _)
-         [(DataAlt _, params, sel_expr)] <- strip expr
+         [(DataAlt _, params, sel_expr, _)] <- strip expr
   , StgApp selectee [{-no args-}] <- strip sel_expr
   , the_fv == scrutinee                -- Scrutinee is the only free variable
 
diff --git a/compiler/codeGen/StgCmmExpr.hs b/compiler/codeGen/StgCmmExpr.hs
index 3fcc935121..ab0e6d0c2a 100644
--- a/compiler/codeGen/StgCmmExpr.hs
+++ b/compiler/codeGen/StgCmmExpr.hs
@@ -31,6 +31,7 @@ import StgCmmClosure
 
 import StgSyn
 
+import BasicTypes (BranchWeight)
 import MkGraph
 import BlockId
 import Cmm
@@ -379,7 +380,7 @@ calls to nonVoidIds in various places.  So we must not look up
 
 cgCase (StgApp v []) _ (PrimAlt _) alts
   | isVoidRep (idPrimRep v)  -- See Note [Scrutinising VoidRep]
-  , [(DEFAULT, _, rhs)] <- alts
+  , [(DEFAULT, _, rhs, _)] <- alts
   = cgExpr rhs
 
 {- Note [Dodgy unsafeCoerce 1]
@@ -561,7 +562,7 @@ chooseReturnBndrs :: Id -> AltType -> [StgAlt] -> [NonVoid Id]
 chooseReturnBndrs bndr (PrimAlt _) _alts
   = assertNonVoidIds [bndr]
 
-chooseReturnBndrs _bndr (MultiValAlt n) [(_, ids, _)]
+chooseReturnBndrs _bndr (MultiValAlt n) [(_, ids, _, _)]
   = ASSERT2(ids `lengthIs` n, ppr n $$ ppr ids $$ ppr _bndr)
     assertNonVoidIds ids     -- 'bndr' is not assigned!
 
@@ -578,10 +579,10 @@ chooseReturnBndrs _ _ _ = panic "chooseReturnBndrs"
 cgAlts :: (GcPlan,ReturnKind) -> NonVoid Id -> AltType -> [StgAlt]
        -> FCode ReturnKind
 -- At this point the result of the case are in the binders
-cgAlts gc_plan _bndr PolyAlt [(_, _, rhs)]
+cgAlts gc_plan _bndr PolyAlt [(_, _, rhs, _)]
   = maybeAltHeapCheck gc_plan (cgExpr rhs)
 
-cgAlts gc_plan _bndr (MultiValAlt _) [(_, _, rhs)]
+cgAlts gc_plan _bndr (MultiValAlt _) [(_, _, rhs, _)]
   = maybeAltHeapCheck gc_plan (cgExpr rhs)
         -- Here bndrs are *already* in scope, so don't rebind them
 
@@ -591,13 +592,13 @@ cgAlts gc_plan bndr (PrimAlt _) alts
         ; tagged_cmms <- cgAltRhss gc_plan bndr alts
 
         ; let bndr_reg = CmmLocal (idToReg dflags bndr)
-              (DEFAULT,deflt) = head tagged_cmms
+              (DEFAULT,deflt,f) = head tagged_cmms
                 -- PrimAlts always have a DEFAULT case
                 -- and it always comes first
 
-              tagged_cmms' = [(lit,code)
-                             | (LitAlt lit, code) <- tagged_cmms]
-        ; emitCmmLitSwitch (CmmReg bndr_reg) tagged_cmms' deflt
+              tagged_cmms' = [(lit,code,f)
+                             | (LitAlt lit, code,f) <- tagged_cmms]
+        ; emitCmmLitSwitch (CmmReg bndr_reg) tagged_cmms' (deflt,f)
         ; return AssignedDirectly }
 
 cgAlts gc_plan bndr (AlgAlt tycon) alts
@@ -613,7 +614,7 @@ cgAlts gc_plan bndr (AlgAlt tycon) alts
           then do
                 let   -- Yes, bndr_reg has constr. tag in ls bits
                    tag_expr = cmmConstrTag1 dflags (CmmReg bndr_reg)
-                   branches' = [(tag+1,branch) | (tag,branch) <- branches]
+                   branches' = [(tag+1,branch,f) | (tag,branch,f) <- branches]
                 emitSwitch tag_expr branches' mb_deflt 1 fam_sz
 
            else -- No, get tag from info table
@@ -651,18 +652,18 @@ cgAlts _ _ _ _ = panic "cgAlts"
 
 -------------------
 cgAlgAltRhss :: (GcPlan,ReturnKind) -> NonVoid Id -> [StgAlt]
-             -> FCode ( Maybe CmmAGraphScoped
-                      , [(ConTagZ, CmmAGraphScoped)] )
+             -> FCode ( Maybe (CmmAGraphScoped, BranchWeight)
+                      , [(ConTagZ, CmmAGraphScoped, BranchWeight)] )
 cgAlgAltRhss gc_plan bndr alts
   = do { tagged_cmms <- cgAltRhss gc_plan bndr alts
 
        ; let { mb_deflt = case tagged_cmms of
-                           ((DEFAULT,rhs) : _) -> Just rhs
+                           ((DEFAULT,rhs,f) : _) -> Just (rhs,f)
                            _other              -> Nothing
                             -- DEFAULT is always first, if present
 
-              ; branches = [ (dataConTagZ con, cmm)
-                           | (DataAlt con, cmm) <- tagged_cmms ]
+              ; branches = [ (dataConTagZ con, cmm, f)
+                           | (DataAlt con, cmm, f) <- tagged_cmms ]
               }
 
        ; return (mb_deflt, branches)
@@ -671,20 +672,20 @@ cgAlgAltRhss gc_plan bndr alts
 
 -------------------
 cgAltRhss :: (GcPlan,ReturnKind) -> NonVoid Id -> [StgAlt]
-          -> FCode [(AltCon, CmmAGraphScoped)]
+          -> FCode [(AltCon, CmmAGraphScoped,BranchWeight)]
 cgAltRhss gc_plan bndr alts = do
   dflags <- getDynFlags
   let
     base_reg = idToReg dflags bndr
-    cg_alt :: StgAlt -> FCode (AltCon, CmmAGraphScoped)
-    cg_alt (con, bndrs, rhs)
-      = getCodeScoped             $
-        maybeAltHeapCheck gc_plan $
-        do { _ <- bindConArgs con base_reg (assertNonVoidIds bndrs)
-                    -- alt binders are always non-void,
-                    -- see Note [Post-unarisation invariants] in UnariseStg
-           ; _ <- cgExpr rhs
-           ; return con }
+    cg_alt :: StgAlt -> FCode (AltCon, CmmAGraphScoped, BranchWeight)
+    cg_alt (con, bndrs, rhs, freq) = do
+      (i,c) <- getCodeScoped $ maybeAltHeapCheck gc_plan $
+          do { _ <- bindConArgs con base_reg (assertNonVoidIds bndrs)
+                      -- alt binders are always non-void,
+                      -- see Note [Post-unarisation invariants] in UnariseStg
+            ; _ <- cgExpr rhs
+            ; return con }
+      return (i,c,freq)
   forkAlts (map cg_alt alts)
 
 maybeAltHeapCheck :: (GcPlan,ReturnKind) -> FCode a -> FCode a
diff --git a/compiler/codeGen/StgCmmMonad.hs b/compiler/codeGen/StgCmmMonad.hs
index 7c3864296c..c3baa7bb70 100644
--- a/compiler/codeGen/StgCmmMonad.hs
+++ b/compiler/codeGen/StgCmmMonad.hs
@@ -799,8 +799,10 @@ getCmm code
         ; return (fromOL (cgs_tops state2)) }
 
 
-mkCmmIfThenElse :: CmmExpr -> CmmAGraph -> CmmAGraph -> FCode CmmAGraph
-mkCmmIfThenElse e tbranch fbranch = mkCmmIfThenElse' e tbranch fbranch Nothing
+mkCmmIfThenElse :: CmmExpr -> CmmAGraph -> CmmAGraph -> Maybe Bool
+                -> FCode CmmAGraph
+mkCmmIfThenElse e tbranch fbranch likely
+  = mkCmmIfThenElse' e tbranch fbranch likely
 
 mkCmmIfThenElse' :: CmmExpr -> CmmAGraph -> CmmAGraph
                  -> Maybe Bool -> FCode CmmAGraph
diff --git a/compiler/codeGen/StgCmmPrim.hs b/compiler/codeGen/StgCmmPrim.hs
index b5cd267c6b..fa2e7d2b6c 100644
--- a/compiler/codeGen/StgCmmPrim.hs
+++ b/compiler/codeGen/StgCmmPrim.hs
@@ -1127,9 +1127,8 @@ genericFabsOp w [res_r] [aa]
       let g3 = catAGraphs [mkAssign res_t aa,
                            mkAssign (CmmLocal res_r) (neg (CmmReg res_t))]
 
-      g4 <- mkCmmIfThenElse (gt aa zero) g2 g3
-
-      emit =<< mkCmmIfThenElse (eq aa zero) g1 g4
+      g4 <- mkCmmIfThenElse (gt aa zero) g2 g3 Nothing
+      emit =<< mkCmmIfThenElse (eq aa zero) g1 g4 Nothing
 
 genericFabsOp _ _ _ = panic "genericFabsOp"
 
@@ -1821,14 +1820,17 @@ doCopyMutableByteArrayOp = emitCopyByteArray copy
   where
     -- The only time the memory might overlap is when the two arrays
     -- we were provided are the same array!
-    -- TODO: Optimize branch for common case of no aliasing.
+    -- The common case is no aliasing so we set the likly value to `Just False`.
     copy src dst dst_p src_p bytes = do
         dflags <- getDynFlags
         [moveCall, cpyCall] <- forkAlts [
             getCode $ emitMemmoveCall dst_p src_p bytes 1,
             getCode $ emitMemcpyCall  dst_p src_p bytes 1
             ]
-        emit =<< mkCmmIfThenElse (cmmEqWord dflags src dst) moveCall cpyCall
+        emit =<< mkCmmIfThenElse
+                    (cmmEqWord dflags src dst)
+                    moveCall cpyCall
+                    (Just False)
 
 emitCopyByteArray :: (CmmExpr -> CmmExpr -> CmmExpr -> CmmExpr -> CmmExpr
                       -> FCode ())
@@ -1965,7 +1967,8 @@ doCopyMutableArrayOp = emitCopyArray copy
   where
     -- The only time the memory might overlap is when the two arrays
     -- we were provided are the same array!
-    -- TODO: Optimize branch for common case of no aliasing.
+    -- Optimize branch for common case of no aliasing by setting likely
+    -- to `Just False`.
     copy src dst dst_p src_p bytes = do
         dflags <- getDynFlags
         [moveCall, cpyCall] <- forkAlts [
@@ -1974,7 +1977,10 @@ doCopyMutableArrayOp = emitCopyArray copy
             getCode $ emitMemcpyCall  dst_p src_p (mkIntExpr dflags bytes)
             (wORD_SIZE dflags)
             ]
-        emit =<< mkCmmIfThenElse (cmmEqWord dflags src dst) moveCall cpyCall
+        emit =<< mkCmmIfThenElse
+                    (cmmEqWord dflags src dst)
+                    moveCall cpyCall
+                    (Just False)
 
 emitCopyArray :: (CmmExpr -> CmmExpr -> CmmExpr -> CmmExpr -> ByteOff
                   -> FCode ())  -- ^ copy function
@@ -2028,7 +2034,8 @@ doCopySmallMutableArrayOp = emitCopySmallArray copy
   where
     -- The only time the memory might overlap is when the two arrays
     -- we were provided are the same array!
-    -- TODO: Optimize branch for common case of no aliasing.
+    -- Optimize branch for common case of no aliasing by setting likelyhood
+    -- to `Just False`.
     copy src dst dst_p src_p bytes = do
         dflags <- getDynFlags
         [moveCall, cpyCall] <- forkAlts
@@ -2037,7 +2044,10 @@ doCopySmallMutableArrayOp = emitCopySmallArray copy
             , getCode $ emitMemcpyCall  dst_p src_p (mkIntExpr dflags bytes)
               (wORD_SIZE dflags)
             ]
-        emit =<< mkCmmIfThenElse (cmmEqWord dflags src dst) moveCall cpyCall
+        emit =<< mkCmmIfThenElse
+                    (cmmEqWord dflags src dst)
+                    moveCall cpyCall
+                    (Just False)
 
 emitCopySmallArray :: (CmmExpr -> CmmExpr -> CmmExpr -> CmmExpr -> ByteOff
                        -> FCode ())  -- ^ copy function
diff --git a/compiler/codeGen/StgCmmProf.hs b/compiler/codeGen/StgCmmProf.hs
index a0bca5d661..f490be5c31 100644
--- a/compiler/codeGen/StgCmmProf.hs
+++ b/compiler/codeGen/StgCmmProf.hs
@@ -352,6 +352,7 @@ ldvEnter cl_ptr = do
         emit =<< mkCmmIfThenElse (CmmMachOp (mo_wordUGt dflags) [loadEra dflags, CmmLit (zeroCLit dflags)])
                      (mkStore ldv_wd new_ldv_wd)
                      mkNop
+                     Nothing
 
 loadEra :: DynFlags -> CmmExpr
 loadEra dflags = CmmMachOp (MO_UU_Conv (cIntWidth dflags) (wordWidth dflags))
diff --git a/compiler/codeGen/StgCmmUtils.hs b/compiler/codeGen/StgCmmUtils.hs
index 94013f5c6d..68949bf190 100644
--- a/compiler/codeGen/StgCmmUtils.hs
+++ b/compiler/codeGen/StgCmmUtils.hs
@@ -55,6 +55,7 @@ import CLabel
 import CmmUtils
 import CmmSwitch
 
+import BasicTypes (BranchWeight)
 import ForeignCall
 import IdInfo
 import Type
@@ -74,8 +75,6 @@ import RepType
 import qualified Data.ByteString as BS
 import qualified Data.Map as M
 import Data.Char
-import Data.List
-import Data.Ord
 import Data.Word
 
 
@@ -448,16 +447,16 @@ unscramble dflags vertices = mapM_ do_component components
 
 
 emitSwitch :: CmmExpr                      -- Tag to switch on
-           -> [(ConTagZ, CmmAGraphScoped)] -- Tagged branches
-           -> Maybe CmmAGraphScoped        -- Default branch (if any)
+           -> [(ConTagZ, CmmAGraphScoped, BranchWeight)] -- Tagged branches
+           -> Maybe (CmmAGraphScoped, BranchWeight)  -- Default branch (if any)
            -> ConTagZ -> ConTagZ           -- Min and Max possible values;
                                            -- behaviour outside this range is
                                            -- undefined
            -> FCode ()
 
 -- First, two rather common cases in which there is no work to do
-emitSwitch _ []         (Just code) _ _ = emit (fst code)
-emitSwitch _ [(_,code)] Nothing     _ _ = emit (fst code)
+emitSwitch _ []           (Just code) _ _ = emit (fst $ fst code)
+emitSwitch _ [(_,code,_)] Nothing     _ _ = emit (fst code)
 
 -- Right, off we go
 emitSwitch tag_expr branches mb_deflt lo_tag hi_tag = do
@@ -467,7 +466,8 @@ emitSwitch tag_expr branches mb_deflt lo_tag hi_tag = do
     tag_expr'     <- assignTemp' tag_expr
 
     -- Sort the branches before calling mk_discrete_switch
-    let branches_lbls' = [ (fromIntegral i, l) | (i,l) <- sortBy (comparing fst) branches_lbls ]
+    let branches_lbls' = [ (fromIntegral i, l, f)
+                         | (i,l,f) <- sortWith fstOf3 branches_lbls ]
     let range = (fromIntegral lo_tag, fromIntegral hi_tag)
 
     emit $ mk_discrete_switch False tag_expr' branches_lbls' mb_deflt_lbl range
@@ -476,19 +476,19 @@ emitSwitch tag_expr branches mb_deflt lo_tag hi_tag = do
 
 mk_discrete_switch :: Bool -- ^ Use signed comparisons
           -> CmmExpr
-          -> [(Integer, BlockId)]
-          -> Maybe BlockId
+          -> [(Integer, BlockId, BranchWeight)]
+          -> Maybe (BlockId, BranchWeight)
           -> (Integer, Integer)
           -> CmmAGraph
 
 -- SINGLETON TAG RANGE: no case analysis to do
-mk_discrete_switch _ _tag_expr [(tag, lbl)] _ (lo_tag, hi_tag)
+mk_discrete_switch _ _tag_expr [(tag, lbl, _f)] _ (lo_tag, hi_tag)
   | lo_tag == hi_tag
   = ASSERT( tag == lo_tag )
     mkBranch lbl
 
 -- SINGLETON BRANCH, NO DEFAULT: no case analysis to do
-mk_discrete_switch _ _tag_expr [(_tag,lbl)] Nothing _
+mk_discrete_switch _ _tag_expr [(_tag,lbl,_)] Nothing _
   = mkBranch lbl
         -- The simplifier might have eliminated a case
         --       so we may have e.g. case xs of
@@ -499,25 +499,17 @@ mk_discrete_switch _ _tag_expr [(_tag,lbl)] Nothing _
 -- SOMETHING MORE COMPLICATED: defer to CmmImplementSwitchPlans
 -- See Note [Cmm Switches, the general plan] in CmmSwitch
 mk_discrete_switch signed tag_expr branches mb_deflt range
-  = mkSwitch tag_expr $ mkSwitchTargets signed range mb_deflt (M.fromList branches)
-
-divideBranches :: Ord a => [(a,b)] -> ([(a,b)], a, [(a,b)])
-divideBranches branches = (lo_branches, mid, hi_branches)
-  where
-    -- 2 branches => n_branches `div` 2 = 1
-    --            => branches !! 1 give the *second* tag
-    -- There are always at least 2 branches here
-    (mid,_) = branches !! (length branches `div` 2)
-    (lo_branches, hi_branches) = span is_lo branches
-    is_lo (t,_) = t < mid
+  = mkSwitch tag_expr $
+      mkSwitchTargets signed range mb_deflt
+        (M.fromList $ map (\(i,e,f)-> (i,(e,f))) branches)
 
 --------------
 emitCmmLitSwitch :: CmmExpr                    -- Tag to switch on
-               -> [(Literal, CmmAGraphScoped)] -- Tagged branches
-               -> CmmAGraphScoped              -- Default branch (always)
+               -> [(Literal, CmmAGraphScoped, BranchWeight)] -- Tagged branches
+               -> (CmmAGraphScoped, BranchWeight)    -- Default branch (always)
                -> FCode ()                     -- Emit the code
-emitCmmLitSwitch _scrut []       deflt = emit $ fst deflt
-emitCmmLitSwitch scrut  branches deflt = do
+emitCmmLitSwitch _scrut []       (deflt,_dfreq) = emit $ fst deflt
+emitCmmLitSwitch scrut  branches (deflt,dfreq) = do
     scrut' <- assignTemp' scrut
     join_lbl <- newBlockId
     deflt_lbl <- label_code join_lbl deflt
@@ -529,20 +521,22 @@ emitCmmLitSwitch scrut  branches deflt = do
 
     -- We find the necessary type information in the literals in the branches
     let signed = case head branches of
-                    (MachInt _, _) ->   True
-                    (MachInt64 _, _) -> True
+                    (MachInt _, _, _) ->   True
+                    (MachInt64 _, _, _) -> True
                     _ -> False
 
     let range | signed    = (tARGET_MIN_INT dflags, tARGET_MAX_INT dflags)
               | otherwise = (0, tARGET_MAX_WORD dflags)
 
     if isFloatType cmm_ty
-    then emit =<< mk_float_switch rep scrut' deflt_lbl noBound branches_lbls
+    then emit =<< mk_float_switch rep scrut'
+                    (deflt_lbl, dfreq) noBound
+                    branches_lbls
     else emit $ mk_discrete_switch
         signed
         scrut'
-        [(litValue lit,l) | (lit,l) <- branches_lbls]
-        (Just deflt_lbl)
+        [(litValue lit,l,f) | (lit,l,f) <- branches_lbls]
+        (Just (deflt_lbl, dfreq))
         range
     emitLabel join_lbl
 
@@ -552,11 +546,30 @@ type LitBound = (Maybe Literal, Maybe Literal)
 noBound :: LitBound
 noBound = (Nothing, Nothing)
 
-mk_float_switch :: Width -> CmmExpr -> BlockId
+{- TODO:
+  Currently this generates a binary search tree for the given value.
+
+  Given we have branch weights we would ideally balance the tree
+  by weight instead.
+
+  Eg. given (lit,weight) of [(0,1),(1,1),(2,1),(3,99)] we want to split the
+  list into [(0,1),(1,1),(2,1)] and [(3,99)].
+
+  Things to consider:
+    * Does it make a difference often enough to be worth the complexity
+      and increase in compile time.
+    * Negative weights have to be rounded up to zero,
+      otherwise they would distort the results.
+    * How should entries with no information be treated?
+      -> Probably good enough to use the default value.
+    * If implemented should this only apply when optimizations are
+      active?
+-}
+mk_float_switch :: Width -> CmmExpr -> (BlockId, BranchWeight)
               -> LitBound
-              -> [(Literal,BlockId)]
+              -> [(Literal,BlockId,BranchWeight)]
               -> FCode CmmAGraph
-mk_float_switch rep scrut deflt _bounds [(lit,blk)]
+mk_float_switch rep scrut (deflt, _dfrq) _bounds [(lit,blk,_frq)]
   = do dflags <- getDynFlags
        return $ mkCbranch (cond dflags) deflt blk Nothing
   where
@@ -565,17 +578,32 @@ mk_float_switch rep scrut deflt _bounds [(lit,blk)]
         cmm_lit = mkSimpleLit dflags lit
         ne      = MO_F_Ne rep
 
-mk_float_switch rep scrut deflt_blk_id (lo_bound, hi_bound) branches
+mk_float_switch rep scrut (deflt_blk_id,dfreq) (lo_bound, hi_bound) branches
   = do dflags <- getDynFlags
-       lo_blk <- mk_float_switch rep scrut deflt_blk_id bounds_lo lo_branches
-       hi_blk <- mk_float_switch rep scrut deflt_blk_id bounds_hi hi_branches
-       mkCmmIfThenElse (cond dflags) lo_blk hi_blk
+       lo_blk <- mk_float_switch
+                    rep scrut (deflt_blk_id,dfreq)
+                    bounds_lo lo_branches
+       hi_blk <- mk_float_switch
+                    rep scrut
+                    (deflt_blk_id,dfreq) bounds_hi hi_branches
+       mkCmmIfThenElse (cond dflags) lo_blk hi_blk Nothing
   where
+
     (lo_branches, mid_lit, hi_branches) = divideBranches branches
 
     bounds_lo = (lo_bound, Just mid_lit)
     bounds_hi = (Just mid_lit, hi_bound)
 
+    divideBranches :: Ord a => [(a,b,c)] -> ([(a,b,c)], a, [(a,b,c)])
+    divideBranches branches = (lo_branches, mid, hi_branches)
+      where
+        -- 2 branches => n_branches `div` 2 = 1
+        --            => branches !! 1 give the *second* tag
+        -- There are always at least 2 branches here
+        (mid,_,_) = branches !! (length branches `div` 2)
+        (lo_branches, hi_branches) = span is_lo branches
+        is_lo (t,_,_) = t < mid
+
     cond dflags = CmmMachOp lt [scrut, CmmLit cmm_lit]
       where
         cmm_lit = mkSimpleLit dflags mid_lit
@@ -583,21 +611,23 @@ mk_float_switch rep scrut deflt_blk_id (lo_bound, hi_bound) branches
 
 
 --------------
-label_default :: BlockId -> Maybe CmmAGraphScoped -> FCode (Maybe BlockId)
+label_default :: BlockId -> Maybe (CmmAGraphScoped, BranchWeight)
+              -> FCode (Maybe (BlockId, BranchWeight))
 label_default _ Nothing
   = return Nothing
-label_default join_lbl (Just code)
+label_default join_lbl (Just (code,f))
   = do lbl <- label_code join_lbl code
-       return (Just lbl)
+       return (Just (lbl,f))
 
 --------------
-label_branches :: BlockId -> [(a,CmmAGraphScoped)] -> FCode [(a,BlockId)]
+label_branches :: BlockId -> [(a,CmmAGraphScoped, BranchWeight)]
+               -> FCode [(a,BlockId,BranchWeight)]
 label_branches _join_lbl []
   = return []
-label_branches join_lbl ((tag,code):branches)
+label_branches join_lbl ((tag,code,freq):branches)
   = do lbl <- label_code join_lbl code
        branches' <- label_branches join_lbl branches
-       return ((tag,lbl):branches')
+       return ((tag,lbl,freq):branches')
 
 --------------
 label_code :: BlockId -> CmmAGraphScoped -> FCode BlockId