Allow multiple entry points when allocating recursive groups (#9303)

Summary:
In this example we ended up with some code that was only reachable via
an info table, because a branch had been optimised away by the native
code generator.  The register allocator then got confused because it
was only considering the first block of the proc to be an entry point,
when actually any of the info tables are entry points.

Test Plan: validate

Reviewers: simonpj, austin

Subscribers: simonmar, relrod, carter

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

2 files changed, 39 insertions, 35 deletions

diff --git a/compiler/nativeGen/RegAlloc/Linear/Main.hs b/compiler/nativeGen/RegAlloc/Linear/Main.hs
index ee43d25aa3..35416926ed 100644
--- a/compiler/nativeGen/RegAlloc/Linear/Main.hs
+++ b/compiler/nativeGen/RegAlloc/Linear/Main.hs
@@ -158,11 +158,11 @@ regAlloc _ (CmmProc (LiveInfo info _ _ _) lbl live [])
                  , Nothing )
 
 regAlloc dflags (CmmProc static lbl live sccs)
-        | LiveInfo info (Just first_id) (Just block_live) _     <- static
+        | LiveInfo info entry_ids@(first_id:_) (Just block_live) _ <- static
         = do
                 -- do register allocation on each component.
                 (final_blocks, stats, stack_use)
-                        <- linearRegAlloc dflags first_id block_live sccs
+                        <- linearRegAlloc dflags entry_ids block_live sccs
 
                 -- make sure the block that was first in the input list
                 --      stays at the front of the output
@@ -196,18 +196,18 @@ regAlloc _ (CmmProc _ _ _ _)
 linearRegAlloc
         :: (Outputable instr, Instruction instr)
         => DynFlags
-        -> BlockId                      -- ^ the first block
+        -> [BlockId]                    -- ^ entry points
         -> BlockMap RegSet              -- ^ live regs on entry to each basic block
         -> [SCC (LiveBasicBlock instr)] -- ^ instructions annotated with "deaths"
         -> UniqSM ([NatBasicBlock instr], RegAllocStats, Int)
 
-linearRegAlloc dflags first_id block_live sccs
+linearRegAlloc dflags entry_ids block_live sccs
  = let platform = targetPlatform dflags
    in case platformArch platform of
-      ArchX86       -> linearRegAlloc' dflags (frInitFreeRegs platform :: X86.FreeRegs)    first_id block_live sccs
-      ArchX86_64    -> linearRegAlloc' dflags (frInitFreeRegs platform :: X86_64.FreeRegs) first_id block_live sccs
-      ArchSPARC     -> linearRegAlloc' dflags (frInitFreeRegs platform :: SPARC.FreeRegs)  first_id block_live sccs
-      ArchPPC       -> linearRegAlloc' dflags (frInitFreeRegs platform :: PPC.FreeRegs)    first_id block_live sccs
+      ArchX86       -> linearRegAlloc' dflags (frInitFreeRegs platform :: X86.FreeRegs)    entry_ids block_live sccs
+      ArchX86_64    -> linearRegAlloc' dflags (frInitFreeRegs platform :: X86_64.FreeRegs) entry_ids block_live sccs
+      ArchSPARC     -> linearRegAlloc' dflags (frInitFreeRegs platform :: SPARC.FreeRegs)  entry_ids block_live sccs
+      ArchPPC       -> linearRegAlloc' dflags (frInitFreeRegs platform :: PPC.FreeRegs)    entry_ids block_live sccs
       ArchARM _ _ _ -> panic "linearRegAlloc ArchARM"
       ArchARM64     -> panic "linearRegAlloc ArchARM64"
       ArchPPC_64    -> panic "linearRegAlloc ArchPPC_64"
@@ -221,21 +221,21 @@ linearRegAlloc'
         :: (FR freeRegs, Outputable instr, Instruction instr)
         => DynFlags
         -> freeRegs
-        -> BlockId                      -- ^ the first block
+        -> [BlockId]                    -- ^ entry points
         -> BlockMap RegSet              -- ^ live regs on entry to each basic block
         -> [SCC (LiveBasicBlock instr)] -- ^ instructions annotated with "deaths"
         -> UniqSM ([NatBasicBlock instr], RegAllocStats, Int)
 
