Simple register allocation, see "Essentials of Compilation" for more https://jeapostrophe.github.io/courses/2017/spring/406/notes/book.pdf

coloring-3reg.svg

example analysis

aseipp@ubuntu:~/t$ ghci RegAlloc1.hs 
GHCi, version 8.0.2: http://www.haskell.org/ghc/  :? for help
[1 of 1] Compiling RegAlloc1        ( RegAlloc1.hs, interpreted )
Ok, modules loaded: RegAlloc1.
*RegAlloc1> printLive program1 
    movq $1, v.1    | {v}
    movq $46, w.2   | {v,w}
    movq v.1, x.3   | {w,x}
    addq $7, x.3    | {w,x}
    movq x.3, y.4   | {w,x,y}
    addq $4, y.4    | {w,x,y}
    movq x.3, z.5   | {w,y,z}
    addq w.2, z.5   | {y,z}
    movq y.4, t1.6  | {z,t1}
    negq t1.6       | {z,t1}
    movq z.5, t2.7  | {t1,t2}
    addq t1.6, t2.7 | {t2}
    movq t2.7, %rax | {}
*RegAlloc1> let (l, i, c) = (liveness program1, interference program1 l, colorGraph 3 i) in writeDotColored "coloring.dot" c
*RegAlloc1> 

RegAlloc1.hs

