Do aggressive register coalescing

Conservative and iterative coalescing come next.

compiler/nativeGen/GraphColor.hs

@@ -41,25 +41,36 @@ import Data.List
 --	the stack (ie in reverse order) and assigning them colors different to their neighbors.
 --
 colorGraph
-	:: ( Uniquable  k, Uniquable cls,  Uniquable  color, Eq color
+	:: ( Uniquable  k, Uniquable cls,  Uniquable  color
+	   , Eq color, Eq cls, Ord k
 	   , Outputable k, Outputable cls, Outputable color)
 	=> UniqFM (UniqSet color)	-- ^ map of (node class -> set of colors available for this class).
 	-> Triv   k cls color 		-- ^ fn to decide whether a node is trivially colorable.
 	-> (Graph k cls color -> k)	-- ^ fn to choose a node to potentially leave uncolored if nothing is trivially colorable.
 	-> Graph  k cls color 		-- ^ the graph to color.
-	-> ( Graph k cls color 		-- ^ the colored graph.
-	   , UniqSet k )		-- ^ the set of nodes that we couldn't find a color for.
+
+	-> ( Graph k cls color 		-- the colored graph.
+	   , UniqSet k			-- the set of nodes that we couldn't find a color for.
+	   , UniqFM  k )		-- map of regs (r1 -> r2) that were coaleced
+	   				--	 r1 should be replaced by r2 in the source
 
 colorGraph colors triv spill graph0
- = let	-- run the scanner to slurp out all the trivially colorable nodes
-  	(ksTriv, ksProblems)	= colorScan colors triv spill [] emptyUniqSet graph0
+ = let
+ 	-- do aggressive coalesing on the graph
+ 	(graph_coalesced, rsCoalesce)
+		= coalesceGraph graph0
+
+ 	-- run the scanner to slurp out all the trivially colorable nodes
+  	(ksTriv, ksProblems)
+		= colorScan colors triv spill [] emptyUniqSet graph_coalesced
  
 	-- color the trivially colorable nodes
-	(graph1, ksNoTriv)	= assignColors colors graph0 ksTriv
+	(graph_triv, ksNoTriv)
+		= assignColors colors graph_coalesced ksTriv
 
  	-- try and color the problem nodes
-	(graph2, ksNoColor)	= assignColors colors graph1 (uniqSetToList ksProblems)
-	
+	(graph_prob, ksNoColor)	= assignColors colors graph_triv (uniqSetToList ksProblems)
+
 	-- if the trivially colorable nodes didn't color then something is wrong
 	--	with the provided triv function.
    in	if not $ null ksNoTriv
@@ -69,8 +80,10 @@ colorGraph colors triv spill graph0
 			$$ text "ksNoTriv  = " <> ppr ksNoTriv
 			$$ empty
 			$$ dotGraph (\x -> text "white") triv graph1) -}
-	 else	(graph2, mkUniqSet ksNoColor)
-	
+
+	 else	( graph_prob
+	 	, mkUniqSet ksNoColor
+		, listToUFM rsCoalesce)
 	
 colorScan colors triv spill safe prob graph
 
@@ -95,7 +108,6 @@ colorScan colors triv spill safe prob graph
 		safe (addOneToUniqSet prob k) (delNode k graph)
 		
 
-
 -- | Try to assign a color to all these nodes.
 
 assignColors 

compiler/nativeGen/GraphOps.hs

@@ -17,8 +17,10 @@ module GraphOps (
 	addCoalesce,	delCoalesce,	
 	addExclusion,	
 	addPreference,
+	coalesceGraph,
+	coalesceNodes,
 	setColor,
-	verify,
+	validateGraph,
 	slurpNodeConflictCount
 )
 where
@@ -91,11 +93,11 @@ delNode k graph
  = let	Just node	= lookupNode graph k
 
 	-- delete conflict edges from other nodes to this one.
-	graph1		= foldl' (\g k1 -> delConflict k1 k g) graph 
+	graph1		= foldl' (\g k1 -> let Just g' = delConflict k1 k g in g') graph
 			$ uniqSetToList (nodeConflicts node)
 	
 	-- delete coalesce edge from other nodes to this one.
-	graph2		= foldl' (\g k1 -> delCoalesce k1 k g) graph1 
+	graph2		= foldl' (\g k1 -> let Just g' = delCoalesce k1 k g in g') graph1
 			$ uniqSetToList (nodeCoalesce node)
 	
 	-- delete the node
