GraphColor.hs 15 KB
Newer Older
1
-- | Graph Coloring.
2 3
--      This is a generic graph coloring library, abstracted over the type of
--      the node keys, nodes and colors.
4
--
5

6 7 8 9 10
module GraphColor (
        module GraphBase,
        module GraphOps,
        module GraphPpr,
        colorGraph
11 12 13 14
)

where

15 16
import GhcPrelude

17 18 19 20 21
import GraphBase
import GraphOps
import GraphPpr

import Unique
22
import UniqFM
23
import UniqSet
24
import Outputable
25 26 27

import Data.Maybe
import Data.List
28

29 30

-- | Try to color a graph with this set of colors.
31 32 33 34 35
--      Uses Chaitin's algorithm to color the graph.
--      The graph is scanned for nodes which are deamed 'trivially colorable'. These nodes
--      are pushed onto a stack and removed from the graph.
--      Once this process is complete the graph can be colored by removing nodes from
--      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 cls, Ord k
40 41 42 43 44 45 46 47 48 49
           , Outputable k, Outputable cls, Outputable color)
        => Bool                         -- ^ whether to do iterative coalescing
        -> Int                          -- ^ how many times we've tried to color this graph so far.
        -> 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.
Gabor Greif's avatar
Gabor Greif committed
50
           , UniqFM  k )                -- map of regs (r1 -> r2) that were coalesced
51
                                        --       r1 should be replaced by r2 in the source
52

53
colorGraph iterative spinCount colors triv spill graph0
54
 = let
55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77
        -- If we're not doing iterative coalescing then do an aggressive coalescing first time
        --      around and then conservative coalescing for subsequent passes.
        --
        --      Aggressive coalescing is a quick way to get rid of many reg-reg moves. However, if
        --      there is a lot of register pressure and we do it on every round then it can make the
        --      graph less colorable and prevent the algorithm from converging in a sensible number
        --      of cycles.
        --
        (graph_coalesced, kksCoalesce1)
         = if iterative
                then (graph0, [])
                else if spinCount == 0
                        then coalesceGraph True  triv graph0
                        else coalesceGraph False triv graph0

        -- run the scanner to slurp out all the trivially colorable nodes
        --      (and do coalescing if iterative coalescing is enabled)
        (ksTriv, ksProblems, kksCoalesce2)
                = colorScan iterative triv spill graph_coalesced

        -- If iterative coalescing is enabled, the scanner will coalesce the graph as does its business.
        --      We need to apply all the coalescences found by the scanner to the original
        --      graph before doing assignColors.
78
        --
79
        --      Because we've got the whole, non-pruned graph here we turn on aggressive coalescing
80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114
        --      to force all the (conservative) coalescences found during scanning.
        --
        (graph_scan_coalesced, _)
                = mapAccumL (coalesceNodes True triv) graph_coalesced kksCoalesce2

        -- color the trivially colorable nodes
        --      during scanning, keys of triv nodes were added to the front of the list as they were found
        --      this colors them in the reverse order, as required by the algorithm.
        (graph_triv, ksNoTriv)
                = assignColors colors graph_scan_coalesced ksTriv

        -- try and color the problem nodes
        --      problem nodes are the ones that were left uncolored because they weren't triv.
        --      theres a change we can color them here anyway.
        (graph_prob, ksNoColor)
                = assignColors colors graph_triv ksProblems

        -- if the trivially colorable nodes didn't color then something is probably wrong
        --      with the provided triv function.
        --
   in   if not $ null ksNoTriv
         then   pprPanic "colorGraph: trivially colorable nodes didn't color!" -- empty
                        (  empty
                        $$ text "ksTriv    = " <> ppr ksTriv
                        $$ text "ksNoTriv  = " <> ppr ksNoTriv
                        $$ text "colors    = " <> ppr colors
                        $$ empty
                        $$ dotGraph (\_ -> text "white") triv graph_triv)

         else   ( graph_prob
                , mkUniqSet ksNoColor   -- the nodes that didn't color (spills)
                , if iterative
                        then (listToUFM kksCoalesce2)
                        else (listToUFM kksCoalesce1))

115 116

