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