@@ -104,19 +106,24 @@ delNode k graph
   in	graph3
 		
 
--- | Modify a node in the graph
+-- | Modify a node in the graph.
+--	returns Nothing if the node isn't present.
+--
 modNode :: Uniquable k
 	=> (Node k cls color -> Node k cls color) 
-	-> k -> Graph k cls color -> Graph k cls color
+	-> k -> Graph k cls color -> Maybe (Graph k cls color)
 
 modNode f k graph
- = case getNode graph k of
- 	Node{} -> graphMapModify
+ = case lookupNode graph k of
+ 	Just Node{}
+	 -> Just
+	 $  graphMapModify
 		 (\fm	-> let	Just node	= lookupUFM fm k
 			   	node'		= f node
 			   in	addToUFM fm k node') 
 		graph
 
+	Nothing	-> Nothing
 
 -- | Get the size of the graph, O(n)
 size	:: Uniquable k 
@@ -157,10 +164,11 @@ addConflict (u1, c1) (u2, c2)
 
  
 -- | Delete a conflict edge. k1 -> k2
+--	returns Nothing if the node isn't in the graph
 delConflict 
 	:: Uniquable k
 	=> k -> k
-	-> Graph k cls color -> Graph k cls color
+	-> Graph k cls color -> Maybe (Graph k cls color)
 	
 delConflict k1 k2
 	= modNode
@@ -237,7 +245,7 @@ addCoalesce (u1, c1) (u2, c2)
 delCoalesce
 	:: Uniquable k
 	=> k -> k 
-	-> Graph k cls color	-> Graph k cls color
+	-> Graph k cls color	-> Maybe (Graph k cls color)
 
 delCoalesce k1 k2
 	= modNode (\node -> node { nodeCoalesce = delOneFromUniqSet (nodeCoalesce node) k2 })
@@ -260,15 +268,129 @@ addPreference (u, c) color
 		(newNode u c)  { nodePreference = [color] }
 		u
 
+
+-- | Do agressive coalescing on this graph.
+--	returns	the new graph and the list of pairs of nodes that got coaleced together.
+--	for each pair, the resulting node will have the least key and be second in the pair.
+--
+coalesceGraph
+	:: (Uniquable k, Ord k, Eq cls, Outputable k)
+	=>  Graph k cls color
+	-> (Graph k cls color, [(k, k)])
+
+coalesceGraph graph
+ = let
+ 	-- find all the nodes that have coalescence edges
+	cNodes	= filter (\node -> not $ isEmptyUniqSet (nodeCoalesce node))
+		$ eltsUFM $ graphMap graph
+
+	-- build a list of pairs of keys for node's we'll try and coalesce
+	--	every pair of nodes will appear twice in this list
+	--	ie [(k1, k2), (k2, k1) ... ]
+	--	This is ok, GrapOps.coalesceNodes handles this and it's convenient for
+	--	build a list of what nodes get coalesced together for later on.
+	--
+	cList	= [ (nodeId node1, k2)
+			| node1	<- cNodes
+			, k2	<- uniqSetToList $ nodeCoalesce node1 ]
+
+	-- do the coalescing, returning the new graph and a list of pairs of keys
+	--	that got coalesced together.
+	(graph', mPairs)
+		= mapAccumL coalesceNodes graph cList
+
+   in	(graph', catMaybes mPairs)
+
+
+-- | Coalesce this pair of nodes unconditionally / agressively.
+--	The resulting node is the one with the least key.
+--
+--	returns: Just 	 the pair of keys if the nodes were coalesced
+--			 the second element of the pair being the least one
+--
+--		 Nothing if either of the nodes weren't in the graph
+
+coalesceNodes
+	:: (Uniquable k, Ord k, Eq cls, Outputable k)
+	=> Graph k cls color
+	-> (k, k)		-- ^ keys of the nodes to be coalesced
+	-> (Graph k cls color, Maybe (k, k))
+
+coalesceNodes graph (k1, k2)
+	| (kMin, kMax)	<- if k1 < k2
+				then (k1, k2)
+				else (k2, k1)
+
+	-- nodes must be in the graph
+	, Just nMin	<- lookupNode graph kMin
+	, Just nMax	<- lookupNode graph kMax
+
+	-- can't coalesce conflicting nodes
+	, not $ elementOfUniqSet kMin (nodeConflicts nMax)
+	, not $ elementOfUniqSet kMax (nodeConflicts nMin)
+
+	= coalesceNodes' graph kMin kMax nMin nMax
+
+
+
+	-- one of the nodes wasn't in the graph anymore
+	| otherwise
+	= (graph, Nothing)
+
+coalesceNodes' graph kMin kMax nMin nMax
+
+	-- sanity checks
+	| nodeClass nMin /= nodeClass nMax
+	= error "GraphOps.coalesceNodes: can't coalesce nodes of different classes."
+
+	| not (isNothing (nodeColor nMin) && isNothing (nodeColor nMax))
+	= error "GraphOps.coalesceNodes: can't coalesce colored nodes."
+
+	---
+	| otherwise
+	= let
+		-- the new node gets all the edges from its two components
+		node	=
+		 Node	{ nodeId		= kMin
+			, nodeClass		= nodeClass nMin
+			, nodeColor		= Nothing
+
+			-- nodes don't conflict with themselves..
+			, nodeConflicts
+				= (unionUniqSets (nodeConflicts nMin) (nodeConflicts nMax))
+					`delOneFromUniqSet` kMin
+					`delOneFromUniqSet` kMax
+
+			, nodeExclusions	= unionUniqSets (nodeExclusions nMin) (nodeExclusions nMax)
+			, nodePreference	= nodePreference nMin ++ nodePreference nMax
+
+			-- nodes don't coalesce with themselves..
+			, nodeCoalesce
+				= (unionUniqSets (nodeCoalesce nMin) (nodeCoalesce nMax))
+					`delOneFromUniqSet` kMin
+					`delOneFromUniqSet` kMax
+			}
+
+		-- delete the old nodes from the graph and add the new one
+	 	graph'	= addNode kMin node
+	  		$ delNode kMin
+			$ delNode kMax
+			$ graph
+
+	  in	(graph', Just (kMax, kMin))
+
 		
--- | Verify the internal structure of a graph
+-- | validate the internal structure of a graph
 --	all its edges should point to valid nodes
+--	if they don't then throw an error
 --
-verify 	:: Uniquable k 
-	=> Graph k cls color
-	-> Bool
+validateGraph
+	:: (Uniquable k, Outputable k)
+	=> SDoc
+	-> Graph k cls color
+	-> Graph k cls color
 
-verify graph
+validateGraph doc graph
  = let	edges	= unionUniqSets
  			(unionManyUniqSets
  				(map nodeConflicts $ eltsUFM $ graphMap graph))
@@ -280,8 +402,12 @@ verify graph
 	badEdges = minusUniqSet edges nodes
 	
   in	if isEmptyUniqSet badEdges 
-  	 then 	True
-	 else	False
+  	 then 	graph
+	 else	pprPanic "GraphOps.validateGraph"
+	 	( text  "-- bad edges"
+		$$ vcat (map ppr $ uniqSetToList badEdges)
+		$$ text "----------------------------"
+		$$ doc)
 
 
 -- | Slurp out a map of how many nodes had a certain number of conflict neighbours

compiler/nativeGen/RegAllocColor.hs

@@ -106,14 +106,22 @@ regAlloc_spin dump (spinCount :: Int) triv regsFree slotsFree debug_codeGraphs c
 	let spill	= chooseSpill_maxLife fmLife
 	
 	-- try and color the graph 
-	let (graph_colored, rsSpill)	
+	let (graph_colored, rsSpill, rmCoalesce)
 			= Color.colorGraph regsFree triv spill graph
 
+	-- rewrite regs in the code that have been coalesced
+	let patchF reg	= case lookupUFM rmCoalesce reg of
+				Just reg'	-> reg'
+				Nothing		-> reg
+	let code_coalesced
+			= map (patchEraseLive patchF) code
+
+
 	-- see if we've found a coloring
 	if isEmptyUniqSet rsSpill
 	 then do
 		-- patch the registers using the info in the graph
-	 	let code_patched	= map (patchRegsFromGraph graph_colored) code
+	 	let code_patched	= map (patchRegsFromGraph graph_colored) code_coalesced
 
 		-- clean out unneeded SPILL/RELOADs
 		let code_spillclean	= map cleanSpills code_patched
@@ -129,6 +137,7 @@ regAlloc_spin dump (spinCount :: Int) triv regsFree slotsFree debug_codeGraphs c
 		let stat		=
 			RegAllocStatsColored
 			{ raGraph	= graph_colored
+			, raCoalesced	= rmCoalesce
 			, raPatched	= code_patched
 			, raSpillClean	= code_spillclean
 			, raFinal	= code_final
@@ -143,7 +152,7 @@ regAlloc_spin dump (spinCount :: Int) triv regsFree slotsFree debug_codeGraphs c
 	 else do
 	 	-- spill the uncolored regs
 		(code_spilled, slotsFree', spillStats)
-			<- regSpill code slotsFree rsSpill
+			<- regSpill code_coalesced slotsFree rsSpill
 			
 		-- recalculate liveness
 		let code_nat	= map stripLive code_spilled
@@ -153,6 +162,7 @@ regAlloc_spin dump (spinCount :: Int) triv regsFree slotsFree debug_codeGraphs c
 		let stat	=
 			RegAllocStatsSpill
 			{ raGraph	= graph_colored
+			, raCoalesced	= rmCoalesce
 			, raSpillStats	= spillStats
 			, raLifetimes	= fmLife
 			, raSpilled	= code_spilled }

compiler/nativeGen/RegAllocStats.hs

@@ -52,6 +52,7 @@ data RegAllocStats
 	-- a spill stage
 	| RegAllocStatsSpill
 	{ raGraph	:: Color.Graph Reg RegClass Reg	-- ^ the partially colored graph
+	, raCoalesced	:: UniqFM Reg			-- ^ the regs that were coaleced
 	, raSpillStats	:: SpillStats 			-- ^ spiller stats
 	, raLifetimes	:: UniqFM (Reg, Int) 		-- ^ number of instrs each reg lives for
 	, raSpilled	:: [LiveCmmTop] }		-- ^ code with spill instructions added
@@ -59,6 +60,7 @@ data RegAllocStats
 	-- a successful coloring
 	| RegAllocStatsColored
 	{ raGraph	:: Color.Graph Reg RegClass Reg -- ^ the colored graph
+	, raCoalesced	:: UniqFM Reg			-- ^ the regs that were coaleced
 	, raPatched	:: [LiveCmmTop] 		-- ^ code with vregs replaced by hregs
 	, raSpillClean  :: [LiveCmmTop]			-- ^ code with unneeded spill/reloads cleaned out
 	, raFinal	:: [NatCmmTop] 			-- ^ final code
@@ -74,28 +76,49 @@ instance Outputable RegAllocStats where
 	$$ text "#  Initial register conflict graph."
 	$$ Color.dotGraph regDotColor trivColorable (raGraph s)
 
+
  ppr (s@RegAllocStatsSpill{})
  	=  text "#  Spill"
+
 	$$ text "#  Register conflict graph."
 	$$ Color.dotGraph regDotColor trivColorable (raGraph s)
 	$$ text ""
+
+	$$ (if (not $ isNullUFM $ raCoalesced s)
+		then 	text "#  Registers coalesced."
+			$$ (vcat $ map ppr $ ufmToList $ raCoalesced s)
+			$$ text ""
+		else empty)
+
 	$$ text "#  Spills inserted."
 	$$ ppr (raSpillStats s)
 	$$ text ""
+
 	$$ text "#  Code with spills inserted."
 	$$ (ppr (raSpilled s))
 
+
  ppr (s@RegAllocStatsColored { raSRMs = (spills, reloads, moves) })
  	=  text "#  Colored"
+
 	$$ text "#  Register conflict graph."
 	$$ Color.dotGraph regDotColor trivColorable (raGraph s)
 	$$ text ""
+
+	$$ (if (not $ isNullUFM $ raCoalesced s)
+		then 	text "#  Registers coalesced."
+			$$ (vcat $ map ppr $ ufmToList $ raCoalesced s)
+			$$ text ""
+		else empty)
+
 	$$ text "#  Native code after register allocation."
 	$$ ppr (raPatched s)
 	$$ text ""
+
 	$$ text "#  Clean out unneeded spill/reloads."
 	$$ ppr (raSpillClean s)
 	$$ text ""
+
 	$$ text "#  Final code, after rewriting spill/rewrite pseudo instrs."
 	$$ ppr (raFinal s)
 	$$ text ""