Allow resizing the stack for the graph allocator.

The graph allocator now dynamically resizes the number of stack
slots when running into the limit.

This fixes #8657.

Also loop membership of basic blocks is now available
in the register allocator for cost heuristics.

6 files changed, 105 insertions, 36 deletions

diff --git a/compiler/nativeGen/AsmCodeGen.hs b/compiler/nativeGen/AsmCodeGen.hs
index 956528bf81..8c62a15429 100644
--- a/compiler/nativeGen/AsmCodeGen.hs
+++ b/compiler/nativeGen/AsmCodeGen.hs
@@ -608,14 +608,26 @@ cmmNativeGen dflags this_mod modLoc ncgImpl us fileIds dbgMap cmm count
                         $ allocatableRegs ncgImpl
 
                 -- do the graph coloring register allocation
-                let ((alloced, regAllocStats), usAlloc)
+                let ((alloced, maybe_more_stack, regAllocStats), usAlloc)
                         = {-# SCC "RegAlloc-color" #-}
                           initUs usLive
                           $ Color.regAlloc
                                 dflags
                                 alloc_regs
                                 (mkUniqSet [0 .. maxSpillSlots ncgImpl])
+                                (maxSpillSlots ncgImpl)
                                 withLiveness
+                                livenessCfg
+
+                let ((alloced', stack_updt_blks), usAlloc')
+                        = initUs usAlloc $
+                                case maybe_more_stack of
+                                Nothing     -> return (alloced, [])
+                                Just amount -> do
+                                    (alloced',stack_updt_blks) <- unzip <$>
+                                                (mapM ((ncgAllocMoreStack ncgImpl) amount) alloced)
+                                    return (alloced', concat stack_updt_blks )
+
 
                 -- dump out what happened during register allocation
                 dumpIfSet_dyn dflags
@@ -637,10 +649,10 @@ cmmNativeGen dflags this_mod modLoc ncgImpl us fileIds dbgMap cmm count
                 -- force evaluation of the Maybe to avoid space leak
                 mPprStats `seq` return ()
 
-                return  ( alloced, usAlloc
+                return  ( alloced', usAlloc'
                         , mPprStats
                         , Nothing
-                        , [], [])
+                        , [], stack_updt_blks)
 
           else do
                 -- do linear register allocation
diff --git a/compiler/nativeGen/CFG.hs b/compiler/nativeGen/CFG.hs
index b19db02b13..155e5bcac4 100644
--- a/compiler/nativeGen/CFG.hs
+++ b/compiler/nativeGen/CFG.hs
@@ -24,6 +24,7 @@ module CFG
     , getSuccEdgesSorted, weightedEdgeList
     , getEdgeInfo
     , getCfgNodes, hasNode
+    , loopMembers
 
     --Construction/Misc
     , getCfg, getCfgProc, pprEdgeWeights, sanityCheckCfg
@@ -636,3 +637,20 @@ optimizeCFG weights (CmmProc info _lab _live graph) cfg =
           | CmmSource (CmmBranch {}) <- source = True
           | CmmSource (CmmCondBranch {}) <- source = True
           | otherwise = False
+
+-- | Determine loop membership of blocks based on SCC analysis
+--   Ideally we would replace this with a variant giving us loop
+--   levels instead but the SCC code will do for now.
+loopMembers :: CFG -> LabelMap Bool
+loopMembers cfg =
+    foldl' (flip setLevel) mapEmpty sccs
+  where
+    mkNode :: BlockId -> Node BlockId BlockId
+    mkNode bid = DigraphNode bid bid (getSuccessors cfg bid)
+    nodes = map mkNode (setElems $ getCfgNodes cfg)
+
+    sccs = stronglyConnCompFromEdgedVerticesOrd nodes
+
+    setLevel :: SCC BlockId -> LabelMap Bool -> LabelMap Bool
+    setLevel (AcyclicSCC bid) m = mapInsert bid False m
+    setLevel (CyclicSCC bids) m = foldl' (\m k -> mapInsert k True m) m bids
diff --git a/compiler/nativeGen/RegAlloc/Graph/Main.hs b/compiler/nativeGen/RegAlloc/Graph/Main.hs
index 4c17d930ea..146f88a8b6 100644
--- a/compiler/nativeGen/RegAlloc/Graph/Main.hs
+++ b/compiler/nativeGen/RegAlloc/Graph/Main.hs
@@ -26,6 +26,7 @@ import UniqFM
 import UniqSet
 import UniqSupply
 import Util (seqList)
+import CFG
 
 import Data.Maybe
 import Control.Monad
@@ -46,12 +47,15 @@ regAlloc
         => DynFlags
         -> UniqFM (UniqSet RealReg)     -- ^ registers we can use for allocation
         -> UniqSet Int                  -- ^ set of available spill slots.
+        -> Int                          -- ^ current number of spill slots
         -> [LiveCmmDecl statics instr]  -- ^ code annotated with liveness information.
-        -> UniqSM ( [NatCmmDecl statics instr], [RegAllocStats statics instr] )
-           -- ^ code with registers allocated and stats for each stage of
-           -- allocation
+        -> Maybe CFG                    -- ^ CFG of basic blocks if available
+        -> UniqSM ( [NatCmmDecl statics instr]
+                  , Maybe Int, [RegAllocStats statics instr] )
+           -- ^ code with registers allocated, additional stacks required
+           -- and stats for each stage of allocation
 
-regAlloc dflags regsFree slotsFree code
+regAlloc dflags regsFree slotsFree slotsCount code cfg
  = do
         -- TODO: the regClass function is currently hard coded to the default
         --       target architecture. Would prefer to determine this from dflags.
@@ -61,12 +65,19 @@ regAlloc dflags regsFree slotsFree code
                         (targetVirtualRegSqueeze platform)
                         (targetRealRegSqueeze platform)
 
-        (code_final, debug_codeGraphs, _)
+        (code_final, debug_codeGraphs, slotsCount', _)
                 <- regAlloc_spin dflags 0
                         triv
-                        regsFree slotsFree [] code
+                        regsFree slotsFree slotsCount [] code cfg
+
+        let needStack
+                | slotsCount == slotsCount'
+                = Nothing
+                | otherwise
+                = Just slotsCount'
 
         return  ( code_final
+                , needStack
                 , reverse debug_codeGraphs )
 
 
@@ -88,13 +99,16 @@ regAlloc_spin
                 --   colourable.
         -> UniqFM (UniqSet RealReg)      -- ^ Free registers that we can allocate.
         -> UniqSet Int                   -- ^ Free stack slots that we can use.
+        -> Int                           -- ^ Number of spill slots in use
         -> [RegAllocStats statics instr] -- ^ Current regalloc stats to add to.
         -> [LiveCmmDecl statics instr]   -- ^ Liveness annotated code to allocate.
+        -> Maybe CFG
         -> UniqSM ( [NatCmmDecl statics instr]
                   , [RegAllocStats statics instr]
+                  , Int                  -- Slots in use
                   , Color.Graph VirtualReg RegClass RealReg)
 
-regAlloc_spin dflags spinCount triv regsFree slotsFree debug_codeGraphs code
+regAlloc_spin dflags spinCount triv regsFree slotsFree slotsCount debug_codeGraphs code cfg
  = do
         let platform = targetPlatform dflags
 
@@ -134,7 +148,7 @@ regAlloc_spin dflags spinCount triv regsFree slotsFree debug_codeGraphs code
         -- This is a lazy binding, so the map will only be computed if we
         -- actually have to spill to the stack.
         let spillCosts  = foldl' plusSpillCostInfo zeroSpillCostInfo
-                        $ map (slurpSpillCostInfo platform) code
+                        $ map (slurpSpillCostInfo platform cfg) code
 
         -- The function to choose regs to leave uncolored.
         let spill       = chooseSpill spillCosts
@@ -227,6 +241,7 @@ regAlloc_spin dflags spinCount triv regsFree slotsFree debug_codeGraphs code
 
                 return  ( code_final
                         , statList
+                        , slotsCount
                         , graph_colored_lint)
 
          -- Coloring was unsuccessful. We need to spill some register to the
@@ -241,8 +256,8 @@ regAlloc_spin dflags spinCount triv regsFree slotsFree debug_codeGraphs code
                                 else graph_colored
 
                 -- Spill uncolored regs to the stack.
-                (code_spilled, slotsFree', spillStats)
-                        <- regSpill platform code_coalesced slotsFree rsSpill
+                (code_spilled, slotsFree', slotsCount', spillStats)
+                        <- regSpill platform code_coalesced slotsFree slotsCount rsSpill
 
                 -- Recalculate liveness information.
                 -- NOTE: we have to reverse the SCCs here to get them back into
@@ -273,8 +288,7 @@ regAlloc_spin dflags spinCount triv regsFree slotsFree debug_codeGraphs code
                 seqList statList (return ())
 
                 regAlloc_spin dflags (spinCount + 1) triv regsFree slotsFree'
-                        statList
-                        code_relive
+                              slotsCount' statList code_relive cfg
 
 
 -- | Build a graph from the liveness and coalesce information in this code.
diff --git a/compiler/nativeGen/RegAlloc/Graph/Spill.hs b/compiler/nativeGen/RegAlloc/Graph/Spill.hs
index bce24bdd3c..2e1879926e 100644
--- a/compiler/nativeGen/RegAlloc/Graph/Spill.hs
+++ b/compiler/nativeGen/RegAlloc/Graph/Spill.hs
@@ -33,6 +33,9 @@ import qualified Data.IntSet    as IntSet
 
 -- | Spill all these virtual regs to stack slots.
 --
+--   Bumps the number of required stack slots if required.
+--
+--
 --   TODO: See if we can split some of the live ranges instead of just globally
 --         spilling the virtual reg. This might make the spill cleaner's job easier.
 --
@@ -45,20 +48,22 @@ regSpill
         => Platform
         -> [LiveCmmDecl statics instr]  -- ^ the code
         -> UniqSet Int                  -- ^ available stack slots
+        -> Int                          -- ^ current number of spill slots.
         -> UniqSet VirtualReg           -- ^ the regs to spill
         -> UniqSM
             ([LiveCmmDecl statics instr]
                  -- code with SPILL and RELOAD meta instructions added.
             , UniqSet Int               -- left over slots
+            , Int                       -- slot count in use now.
             , SpillStats )              -- stats about what happened during spilling
 
-regSpill platform code slotsFree regs
+regSpill platform code slotsFree slotCount regs
 
         -- Not enough slots to spill these regs.
         | sizeUniqSet slotsFree < sizeUniqSet regs
-        = pprPanic "regSpill: out of spill slots!"
-                (  text "   regs to spill = " <> ppr (sizeUniqSet regs)
-                $$ text "   slots left    = " <> ppr (sizeUniqSet slotsFree))
+        = -- pprTrace "Bumping slot count:" (ppr slotCount <> text " -> " <> ppr (slotCount+512)) $
+          let slotsFree' = (addListToUniqSet slotsFree [slotCount+1 .. slotCount+512])
+          in regSpill platform code slotsFree' (slotCount+512) regs
 
         | otherwise
         = do
@@ -80,6 +85,7 @@ regSpill platform code slotsFree regs
 
                 return  ( code'
                         , minusUniqSet slotsFree (mkUniqSet slots)
+                        , slotCount
                         , makeSpillStats state')
 
 
diff --git a/compiler/nativeGen/RegAlloc/Graph/SpillCost.hs b/compiler/nativeGen/RegAlloc/Graph/SpillCost.hs
index f603b609df..4d5f44a8d3 100644
--- a/compiler/nativeGen/RegAlloc/Graph/SpillCost.hs
+++ b/compiler/nativeGen/RegAlloc/Graph/SpillCost.hs
@@ -1,4 +1,4 @@
-
+{-# LANGUAGE ScopedTypeVariables #-}
 module RegAlloc.Graph.SpillCost (
         SpillCostRecord,
         plusSpillCostRecord,
@@ -30,9 +30,11 @@ import Digraph          (flattenSCCs)
 import Outputable
 import Platform
 import State
+import CFG
 
 import Data.List        (nub, minimumBy)
 import Data.Maybe
+import Control.Monad (join)
 
 
 -- | Records the expected cost to spill some regster.
@@ -47,6 +49,10 @@ type SpillCostRecord
 type SpillCostInfo
         = UniqFM SpillCostRecord
 
+-- | Block membership in a loop
+type LoopMember = Bool
+
+type SpillCostState = State (UniqFM SpillCostRecord) ()
 
 -- | An empty map of spill costs.
 zeroSpillCostInfo :: SpillCostInfo
@@ -71,12 +77,13 @@ plusSpillCostRecord (r1, a1, b1, c1) (r2, a2, b2, c2)
 --   For each vreg, the number of times it was written to, read from,
 --   and the number of instructions it was live on entry to (lifetime)
 --
-slurpSpillCostInfo :: (Outputable instr, Instruction instr)
+slurpSpillCostInfo :: forall instr statics. (Outputable instr, Instruction instr)
                    => Platform
+                   -> Maybe CFG
                    -> LiveCmmDecl statics instr
                    -> SpillCostInfo
 
-slurpSpillCostInfo platform cmm
+slurpSpillCostInfo platform cfg cmm
         = execState (countCmm cmm) zeroSpillCostInfo
  where
         countCmm CmmData{}              = return ()
@@ -90,35 +97,36 @@ slurpSpillCostInfo platform cmm
                 | LiveInfo _ _ (Just blockLive) _ <- info
                 , Just rsLiveEntry  <- mapLookup blockId blockLive
                 , rsLiveEntry_virt  <- takeVirtuals rsLiveEntry
-                = countLIs rsLiveEntry_virt instrs
+                = countLIs (loopMember blockId) rsLiveEntry_virt instrs
 
                 | otherwise
                 = error "RegAlloc.SpillCost.slurpSpillCostInfo: bad block"
 
-        countLIs _      []
+        countLIs :: LoopMember -> UniqSet VirtualReg -> [LiveInstr instr] -> SpillCostState
+        countLIs _      _      []
                 = return ()
 
         -- Skip over comment and delta pseudo instrs.
-        countLIs rsLive (LiveInstr instr Nothing : lis)
+        countLIs inLoop rsLive (LiveInstr instr Nothing : lis)
                 | isMetaInstr instr
-                = countLIs rsLive lis
+                = countLIs inLoop rsLive lis
 
                 | otherwise
                 = pprPanic "RegSpillCost.slurpSpillCostInfo"
                 $ text "no liveness information on instruction " <> ppr instr
 
-        countLIs rsLiveEntry (LiveInstr instr (Just live) : lis)
+        countLIs inLoop rsLiveEntry (LiveInstr instr (Just live) : lis)
          = do
                 -- Increment the lifetime counts for regs live on entry to this instr.
-                mapM_ incLifetime $ nonDetEltsUniqSet rsLiveEntry
+                mapM_ (incLifetime (loopCount inLoop)) $ nonDetEltsUniqSet rsLiveEntry
                     -- This is non-deterministic but we do not
                     -- currently support deterministic code-generation.
                     -- See Note [Unique Determinism and code generation]
 
                 -- Increment counts for what regs were read/written from.
                 let (RU read written)   = regUsageOfInstr platform instr
-                mapM_ incUses   $ catMaybes $ map takeVirtualReg $ nub read
-                mapM_ incDefs   $ catMaybes $ map takeVirtualReg $ nub written
+                mapM_ (incUses (loopCount inLoop)) $ catMaybes $ map takeVirtualReg $ nub read
+                mapM_ (incDefs (loopCount inLoop)) $ catMaybes $ map takeVirtualReg $ nub written
 
                 -- Compute liveness for entry to next instruction.
                 let liveDieRead_virt    = takeVirtuals (liveDieRead  live)
@@ -132,12 +140,21 @@ slurpSpillCostInfo platform cmm
                         = (rsLiveAcross `unionUniqSets` liveBorn_virt)
                                         `minusUniqSet`  liveDieWrite_virt
 
-                countLIs rsLiveNext lis
-
-        incDefs     reg = modify $ \s -> addToUFM_C plusSpillCostRecord s reg (reg, 1, 0, 0)
-        incUses     reg = modify $ \s -> addToUFM_C plusSpillCostRecord s reg (reg, 0, 1, 0)
-        incLifetime reg = modify $ \s -> addToUFM_C plusSpillCostRecord s reg (reg, 0, 0, 1)
-
+                countLIs inLoop rsLiveNext lis
+
+        loopCount inLoop
+          | inLoop = 10
+          | otherwise = 1
+        incDefs     count reg = modify $ \s -> addToUFM_C plusSpillCostRecord s reg (reg, count, 0, 0)
+        incUses     count reg = modify $ \s -> addToUFM_C plusSpillCostRecord s reg (reg, 0, count, 0)
+        incLifetime count reg = modify $ \s -> addToUFM_C plusSpillCostRecord s reg (reg, 0, 0, count)
+
+        loopBlocks = CFG.loopMembers <$> cfg
+        loopMember bid
+          | Just isMember <- join (mapLookup bid <$> loopBlocks)
+          = isMember
+          | otherwise
+          = False
 
 -- | Take all the virtual registers from this set.
 takeVirtuals :: UniqSet Reg -> UniqSet VirtualReg
diff --git a/compiler/nativeGen/X86/Instr.hs b/compiler/nativeGen/X86/Instr.hs
index 5e790e481e..4717ec10d6 100644
--- a/compiler/nativeGen/X86/Instr.hs
+++ b/compiler/nativeGen/X86/Instr.hs
@@ -1063,6 +1063,8 @@ is_G_instr instr
 -- Otherwise, we would repeat the $rsp adjustment for each branch to
 -- L.
 --
+-- Returns a list of (L,Lnew) pairs.
+--
 allocMoreStack
   :: Platform
   -> Int