{-# OPTIONS_GHC -Wall #-}

{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE PartialTypeSignatures #-}
module RegAlloc1
  ( -- * Types
    Var
  , Reg(..)
  , Arg(..)
  , Instr(..)
  , Program(..)

    -- * Liveness analysis
  , Live
  , liveness
  , printLive

    -- * Interference
  , Interference
  , interference

    -- * Graph coloring
  , colorGraph

    -- * Example programs
  , program1

    -- * Utilities
    -- ** GraphViz
  , writeDot
  , writeDotInterference
  , writeDotColored

    -- ** x86 Registers
  , rax, rdi, rsi, rdx, rcx, r8, r9
  ) where

import           Control.Monad

import           Data.Data
import           Data.Function
import           Data.List
import           Data.Monoid
import           Data.String

-- containers
import           Data.Map            ( Map )
import qualified Data.Map            as Map
import           Data.Set            ( Set )
import qualified Data.Set            as Set

-- fgl, graphviz
import           Data.Graph.Inductive.Graph
import           Data.Graph.Inductive.PatriciaTree
import qualified Data.GraphViz                       as GV
import qualified Data.GraphViz.Attributes.Complete   as GV
import qualified Data.GraphViz.Commands.IO           as GV

--------------------------------------------------------------------------------
-- basics

data Var = VarU { _varI :: !Int, _varS :: String }
  deriving (Eq, Ord, Data, Typeable)

instance Show Var where
  show (VarU i s) = s <> ('.':show i)

data Reg = Rax | Rdi | Rsi | Rdx | Rcx | R8 | R9
  deriving (Eq, Ord, Data, Typeable)

instance Show Reg where
  show Rax = "%rax"
  show Rdi = "%rdi"
  show Rsi = "%rsi"
  show Rdx = "%rdx"
  show Rcx = "%rcx"
  show R8  = "%r8"
  show R9  = "%r9"

data Arg = LitI Integer | Var Var | Reg Reg
  deriving (Eq, Ord, Data, Typeable)

instance Show Arg where
  show (LitI i) = '$':show i
  show (Var v)  = show v
  show (Reg r)  = show r

data Instr
  = Movq Arg Arg
  | Addq Arg Arg
  | Negq Arg
  deriving (Eq, Ord, Data, Typeable)

instance Show Instr where
  show (Movq a d) = "movq" <+> show a <> "," <+> show d
  show (Addq a d) = "addq" <+> show a <> "," <+> show d
  show (Negq a)   = "negq" <+> show a

data Program = Program String [Instr]
  deriving (Eq, Ord, Data, Typeable)

instance Show Program where
  show (Program name instrs) =
    let fullName = '_':name
        indent   = ("    "++)

    in   ".global" <+> fullName
    <++> fullName <> ":"
    <++> unlines (map (indent . show) instrs)

--
-- utilities
--

rax, rdi, rsi, rdx, rcx, r8, r9 :: Arg
(rax, rdi, rsi, rdx, rcx, r8, r9)
  = (Reg Rax, Reg Rdi, Reg Rsi, Reg Rdx, Reg Rcx, Reg R8, Reg R9)

(<+>) :: (Monoid m, IsString m) => m -> m -> m
(<+>) a b  = a <> " "  <> b

(<++>) :: (Monoid m, IsString m) => m -> m -> m
(<++>) a b = a <> "\n" <> b

--
-- useful but stupid instances for writing clear examples
--

instance Num Arg where
  fromInteger = LitI

  -- warnings, lol
  (+)    = error "lol"
  (*)    = error "lol"
  abs    = error "lol"
  signum = error "lol"
  negate = error "lol"

--------------------------------------------------------------------------------
-- example program

program1 :: Program
program1 = Program "prog1"
  [ Movq 1  v
  , Movq 46 w
  , Movq v  x
  , Addq 7  x
  , Movq x  y
  , Addq 4  y
  , Movq x  z
  , Addq w  z
  , Movq y  t1
  , Negq t1
  , Movq z  t2
  , Addq t1 t2
  , Movq t2 rax
  ]
  where
    v  = Var $ VarU 1 "v"
    w  = Var $ VarU 2 "w"
    x  = Var $ VarU 3 "x"
    y  = Var $ VarU 4 "y"
    z  = Var $ VarU 5 "z"
    t1 = Var $ VarU 6 "t1"
    t2 = Var $ VarU 7 "t2"

--------------------------------------------------------------------------------
-- liveness analysis

type Live = Set Var

data LiveSet = Read | Write

getVar :: Arg -> Set Var
getVar (Var a) = Set.singleton a
getVar _       = Set.empty

need :: LiveSet -> Instr -> Live
-- the variables that are read in an instruction
need Read (Movq a _) = getVar a
need Read (Addq a d) = getVar a `Set.union` getVar d
need Read (Negq a)   = getVar a

-- the variables that are written in an instruction
need Write (Movq _ d) = getVar d
need Write (Addq _ d) = getVar d
need Write (Negq a)   = getVar a

-- | Perform a liveness analysis on the input program and determine the set of
-- live variables for every point in the program.
liveness :: Program -> [Live]
liveness (Program _ instrs)
  = reverse $ Set.empty : go input (k end Set.empty)
  where
    (end:input) = reverse instrs

    k x after = (after `Set.difference` w) `Set.union` r
      where (r, w) = (need Read x, need Write x)

    go []     _     = []
    go (x:xs) after = after : go xs (k x after)

-- | Print a program and its liveness information. Useful for debugging.
printLive :: Program -> IO ()
printLive prog@(Program _ instrs)
  = void
  $ flip traverse (zip instrs $ liveness prog)
  $ \(instr, lv) -> do
    let msg    = "    " ++ show instr
        buffer = replicate (20 - length msg) ' '
        info   = intercalate "," $ map _varS (Set.toList lv)
    putStrLn (msg ++ buffer ++ "| {" ++ info ++ "}")

--------------------------------------------------------------------------------
-- interference graphs

-- | An interference graph for a given @'Program'@. Nodes in the graph (which
-- are @Int@s) are labeled with the original variable they represented. Edges
-- are undirected, and have no labels as they convey no meaning.
type Interference = Gr Var ()

-- | Build an interference graph for a given program, based on the results of a
-- liveness analysis. Two variables in the input program interfere -- they are
-- both live at the same time -- if those variables have an edge between them in
-- the resulting graph.
interference :: Program
             -- ^ the input program
             -> [Live]
             -- ^ results of a liveness analysis on the input program
             -> Interference
             -- ^ interference graph
interference (Program _ instrs) live = foldr loop empty (zip instrs live)
  where
    -- create an edge between two variables. does not update the graph if the
    -- edge already exists.
    mkIEdge (VarU a _) (VarU b _) gr
      | hasEdge gr (a, b) = gr
      | otherwise         = insEdge (a, b, ()) gr

    -- insert a node if it doesn't exist already. does not update the graph if
    -- the given node already exists.
    mkINode v@(VarU i _) gr
      | gelem i gr = gr
      | otherwise  = insNode (i, v) gr

    -- insert a list of nodes, all at once.
    mkINodes vs gr = foldr mkINode gr vs

    -- mark two variables as interfering with each other, e.g. they are both
    -- live at the same time. this is represented as a single edge between both
    -- nodes in the graph.
    interferes :: Var -> Var -> Interference -> Interference
    interferes v1 v2 gr
      = mkIEdge v1 v2     -- create an edge between them
      $ mkINodes [v1, v2] -- insert both nodes
      $ gr

    -- mark a variable @d@ as interfering with all of the variables @v@ in given
    -- live variable set @lv@, iff @v@ passes a given predicate. this
    -- essentially creates a one-to-many mapping between @d@ and the elements of
    -- set @lv@ in the graph.
    insertP :: (Var -> Bool)
            -- ^ predicate function. for every variable @v \elem vs@ in
            -- the live set @vs@, if @v@ passes this predicate, then the two
            -- variables @(d, v)@ are marked as interfering, i.e. an edge is
            -- created between them.
            -> (Var, Live)
            -- ^ @(d, lv)@ -- the variable @d@ and live variable set @lv@
            -> Interference
            -- ^ input graph
            -> Interference
            -- ^ resulting graph
    insertP predicate (d, lv) gr =
        let k :: Var -> Interference -> Interference
            k v g | predicate v = interferes d v g
                  | otherwise   = g
        in Set.foldr k gr lv

    -- mark all of the interferences for a given instruction and its live
    -- variable set.
    loop :: (Instr, Live)
         -- ^ input instruction, and live variable set for that instruction
         -> Interference
         -- ^ input graph
         -> Interference
         -- ^ output graph
    loop x gr = case x of
      -- for move instructions, we mark @d@ as interfering with @v \elem lv@,
      -- iff @v /= d@ (we don't associate @d@ with itself) and @d /= a@ (we
      -- don't associate @d@ with the source @a@, since @a@ may be dead after
      -- this, or it may not be a variable at all)
      (Movq a (Var d), lv) -> insertP (\v -> v /= d && (Var v) /= a) (d, lv) gr

      -- for arithmetic instructions, we mark @d@ as interfering with
      -- @v \elem lv@, iff @v /= d@ (we don't associate @d@ with itself)
      (Addq _ (Var d), lv) -> insertP (\v -> v /= d)                 (d, lv) gr
      (Negq   (Var d), lv) -> insertP (\v -> v /= d)                 (d, lv) gr

      -- otherwise, return
      _ -> gr

writeDotInterference :: FilePath -> Interference -> IO ()
writeDotInterference = writeDot $ \l -> [ GV.toLabel (_varS l) ]

--------------------------------------------------------------------------------
-- Graph coloring via DSATUR

type GColor  = (Int, GV.Color)
type Colored = Gr (Var, Maybe GColor) ()

-- see https://stackoverflow.com/a/13781114
allColors :: [GV.Color]
allColors = do
  i <- [ 2** k | k <- [0..] ]
  j <- enumFromThenTo 1 3 i
  v <- [ 8 / 10, 5 / 10 ]
  return $ GV.HSV (j / i) (6 / 10) v

color :: Int -> Node -> GColor
color limit i = zip [0..] allColors !! (i `mod` limit)

saturation :: Graph gr => Int -> gr a b -> Node -> [Node]
saturation limit gr u
  = map fst $ ordNubBy id ((==) `on` fst)
      [ color limit v
      | v <- neighbors gr u
      ]

sortForColoring :: Graph gr => Int -> gr a b -> [Node]
sortForColoring limit gr
  = sortBy (flip compare `on` length . saturation limit gr)
  $ nodes gr

getMinColor :: Graph gr => Int -> gr a b -> Node -> Maybe GColor
getMinColor limit gr u
  = let taken      = saturation limit gr u -- find all the colors that are taken
        everyColor = [ 0 .. limit-1 ]      -- set of all possible colors
        getColor   = color limit
    in case everyColor \\ taken of          -- find non-taken colors
         [] -> Nothing                      -- none available; stack slot
         vs -> Just (getColor $ minimum vs) -- otherwise, get min color

buildColorMap :: Int -> Interference -> Map Node (Maybe GColor)
buildColorMap limit gr
  = foldr (\n -> Map.insert n (getMinColor limit gr n)) Map.empty
  $ sortForColoring limit gr

colorGraph :: Int -> Interference -> Colored
colorGraph limit gr = nmap (\v@(VarU i _) -> (v, get i)) gr
  where
    get       = maybe (error "colorGraph: impossible!") id . findMap
    findMap x = Map.lookup x (buildColorMap limit gr)

writeDotColored :: FilePath -> Colored -> IO ()
writeDotColored = writeDot $ \(l, c) ->
  case c of
    -- no coloring; must be a stack slot
    Nothing      -> [ GV.toLabel $ _varS l <+> "(stack)"
                    , GV.style GV.dashed
                    ]
    Just (_, cl) -> [ GV.toLabel $ _varS l
                    , GV.Color $ GV.toColorList [ cl ]
                    , GV.style GV.solid
                    ]

--------------------------------------------------------------------------------
-- GraphViz utilities

writeDot :: Graph gr
         => (a -> GV.Attributes)
         -> FilePath
         -> gr a b
         -> IO ()
writeDot markup fp
  = GV.writeDotFile fp
  . GV.graphToDot GV.nonClusteredParams
      { GV.globalAttributes = [ noDir ]   -- undirected edges
      , GV.fmtEdge = \(_, _, _l) -> []    -- edges have no label
      , GV.fmtNode = \(_, l) -> markup l  -- label nodes by name
      }
  where
    noDir = GV.EdgeAttrs [ GV.Dir GV.NoDir ]

ordNubBy :: (Ord b) => (a -> b) -> (a -> a -> Bool) -> [a] -> [a]
ordNubBy p f l = go Map.empty l
  where
    go _ []     = []
    go m (x:xs) =
      let b = p x
      in case b `Map.lookup` m of
           Nothing     -> x : go (Map.insert b [x] m) xs
           Just bucket
             | elem_by f x bucket -> go m xs
             | otherwise          -> x : go (Map.insert b (x:bucket) m) xs

    -- From the Data.List source code.
    elem_by :: (a -> a -> Bool) -> a -> [a] -> Bool
    elem_by _  _ []     = False
    elem_by eq y (x:xs) = y `eq` x || elem_by eq y xs