-- | Scan through the conflict graph separating out trivially colorable and
117
--      potentially uncolorable (problem) nodes.
118
--
Krzysztof Gogolewski's avatar
Krzysztof Gogolewski committed
119
--      Checking whether a node is trivially colorable or not is a reasonably expensive operation,
120 121
--      so after a triv node is found and removed from the graph it's no good to return to the 'start'
--      of the graph and recheck a bunch of nodes that will probably still be non-trivially colorable.
122
--
123 124 125 126
--      To ward against this, during each pass through the graph we collect up a list of triv nodes
--      that were found, and only remove them once we've finished the pass. The more nodes we can delete
--      at once the more likely it is that nodes we've already checked will become trivially colorable
--      for the next pass.
127
--
128 129 130
--      TODO:   add work lists to finding triv nodes is easier.
--              If we've just scanned the graph, and removed triv nodes, then the only
--              nodes that we need to rescan are the ones we've removed edges from.
131

132
colorScan
133 134 135 136 137 138 139
        :: ( Uniquable k, Uniquable cls, Uniquable color
           , Ord k,       Eq cls
           , Outputable k, Outputable cls)
        => Bool                         -- ^ whether to do iterative coalescing
        -> 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 scan
140

141
        -> ([k], [k], [(k, k)])         --  triv colorable nodes, problem nodes, pairs of nodes to coalesce
142

143
colorScan iterative triv spill graph
144
        = colorScan_spin iterative triv spill graph [] [] []
145

146 147 148 149 150 151 152 153 154 155 156 157 158
colorScan_spin
        :: ( Uniquable k, Uniquable cls, Uniquable color
           , Ord k,       Eq cls
           , Outputable k, Outputable cls)
        => Bool
        -> Triv k cls color
        -> (Graph k cls color -> k)
        -> Graph k cls color
        -> [k]
        -> [k]
        -> [(k, k)]
        -> ([k], [k], [(k, k)])