-linearRegAlloc' dflags initFreeRegs first_id block_live sccs
+linearRegAlloc' dflags initFreeRegs entry_ids block_live sccs
  = do   us      <- getUs
         let (_, stack, stats, blocks) =
                 runR dflags emptyBlockMap initFreeRegs emptyRegMap (emptyStackMap dflags) us
-                    $ linearRA_SCCs first_id block_live [] sccs
+                    $ linearRA_SCCs entry_ids block_live [] sccs
         return  (blocks, stats, getStackUse stack)
 
 
 linearRA_SCCs :: (FR freeRegs, Instruction instr, Outputable instr)
-              => BlockId
+              => [BlockId]
               -> BlockMap RegSet
               -> [NatBasicBlock instr]
               -> [SCC (LiveBasicBlock instr)]
@@ -244,16 +244,16 @@ linearRA_SCCs :: (FR freeRegs, Instruction instr, Outputable instr)
 linearRA_SCCs _ _ blocksAcc []
         = return $ reverse blocksAcc
 
-linearRA_SCCs first_id block_live blocksAcc (AcyclicSCC block : sccs)
+linearRA_SCCs entry_ids block_live blocksAcc (AcyclicSCC block : sccs)
  = do   blocks' <- processBlock block_live block
-        linearRA_SCCs first_id block_live
+        linearRA_SCCs entry_ids block_live
                 ((reverse blocks') ++ blocksAcc)
                 sccs
 
-linearRA_SCCs first_id block_live blocksAcc (CyclicSCC blocks : sccs)
+linearRA_SCCs entry_ids block_live blocksAcc (CyclicSCC blocks : sccs)
  = do
-        blockss' <- process first_id block_live blocks [] (return []) False
-        linearRA_SCCs first_id block_live
+        blockss' <- process entry_ids block_live blocks [] (return []) False
+        linearRA_SCCs entry_ids block_live
                 (reverse (concat blockss') ++ blocksAcc)
                 sccs
 
@@ -270,7 +270,7 @@ linearRA_SCCs first_id block_live blocksAcc (CyclicSCC blocks : sccs)
 -}
 
 process :: (FR freeRegs, Instruction instr, Outputable instr)
-        => BlockId
+        => [BlockId]
         -> BlockMap RegSet
         -> [GenBasicBlock (LiveInstr instr)]
         -> [GenBasicBlock (LiveInstr instr)]
@@ -281,7 +281,7 @@ process :: (FR freeRegs, Instruction instr, Outputable instr)
 process _ _ [] []         accum _
         = return $ reverse accum
 
-process first_id block_live [] next_round accum madeProgress
+process entry_ids block_live [] next_round accum madeProgress
         | not madeProgress
 
           {- BUGS: There are so many unreachable blocks in the code the warnings are overwhelming.
@@ -291,22 +291,22 @@ process first_id block_live [] next_round accum madeProgress
         = return $ reverse accum
 
         | otherwise
-        = process first_id block_live
+        = process entry_ids block_live
                   next_round [] accum False
 
-process first_id block_live (b@(BasicBlock id _) : blocks)
+process entry_ids block_live (b@(BasicBlock id _) : blocks)
         next_round accum madeProgress
  = do
         block_assig <- getBlockAssigR
 
         if isJust (mapLookup id block_assig)
-             || id == first_id
+             || id `elem` entry_ids
          then do
                 b'  <- processBlock block_live b
-                process first_id block_live blocks
+                process entry_ids block_live blocks
                         next_round (b' : accum) True
 
-         else   process first_id block_live blocks
+         else   process entry_ids block_live blocks
                         (b : next_round) accum madeProgress
 
 
diff --git a/compiler/nativeGen/RegAlloc/Liveness.hs b/compiler/nativeGen/RegAlloc/Liveness.hs
index 1cb6dc8268..d7fd8bdcb4 100644
--- a/compiler/nativeGen/RegAlloc/Liveness.hs
+++ b/compiler/nativeGen/RegAlloc/Liveness.hs
@@ -169,10 +169,11 @@ data Liveness
 -- | Stash regs live on entry to each basic block in the info part of the cmm code.
 data LiveInfo
         = LiveInfo
-                (BlockEnv CmmStatics)                   -- cmm info table static stuff
-                (Maybe BlockId)                         -- id of the first block
-                (Maybe (BlockMap RegSet))               -- argument locals live on entry to this block
-                (Map BlockId (Set Int))                 -- stack slots live on entry to this block
+                (BlockEnv CmmStatics)     -- cmm info table static stuff
+                [BlockId]                 -- entry points (first one is the
+                                          -- entry point for the proc).
+                (Maybe (BlockMap RegSet)) -- argument locals live on entry to this block
+                (Map BlockId (Set Int))   -- stack slots live on entry to this block
 
 
 -- | A basic block with liveness information.
@@ -223,9 +224,9 @@ instance Outputable instr
                  | otherwise            = name <> (hcat $ punctuate space $ map ppr $ uniqSetToList regs)
 
 instance Outputable LiveInfo where
-    ppr (LiveInfo mb_static firstId liveVRegsOnEntry liveSlotsOnEntry)
+    ppr (LiveInfo mb_static entryIds liveVRegsOnEntry liveSlotsOnEntry)
         =  (ppr mb_static)
-        $$ text "# firstId          = " <> ppr firstId
+        $$ text "# entryIds         = " <> ppr entryIds
         $$ text "# liveVRegsOnEntry = " <> ppr liveVRegsOnEntry
         $$ text "# liveSlotsOnEntry = " <> text (show liveSlotsOnEntry)
 
@@ -480,7 +481,7 @@ stripLive dflags live
  where  stripCmm :: (Outputable statics, Outputable instr, Instruction instr)
                  => LiveCmmDecl statics instr -> NatCmmDecl statics instr
         stripCmm (CmmData sec ds)       = CmmData sec ds
-        stripCmm (CmmProc (LiveInfo info (Just first_id) _ _) label live sccs)
+        stripCmm (CmmProc (LiveInfo info (first_id:_) _ _) label live sccs)
          = let  final_blocks    = flattenSCCs sccs
 
                 -- make sure the block that was first in the input list
@@ -493,7 +494,7 @@ stripLive dflags live
                           (ListGraph $ map (stripLiveBlock dflags) $ first' : rest')
 
         -- procs used for stg_split_markers don't contain any blocks, and have no first_id.
-        stripCmm (CmmProc (LiveInfo info Nothing _ _) label live [])
+        stripCmm (CmmProc (LiveInfo info [] _ _) label live [])
          =      CmmProc info label live (ListGraph [])
 
         -- If the proc has blocks but we don't know what the first one was, then we're dead.
@@ -641,16 +642,19 @@ natCmmTopToLive (CmmData i d)
         = CmmData i d
 
 natCmmTopToLive (CmmProc info lbl live (ListGraph []))
-        = CmmProc (LiveInfo info Nothing Nothing Map.empty) lbl live []
+        = CmmProc (LiveInfo info [] Nothing Map.empty) lbl live []
 
 natCmmTopToLive proc@(CmmProc info lbl live (ListGraph blocks@(first : _)))
  = let  first_id        = blockId first
-        sccs            = sccBlocks blocks (entryBlocks proc)
+        all_entry_ids   = entryBlocks proc
+        sccs            = sccBlocks blocks all_entry_ids
+        entry_ids       = filter (/= first_id) all_entry_ids
         sccsLive        = map (fmap (\(BasicBlock l instrs) ->
                                         BasicBlock l (map (\i -> LiveInstr (Instr i) Nothing) instrs)))
                         $ sccs
 
-   in   CmmProc (LiveInfo info (Just first_id) Nothing Map.empty) lbl live sccsLive
+   in   CmmProc (LiveInfo info (first_id : entry_ids) Nothing Map.empty)
+                lbl live sccsLive
 
 
 --