Add iterative coalescing to graph coloring allocator
[ghc-hetmet.git] / compiler / nativeGen / GraphOps.hs
1 -- | Basic operations on graphs.
2 --
3 --      TODO: refine coalescing crieteria
4
5 {-# OPTIONS -fno-warn-missing-signatures #-}
6
7 module GraphOps (
8         addNode,        delNode,        getNode,        lookupNode,     modNode,
9         size,
10         union,
11         addConflict,    delConflict,    addConflicts,
12         addCoalesce,    delCoalesce,    
13         addExclusion,   
14         addPreference,
15         coalesceNodes,  coalesceGraph,
16         freezeNode,     freezeOneInGraph, freezeAllInGraph,
17         scanGraph,
18         setColor,
19         validateGraph,
20         slurpNodeConflictCount
21 )
22 where
23
24 import GraphBase
25
26 import Outputable
27 import Unique
28 import UniqSet
29 import UniqFM
30
31 import Data.List        hiding (union)
32 import Data.Maybe
33
34 -- | Lookup a node from the graph.
35 lookupNode 
36         :: Uniquable k
37         => Graph k cls color
38         -> k -> Maybe (Node  k cls color)
39
40 lookupNode graph k      
41         = lookupUFM (graphMap graph) k
42
43
44 -- | Get a node from the graph, throwing an error if it's not there
45 getNode
46         :: Uniquable k
47         => Graph k cls color
48         -> k -> Node k cls color
49
50 getNode graph k
51  = case lookupUFM (graphMap graph) k of
52         Just node       -> node
53         Nothing         -> panic "ColorOps.getNode: not found" 
54
55
56 -- | Add a node to the graph, linking up its edges
57 addNode :: Uniquable k
58         => k -> Node k cls color 
59         -> Graph k cls color -> Graph k cls color
60         
61 addNode k node graph
62  = let  
63         -- add back conflict edges from other nodes to this one
64         map_conflict    
65                 = foldUniqSet 
66                         (adjustUFM (\n -> n { nodeConflicts = addOneToUniqSet (nodeConflicts n) k}))
67                         (graphMap graph)
68                         (nodeConflicts node)
69                         
70         -- add back coalesce edges from other nodes to this one
71         map_coalesce
72                 = foldUniqSet
73                         (adjustUFM (\n -> n { nodeCoalesce = addOneToUniqSet (nodeCoalesce n) k}))
74                         map_conflict
75                         (nodeCoalesce node)
76         
77   in    graph
78         { graphMap      = addToUFM map_coalesce k node}
79                 
80
81
82 -- | Delete a node and all its edges from the graph.
83 --      Throws an error if it's not there.
84 delNode :: Uniquable k
85         => k -> Graph k cls color -> Graph k cls color
86
87 delNode k graph
88  = let  Just node       = lookupNode graph k
89
90         -- delete conflict edges from other nodes to this one.
91         graph1          = foldl' (\g k1 -> let Just g' = delConflict k1 k g in g') graph
92                         $ uniqSetToList (nodeConflicts node)
93         
94         -- delete coalesce edge from other nodes to this one.
95         graph2          = foldl' (\g k1 -> let Just g' = delCoalesce k1 k g in g') graph1
96                         $ uniqSetToList (nodeCoalesce node)
97         
98         -- delete the node
99         graph3          = graphMapModify (\fm -> delFromUFM fm k) graph2
100         
101   in    graph3
102                 
103
104 -- | Modify a node in the graph.
105 --      returns Nothing if the node isn't present.
106 --
107 modNode :: Uniquable k
108         => (Node k cls color -> Node k cls color) 
109         -> k -> Graph k cls color -> Maybe (Graph k cls color)
110
111 modNode f k graph
112  = case lookupNode graph k of
113         Just Node{}
114          -> Just
115          $  graphMapModify
116                  (\fm   -> let  Just node       = lookupUFM fm k
117                                 node'           = f node
118                            in   addToUFM fm k node') 
119                 graph
120
121         Nothing -> Nothing
122
123
124 -- | Get the size of the graph, O(n)
125 size    :: Uniquable k 
126         => Graph k cls color -> Int
127         
128 size graph      
129         = sizeUFM $ graphMap graph
130         
131
132 -- | Union two graphs together.
133 union   :: Uniquable k
134         => Graph k cls color -> Graph k cls color -> Graph k cls color
135         
136 union   graph1 graph2
137         = Graph 
138         { graphMap              = plusUFM (graphMap graph1) (graphMap graph2) }
139
140
141 -- | Add a conflict between nodes to the graph, creating the nodes required.
142 --      Conflicts are virtual regs which need to be colored differently.
143 addConflict
144         :: Uniquable k
145         => (k, cls) -> (k, cls) 
146         -> Graph k cls color -> Graph k cls color
147
148 addConflict (u1, c1) (u2, c2)
149  = let  addNeighbor u c u'
150                 = adjustWithDefaultUFM
151                         (\node -> node { nodeConflicts = addOneToUniqSet (nodeConflicts node) u' })
152                         (newNode u c)  { nodeConflicts = unitUniqSet u' }
153                         u
154         
155    in   graphMapModify
156         ( addNeighbor u1 c1 u2 
157         . addNeighbor u2 c2 u1)
158
159  
160 -- | Delete a conflict edge. k1 -> k2
161 --      returns Nothing if the node isn't in the graph
162 delConflict 
163         :: Uniquable k
164         => k -> k
165         -> Graph k cls color -> Maybe (Graph k cls color)
166         
167 delConflict k1 k2
168         = modNode
169                 (\node -> node { nodeConflicts = delOneFromUniqSet (nodeConflicts node) k2 })
170                 k1
171
172
173 -- | Add some conflicts to the graph, creating nodes if required.
174 --      All the nodes in the set are taken to conflict with each other.
175 addConflicts
176         :: Uniquable k
177         => UniqSet k -> (k -> cls)
178         -> Graph k cls color -> Graph k cls color
179         
180 addConflicts conflicts getClass
181
182         -- just a single node, but no conflicts, create the node anyway.
183         | (u : [])      <- uniqSetToList conflicts
184         = graphMapModify 
185         $ adjustWithDefaultUFM 
186                 id
187                 (newNode u (getClass u)) 
188                 u
189
190         | otherwise
191         = graphMapModify
192         $ (\fm -> foldl' (\g u  -> addConflictSet1 u getClass conflicts g) fm
193                 $ uniqSetToList conflicts)
194
195
196 addConflictSet1 u getClass set 
197  = case delOneFromUniqSet set u of
198     set' -> adjustWithDefaultUFM
199                 (\node -> node                  { nodeConflicts = unionUniqSets set' (nodeConflicts node) } )
200                 (newNode u (getClass u))        { nodeConflicts = set' }
201                 u
202
203
204 -- | Add an exclusion to the graph, creating nodes if required.
205 --      These are extra colors that the node cannot use.
206 addExclusion
207         :: (Uniquable k, Uniquable color)
208         => k -> (k -> cls) -> color 
209         -> Graph k cls color -> Graph k cls color
210         
211 addExclusion u getClass color 
212         = graphMapModify
213         $ adjustWithDefaultUFM 
214                 (\node -> node                  { nodeExclusions = addOneToUniqSet (nodeExclusions node) color })
215                 (newNode u (getClass u))        { nodeExclusions = unitUniqSet color }
216                 u
217
218
219 -- | Add a coalescence edge to the graph, creating nodes if requried.
220 --      It is considered adventageous to assign the same color to nodes in a coalesence.
221 addCoalesce 
222         :: Uniquable k
223         => (k, cls) -> (k, cls) 
224         -> Graph k cls color -> Graph k cls color
225         
226 addCoalesce (u1, c1) (u2, c2) 
227  = let  addCoalesce u c u'
228          =      adjustWithDefaultUFM
229                         (\node -> node { nodeCoalesce = addOneToUniqSet (nodeCoalesce node) u' })
230                         (newNode u c)  { nodeCoalesce = unitUniqSet u' }
231                         u
232                         
233    in   graphMapModify
234         ( addCoalesce u1 c1 u2
235         . addCoalesce u2 c2 u1)
236
237
238 -- | Delete a coalescence edge (k1 -> k2) from the graph.
239 delCoalesce
240         :: Uniquable k
241         => k -> k 
242         -> Graph k cls color    -> Maybe (Graph k cls color)
243
244 delCoalesce k1 k2
245         = modNode (\node -> node { nodeCoalesce = delOneFromUniqSet (nodeCoalesce node) k2 })
246                 k1
247
248
249 -- | Add a color preference to the graph, creating nodes if required.
250 --      The most recently added preference is the most prefered.
251 --      The algorithm tries to assign a node it's prefered color if possible.
252 --
253 addPreference 
254         :: Uniquable k
255         => (k, cls) -> color
256         -> Graph k cls color -> Graph k cls color
257         
258 addPreference (u, c) color 
259         = graphMapModify
260         $ adjustWithDefaultUFM 
261                 (\node -> node { nodePreference = color : (nodePreference node) })
262                 (newNode u c)  { nodePreference = [color] }
263                 u
264
265
266 -- | Do agressive coalescing on this graph.
267 --      returns the new graph and the list of pairs of nodes that got coaleced together.
268 --      for each pair, the resulting node will have the least key and be second in the pair.
269 --
270 coalesceGraph
271         :: (Uniquable k, Ord k, Eq cls, Outputable k)
272         => Bool                 -- ^ If True, coalesce nodes even if this might make the graph
273                                 --      less colorable (aggressive coalescing)
274         -> Triv k cls color
275         -> Graph k cls color
276         -> (Graph k cls color, [(k, k)])
277
278 coalesceGraph aggressive triv graph
279         = coalesceGraph' aggressive triv graph []
280
281 coalesceGraph' aggressive triv graph kkPairsAcc
282  = let
283         -- find all the nodes that have coalescence edges
284         cNodes  = filter (\node -> not $ isEmptyUniqSet (nodeCoalesce node))
285                 $ eltsUFM $ graphMap graph
286
287         -- build a list of pairs of keys for node's we'll try and coalesce
288         --      every pair of nodes will appear twice in this list
289         --      ie [(k1, k2), (k2, k1) ... ]
290         --      This is ok, GrapOps.coalesceNodes handles this and it's convenient for
291         --      build a list of what nodes get coalesced together for later on.
292         --
293         cList   = [ (nodeId node1, k2)
294                         | node1 <- cNodes
295                         , k2    <- uniqSetToList $ nodeCoalesce node1 ]
296
297         -- do the coalescing, returning the new graph and a list of pairs of keys
298         --      that got coalesced together.
299         (graph', mPairs)
300                 = mapAccumL (coalesceNodes aggressive triv) graph cList
301
302         -- keep running until there are no more coalesces can be found
303    in   case catMaybes mPairs of
304          []     -> (graph', kkPairsAcc)
305          pairs  -> coalesceGraph' aggressive triv graph' (pairs ++ kkPairsAcc)
306
307
308 -- | Coalesce this pair of nodes unconditionally / agressively.
309 --      The resulting node is the one with the least key.
310 --
311 --      returns: Just    the pair of keys if the nodes were coalesced
312 --                       the second element of the pair being the least one
313 --
314 --               Nothing if either of the nodes weren't in the graph
315
316 coalesceNodes
317         :: (Uniquable k, Ord k, Eq cls, Outputable k)
318         => Bool                 -- ^ If True, coalesce nodes even if this might make the graph
319                                 --      less colorable (aggressive coalescing)
320         -> Triv  k cls color
321         -> Graph k cls color
322         -> (k, k)               -- ^ keys of the nodes to be coalesced
323         -> (Graph k cls color, Maybe (k, k))
324
325 coalesceNodes aggressive triv graph (k1, k2)
326         | (kMin, kMax)  <- if k1 < k2
327                                 then (k1, k2)
328                                 else (k2, k1)
329
330         -- the nodes being coalesced must be in the graph
331         , Just nMin     <- lookupNode graph kMin
332         , Just nMax     <- lookupNode graph kMax
333
334         -- can't coalesce conflicting modes
335         , not $ elementOfUniqSet kMin (nodeConflicts nMax)
336         , not $ elementOfUniqSet kMax (nodeConflicts nMin)
337
338         = coalesceNodes_merge aggressive triv graph kMin kMax nMin nMax
339
340         -- don't do the coalescing after all
341         | otherwise
342         = (graph, Nothing)
343
344 coalesceNodes_merge aggressive triv graph kMin kMax nMin nMax
345
346         -- sanity checks
347         | nodeClass nMin /= nodeClass nMax
348         = error "GraphOps.coalesceNodes: can't coalesce nodes of different classes."
349
350         | not (isNothing (nodeColor nMin) && isNothing (nodeColor nMax))
351         = error "GraphOps.coalesceNodes: can't coalesce colored nodes."
352
353         ---
354         | otherwise
355         = let
356                 -- the new node gets all the edges from its two components
357                 node    =
358                  Node   { nodeId                = kMin
359                         , nodeClass             = nodeClass nMin
360                         , nodeColor             = Nothing
361
362                         -- nodes don't conflict with themselves..
363                         , nodeConflicts
364                                 = (unionUniqSets (nodeConflicts nMin) (nodeConflicts nMax))
365                                         `delOneFromUniqSet` kMin
366                                         `delOneFromUniqSet` kMax
367
368                         , nodeExclusions        = unionUniqSets (nodeExclusions nMin) (nodeExclusions nMax)
369                         , nodePreference        = nodePreference nMin ++ nodePreference nMax
370
371                         -- nodes don't coalesce with themselves..
372                         , nodeCoalesce
373                                 = (unionUniqSets (nodeCoalesce nMin) (nodeCoalesce nMax))
374                                         `delOneFromUniqSet` kMin
375                                         `delOneFromUniqSet` kMax
376                         }
377
378           in    coalesceNodes_check aggressive triv graph kMin kMax node
379
380 coalesceNodes_check aggressive triv graph kMin kMax node
381
382         -- Unless we're coalescing aggressively, if the result node is not trivially
383         --      colorable then don't do the coalescing.
384         | not aggressive
385         , not $ triv (nodeClass node) (nodeConflicts node) (nodeExclusions node)
386         = (graph, Nothing)
387
388         | otherwise
389         = let -- delete the old nodes from the graph and add the new one
390                 graph'  = addNode kMin node
391                         $ delNode kMin
392                         $ delNode kMax
393                         $ graph
394
395           in    (graph', Just (kMax, kMin))
396
397
398 -- | Freeze a node
399 --      This is for the iterative coalescer.
400 --      By freezing a node we give up on ever coalescing it.
401 --      Move all its coalesce edges into the frozen set - and update
402 --      back edges from other nodes.
403 --
404 freezeNode
405         :: Uniquable k
406         => k                    -- ^ key of the node to freeze
407         -> Graph k cls color    -- ^ the graph
408         -> Graph k cls color    -- ^ graph with that node frozen
409
410 freezeNode k
411   = graphMapModify
412   $ \fm ->
413     let
414         -- freeze all the edges in the node to be frozen
415         Just node = lookupUFM fm k
416         node'   = node
417                 { nodeCoalesce          = emptyUniqSet }
418
419         fm1     = addToUFM fm k node'
420
421         -- update back edges pointing to this node
422         freezeEdge k node
423          = if elementOfUniqSet k (nodeCoalesce node)
424                 then node
425                         { nodeCoalesce          = delOneFromUniqSet (nodeCoalesce node) k }
426                 else    panic "GraphOps.freezeNode: edge to freeze wasn't in the coalesce set"
427
428         fm2     = foldUniqSet (adjustUFM (freezeEdge k)) fm1
429                         $ nodeCoalesce node
430
431     in  fm2
432
433
434 -- | Freeze one node in the graph
435 --      This if for the iterative coalescer.
436 --      Look for a move related node of low degree and freeze it.
437 --
438 --      We probably don't need to scan the whole graph looking for the node of absolute
439 --      lowest degree. Just sample the first few and choose the one with the lowest 
440 --      degree out of those. Also, we don't make any distinction between conflicts of different
441 --      classes.. this is just a heuristic, after all.
442 --
443 --      IDEA:   freezing a node might free it up for Simplify.. would be good to check for triv
444 --              right here, and add it to a worklist if known triv/non-move nodes.
445 --
446 freezeOneInGraph
447         :: (Uniquable k, Outputable k)
448         => Graph k cls color
449         -> ( Graph k cls color          -- the new graph
450            , Bool )                     -- whether we found a node to freeze
451
452 freezeOneInGraph graph
453  = let  compareNodeDegree n1 n2
454                 = compare (sizeUniqSet $ nodeConflicts n1) (sizeUniqSet $ nodeConflicts n2)
455
456         candidates
457                 = sortBy compareNodeDegree
458                 $ take 5        -- 5 isn't special, it's just a small number.
459                 $ scanGraph (\node -> not $ isEmptyUniqSet (nodeCoalesce node)) graph
460
461    in   case candidates of
462
463          -- there wasn't anything available to freeze
464          []     -> (graph, False)
465
466          -- we found something to freeze
467          (n : _)
468           -> ( freezeNode (nodeId n) graph
469              , True)
470
471
472 -- | Freeze all the nodes in the graph
473 --      for debugging the iterative allocator.
474 --
475 freezeAllInGraph
476         :: (Uniquable k, Outputable k)
477         => Graph k cls color
478         -> Graph k cls color
479
480 freezeAllInGraph graph
481         = foldr freezeNode graph
482                 $ map nodeId
483                 $ eltsUFM $ graphMap graph
484
485
486 -- | Find all the nodes in the graph that meet some criteria
487 --
488 scanGraph
489         :: Uniquable k
490         => (Node k cls color -> Bool)
491         -> Graph k cls color
492         -> [Node k cls color]
493
494 scanGraph match graph
495         = filter match $ eltsUFM $ graphMap graph
496
497
498 -- | validate the internal structure of a graph
499 --      all its edges should point to valid nodes
500 --      if they don't then throw an error
501 --
502 validateGraph
503         :: (Uniquable k, Outputable k)
504         => SDoc
505         -> Graph k cls color
506         -> Graph k cls color
507
508 validateGraph doc graph
509  = let  edges   = unionManyUniqSets
510                         (  (map nodeConflicts       $ eltsUFM $ graphMap graph)
511                         ++ (map nodeCoalesce        $ eltsUFM $ graphMap graph))
512
513         nodes   = mkUniqSet $ map nodeId $ eltsUFM $ graphMap graph
514         
515         badEdges = minusUniqSet edges nodes
516         
517   in    if isEmptyUniqSet badEdges 
518          then   graph
519          else   pprPanic "GraphOps.validateGraph"
520                 ( text  "-- bad edges"
521                 $$ vcat (map ppr $ uniqSetToList badEdges)
522                 $$ text "----------------------------"
523                 $$ doc)
524
525
526 -- | Slurp out a map of how many nodes had a certain number of conflict neighbours
527
528 slurpNodeConflictCount
529         :: Uniquable k
530         => Graph k cls color
531         -> UniqFM (Int, Int)    -- ^ (conflict neighbours, num nodes with that many conflicts)
532
533 slurpNodeConflictCount graph
534         = addListToUFM_C
535                 (\(c1, n1) (_, n2) -> (c1, n1 + n2))
536                 emptyUFM
537         $ map   (\node
538                   -> let count  = sizeUniqSet $ nodeConflicts node
539                      in  (count, (count, 1)))
540         $ eltsUFM
541         $ graphMap graph
542
543
544 -- | Set the color of a certain node
545 setColor 
546         :: Uniquable k
547         => k -> color
548         -> Graph k cls color -> Graph k cls color
549         
550 setColor u color
551         = graphMapModify
552         $ adjustUFM
553                 (\n -> n { nodeColor = Just color })
554                 u 
555         
556
557 {-# INLINE      adjustWithDefaultUFM #-}
558 adjustWithDefaultUFM 
559         :: Uniquable k 
560         => (a -> a) -> a -> k 
561         -> UniqFM a -> UniqFM a
562
563 adjustWithDefaultUFM f def k map
564         = addToUFM_C 
565                 (\old _ -> f old)
566                 map
567                 k def
568                 
569 {-# INLINE adjustUFM #-}
570 adjustUFM 
571         :: Uniquable k
572         => (a -> a)
573         -> k -> UniqFM a -> UniqFM a
574
575 adjustUFM f k map
576  = case lookupUFM map k of
577         Nothing -> map
578         Just a  -> addToUFM map k (f a)
579