{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE TemplateHaskell #-}
{-# OPTIONS_GHC -Wno-partial-type-signatures #-}
module RegAlloc (Operation (..), RegCount, allocRegs, allocRegs', colorize) where
import Prelude hiding (id, (.))
import Control.Applicative
import Control.Category
import Control.Lens.TH (mkLens)
import Control.Monad (guard, join)
import Control.Monad.Except
import Control.Monad.State
import qualified Control.Monad.State.Lens as ML
import Control.Monad.Writer
import Data.Bool (bool)
import Data.Foldable (find, toList, traverse_)
import Data.IntMap (IntMap)
import qualified Data.IntMap as IM
import Data.IntSet (IntSet)
import qualified Data.IntSet as IS
import qualified Data.List as List
import Data.Ord (Down (..))
import Data.Peano
import Data.Traversable (for)
import qualified Lens.Micro.Mtl as ML
import Util hiding ((∈))
import RegAlloc.Interference (Interferences, Node (..), Operands, UGraph, interferences, interferes, (!), (∈))
import qualified RegAlloc.Interference as If
import qualified RegAlloc.Nodes as Nodes
import qualified RegAlloc.UGraph as UGr
data St = St
{ _degree :: !Int
, _ifs :: !Interferences
, _moves :: !Moves
}
deriving (Eq, Show)
type Moves = UGraph
$(mkLens (dropWhile (== '_')) ''St)
allocRegs :: (Traversable f) => RegCount -> f Operation -> Except Interferences (f Int)
allocRegs deg insns = do
colors <- (allocRegs' deg ifm >=> uncurry (colorize deg ifm)) moves
for (count insns) \ (k, _) -> maybe (throwError ifm) pure $ colors IM.!? k
where
ifm = interferences insns'
insns' = (\ case NonMove xs -> xs; Move x -> Nodes.fromList [x]) <$> insns
moves =
UGr.insertEdges [(k', k :: Int) | (k, Move k') <- toList (count insns)] (UGr.empty deg)
allocRegs' :: RegCount -> Interferences -> Moves -> Except Interferences ([Op], Colors)
allocRegs' deg ifm theMoves =
execWriterT . flip evalStateT (St { _degree = 1, _ifs = ifm, _moves = theMoves }) $
whileM (untilFixpointBy (==) (simplifyAndCoalescePhase >> freezePhase) >>
not . If.null <$> ML.gets ifs) potentialSpillPhase
where
simplifyAndCoalescePhase = doWhileM bumpDegree do
St { _degree = deg } <- get
untilFixpointBy (==) do
_ <- untilFixpointBy (==) simplify
St { _ifs = theIfs, _moves = theMoves } <- get
let isCoalescibleMove (k, k') =
not (interferes k' k theIfs) && coalescible deg k' k theIfs
for (find isCoalescibleMove (UGr.edges theMoves)) \ (k, k') -> coalesce1 k' k
freezePhase = do
St { _degree = deg, _ifs = theIfs, _moves = theMoves } <- get
case [k | k <- [0..deg-1]
, k' <- Nodes.toList (theMoves ! k)
, not $ If.interferes k' k theIfs
, Nodes.size (theIfs ! k) < deg] of
[] -> pure ()
k:_ -> () <$ ML.puts moves (UGr.deleteNode k theMoves)
potentialSpillPhase = do
k <- ML.zoom ifs potentialSpill
deleteNode k
tell ([Select k], IM.empty)
bumpDegree = compare deg <$> ML.gets degree >>= \ case
GT -> True <$ ML.modify degree (+1)
_ -> pure False
colorize
:: (MonadError Interferences m, Foldable t)
=> RegCount -> Interferences -> t Op -> Colors -> m Colors
colorize deg ifm = execStateT . traverse_ \ case
Select k -> do
colors <- get
let nbrs = ifm ! k
nbrColors = IS.fromList
[c | nbr <- Nodes.toList nbrs
, Just c <- [case nbr of
Node k -> colors IM.!? k
Precolored c -> Just c]]
go k $ find (`IS.notMember` nbrColors) [0..deg-1]
Coalesce k k' -> go k =<< gets (IM.!? k')
where go k = maybe (throwError ifm) (modify . IM.insert k)
data Op = Select !Int | Coalesce !Int !Int
deriving (Show)
deleteNode :: MonadState St m => Int -> m ()
deleteNode k = traverse_ ($ UGr.deleteNode k) [ML.modify ifs, ML.modify moves]
deleteNodes :: MonadState St m => IntSet -> m ()
deleteNodes ks = traverse_ ($ UGr.deleteNodes ks) [ML.modify ifs, ML.modify moves]
coalesce1 :: (MonadState St m, MonadWriter ([Op], Colors) m) => Node -> Int -> m ()
coalesce1 k' k = do
traverse_ ($ UGr.coalesce k' k) [ML.modify ifs, ML.modify moves]
tell case k' of
Node k' -> ([Coalesce k k'], IM.empty)
Precolored c -> ([], IM.singleton k c)
simplify :: (MonadState St m, MonadWriter ([Op], Colors) m) => m ()
simplify =
concatMap IS.toList <$> untilFixpointBy (==) simplify1 >>=
tell . flip (,) IM.empty . fmap Select
simplify1 :: MonadState St m => m IntSet
simplify1 = do
St { _degree = deg, _ifs = theIfs, _moves = theMoves } <- get
let moveRelateds = foldMap IS.fromList [[k, k'] | (k, Node k') <- UGr.edges theMoves]
delenda = IS.fromAscList
[k | (k, ks) <- If.toAscList theIfs
, IS.notMember k moveRelateds && k >= 0 && Nodes.size ks < deg]
delenda <$ deleteNodes delenda
untilFixpointBy :: (MonadState s m) => (s -> s -> Bool) -> m a -> m [a]
untilFixpointBy eq x = go [] where
go as = do
s <- get
a <- x
t <- get
bool (go . (a:)) pure (eq s t) as
whileM :: Monad m => m Bool -> m a -> m [a]
whileM = compose2 whileJust (fmap guard) pure
doWhileM :: Monad m => m Bool -> m a -> m [a]
doWhileM p = liftA2 (:) <*> whileM p
coalescible :: RegCount -> Node -> Int -> Interferences -> Bool
coalescible n = join \ case
Node _ -> briggs n
Precolored _ -> george n
briggs, george :: RegCount -> Node -> Int -> Interferences -> Bool
briggs n a b ifm = (fromIntegral n :: Peano) > flip count' (Nodes.toList (ifm' ! b)) \ case
Node c -> Nodes.size (ifm' ! c) > n
Precolored _ -> True
where
ifm' = UGr.coalesce a b ifm
count' f = List.genericLength . filter f
george n a b ifm = flip all (Nodes.toList (ifm ! b)) \ case
Node c -> let nbrs = ifm ! c
in Nodes.size nbrs < n || a ∈ nbrs
c@(Precolored _) -> all (c ∈) aNbrs'
where
aNbrs' = case a of
Precolored _ -> Nothing
Node a -> Just (ifm ! a)
potentialSpill :: (MonadState Interferences m, MonadError Interferences m) => m Int
potentialSpill = List.sortOn (Down . Nodes.size . snd) . If.toAscList <$> get >>= \ case
[] -> get >>= throwError
(k, _):_ -> pure k
data Operation = Move !Node | NonMove !Operands
type RegCount = Int
type Colors = IntMap Int