RegAlloc: Track slot liveness over jumps in spill cleaner

compiler/nativeGen/RegAlloc/Graph/Spill.hs

-
 {-# OPTIONS -fno-warn-missing-signatures #-}
 
+-- | When there aren't enough registers to hold all the vregs we have to spill some of those
+--   vregs to slots on the stack. This module is used modify the code to use those slots.
+--
 module RegAlloc.Graph.Spill (
 	regSpill,
 	SpillStats(..),
 	accSpillSL
 )
-
 where
-
 import RegAlloc.Liveness
 import Instruction
 import Reg
-import Cmm
+import Cmm	hiding (RegSet)
+import BlockId
 
 import State
 import Unique
@@ -22,15 +23,21 @@ import UniqSupply
 import Outputable
 
 import Data.List
+import Data.Maybe
+import Data.Map			(Map)
+import Data.Set			(Set)
+import qualified Data.Map	as Map
+import qualified Data.Set	as Set
 
 
--- | Spill all these virtual regs to memory
---	TODO: 	see if we can split some of the live ranges instead of just globally
---		spilling the virtual reg.
+-- | Spill all these virtual regs to stack slots.
+-- 
+--   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.
 --
---	TODO:	On ciscy x86 and x86_64 we don't nessesarally have to add a mov instruction
---		when making spills. If an instr is using a spilled virtual we may be able to
---		address the spill slot directly.
+--   TODO: On CISCy x86 and x86_64 we don't nessesarally have to add a mov instruction
+--         when making spills. If an instr is using a spilled virtual we may be able to
+--         address the spill slot directly.
 --
 regSpill
 	:: Instruction instr
@@ -38,7 +45,7 @@ regSpill
 	-> UniqSet Int			-- ^ available stack slots
 	-> UniqSet VirtualReg		-- ^ the regs to spill
 	-> UniqSM
-		([LiveCmmTop instr]	-- code will spill instructions
+		([LiveCmmTop instr]	-- code with SPILL and RELOAD meta instructions added.
 		, UniqSet Int		-- left over slots
 		, SpillStats )		-- stats about what happened during spilling
 
@@ -62,7 +69,7 @@ regSpill code slotsFree regs
 
 		-- run the spiller on all the blocks
 		let (code', state')	=
-			runState (mapM (mapBlockTopM (regSpill_block regSlotMap)) code)
+			runState (mapM (regSpill_top regSlotMap) code)
 				 (initSpillS us)
 
 		return	( code'
@@ -70,15 +77,84 @@ regSpill code slotsFree regs
 			, makeSpillStats state')
 
 
+-- | Spill some registers to stack slots in a top-level thing.
+regSpill_top 
+	:: Instruction instr
+	=> RegMap Int 			-- ^ map of vregs to slots they're being spilled to.
+	-> LiveCmmTop instr		-- ^ the top level thing.
+	-> SpillM (LiveCmmTop instr)
+	
+regSpill_top regSlotMap cmm
+ = case cmm of
+	CmmData{}				
+	 -> return cmm
+
+	CmmProc info label params sccs
+	 |  LiveInfo static firstId mLiveVRegsOnEntry liveSlotsOnEntry <- info
+	 -> do	
+		-- We should only passed Cmms with the liveness maps filled in,  but we'll
+		-- create empty ones if they're not there just in case.
+		let liveVRegsOnEntry	= fromMaybe emptyBlockEnv mLiveVRegsOnEntry
+		
+		-- The liveVRegsOnEntry contains the set of vregs that are live on entry to
+		-- each basic block. If we spill one of those vregs we remove it from that
+		-- set and add the corresponding slot number to the liveSlotsOnEntry set.
+		-- The spill cleaner needs this information to erase unneeded spill and 
+		-- reload instructions after we've done a successful allocation.
+		let liveSlotsOnEntry' :: Map BlockId (Set Int)
+		    liveSlotsOnEntry'
+			= foldBlockEnv patchLiveSlot liveSlotsOnEntry liveVRegsOnEntry
+
+		let info'
+			= LiveInfo static firstId
+				(Just liveVRegsOnEntry)
+				liveSlotsOnEntry'
+					
+		-- Apply the spiller to all the basic blocks in the CmmProc.
+		sccs'		<- mapM (mapSCCM (regSpill_block regSlotMap)) sccs
+
+		return	$ CmmProc info' label params sccs'
+
+ where	-- | Given a BlockId and the set of registers live in it, 
+	--   if registers in this block are being spilled to stack slots, 
+	--   then record the fact that these slots are now live in those blocks
+	--   in the given slotmap.
+	patchLiveSlot :: BlockId -> RegSet -> Map BlockId (Set Int) -> Map BlockId (Set Int)
+	patchLiveSlot blockId regsLive slotMap
+	 = let	curSlotsLive	= fromMaybe Set.empty
+				$ Map.lookup blockId slotMap
+
+		moreSlotsLive	= Set.fromList
+				$ catMaybes 
+				$ map (lookupUFM regSlotMap)
+				$ uniqSetToList regsLive
+		
+	  	slotMap'	= Map.insert blockId (Set.union curSlotsLive moreSlotsLive) slotMap
+
+	   in	slotMap'
+
+
+
+-- | Spill some registers to stack slots in a basic block.
+regSpill_block
+	:: Instruction instr
+	=> UniqFM Int		-- ^ map of vregs to slots they're being spilled to.
+	-> LiveBasicBlock instr	
+	-> SpillM (LiveBasicBlock instr)
+	
 regSpill_block regSlotMap (BasicBlock i instrs)
  = do	instrss'	<- mapM (regSpill_instr regSlotMap) instrs
  	return	$ BasicBlock i (concat instrss')
 
 
+-- | Spill some registers to stack slots in a single instruction.  If the instruction
+--   uses registers that need to be spilled, then it is prefixed (or postfixed) with
+--   the appropriate RELOAD or SPILL meta instructions.
 regSpill_instr
 	:: Instruction instr
-	=> UniqFM Int 
-	-> LiveInstr instr -> SpillM [LiveInstr instr]
+	=> UniqFM Int 		-- ^ map of vregs to slots they're being spilled to.
+	-> LiveInstr instr
+	-> SpillM [LiveInstr instr]
 
 regSpill_instr _ li@(LiveInstr _ Nothing)
  = do	return [li]
@@ -174,7 +250,7 @@ spillModify regSlotMap instr reg
 
 
 
--- | rewrite uses of this virtual reg in an instr to use a different virtual reg
+-- | Rewrite uses of this virtual reg in an instr to use a different virtual reg
 patchInstr 
 	:: Instruction instr
 	=> Reg -> instr -> SpillM (instr, Reg)
@@ -198,13 +274,14 @@ patchReg1 old new instr
    in	patchRegsOfInstr instr patchF
 
 
-------------------------------------------------------
--- Spiller monad
-
+-- Spiller monad --------------------------------------------------------------
 data SpillS
 	= SpillS
-	{ stateUS	:: UniqSupply
-	, stateSpillSL	:: UniqFM (Reg, Int, Int) } -- ^ spilled reg vs number of times vreg was loaded, stored
+	{ -- | unique supply for generating fresh vregs.
+	  stateUS	:: UniqSupply
+	
+	  -- | spilled vreg vs the number of times it was loaded, stored 
+	, stateSpillSL	:: UniqFM (Reg, Int, Int) }
 
 initSpillS uniqueSupply
 	= SpillS
@@ -226,9 +303,7 @@ accSpillSL (r1, s1, l1) (_, s2, l2)
 	= (r1, s1 + s2, l1 + l2)
 
 
-----------------------------------------------------
--- Spiller stats
-
+-- Spiller stats --------------------------------------------------------------
 data SpillStats
 	= SpillStats
 	{ spillStoreLoad	:: UniqFM (Reg, Int, Int) }

compiler/nativeGen/RegAlloc/Graph/SpillClean.hs

@@ -23,7 +23,6 @@
 --	This also works if the reloads in B1\/B2 were spills instead, because
 --	spilling %r1 to a slot makes that slot have the same value as %r1.
 --
-
 module RegAlloc.Graph.SpillClean (
 	cleanSpills
 )
@@ -42,7 +41,13 @@ import State
 import Outputable
 import Util
 
-import Data.List        ( find, nub )
+import Data.List
+import Data.Maybe
+import Data.Map			(Map)
+import Data.Set			(Set)
+import qualified Data.Map	as Map
+import qualified Data.Set	as Set
+
 
 --
 type Slot = Int
@@ -84,8 +89,8 @@ cleanSpin spinCount code
 		, sReloadedBy		= emptyUFM }
 
  	code_forward	<- mapBlockTopM cleanBlockForward  code
-	code_backward	<- mapBlockTopM cleanBlockBackward code_forward
-
+	code_backward	<- cleanTopBackward code_forward
+	
 	-- During the cleaning of each block we collected information about what regs
 	--	were valid across each jump. Based on this, work out whether it will be
 	--	safe to erase reloads after join points for the next pass.
@@ -125,17 +130,6 @@ cleanBlockForward (BasicBlock blockId instrs)
 	return	$ BasicBlock blockId instrs_reload
 
 
-cleanBlockBackward 
-	:: Instruction instr
-	=> LiveBasicBlock instr 
-	-> CleanM (LiveBasicBlock instr)
-
-cleanBlockBackward (BasicBlock blockId instrs)
- = do	instrs_spill	<- cleanBackward  emptyUniqSet  [] instrs
-	return	$ BasicBlock blockId instrs_spill
-
-
-
 
 -- | Clean out unneeded reload instructions.
 --	Walking forwards across the code
@@ -286,27 +280,59 @@ cleanReload _ _ _
 -- TODO: This is mostly inter-block
 --	 we should really be updating the noReloads set as we cross jumps also.
 --
+-- TODO: generate noReloads from liveSlotsOnEntry
+-- 
+cleanTopBackward
+	:: Instruction instr
+	=> LiveCmmTop instr
+	-> CleanM (LiveCmmTop instr)
+
+cleanTopBackward cmm
+ = case cmm of
+	CmmData{}
+	 -> return cmm
+	
+	CmmProc info label params sccs
+	 | LiveInfo _ _ _ liveSlotsOnEntry <- info
+	 -> do	sccs'	<- mapM (mapSCCM (cleanBlockBackward liveSlotsOnEntry)) sccs
+		return	$ CmmProc info label params sccs' 
+
+
+cleanBlockBackward 
+	:: Instruction instr
+	=> Map BlockId (Set Int)
+	-> LiveBasicBlock instr 
+	-> CleanM (LiveBasicBlock instr)
+
+cleanBlockBackward liveSlotsOnEntry (BasicBlock blockId instrs)
+ = do	instrs_spill	<- cleanBackward liveSlotsOnEntry  emptyUniqSet  [] instrs
+	return	$ BasicBlock blockId instrs_spill
+
+
+
 cleanBackward
-	:: UniqSet Int 			-- ^ slots that have been spilled, but not reloaded from
+	:: Instruction instr
+	=> Map BlockId (Set Int)	-- ^ Slots live on entry to each block
+	-> UniqSet Int 			-- ^ slots that have been spilled, but not reloaded from
 	-> [LiveInstr instr]		-- ^ acc
 	-> [LiveInstr instr]		-- ^ instrs to clean (in forwards order)
 	-> CleanM [LiveInstr instr]	-- ^ cleaned instrs  (in backwards order)
 
 
-cleanBackward noReloads acc lis
+cleanBackward liveSlotsOnEntry noReloads acc lis
  = do	reloadedBy	<- gets sReloadedBy
- 	cleanBackward' reloadedBy noReloads acc lis
+ 	cleanBackward' liveSlotsOnEntry reloadedBy noReloads acc lis
 
-cleanBackward' _ _      acc []
+cleanBackward' _ _ _      acc []
 	= return  acc
 
-cleanBackward' reloadedBy noReloads acc (li : instrs)
+cleanBackward' liveSlotsOnEntry reloadedBy noReloads acc (li : instrs)
 
 	-- if nothing ever reloads from this slot then we don't need the spill
 	| LiveInstr (SPILL _ slot) _	<- li
 	, Nothing	<- lookupUFM reloadedBy (SSlot slot)
 	= do	modify $ \s -> s { sCleanedSpillsAcc = sCleanedSpillsAcc s + 1 }
-		cleanBackward noReloads acc instrs
+		cleanBackward liveSlotsOnEntry noReloads acc instrs
 
 	| LiveInstr (SPILL _ slot) _	<- li
 	= if elementOfUniqSet slot noReloads
@@ -314,21 +340,39 @@ cleanBackward' reloadedBy noReloads acc (li : instrs)
 	   -- we can erase this spill because the slot won't be read until after the next one
 	   then do
 		modify $ \s -> s { sCleanedSpillsAcc = sCleanedSpillsAcc s + 1 }
-	   	cleanBackward noReloads acc instrs
+	   	cleanBackward liveSlotsOnEntry noReloads acc instrs
 
 	   else do
 		-- this slot is being spilled to, but we haven't seen any reloads yet.
 		let noReloads'	= addOneToUniqSet noReloads slot
-	   	cleanBackward noReloads' (li : acc) instrs
+	   	cleanBackward liveSlotsOnEntry noReloads' (li : acc) instrs
 
 	-- if we reload from a slot then it's no longer unused
 	| LiveInstr (RELOAD slot _) _	<- li
 	, noReloads'		<- delOneFromUniqSet noReloads slot
-	= cleanBackward noReloads' (li : acc) instrs
+	= cleanBackward liveSlotsOnEntry noReloads' (li : acc) instrs
+
+	-- If a slot is live in a jump target then assume it's reloaded there.
+	-- TODO: A real dataflow analysis would do a better job here.
+	--       If the target block _ever_ used the slot then we assume it always does,
+	--       but if those reloads are cleaned the slot liveness map doesn't get updated.
+	| LiveInstr instr _ 	<- li
+	, targets		<- jumpDestsOfInstr instr
+	= do	
+		let slotsReloadedByTargets
+				= Set.unions
+				$ catMaybes
+				$ map (flip Map.lookup liveSlotsOnEntry) 
+				$ targets
+		
+		let noReloads'	= foldl' delOneFromUniqSet noReloads 
+				$ Set.toList slotsReloadedByTargets
+		
+		cleanBackward liveSlotsOnEntry noReloads' (li : acc) instrs
 
 	-- some other instruction
 	| otherwise
-	= cleanBackward noReloads (li : acc) instrs
+	= cleanBackward liveSlotsOnEntry noReloads (li : acc) instrs
 
 
 -- collateJoinPoints:

compiler/nativeGen/RegAlloc/Graph/SpillCost.hs

@@ -78,7 +78,7 @@ slurpSpillCostInfo cmm
 	-- lookup the regs that are live on entry to this block in
 	--	the info table from the CmmProc
  	countBlock info (BasicBlock blockId instrs)
-		| LiveInfo _ _ (Just blockLive)	<- info
+		| LiveInfo _ _ (Just blockLive) _ <- info
 		, Just rsLiveEntry  <- lookupBlockEnv blockLive blockId
 		, rsLiveEntry_virt  <- takeVirtuals rsLiveEntry
 		= countLIs rsLiveEntry_virt instrs

compiler/nativeGen/RegAlloc/Linear/Main.hs

@@ -132,12 +132,12 @@ regAlloc (CmmData sec d)
 		( CmmData sec d
 		, Nothing )
 	
-regAlloc (CmmProc (LiveInfo info _ _) lbl params [])
+regAlloc (CmmProc (LiveInfo info _ _ _) lbl params [])
 	= return ( CmmProc info lbl params (ListGraph [])
 		 , Nothing )
 	
 regAlloc (CmmProc static lbl params sccs)
-	| LiveInfo info (Just first_id) (Just block_live)	<- static
+	| LiveInfo info (Just first_id) (Just block_live) _	<- static
 	= do	
  		-- do register allocation on each component.
 		(final_blocks, stats)

compiler/nativeGen/RegAlloc/Liveness.hs

@@ -18,7 +18,7 @@ module RegAlloc.Liveness (
 	LiveInfo (..),
 	LiveBasicBlock,
 
-	mapBlockTop, 	mapBlockTopM,
+	mapBlockTop, 	mapBlockTopM, 	mapSCCM,
 	mapGenBlockTop,	mapGenBlockTopM,
 	stripLive,
 	stripLiveBlock,
@@ -31,8 +31,6 @@ module RegAlloc.Liveness (
 	regLiveness,
 	natCmmTopToLive
   ) where
-
-
 import Reg
 import Instruction
 
@@ -52,6 +50,9 @@ import FastString
 
 import Data.List
 import Data.Maybe
+import Data.Map			(Map)
+import Data.Set			(Set)
+import qualified Data.Map	as Map
 
 -----------------------------------------------------------------------------
 type RegSet = UniqSet Reg
@@ -160,9 +161,11 @@ data Liveness
 -- | Stash regs live on entry to each basic block in the info part of the cmm code.
 data LiveInfo
 	= LiveInfo
-		[CmmStatic]			-- cmm static stuff
-		(Maybe BlockId)			-- id of the first block
-		(Maybe (BlockMap RegSet))	-- argument locals live on entry to this block
+		[CmmStatic]				-- cmm 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
+
 
 -- | A basic block with liveness information.
 type LiveBasicBlock instr
@@ -212,10 +215,11 @@ instance Outputable instr
 		 | otherwise		= name <> (hcat $ punctuate space $ map ppr $ uniqSetToList regs)
 
 instance Outputable LiveInfo where
-	ppr (LiveInfo static firstId liveOnEntry)
+	ppr (LiveInfo static firstId liveVRegsOnEntry liveSlotsOnEntry)
 		=  (vcat $ map ppr static)
-		$$ text "# firstId     = " <> ppr firstId
-		$$ text "# liveOnEntry = " <> ppr liveOnEntry
+		$$ text "# firstId          = " <> ppr firstId
+		$$ text "# liveVRegsOnEntry = " <> ppr liveVRegsOnEntry
+		$$ text "# liveSlotsOnEntry = " <> text (show liveSlotsOnEntry)
 
 
 
@@ -299,9 +303,9 @@ slurpConflicts live
 		= foldl'  (slurpBlock info) rs bs
 
 	slurpBlock info rs (BasicBlock blockId instrs)	
-		| LiveInfo _ _ (Just blockLive)	<- info
-		, Just rsLiveEntry		<- lookupBlockEnv blockLive blockId
-		, (conflicts, moves)		<- slurpLIs rsLiveEntry rs instrs
+		| LiveInfo _ _ (Just blockLive)	_ <- info
+		, Just rsLiveEntry		  <- lookupBlockEnv blockLive blockId
+		, (conflicts, moves)		  <- slurpLIs rsLiveEntry rs instrs
 		= (consBag rsLiveEntry conflicts, moves)
 
 		| otherwise
@@ -466,7 +470,7 @@ stripLive live
 
  where	stripCmm (CmmData sec ds)	= CmmData sec ds
 
- 	stripCmm (CmmProc (LiveInfo info (Just first_id) _) label params sccs)
+ 	stripCmm (CmmProc (LiveInfo info (Just first_id) _ _) label params sccs)
 	 = let	final_blocks	= flattenSCCs sccs
 		
 		-- make sure the block that was first in the input list
@@ -479,7 +483,7 @@ stripLive live
                           (ListGraph $ map stripLiveBlock $ first' : rest')
 
 	-- procs used for stg_split_markers don't contain any blocks, and have no first_id.
-	stripCmm (CmmProc (LiveInfo info Nothing _) label params [])
+	stripCmm (CmmProc (LiveInfo info Nothing _ _) label params [])
 	 =	CmmProc info label params (ListGraph [])
 
 	-- If the proc has blocks but we don't know what the first one was, then we're dead.
@@ -540,7 +544,6 @@ eraseDeltasLive cmm
 -- | Patch the registers in this code according to this register mapping.
 --	also erase reg -> reg moves when the reg is the same.
 --	also erase reg -> reg moves when the destination dies in this instr.
-
 patchEraseLive
 	:: Instruction instr
 	=> (Reg -> Reg)
@@ -552,12 +555,12 @@ patchEraseLive patchF cmm
 	patchCmm cmm@CmmData{}	= cmm
 
 	patchCmm (CmmProc info label params sccs)
-	 | LiveInfo static id (Just blockMap)	<- info
+	 | LiveInfo static id (Just blockMap) mLiveSlots <- info
 	 = let 	
 	 	patchRegSet set	= mkUniqSet $ map patchF $ uniqSetToList set
 		blockMap'	= mapBlockEnv patchRegSet blockMap
 
-		info'		= LiveInfo static id (Just blockMap')
+		info'		= LiveInfo static id (Just blockMap') mLiveSlots
 	   in	CmmProc info' label params $ map patchSCC sccs
 
 	 | otherwise
@@ -628,7 +631,7 @@ natCmmTopToLive (CmmData i d)
 	= CmmData i d
 
 natCmmTopToLive (CmmProc info lbl params (ListGraph []))
-	= CmmProc (LiveInfo info Nothing Nothing)
+	= CmmProc (LiveInfo info Nothing Nothing Map.empty)
 		  lbl params []
 
 natCmmTopToLive (CmmProc info lbl params (ListGraph blocks@(first : _)))
@@ -638,7 +641,7 @@ natCmmTopToLive (CmmProc info lbl params (ListGraph blocks@(first : _)))
 					BasicBlock l (map (\i -> LiveInstr (Instr i) Nothing) instrs)))
 			$ sccs
 				
-   in	CmmProc (LiveInfo info (Just first_id) Nothing)
+   in	CmmProc (LiveInfo info (Just first_id) Nothing Map.empty)
 		lbl params sccsLive
 
 
@@ -668,16 +671,16 @@ regLiveness (CmmData i d)
 	= returnUs $ CmmData i d
 
 regLiveness (CmmProc info lbl params [])
-	| LiveInfo static mFirst _	<- info
+	| LiveInfo static mFirst _ _	<- info
 	= returnUs $ CmmProc
-			(LiveInfo static mFirst (Just emptyBlockEnv))
+			(LiveInfo static mFirst (Just emptyBlockEnv) Map.empty)
 			lbl params []
 
 regLiveness (CmmProc info lbl params sccs)
-	| LiveInfo static mFirst _	<- info
+	| LiveInfo static mFirst _ liveSlotsOnEntry	<- info
 	= let	(ann_sccs, block_live)	= computeLiveness sccs
 
-   	  in	returnUs $ CmmProc (LiveInfo static mFirst (Just block_live))
+   	  in	returnUs $ CmmProc (LiveInfo static mFirst (Just block_live) liveSlotsOnEntry)
 			   lbl params ann_sccs