module Futhark.Optimise.MemoryBlockMerging.GreedyColoring (colorGraph, Coloring) where
import Data.Function ((&))
import qualified Data.Map as M
import Data.Maybe (fromMaybe)
import qualified Data.Set as S
import qualified Futhark.Analysis.Interference as Interference
type Coloring a = M.Map a Int
type Neighbors a = M.Map a (S.Set a)
neighbors :: Ord a => Interference.Graph a -> Neighbors a
neighbors :: Graph a -> Neighbors a
neighbors =
((a, a) -> Neighbors a -> Neighbors a)
-> Neighbors a -> Graph a -> Neighbors a
forall a b. (a -> b -> b) -> b -> Set a -> b
S.foldr
( \(a
x, a
y) Neighbors a
acc ->
Neighbors a
acc
Neighbors a -> (Neighbors a -> Neighbors a) -> Neighbors a
forall a b. a -> (a -> b) -> b
& (Set a -> Set a -> Set a)
-> a -> Set a -> Neighbors a -> Neighbors a
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
S.union a
x (a -> Set a
forall a. a -> Set a
S.singleton a
y)
Neighbors a -> (Neighbors a -> Neighbors a) -> Neighbors a
forall a b. a -> (a -> b) -> b
& (Set a -> Set a -> Set a)
-> a -> Set a -> Neighbors a -> Neighbors a
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
S.union a
y (a -> Set a
forall a. a -> Set a
S.singleton a
x)
)
Neighbors a
forall k a. Map k a
M.empty
firstAvailable :: Eq space => M.Map Int space -> S.Set Int -> Int -> space -> (M.Map Int space, Int)
firstAvailable :: Map Int space -> Set Int -> Int -> space -> (Map Int space, Int)
firstAvailable Map Int space
spaces Set Int
xs Int
i space
sp =
case (Int
i Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Int
xs, Map Int space
spaces Map Int space -> Int -> Maybe space
forall k a. Ord k => Map k a -> k -> Maybe a
M.!? Int
i) of
(Bool
False, Just space
sp') | space
sp' space -> space -> Bool
forall a. Eq a => a -> a -> Bool
== space
sp -> (Map Int space
spaces, Int
i)
(Bool
False, Maybe space
Nothing) -> (Int -> space -> Map Int space -> Map Int space
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
i space
sp Map Int space
spaces, Int
i)
(Bool, Maybe space)
_ -> Map Int space -> Set Int -> Int -> space -> (Map Int space, Int)
forall space.
Eq space =>
Map Int space -> Set Int -> Int -> space -> (Map Int space, Int)
firstAvailable Map Int space
spaces Set Int
xs (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) space
sp
colorNode ::
(Ord a, Eq space) =>
Neighbors a ->
(a, space) ->
(M.Map Int space, Coloring a) ->
(M.Map Int space, Coloring a)
colorNode :: Neighbors a
-> (a, space)
-> (Map Int space, Coloring a)
-> (Map Int space, Coloring a)
colorNode Neighbors a
nbs (a
x, space
sp) (Map Int space
spaces, Coloring a
coloring) =
let nb_colors :: Set Int
nb_colors =
(a -> Set Int) -> Set a -> Set Int
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Set Int -> (Int -> Set Int) -> Maybe Int -> Set Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set Int
forall a. Set a
S.empty Int -> Set Int
forall a. a -> Set a
S.singleton (Maybe Int -> Set Int) -> (a -> Maybe Int) -> a -> Set Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Coloring a
coloring Coloring a -> a -> Maybe Int
forall k a. Ord k => Map k a -> k -> Maybe a
M.!?)) (Set a -> Set Int) -> Set a -> Set Int
forall a b. (a -> b) -> a -> b
$
Set a -> Maybe (Set a) -> Set a
forall a. a -> Maybe a -> a
fromMaybe Set a
forall a. Monoid a => a
mempty (Neighbors a
nbs Neighbors a -> a -> Maybe (Set a)
forall k a. Ord k => Map k a -> k -> Maybe a
M.!? a
x)
(Map Int space
spaces', Int
color) = Map Int space -> Set Int -> Int -> space -> (Map Int space, Int)
forall space.
Eq space =>
Map Int space -> Set Int -> Int -> space -> (Map Int space, Int)
firstAvailable Map Int space
spaces Set Int
nb_colors Int
0 space
sp
in (Map Int space
spaces', a -> Int -> Coloring a -> Coloring a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
x Int
color Coloring a
coloring)
colorGraph ::
(Ord a, Ord space) =>
M.Map a space ->
Interference.Graph a ->
(M.Map Int space, Coloring a)
colorGraph :: Map a space -> Graph a -> (Map Int space, Coloring a)
colorGraph Map a space
spaces Graph a
graph =
let nodes :: Set (a, space)
nodes = [(a, space)] -> Set (a, space)
forall a. Ord a => [a] -> Set a
S.fromList ([(a, space)] -> Set (a, space)) -> [(a, space)] -> Set (a, space)
forall a b. (a -> b) -> a -> b
$ Map a space -> [(a, space)]
forall k a. Map k a -> [(k, a)]
M.toList Map a space
spaces
nbs :: Neighbors a
nbs = Graph a -> Neighbors a
forall a. Ord a => Graph a -> Neighbors a
neighbors Graph a
graph
in ((a, space)
-> (Map Int space, Coloring a) -> (Map Int space, Coloring a))
-> (Map Int space, Coloring a)
-> Set (a, space)
-> (Map Int space, Coloring a)
forall a b. (a -> b -> b) -> b -> Set a -> b
S.foldr (Neighbors a
-> (a, space)
-> (Map Int space, Coloring a)
-> (Map Int space, Coloring a)
forall a space.
(Ord a, Eq space) =>
Neighbors a
-> (a, space)
-> (Map Int space, Coloring a)
-> (Map Int space, Coloring a)
colorNode Neighbors a
nbs) (Map Int space, Coloring a)
forall a. Monoid a => a
mempty Set (a, space)
nodes