159
colorScan_spin iterative triv spill graph
160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188
        ksTriv ksSpill kksCoalesce

        -- if the graph is empty then we're done
        | isNullUFM $ graphMap graph
        = (ksTriv, ksSpill, reverse kksCoalesce)

        -- Simplify:
        --      Look for trivially colorable nodes.
        --      If we can find some then remove them from the graph and go back for more.
        --
        | nsTrivFound@(_:_)
                <-  scanGraph   (\node -> triv  (nodeClass node) (nodeConflicts node) (nodeExclusions node)

                                  -- for iterative coalescing we only want non-move related
                                  --    nodes here
                                  && (not iterative || isEmptyUniqSet (nodeCoalesce node)))
                        $ graph

        , ksTrivFound   <- map nodeId nsTrivFound
        , graph2        <- foldr (\k g -> let Just g' = delNode k g
                                          in  g')
                                graph ksTrivFound

        = colorScan_spin iterative triv spill graph2
                (ksTrivFound ++ ksTriv)
                ksSpill
                kksCoalesce

        -- Coalesce:
Gabor Greif's avatar
Gabor Greif committed
189
        --      If we're doing iterative coalescing and no triv nodes are available
190 191 192 193 194 195
        --      then it's time for a coalescing pass.
        | iterative
        = case coalesceGraph False triv graph of

                -- we were able to coalesce something
                --      go back to Simplify and see if this frees up more nodes to be trivially colorable.
196
                (graph2, kksCoalesceFound@(_:_))
197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221
                 -> colorScan_spin iterative triv spill graph2
                        ksTriv ksSpill (reverse kksCoalesceFound ++ kksCoalesce)

                -- Freeze:
                -- nothing could be coalesced (or was triv),
                --      time to choose a node to freeze and give up on ever coalescing it.
                (graph2, [])
                 -> case freezeOneInGraph graph2 of

                        -- we were able to freeze something
                        --      hopefully this will free up something for Simplify
                        (graph3, True)
                         -> colorScan_spin iterative triv spill graph3
                                ksTriv ksSpill kksCoalesce

                        -- we couldn't find something to freeze either
                        --      time for a spill
                        (graph3, False)
                         -> colorScan_spill iterative triv spill graph3
                                ksTriv ksSpill kksCoalesce

        -- spill time
        | otherwise
        = colorScan_spill iterative triv spill graph
                ksTriv ksSpill kksCoalesce
222 223


224 225
-- Select:
-- we couldn't find any triv nodes or things to freeze or coalesce,
226 227
--      and the graph isn't empty yet.. We'll have to choose a spill
--      candidate and leave it uncolored.
228
--
229 230 231 232 233 234 235 236 237 238 239 240 241
colorScan_spill
        :: ( Uniquable k, Uniquable cls, Uniquable color
           , Ord k,       Eq cls
           , Outputable k, Outputable cls)
        => Bool
        -> Triv k cls color
        -> (Graph k cls color -> k)
        -> Graph k cls color
        -> [k]
        -> [k]
        -> [(k, k)]
        -> ([k], [k], [(k, k)])

242
colorScan_spill iterative triv spill graph
243 244 245 246 247 248
        ksTriv ksSpill kksCoalesce

 = let  kSpill          = spill graph
        Just graph'     = delNode kSpill graph
   in   colorScan_spin iterative triv spill graph'
                ksTriv (kSpill : ksSpill) kksCoalesce
249 250 251 252


-- | Try to assign a color to all these nodes.

253 254
assignColors
        :: ( Uniquable k, Uniquable cls, Uniquable color
255
           , Outputable cls)
256 257 258 259 260 261 262 263 264 265 266
        => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
        -> Graph k cls color            -- ^ the graph
        -> [k]                          -- ^ nodes to assign a color to.
        -> ( Graph k cls color          -- the colored graph
           , [k])                       -- the nodes that didn't color.

assignColors colors graph ks
        = assignColors' colors graph [] ks

 where  assignColors' _ graph prob []
                = (graph, prob)
267

268 269
        assignColors' colors graph prob (k:ks)
         = case assignColor colors k graph of
270

271 272
                -- couldn't color this node
                Nothing         -> assignColors' colors graph (k : prob) ks
273

274 275
                -- this node colored ok, so do the rest
                Just graph'     -> assignColors' colors graph' prob ks
276 277


278 279 280
        assignColor colors u graph
                | Just c        <- selectColor colors graph u
                = Just (setColor u c graph)
281

282 283
                | otherwise
                = Nothing
284 285 286 287



-- | Select a color for a certain node
288 289
--      taking into account preferences, neighbors and exclusions.
--      returns Nothing if no color can be assigned to this node.
290 291
--
selectColor
292
        :: ( Uniquable k, Uniquable cls, Uniquable color
293
           , Outputable cls)
294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313
        => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
        -> Graph k cls color            -- ^ the graph
        -> k                            -- ^ key of the node to select a color for.
        -> Maybe color

selectColor colors graph u
 = let  -- lookup the node
        Just node       = lookupNode graph u

        -- lookup the available colors for the class of this node.
        colors_avail
         = case lookupUFM colors (nodeClass node) of
                Nothing -> pprPanic "selectColor: no colors available for class " (ppr (nodeClass node))
                Just cs -> cs

        -- find colors we can't use because they're already being used
        --      by a node that conflicts with this one.
        Just nsConflicts
                        = sequence
                        $ map (lookupNode graph)
David Feuer's avatar
David Feuer committed
314
                        $ nonDetEltsUniqSet
315
                        $ nodeConflicts node
niteria's avatar
niteria committed
316
                        -- See Note [Unique Determinism and code generation]
317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360

        colors_conflict = mkUniqSet
                        $ catMaybes
                        $ map nodeColor nsConflicts

        -- the prefs of our neighbors
        colors_neighbor_prefs
                        = mkUniqSet
                        $ concat $ map nodePreference nsConflicts

        -- colors that are still valid for us
        colors_ok_ex    = minusUniqSet colors_avail (nodeExclusions node)
        colors_ok       = minusUniqSet colors_ok_ex colors_conflict

        -- the colors that we prefer, and are still ok
        colors_ok_pref  = intersectUniqSets
                                (mkUniqSet $ nodePreference node) colors_ok

        -- the colors that we could choose while being nice to our neighbors
        colors_ok_nice  = minusUniqSet
                                colors_ok colors_neighbor_prefs

        -- the best of all possible worlds..
        colors_ok_pref_nice
                        = intersectUniqSets
                                colors_ok_nice colors_ok_pref

        -- make the decision
        chooseColor

                -- everyone is happy, yay!
                | not $ isEmptyUniqSet colors_ok_pref_nice
                , c : _         <- filter (\x -> elementOfUniqSet x colors_ok_pref_nice)
                                        (nodePreference node)
                = Just c

                -- we've got one of our preferences
                | not $ isEmptyUniqSet colors_ok_pref
                , c : _         <- filter (\x -> elementOfUniqSet x colors_ok_pref)
                                        (nodePreference node)
                = Just c

                -- it wasn't a preference, but it was still ok
                | not $ isEmptyUniqSet colors_ok
David Feuer's avatar
David Feuer committed
361
                , c : _         <- nonDetEltsUniqSet colors_ok
niteria's avatar
niteria committed
362
                -- See Note [Unique Determinism and code generation]
363 364 365 366 367 368 369 370
                = Just c

                -- no colors were available for us this time.
                --      looks like we're going around the loop again..
                | otherwise
                = Nothing

   in   chooseColor
371 372 373