compiler/utils/GraphColor.hs

   1
   2 -- | Graph Coloring.
   3 --      This is a generic graph coloring library, abstracted over the type of
   4 --      the node keys, nodes and colors.
   5 --
   6 {-# OPTIONS -fno-warn-missing-signatures #-}
   7
   8 module GraphColor (
   9         module GraphBase,
  10         module GraphOps,
  11         module GraphPpr,
  12         colorGraph
  13 )
  14
  15 where
  16
  17 import GraphBase
  18 import GraphOps
  19 import GraphPpr
  20
  21 import Unique
  22 import UniqFM
  23 import UniqSet
  24 import Outputable
  25
  26 import Data.Maybe
  27 import Data.List
  28
  29
  30 -- | Try to color a graph with this set of colors.
  31 --      Uses Chaitin's algorithm to color the graph.
  32 --      The graph is scanned for nodes which are deamed 'trivially colorable'. These nodes
  33 --      are pushed onto a stack and removed from the graph.
  34 --      Once this process is complete the graph can be colored by removing nodes from
  35 --      the stack (ie in reverse order) and assigning them colors different to their neighbors.
  36 --
  37 colorGraph
  38         :: ( Uniquable  k, Uniquable cls,  Uniquable  color
  39            , Eq color, Eq cls, Ord k
  40            , Outputable k, Outputable cls, Outputable color)
  41         => Bool                         -- ^ whether to do iterative coalescing
  42         -> UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
  43         -> Triv   k cls color           -- ^ fn to decide whether a node is trivially colorable.
  44         -> (Graph k cls color -> k)     -- ^ fn to choose a node to potentially leave uncolored if nothing is trivially colorable.
  45         -> Graph  k cls color           -- ^ the graph to color.
  46
  47         -> ( Graph k cls color          -- the colored graph.
  48            , UniqSet k                  -- the set of nodes that we couldn't find a color for.
  49            , UniqFM  k )                -- map of regs (r1 -> r2) that were coaleced
  50                                         --       r1 should be replaced by r2 in the source
  51
  52 colorGraph iterative colors triv spill graph0
  53  = let
  54         -- if we're not doing iterative coalescing, then just do a single coalescing
  55         --      pass at the front. This won't be as good but should still eat up a
  56         --      lot of the reg-reg moves.
  57         (graph_coalesced, kksCoalesce1)
  58                 = if not iterative
  59                         then coalesceGraph False triv graph0
  60                         else (graph0, [])
  61
  62         -- run the scanner to slurp out all the trivially colorable nodes
  63         --      (and do coalescing if iterative coalescing is enabled)
  64         (ksTriv, ksProblems, kksCoalesce2)
  65                 = colorScan iterative triv spill graph_coalesced
  66
  67         -- If iterative coalescing is enabled, the scanner will coalesce the graph as does its business.
  68         --      We need to apply all the coalescences found by the scanner to the original
  69         --      graph before doing assignColors.
  70         --
  71         --      Because we've got the whole, non-pruned graph here we turn on aggressive coalecing
  72         --      to force all the (conservative) coalescences found during scanning.
  73         --
  74         (graph_scan_coalesced, _)
  75                 = mapAccumL (coalesceNodes True triv) graph_coalesced kksCoalesce2
  76
  77         -- color the trivially colorable nodes
  78         --      during scanning, keys of triv nodes were added to the front of the list as they were found
  79         --      this colors them in the reverse order, as required by the algorithm.
  80         (graph_triv, ksNoTriv)
  81                 = assignColors colors graph_scan_coalesced ksTriv
  82
  83         -- try and color the problem nodes
  84         --      problem nodes are the ones that were left uncolored because they weren't triv.
  85         --      theres a change we can color them here anyway.
  86         (graph_prob, ksNoColor)
  87                 = assignColors colors graph_triv ksProblems
  88
  89         -- if the trivially colorable nodes didn't color then something is probably wrong
  90         --      with the provided triv function.
  91         --
  92    in   if not $ null ksNoTriv
  93          then   pprPanic "colorGraph: trivially colorable nodes didn't color!" empty
  94 {-                      (  empty
  95                         $$ text "ksTriv    = " <> ppr ksTriv
  96                         $$ text "ksNoTriv  = " <> ppr ksNoTriv
  97                         $$ empty
  98                         $$ dotGraph (\x -> text "white") triv graph1) -}
  99
 100          else   ( graph_prob
 101                 , mkUniqSet ksNoColor   -- the nodes that didn't color (spills)
 102                 , if iterative
 103                         then (listToUFM kksCoalesce2)
 104                         else (listToUFM kksCoalesce1))
 105
 106
 107 -- | Scan through the conflict graph separating out trivially colorable and
 108 --      potentially uncolorable (problem) nodes.
 109 --
 110 --      Checking whether a node is trivially colorable or not is a resonably expensive operation,
 111 --      so after a triv node is found and removed from the graph it's no good to return to the 'start'
 112 --      of the graph and recheck a bunch of nodes that will probably still be non-trivially colorable.
 113 --
 114 --      To ward against this, during each pass through the graph we collect up a list of triv nodes
 115 --      that were found, and only remove them once we've finished the pass. The more nodes we can delete
 116 --      at once the more likely it is that nodes we've already checked will become trivially colorable
 117 --      for the next pass.
 118 --
 119 --      TODO:   add work lists to finding triv nodes is easier.
 120 --              If we've just scanned the graph, and removed triv nodes, then the only
 121 --              nodes that we need to rescan are the ones we've removed edges from.
 122
 123 colorScan
 124         :: ( Uniquable k, Uniquable cls, Uniquable color
 125            , Ord k,       Eq cls
 126            , Outputable k, Outputable color)
 127         => Bool                         -- ^ whether to do iterative coalescing
 128         -> Triv k cls color             -- ^ fn to decide whether a node is trivially colorable
 129         -> (Graph k cls color -> k)     -- ^ fn to choose a node to potentially leave uncolored if nothing is trivially colorable.
 130         -> Graph k cls color            -- ^ the graph to scan
 131
 132         -> ([k], [k], [(k, k)])         --  triv colorable nodes, problem nodes, pairs of nodes to coalesce
 133
 134 colorScan iterative triv spill graph
 135         = colorScan_spin iterative triv spill graph [] [] []
 136
 137 colorScan_spin iterative triv spill graph
 138         ksTriv ksSpill kksCoalesce
 139
 140         -- if the graph is empty then we're done
 141         | isNullUFM $ graphMap graph
 142         = (ksTriv, ksSpill, kksCoalesce)
 143
 144         -- Simplify:
 145         --      Look for trivially colorable nodes.
 146         --      If we can find some then remove them from the graph and go back for more.
 147         --
 148         | nsTrivFound@(_:_)
 149                 <-  scanGraph   (\node -> triv  (nodeClass node) (nodeConflicts node) (nodeExclusions node)
 150
 151                                   -- for iterative coalescing we only want non-move related
 152                                   --    nodes here
 153                                   && (not iterative || isEmptyUniqSet (nodeCoalesce node)))
 154                         $ graph
 155
 156         , ksTrivFound   <- map nodeId nsTrivFound
 157         , graph3        <- foldr (\k g -> let Just g' = delNode k g
 158                                           in  g')
 159                                 graph ksTrivFound
 160
 161         = colorScan_spin iterative triv spill graph3
 162                 (ksTrivFound ++ ksTriv)
 163                 ksSpill
 164                 kksCoalesce
 165
 166         -- Coalesce:
 167         --      If we're doing iterative coalescing and no triv nodes are avaliable
 168         --      then it's type for a coalescing pass.
 169         | iterative
 170         = case coalesceGraph False triv graph of
 171
 172                 -- we were able to coalesce something
 173                 --      go back and see if this frees up more nodes to be trivially colorable.
 174                 (graph2, kksCoalesceFound @(_:_))
 175                  -> colorScan_spin iterative triv spill graph2
 176                         ksTriv ksSpill (kksCoalesceFound ++ kksCoalesce)
 177
 178                 -- Freeze:
 179                 -- nothing could be coalesced (or was triv),
 180                 --      time to choose a node to freeze and give up on ever coalescing it.
 181                 (graph2, [])
 182                  -> case freezeOneInGraph graph2 of
 183
 184                         -- we were able to freeze something
 185                         --      hopefully this will free up something for Simplify
 186                         (graph3, True)
 187                          -> colorScan_spin iterative triv spill graph3
 188                                 ksTriv ksSpill kksCoalesce
 189
 190                         -- we couldn't find something to freeze either
 191                         --      time for a spill
 192                         (graph3, False)
 193                          -> colorScan_spill iterative triv spill graph3
 194                                 ksTriv ksSpill kksCoalesce
 195
 196         -- spill time
 197         | otherwise
 198         = colorScan_spill iterative triv spill graph
 199                 ksTriv ksSpill kksCoalesce
 200
 201
 202 -- Select:
 203 -- we couldn't find any triv nodes or things to freeze or coalesce,
 204 --      and the graph isn't empty yet.. We'll have to choose a spill
 205 --      candidate and leave it uncolored.
 206 --
 207 colorScan_spill iterative triv spill graph
 208         ksTriv ksSpill kksCoalesce
 209
 210  = let  kSpill          = spill graph
 211         Just graph'     = delNode kSpill graph
 212    in   colorScan_spin iterative triv spill graph'
 213                 ksTriv (kSpill : ksSpill) kksCoalesce
 214
 215
 216 -- | Try to assign a color to all these nodes.
 217
 218 assignColors
 219         :: ( Uniquable k, Uniquable cls, Uniquable color, Eq color )
 220         => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
 221         -> Graph k cls color            -- ^ the graph
 222         -> [k]                          -- ^ nodes to assign a color to.
 223         -> ( Graph k cls color          -- the colored graph
 224            , [k])                       -- the nodes that didn't color.
 225
 226 assignColors colors graph ks
 227         = assignColors' colors graph [] ks
 228
 229  where  assignColors' _ graph prob []
 230                 = (graph, prob)
 231
 232         assignColors' colors graph prob (k:ks)
 233          = case assignColor colors k graph of
 234
 235                 -- couldn't color this node
 236                 Nothing         -> assignColors' colors graph (k : prob) ks
 237
 238                 -- this node colored ok, so do the rest
 239                 Just graph'     -> assignColors' colors graph' prob ks
 240
 241
 242         assignColor colors u graph
 243                 | Just c        <- selectColor colors graph u
 244                 = Just (setColor u c graph)
 245
 246                 | otherwise
 247                 = Nothing
 248
 249
 250
 251 -- | Select a color for a certain node
 252 --      taking into account preferences, neighbors and exclusions.
 253 --      returns Nothing if no color can be assigned to this node.
 254 --
 255 selectColor
 256         :: ( Uniquable k, Uniquable cls, Uniquable color, Eq color)
 257         => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
 258         -> Graph k cls color            -- ^ the graph
 259         -> k                            -- ^ key of the node to select a color for.
 260         -> Maybe color
 261
 262 selectColor colors graph u
 263  = let  -- lookup the node
 264         Just node       = lookupNode graph u
 265
 266         -- lookup the available colors for the class of this node.
 267         Just colors_avail
 268                         = lookupUFM colors (nodeClass node)
 269
 270         -- find colors we can't use because they're already being used
 271         --      by a node that conflicts with this one.
 272         Just nsConflicts
 273                         = sequence
 274                         $ map (lookupNode graph)
 275                         $ uniqSetToList
 276                         $ nodeConflicts node
 277
 278         colors_conflict = mkUniqSet
 279                         $ catMaybes
 280                         $ map nodeColor nsConflicts
 281
 282         -- the prefs of our neighbors
 283         colors_neighbor_prefs
 284                         = mkUniqSet
 285                         $ concat $ map nodePreference nsConflicts
 286
 287         -- colors that are still valid for us
 288         colors_ok_ex    = minusUniqSet colors_avail (nodeExclusions node)
 289         colors_ok       = minusUniqSet colors_ok_ex colors_conflict
 290
 291         -- the colors that we prefer, and are still ok
 292         colors_ok_pref  = intersectUniqSets
 293                                 (mkUniqSet $ nodePreference node) colors_ok
 294
 295         -- the colors that we could choose while being nice to our neighbors
 296         colors_ok_nice  = minusUniqSet
 297                                 colors_ok colors_neighbor_prefs
 298
 299         -- the best of all possible worlds..
 300         colors_ok_pref_nice
 301                         = intersectUniqSets
 302                                 colors_ok_nice colors_ok_pref
 303
 304         -- make the decision
 305         chooseColor
 306
 307                 -- everyone is happy, yay!
 308                 | not $ isEmptyUniqSet colors_ok_pref_nice
 309                 , c : _         <- filter (\x -> elementOfUniqSet x colors_ok_pref_nice)
 310                                         (nodePreference node)
 311                 = Just c
 312
 313                 -- we've got one of our preferences
 314                 | not $ isEmptyUniqSet colors_ok_pref
 315                 , c : _         <- filter (\x -> elementOfUniqSet x colors_ok_pref)
 316                                         (nodePreference node)
 317                 = Just c
 318
 319                 -- it wasn't a preference, but it was still ok
 320                 | not $ isEmptyUniqSet colors_ok
 321                 , c : _         <- uniqSetToList colors_ok
 322                 = Just c
 323
 324                 -- no colors were available for us this time.
 325                 --      looks like we're going around the loop again..
 326                 | otherwise
 327                 = Nothing
 328
 329    in   chooseColor
 330
 331
 332