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