compiler/nativeGen/RegisterAlloc.hs

   1 -----------------------------------------------------------------------------
   2 --
   3 -- The register allocator
   4 --
   5 -- (c) The University of Glasgow 2004
   6 --
   7 -----------------------------------------------------------------------------
   8
   9 {-
  10 The algorithm is roughly:
  11
  12   1) Compute strongly connected components of the basic block list.
  13
  14   2) Compute liveness (mapping from pseudo register to
  15      point(s) of death?).
  16
  17   3) Walk instructions in each basic block.  We keep track of
  18         (a) Free real registers (a bitmap?)
  19         (b) Current assignment of temporaries to machine registers and/or
  20             spill slots (call this the "assignment").
  21         (c) Partial mapping from basic block ids to a virt-to-loc mapping.
  22             When we first encounter a branch to a basic block,
  23             we fill in its entry in this table with the current mapping.
  24
  25      For each instruction:
  26         (a) For each real register clobbered by this instruction:
  27             If a temporary resides in it,
  28                 If the temporary is live after this instruction,
  29                     Move the temporary to another (non-clobbered & free) reg,
  30                     or spill it to memory.  Mark the temporary as residing
  31                     in both memory and a register if it was spilled (it might
  32                     need to be read by this instruction).
  33             (ToDo: this is wrong for jump instructions?)
  34
  35         (b) For each temporary *read* by the instruction:
  36             If the temporary does not have a real register allocation:
  37                 - Allocate a real register from the free list.  If
  38                   the list is empty:
  39                   - Find a temporary to spill.  Pick one that is
  40                     not used in this instruction (ToDo: not
  41                     used for a while...)
  42                   - generate a spill instruction
  43                 - If the temporary was previously spilled,
  44                   generate an instruction to read the temp from its spill loc.
  45             (optimisation: if we can see that a real register is going to
  46             be used soon, then don't use it for allocation).
  47
  48         (c) Update the current assignment
  49
  50         (d) If the intstruction is a branch:
  51               if the destination block already has a register assignment,
  52                 Generate a new block with fixup code and redirect the
  53                 jump to the new block.
  54               else,
  55                 Update the block id->assignment mapping with the current
  56                 assignment.
  57
  58         (e) Delete all register assignments for temps which are read
  59             (only) and die here.  Update the free register list.
  60
  61         (f) Mark all registers clobbered by this instruction as not free,
  62             and mark temporaries which have been spilled due to clobbering
  63             as in memory (step (a) marks then as in both mem & reg).
  64
  65         (g) For each temporary *written* by this instruction:
  66             Allocate a real register as for (b), spilling something
  67             else if necessary.
  68                 - except when updating the assignment, drop any memory
  69                   locations that the temporary was previously in, since
  70                   they will be no longer valid after this instruction.
  71
  72         (h) Delete all register assignments for temps which are
  73             written and die here (there should rarely be any).  Update
  74             the free register list.
  75
  76         (i) Rewrite the instruction with the new mapping.
  77
  78         (j) For each spilled reg known to be now dead, re-add its stack slot
  79             to the free list.
  80
  81 -}
  82
  83 module RegisterAlloc (
  84         regAlloc
  85   ) where
  86
  87 #include "HsVersions.h"
  88
  89 import PprMach
  90 import MachRegs
  91 import MachInstrs
  92 import RegAllocInfo
  93 import Cmm
  94
  95 import Digraph
  96 import Unique           ( Uniquable(getUnique), Unique )
  97 import UniqSet
  98 import UniqFM
  99 import UniqSupply
 100 import Outputable
 101
 102 #ifndef DEBUG
 103 import Data.Maybe       ( fromJust )
 104 #endif
 105 import Data.List        ( nub, partition, mapAccumL, groupBy )
 106 import Control.Monad    ( when )
 107 import Data.Word
 108 import Data.Bits
 109
 110 -- -----------------------------------------------------------------------------
 111 -- Some useful types
 112
 113 type RegSet = UniqSet Reg
 114
 115 type RegMap a = UniqFM a
 116 emptyRegMap = emptyUFM
 117
 118 type BlockMap a = UniqFM a
 119 emptyBlockMap = emptyUFM
 120
 121 -- A basic block where the isntructions are annotated with the registers
 122 -- which are no longer live in the *next* instruction in this sequence.
 123 -- (NB. if the instruction is a jump, these registers might still be live
 124 -- at the jump target(s) - you have to check the liveness at the destination
 125 -- block to find out).
 126 type AnnBasicBlock
 127         = GenBasicBlock (Instr,
 128                          [Reg],         -- registers read (only) which die
 129                          [Reg])         -- registers written which die
 130
 131 -- -----------------------------------------------------------------------------
 132 -- The free register set
 133
 134 -- This needs to be *efficient*
 135
 136 {- Here's an inefficient 'executable specification' of the FreeRegs data type:
 137 type FreeRegs = [RegNo]
 138
 139 noFreeRegs = 0
 140 releaseReg n f = if n `elem` f then f else (n : f)
 141 initFreeRegs = allocatableRegs
 142 getFreeRegs cls f = filter ( (==cls) . regClass . RealReg ) f
 143 allocateReg f r = filter (/= r) f
 144 -}
 145
 146 #if defined(powerpc_TARGET_ARCH)
 147
 148 -- The PowerPC has 32 integer and 32 floating point registers.
 149 -- This is 32bit PowerPC, so Word64 is inefficient - two Word32s are much
 150 -- better.
 151 -- Note that when getFreeRegs scans for free registers, it starts at register
 152 -- 31 and counts down. This is a hack for the PowerPC - the higher-numbered
 153 -- registers are callee-saves, while the lower regs are caller-saves, so it
 154 -- makes sense to start at the high end.
 155 -- Apart from that, the code does nothing PowerPC-specific, so feel free to
 156 -- add your favourite platform to the #if (if you have 64 registers but only
 157 -- 32-bit words).
 158
 159 data FreeRegs = FreeRegs !Word32 !Word32
 160               deriving( Show )  -- The Show is used in an ASSERT
 161
 162 noFreeRegs :: FreeRegs
 163 noFreeRegs = FreeRegs 0 0
 164
 165 releaseReg :: RegNo -> FreeRegs -> FreeRegs
 166 releaseReg r (FreeRegs g f)
 167     | r > 31    = FreeRegs g (f .|. (1 `shiftL` (fromIntegral r - 32)))
 168     | otherwise = FreeRegs (g .|. (1 `shiftL` fromIntegral r)) f
 169
 170 initFreeRegs :: FreeRegs
 171 initFreeRegs = foldr releaseReg noFreeRegs allocatableRegs
 172
 173 getFreeRegs :: RegClass -> FreeRegs -> [RegNo]  -- lazilly
 174 getFreeRegs cls (FreeRegs g f)
 175     | RcDouble <- cls = go f (0x80000000) 63
 176     | RcInteger <- cls = go g (0x80000000) 31
 177     where
 178         go x 0 i = []
 179         go x m i | x .&. m /= 0 = i : (go x (m `shiftR` 1) $! i-1)
 180                  | otherwise    = go x (m `shiftR` 1) $! i-1
 181
 182 allocateReg :: RegNo -> FreeRegs -> FreeRegs
 183 allocateReg r (FreeRegs g f)
 184     | r > 31    = FreeRegs g (f .&. complement (1 `shiftL` (fromIntegral r - 32)))
 185     | otherwise = FreeRegs (g .&. complement (1 `shiftL` fromIntegral r)) f
 186
 187 #else
 188
 189 -- If we have less than 32 registers, or if we have efficient 64-bit words,
 190 -- we will just use a single bitfield.
 191
 192 #if defined(alpha_TARGET_ARCH)
 193 type FreeRegs = Word64
 194 #else
 195 type FreeRegs = Word32
 196 #endif
 197
 198 noFreeRegs :: FreeRegs
 199 noFreeRegs = 0
 200
 201 releaseReg :: RegNo -> FreeRegs -> FreeRegs
 202 releaseReg n f = f .|. (1 `shiftL` n)
 203
 204 initFreeRegs :: FreeRegs
 205 initFreeRegs = foldr releaseReg noFreeRegs allocatableRegs
 206
 207 getFreeRegs :: RegClass -> FreeRegs -> [RegNo]  -- lazilly
 208 getFreeRegs cls f = go f 0
 209   where go 0 m = []
 210         go n m
 211           | n .&. 1 /= 0 && regClass (RealReg m) == cls
 212           = m : (go (n `shiftR` 1) $! (m+1))
 213           | otherwise
 214           = go (n `shiftR` 1) $! (m+1)
 215         -- ToDo: there's no point looking through all the integer registers
 216         -- in order to find a floating-point one.
 217
 218 allocateReg :: RegNo -> FreeRegs -> FreeRegs
 219 allocateReg r f = f .&. complement (1 `shiftL` fromIntegral r)
 220
 221 #endif
 222
 223 -- -----------------------------------------------------------------------------
 224 -- The assignment of virtual registers to stack slots
 225
 226 -- We have lots of stack slots. Memory-to-memory moves are a pain on most
 227 -- architectures. Therefore, we avoid having to generate memory-to-memory moves
 228 -- by simply giving every virtual register its own stack slot.
 229
 230 -- The StackMap stack map keeps track of virtual register - stack slot
 231 -- associations and of which stack slots are still free. Once it has been
 232 -- associated, a stack slot is never "freed" or removed from the StackMap again,
 233 -- it remains associated until we are done with the current CmmProc.
 234
 235 type StackSlot = Int
 236 data StackMap = StackMap [StackSlot] (UniqFM StackSlot)
 237
 238 emptyStackMap :: StackMap
 239 emptyStackMap = StackMap [0..maxSpillSlots] emptyUFM
 240
 241 getStackSlotFor :: StackMap -> Unique -> (StackMap,Int)
 242 getStackSlotFor fs@(StackMap (freeSlot:stack') reserved) reg =
 243     case lookupUFM reserved reg of
 244         Just slot -> (fs,slot)
 245         Nothing -> (StackMap stack' (addToUFM reserved reg freeSlot), freeSlot)
 246
 247 -- -----------------------------------------------------------------------------
 248 -- Top level of the register allocator
 249
 250 regAlloc :: NatCmmTop -> UniqSM NatCmmTop
 251 regAlloc (CmmData sec d) = returnUs $ CmmData sec d
 252 regAlloc (CmmProc info lbl params [])
 253   = returnUs $ CmmProc info lbl params []  -- no blocks to run the regalloc on
 254 regAlloc (CmmProc info lbl params blocks@(first:rest))
 255   = let
 256         first_id               = blockId first
 257         sccs                   = sccBlocks blocks
 258         (ann_sccs, block_live) = computeLiveness sccs
 259     in  linearRegAlloc block_live ann_sccs `thenUs` \final_blocks ->
 260     let ((first':_),rest')     = partition ((== first_id) . blockId) final_blocks
 261     in returnUs $ -- pprTrace "Liveness" (ppr block_live) $
 262                   CmmProc info lbl params (first':rest')
 263
 264 sccBlocks :: [NatBasicBlock] -> [SCC NatBasicBlock]
 265 sccBlocks blocks = stronglyConnComp graph
 266   where
 267         getOutEdges :: [Instr] -> [BlockId]
 268         getOutEdges instrs = foldr jumpDests [] instrs
 269
 270         graph = [ (block, getUnique id, map getUnique (getOutEdges instrs))
 271                 | block@(BasicBlock id instrs) <- blocks ]
 272
 273
 274 -- -----------------------------------------------------------------------------
 275 -- Computing liveness
 276
 277 computeLiveness
 278    :: [SCC NatBasicBlock]
 279    -> ([SCC AnnBasicBlock],     -- instructions annotated with list of registers
 280                                 -- which are "dead after this instruction".
 281        BlockMap RegSet)         -- blocks annontated with set of live registers
 282                                 -- on entry to the block.
 283
 284   -- NOTE: on entry, the SCCs are in "reverse" order: later blocks may transfer
 285   -- control to earlier ones only.  The SCCs returned are in the *opposite*
 286   -- order, which is exactly what we want for the next pass.
 287
 288 computeLiveness sccs
 289   = livenessSCCs emptyBlockMap [] sccs
 290   where
 291   livenessSCCs
 292          :: BlockMap RegSet
 293          -> [SCC AnnBasicBlock]         -- accum
 294          -> [SCC NatBasicBlock]
 295          -> ([SCC AnnBasicBlock], BlockMap RegSet)
 296
 297   livenessSCCs blockmap done [] = (done, blockmap)
 298   livenessSCCs blockmap done
 299         (AcyclicSCC (BasicBlock block_id instrs) : sccs) =
 300           {- pprTrace "live instrs" (ppr (getUnique block_id) $$
 301                                   vcat (map (\(instr,regs) -> docToSDoc (pprInstr instr) $$ ppr regs) instrs')) $
 302           -}
 303           livenessSCCs blockmap'
 304                 (AcyclicSCC (BasicBlock block_id instrs'):done) sccs
 305         where (live,instrs') = liveness emptyUniqSet blockmap []
 306                                         (reverse instrs)
 307               blockmap' = addToUFM blockmap block_id live
 308
 309   livenessSCCs blockmap done
 310         (CyclicSCC blocks : sccs) =
 311           livenessSCCs blockmap' (CyclicSCC blocks':done) sccs
 312         where (blockmap', blocks')
 313                   = iterateUntilUnchanged linearLiveness equalBlockMaps
 314                                         blockmap blocks
 315
 316               iterateUntilUnchanged
 317                   :: (a -> b -> (a,c)) -> (a -> a -> Bool)
 318                   -> a -> b
 319                   -> (a,c)
 320
 321               iterateUntilUnchanged f eq a b
 322                   = head $
 323                     concatMap tail $
 324                     groupBy (\(a1, _) (a2, _) -> eq a1 a2) $
 325                     iterate (\(a, _) -> f a b) $
 326                     (a, error "RegisterAlloc.livenessSCCs")
 327
 328
 329               linearLiveness :: BlockMap RegSet -> [NatBasicBlock]
 330                              -> (BlockMap RegSet, [AnnBasicBlock])
 331               linearLiveness = mapAccumL processBlock
 332
 333               processBlock blockmap input@(BasicBlock block_id instrs)
 334                   = (blockmap', BasicBlock block_id instrs')
 335                 where (live,instrs') = liveness emptyUniqSet blockmap []
 336                                                 (reverse instrs)
 337                       blockmap' = addToUFM blockmap block_id live
 338
 339                   -- probably the least efficient way to compare two
 340                   -- BlockMaps for equality.
 341               equalBlockMaps a b
 342                   = a' == b'
 343                 where a' = map f $ ufmToList a
 344                       b' = map f $ ufmToList b
 345                       f (key,elt) = (key, uniqSetToList elt)
 346
 347   liveness :: RegSet                    -- live regs
 348            -> BlockMap RegSet           -- live regs on entry to other BBs
 349            -> [(Instr,[Reg],[Reg])]     -- instructions (accum)
 350            -> [Instr]                   -- instructions
 351            -> (RegSet, [(Instr,[Reg],[Reg])])
 352
 353   liveness liveregs blockmap done []  = (liveregs, done)
 354   liveness liveregs blockmap done (instr:instrs)
 355         | not_a_branch = liveness liveregs1 blockmap
 356                                 ((instr,r_dying,w_dying):done) instrs
 357         | otherwise = liveness liveregs_br blockmap
 358                                 ((instr,r_dying_br,w_dying):done) instrs
 359         where
 360               RU read written = regUsage instr
 361
 362               -- registers that were written here are dead going backwards.
 363               -- registers that were read here are live going backwards.
 364               liveregs1 = (liveregs `delListFromUniqSet` written)
 365                                     `addListToUniqSet` read
 366
 367               -- registers that are not live beyond this point, are recorded
 368               --  as dying here.
 369               r_dying  = [ reg | reg <- read, reg `notElem` written,
 370                                  not (elementOfUniqSet reg liveregs) ]
 371
 372               w_dying = [ reg | reg <- written,
 373                                 not (elementOfUniqSet reg liveregs) ]
 374
 375               -- union in the live regs from all the jump destinations of this
 376               -- instruction.
 377               targets = jumpDests instr [] -- where we go from here
 378               not_a_branch = null targets
 379
 380               targetLiveRegs target = case lookupUFM blockmap target of
 381                                         Just ra -> ra
 382                                         Nothing -> emptyBlockMap
 383
 384               live_from_branch = unionManyUniqSets (map targetLiveRegs targets)
 385
 386               liveregs_br = liveregs1 `unionUniqSets` live_from_branch
 387
 388               -- registers that are live only in the branch targets should
 389               -- be listed as dying here.
 390               live_branch_only = live_from_branch `minusUniqSet` liveregs
 391               r_dying_br = uniqSetToList (mkUniqSet r_dying `unionUniqSets`
 392                                           live_branch_only)
 393
 394 -- -----------------------------------------------------------------------------
 395 -- Linear sweep to allocate registers
 396
 397 data Loc = InReg   {-# UNPACK #-} !RegNo
 398          | InMem   {-# UNPACK #-} !Int          -- stack slot
 399          | InBoth  {-# UNPACK #-} !RegNo
 400                    {-# UNPACK #-} !Int          -- stack slot
 401   deriving (Eq, Show, Ord)
 402
 403 {-
 404 A temporary can be marked as living in both a register and memory
 405 (InBoth), for example if it was recently loaded from a spill location.
 406 This makes it cheap to spill (no save instruction required), but we
 407 have to be careful to turn this into InReg if the value in the
 408 register is changed.
 409
 410 This is also useful when a temporary is about to be clobbered.  We
 411 save it in a spill location, but mark it as InBoth because the current
 412 instruction might still want to read it.
 413 -}
 414
 415 #ifdef DEBUG
 416 instance Outputable Loc where
 417   ppr l = text (show l)
 418 #endif
 419
 420 linearRegAlloc
 421    :: BlockMap RegSet           -- live regs on entry to each basic block
 422    -> [SCC AnnBasicBlock]       -- instructions annotated with "deaths"
 423    -> UniqSM [NatBasicBlock]
 424 linearRegAlloc block_live sccs = linearRA_SCCs emptyBlockMap emptyStackMap sccs
 425   where
 426   linearRA_SCCs
 427         :: BlockAssignment
 428         -> StackMap
 429         -> [SCC AnnBasicBlock]
 430         -> UniqSM [NatBasicBlock]
 431   linearRA_SCCs block_assig stack [] = returnUs []
 432   linearRA_SCCs block_assig stack
 433         (AcyclicSCC (BasicBlock id instrs) : sccs)
 434         = getUs `thenUs` \us ->
 435           let
 436             (block_assig',stack',(instrs',fixups)) =
 437                case lookupUFM block_assig id of
 438                     -- no prior info about this block: assume everything is
 439                     -- free and the assignment is empty.
 440                     Nothing ->
 441                         runR block_assig initFreeRegs
 442                                     emptyRegMap stack us $
 443                             linearRA [] [] instrs
 444                     Just (freeregs,assig) ->
 445                        runR block_assig freeregs assig stack us $
 446                             linearRA [] [] instrs
 447           in
 448           linearRA_SCCs block_assig' stack' sccs `thenUs` \moreBlocks ->
 449           returnUs $ BasicBlock id instrs' : fixups ++ moreBlocks
 450
 451   linearRA_SCCs block_assig stack
 452         (CyclicSCC blocks : sccs)
 453         = getUs `thenUs` \us ->
 454           let
 455             ((block_assig', stack', _), blocks') = mapAccumL processBlock
 456                                                        (block_assig, stack, us)
 457                                                        ({-reverse-} blocks)
 458           in
 459           linearRA_SCCs block_assig' stack' sccs `thenUs` \moreBlocks ->
 460           returnUs $ concat blocks' ++ moreBlocks
 461     where
 462         processBlock (block_assig, stack, us0) (BasicBlock id instrs)
 463           = ((block_assig', stack', us'), BasicBlock id instrs' : fixups)
 464           where
 465                 (us, us') = splitUniqSupply us0
 466                 (block_assig',stack',(instrs',fixups)) =
 467                    case lookupUFM block_assig id of
 468                         -- no prior info about this block: assume everything is
 469                         -- free and the assignment is empty.
 470                         Nothing ->
 471                            runR block_assig initFreeRegs
 472                                         emptyRegMap stack us $
 473                                 linearRA [] [] instrs
 474                         Just (freeregs,assig) ->
 475                            runR block_assig freeregs assig stack us $
 476                                 linearRA [] [] instrs
 477
 478   linearRA :: [Instr] -> [NatBasicBlock] -> [(Instr,[Reg],[Reg])]
 479         -> RegM ([Instr], [NatBasicBlock])
 480   linearRA instr_acc fixups [] =
 481     return (reverse instr_acc, fixups)
 482   linearRA instr_acc fixups (instr:instrs) = do
 483     (instr_acc', new_fixups) <- raInsn block_live instr_acc instr
 484     linearRA instr_acc' (new_fixups++fixups) instrs
 485
 486 -- -----------------------------------------------------------------------------
 487 -- Register allocation for a single instruction
 488
 489 type BlockAssignment = BlockMap (FreeRegs, RegMap Loc)
 490
 491 raInsn  :: BlockMap RegSet              -- Live temporaries at each basic block
 492         -> [Instr]                      -- new instructions (accum.)
 493         -> (Instr,[Reg],[Reg])          -- the instruction (with "deaths")
 494         -> RegM (
 495              [Instr],                   -- new instructions
 496              [NatBasicBlock]            -- extra fixup blocks
 497            )
 498
 499 raInsn block_live new_instrs (instr@(DELTA n), _, _) = do
 500     setDeltaR n
 501     return (new_instrs, [])
 502
 503 raInsn block_live new_instrs (instr, r_dying, w_dying) = do
 504     assig    <- getAssigR
 505
 506     -- If we have a reg->reg move between virtual registers, where the
 507     -- src register is not live after this instruction, and the dst
 508     -- register does not already have an assignment,
 509     -- and the source register is assigned to a register, not to a spill slot,
 510     -- then we can eliminate the instruction.
 511     -- (we can't eliminitate it if the source register is on the stack, because
 512     --  we do not want to use one spill slot for different virtual registers)
 513     case isRegRegMove instr of
 514         Just (src,dst)  | src `elem` r_dying,
 515                           isVirtualReg dst,
 516                           not (dst `elemUFM` assig),
 517                           Just (InReg _) <- (lookupUFM assig src) -> do
 518            case src of
 519               RealReg i -> setAssigR (addToUFM assig dst (InReg i))
 520                 -- if src is a fixed reg, then we just map dest to this
 521                 -- reg in the assignment.  src must be an allocatable reg,
 522                 -- otherwise it wouldn't be in r_dying.
 523               _virt -> case lookupUFM assig src of
 524                          Nothing -> panic "raInsn"
 525                          Just loc ->
 526                            setAssigR (addToUFM (delFromUFM assig src) dst loc)
 527
 528            -- we have elimianted this instruction
 529            {-
 530            freeregs <- getFreeRegsR
 531            assig <- getAssigR
 532            pprTrace "raInsn" (text "ELIMINATED: " <> docToSDoc (pprInstr instr) $$ ppr r_dying <+> ppr w_dying $$ text (show freeregs) $$ ppr assig) $ do
 533            -}
 534            return (new_instrs, [])
 535
 536         other -> genRaInsn block_live new_instrs instr r_dying w_dying
 537
 538
 539 genRaInsn block_live new_instrs instr r_dying w_dying =
 540     case regUsage instr              of { RU read written ->
 541     case partition isRealReg written of { (real_written1,virt_written) ->
 542     do
 543     let
 544         real_written = [ r | RealReg r <- real_written1 ]
 545
 546         -- we don't need to do anything with real registers that are
 547         -- only read by this instr.  (the list is typically ~2 elements,
 548         -- so using nub isn't a problem).
 549         virt_read = nub (filter isVirtualReg read)
 550     -- in
 551
 552     -- (a) save any temporaries which will be clobbered by this instruction
 553     clobber_saves <- saveClobberedTemps real_written r_dying
 554
 555     {-
 556     freeregs <- getFreeRegsR
 557     assig <- getAssigR
 558     pprTrace "raInsn" (docToSDoc (pprInstr instr) $$ ppr r_dying <+> ppr w_dying $$ ppr virt_read <+> ppr virt_written $$ text (show freeregs) $$ ppr assig) $ do
 559     -}
 560
 561     -- (b), (c) allocate real regs for all regs read by this instruction.
 562     (r_spills, r_allocd) <-
 563         allocateRegsAndSpill True{-reading-} virt_read [] [] virt_read
 564
 565     -- (d) Update block map for new destinations
 566     -- NB. do this before removing dead regs from the assignment, because
 567     -- these dead regs might in fact be live in the jump targets (they're
 568     -- only dead in the code that follows in the current basic block).
 569     (fixup_blocks, adjusted_instr)
 570         <- joinToTargets block_live [] instr (jumpDests instr [])
 571
 572     -- (e) Delete all register assignments for temps which are read
 573     --     (only) and die here.  Update the free register list.
 574     releaseRegs r_dying
 575
 576     -- (f) Mark regs which are clobbered as unallocatable
 577     clobberRegs real_written
 578
 579     -- (g) Allocate registers for temporaries *written* (only)
 580     (w_spills, w_allocd) <-
 581         allocateRegsAndSpill False{-writing-} virt_written [] [] virt_written
 582
 583     -- (h) Release registers for temps which are written here and not
 584     -- used again.
 585     releaseRegs w_dying
 586
 587     let
 588         -- (i) Patch the instruction
 589         patch_map = listToUFM   [ (t,RealReg r) |
 590                                   (t,r) <- zip virt_read r_allocd
 591                                           ++ zip virt_written w_allocd ]
 592
 593         patched_instr = patchRegs adjusted_instr patchLookup
 594         patchLookup x = case lookupUFM patch_map x of
 595                                 Nothing -> x
 596                                 Just y  -> y
 597     -- in
 598
 599     -- pprTrace "patched" (docToSDoc (pprInstr patched_instr)) $ do
 600
 601     -- (j) free up stack slots for dead spilled regs
 602     -- TODO (can't be bothered right now)
 603
 604     return (patched_instr : w_spills ++ reverse r_spills
 605                  ++ clobber_saves ++ new_instrs,
 606             fixup_blocks)
 607   }}
 608
 609 -- -----------------------------------------------------------------------------
 610 -- releaseRegs
 611
 612 releaseRegs regs = do
 613   assig <- getAssigR
 614   free <- getFreeRegsR
 615   loop assig free regs
 616  where
 617   loop assig free _ | free `seq` False = undefined
 618   loop assig free [] = do setAssigR assig; setFreeRegsR free; return ()
 619   loop assig free (RealReg r : rs) = loop assig (releaseReg r free) rs
 620   loop assig free (r:rs) =
 621      case lookupUFM assig r of
 622         Just (InBoth real _) -> loop (delFromUFM assig r) (releaseReg real free) rs
 623         Just (InReg real) -> loop (delFromUFM assig r) (releaseReg real free) rs
 624         _other            -> loop (delFromUFM assig r) free rs
 625
 626 -- -----------------------------------------------------------------------------
 627 -- Clobber real registers
 628
 629 {-
 630 For each temp in a register that is going to be clobbered:
 631   - if the temp dies after this instruction, do nothing
 632   - otherwise, put it somewhere safe (another reg if possible,
 633     otherwise spill and record InBoth in the assignment).
 634
 635 for allocateRegs on the temps *read*,
 636   - clobbered regs are allocatable.
 637
 638 for allocateRegs on the temps *written*,
 639   - clobbered regs are not allocatable.
 640 -}
 641
 642 saveClobberedTemps
 643    :: [RegNo]              -- real registers clobbered by this instruction
 644    -> [Reg]                -- registers which are no longer live after this insn
 645    -> RegM [Instr]         -- return: instructions to spill any temps that will
 646                            -- be clobbered.
 647
 648 saveClobberedTemps [] _ = return [] -- common case
 649 saveClobberedTemps clobbered dying =  do
 650   assig <- getAssigR
 651   let
 652         to_spill  = [ (temp,reg) | (temp, InReg reg) <- ufmToList assig,
 653                                    reg `elem` clobbered,
 654                                    temp `notElem` map getUnique dying  ]
 655   -- in
 656   (instrs,assig') <- clobber assig [] to_spill
 657   setAssigR assig'
 658   return instrs
 659  where
 660   clobber assig instrs [] = return (instrs,assig)
 661   clobber assig instrs ((temp,reg):rest)
 662     = do
 663         --ToDo: copy it to another register if possible
 664       (spill,slot) <- spillR (RealReg reg) temp
 665       clobber (addToUFM assig temp (InBoth reg slot)) (spill:instrs) rest
 666
 667 clobberRegs :: [RegNo] -> RegM ()
 668 clobberRegs [] = return () -- common case
 669 clobberRegs clobbered = do
 670   freeregs <- getFreeRegsR
 671   setFreeRegsR $! foldr allocateReg freeregs clobbered
 672   assig <- getAssigR
 673   setAssigR $! clobber assig (ufmToList assig)
 674  where
 675     -- if the temp was InReg and clobbered, then we will have
 676     -- saved it in saveClobberedTemps above.  So the only case
 677     -- we have to worry about here is InBoth.  Note that this
 678     -- also catches temps which were loaded up during allocation
 679     -- of read registers, not just those saved in saveClobberedTemps.
 680   clobber assig [] = assig
 681   clobber assig ((temp, InBoth reg slot) : rest)
 682         | reg `elem` clobbered
 683         = clobber (addToUFM assig temp (InMem slot)) rest
 684   clobber assig (entry:rest)
 685         = clobber assig rest
 686
 687 -- -----------------------------------------------------------------------------
 688 -- allocateRegsAndSpill
 689
 690 -- This function does several things:
 691 --   For each temporary referred to by this instruction,
 692 --   we allocate a real register (spilling another temporary if necessary).
 693 --   We load the temporary up from memory if necessary.
 694 --   We also update the register assignment in the process, and
 695 --   the list of free registers and free stack slots.
 696
 697 allocateRegsAndSpill
 698         :: Bool                 -- True <=> reading (load up spilled regs)
 699         -> [Reg]                -- don't push these out
 700         -> [Instr]              -- spill insns
 701         -> [RegNo]              -- real registers allocated (accum.)
 702         -> [Reg]                -- temps to allocate
 703         -> RegM ([Instr], [RegNo])
 704
 705 allocateRegsAndSpill reading keep spills alloc []
 706   = return (spills,reverse alloc)
 707
 708 allocateRegsAndSpill reading keep spills alloc (r:rs) = do
 709   assig <- getAssigR
 710   case lookupUFM assig r of
 711   -- case (1a): already in a register
 712      Just (InReg my_reg) ->
 713         allocateRegsAndSpill reading keep spills (my_reg:alloc) rs
 714
 715   -- case (1b): already in a register (and memory)
 716   -- NB1. if we're writing this register, update its assignemnt to be
 717   -- InReg, because the memory value is no longer valid.
 718   -- NB2. This is why we must process written registers here, even if they
 719   -- are also read by the same instruction.
 720      Just (InBoth my_reg mem) -> do
 721         when (not reading) (setAssigR (addToUFM assig r (InReg my_reg)))
 722         allocateRegsAndSpill reading keep spills (my_reg:alloc) rs
 723
 724   -- Not already in a register, so we need to find a free one...
 725      loc -> do
 726         freeregs <- getFreeRegsR
 727
 728         case getFreeRegs (regClass r) freeregs of
 729
 730         -- case (2): we have a free register
 731           my_reg:_ -> do
 732             spills'   <- do_load reading loc my_reg spills
 733             let new_loc
 734                  | Just (InMem slot) <- loc, reading = InBoth my_reg slot
 735                  | otherwise                         = InReg my_reg
 736             setAssigR (addToUFM assig r $! new_loc)
 737             setFreeRegsR (allocateReg my_reg freeregs)
 738             allocateRegsAndSpill reading keep spills' (my_reg:alloc) rs
 739
 740         -- case (3): we need to push something out to free up a register
 741           [] -> do
 742             let
 743               keep' = map getUnique keep
 744               candidates1 = [ (temp,reg,mem)
 745                             | (temp, InBoth reg mem) <- ufmToList assig,
 746                               temp `notElem` keep', regClass (RealReg reg) == regClass r ]
 747               candidates2 = [ (temp,reg)
 748                             | (temp, InReg reg) <- ufmToList assig,
 749                               temp `notElem` keep', regClass (RealReg reg) == regClass r  ]
 750             -- in
 751             ASSERT2(not (null candidates1 && null candidates2),
 752                     text (show freeregs) <+> ppr r <+> ppr assig) do
 753
 754             case candidates1 of
 755
 756              -- we have a temporary that is in both register and mem,
 757              -- just free up its register for use.
 758              --
 759              (temp,my_reg,slot):_ -> do
 760                 spills' <- do_load reading loc my_reg spills
 761                 let
 762                   assig1  = addToUFM assig temp (InMem slot)
 763                   assig2  = addToUFM assig1 r (InReg my_reg)
 764                 -- in
 765                 setAssigR assig2
 766                 allocateRegsAndSpill reading keep spills' (my_reg:alloc) rs
 767
 768              -- otherwise, we need to spill a temporary that currently
 769              -- resides in a register.
 770              [] -> do
 771                 let
 772                   (temp_to_push_out, my_reg) = myHead "regalloc" candidates2
 773                   -- TODO: plenty of room for optimisation in choosing which temp
 774                   -- to spill.  We just pick the first one that isn't used in
 775                   -- the current instruction for now.
 776                 -- in
 777                 (spill_insn,slot) <- spillR (RealReg my_reg) temp_to_push_out
 778                 let
 779                   assig1  = addToUFM assig temp_to_push_out (InMem slot)
 780                   assig2  = addToUFM assig1 r (InReg my_reg)
 781                 -- in
 782                 setAssigR assig2
 783                 spills' <- do_load reading loc my_reg spills
 784                 allocateRegsAndSpill reading keep (spill_insn:spills')
 785                         (my_reg:alloc) rs
 786   where
 787         -- load up a spilled temporary if we need to
 788         do_load True (Just (InMem slot)) reg spills = do
 789            insn <- loadR (RealReg reg) slot
 790            return (insn : spills)
 791         do_load _ _ _ spills =
 792            return spills
 793
 794 myHead s [] = panic s
 795 myHead s (x:xs) = x
 796
 797 -- -----------------------------------------------------------------------------
 798 -- Joining a jump instruction to its targets
 799
 800 -- The first time we encounter a jump to a particular basic block, we
 801 -- record the assignment of temporaries.  The next time we encounter a
 802 -- jump to the same block, we compare our current assignment to the
 803 -- stored one.  They might be different if spilling has occrred in one
 804 -- branch; so some fixup code will be required to match up the
 805 -- assignments.
 806
 807 joinToTargets
 808         :: BlockMap RegSet
 809         -> [NatBasicBlock]
 810         -> Instr
 811         -> [BlockId]
 812         -> RegM ([NatBasicBlock], Instr)
 813
 814 joinToTargets block_live new_blocks instr []
 815   = return (new_blocks, instr)
 816 joinToTargets block_live new_blocks instr (dest:dests) = do
 817   block_assig <- getBlockAssigR
 818   assig <- getAssigR
 819   let
 820         -- adjust the assignment to remove any registers which are not
 821         -- live on entry to the destination block.
 822         adjusted_assig = filterUFM_Directly still_live assig
 823         still_live uniq _ = uniq `elemUniqSet_Directly` live_set
 824
 825         -- and free up those registers which are now free.
 826         to_free =
 827           [ r | (reg, loc) <- ufmToList assig,
 828                 not (elemUniqSet_Directly reg live_set),
 829                 r <- regsOfLoc loc ]
 830
 831         regsOfLoc (InReg r)    = [r]
 832         regsOfLoc (InBoth r _) = [r]
 833         regsOfLoc (InMem _)    = []
 834   -- in
 835   case lookupUFM block_assig dest of
 836         -- Nothing <=> this is the first time we jumped to this
 837         -- block.
 838         Nothing -> do
 839           freeregs <- getFreeRegsR
 840           let freeregs' = foldr releaseReg freeregs to_free
 841           setBlockAssigR (addToUFM block_assig dest
 842                                 (freeregs',adjusted_assig))
 843           joinToTargets block_live new_blocks instr dests
 844
 845         Just (freeregs,dest_assig)
 846            | ufmToList dest_assig == ufmToList adjusted_assig
 847            -> -- ok, the assignments match
 848              joinToTargets block_live new_blocks instr dests
 849            | otherwise
 850            -> -- need fixup code
 851              do
 852                delta <- getDeltaR
 853                -- Construct a graph of register/spill movements and
 854                -- untangle it component by component.
 855                --
 856                -- We cut some corners by
 857                -- a) not handling cyclic components
 858                -- b) not handling memory-to-memory moves.
 859                --
 860                -- Cyclic components seem to occur only very rarely,
 861                -- and we don't need memory-to-memory moves because we
 862                -- make sure that every temporary always gets its own
 863                -- stack slot.
 864
 865                let graph = [ node | (vreg, src) <- ufmToList adjusted_assig,
 866                                     node <- mkNodes src vreg ]
 867
 868                    sccs = stronglyConnCompR graph
 869
 870                    mkNodes src vreg =
 871                         expandNode vreg src (lookupWithDefaultUFM_Directly
 872                                           dest_assig
 873                                           (panic "RegisterAlloc.joinToTargets")
 874                                           vreg)
 875
 876                 -- The InBoth handling is a little tricky here.  If
 877                 -- the destination is InBoth, then we must ensure that
 878                 -- the value ends up in both locations.  An InBoth
 879                 -- destination must conflict with an InReg or InMem
 880                 -- source, so we expand an InBoth destination as
 881                 -- necessary.  An InBoth source is slightly different:
 882                 -- we only care about the register that the source value
 883                 -- is in, so that we can move it to the destinations.
 884
 885                    expandNode vreg loc@(InReg src) (InBoth dst mem)
 886                         | src == dst = [(vreg, loc, [InMem mem])]
 887                         | otherwise  = [(vreg, loc, [InReg dst, InMem mem])]
 888                    expandNode vreg loc@(InMem src) (InBoth dst mem)
 889                         | src == mem = [(vreg, loc, [InReg dst])]
 890                         | otherwise  = [(vreg, loc, [InReg dst, InMem mem])]
 891                    expandNode vreg loc@(InBoth _ src) (InMem dst)
 892                         | src == dst = [] -- guaranteed to be true
 893                    expandNode vreg loc@(InBoth src _) (InReg dst)
 894                         | src == dst = []
 895                    expandNode vreg loc@(InBoth src _) dst
 896                         = expandNode vreg (InReg src) dst
 897                    expandNode vreg src dst
 898                         | src == dst = []
 899                         | otherwise  = [(vreg, src, [dst])]
 900
 901                 -- we have eliminated any possibility of single-node cylces
 902                 -- in expandNode above.
 903                    handleComponent (AcyclicSCC (vreg,src,dsts))
 904                        = map (makeMove vreg src) dsts
 905                    handleComponent (CyclicSCC things)
 906                        = panic $ "Register Allocator: handleComponent: cyclic"
 907                                  ++ " (workaround: use -fviaC)"
 908
 909                    makeMove vreg (InReg src) (InReg dst)
 910                        = mkRegRegMoveInstr (RealReg src) (RealReg dst)
 911                    makeMove vreg (InMem src) (InReg dst)
 912                        = mkLoadInstr (RealReg dst) delta src
 913                    makeMove vreg (InReg src) (InMem dst)
 914                        = mkSpillInstr (RealReg src) delta dst
 915                    makeMove vreg src dst
 916                        = panic $ "makeMove " ++ show vreg ++ " (" ++ show src ++ ") ("
 917                                  ++ show dst ++ ")"
 918                                  ++ " (workaround: use -fviaC)"
 919
 920                block_id <- getUniqueR
 921                let block = BasicBlock (BlockId block_id) $
 922                        concatMap handleComponent sccs ++ mkBranchInstr dest
 923                let instr' = patchJump instr dest (BlockId block_id)
 924                joinToTargets block_live (block : new_blocks) instr' dests
 925   where
 926         live_set = lookItUp "joinToTargets" block_live dest
 927
 928 -- -----------------------------------------------------------------------------
 929 -- The register allocator's monad.
 930
 931 -- Here we keep all the state that the register allocator keeps track
 932 -- of as it walks the instructions in a basic block.
 933
 934 data RA_State
 935   = RA_State {
 936         ra_blockassig :: BlockAssignment,
 937                 -- The current mapping from basic blocks to
 938                 -- the register assignments at the beginning of that block.
 939         ra_freeregs   :: {-#UNPACK#-}!FreeRegs, -- free machine registers
 940         ra_assig      :: RegMap Loc,    -- assignment of temps to locations
 941         ra_delta      :: Int,           -- current stack delta
 942         ra_stack      :: StackMap,      -- free stack slots for spilling
 943         ra_us         :: UniqSupply     -- unique supply for generating names
 944                                         -- for fixup blocks.
 945   }
 946
 947 newtype RegM a = RegM { unReg :: RA_State -> (# RA_State, a #) }
 948
 949 instance Monad RegM where
 950   m >>= k   =  RegM $ \s -> case unReg m s of { (# s, a #) -> unReg (k a) s }
 951   return a  =  RegM $ \s -> (# s, a #)
 952
 953 runR :: BlockAssignment -> FreeRegs -> RegMap Loc -> StackMap -> UniqSupply
 954   -> RegM a -> (BlockAssignment, StackMap, a)
 955 runR block_assig freeregs assig stack us thing =
 956   case unReg thing (RA_State{ ra_blockassig=block_assig, ra_freeregs=freeregs,
 957                         ra_assig=assig, ra_delta=0{-???-}, ra_stack=stack,
 958                         ra_us = us }) of
 959         (# RA_State{ ra_blockassig=block_assig, ra_stack=stack' }, returned_thing #)
 960                 -> (block_assig, stack', returned_thing)
 961
 962 spillR :: Reg -> Unique -> RegM (Instr, Int)
 963 spillR reg temp = RegM $ \ s@RA_State{ra_delta=delta, ra_stack=stack} ->
 964   let (stack',slot) = getStackSlotFor stack temp
 965       instr  = mkSpillInstr reg delta slot
 966   in
 967   (# s{ra_stack=stack'}, (instr,slot) #)
 968
 969 loadR :: Reg -> Int -> RegM Instr
 970 loadR reg slot = RegM $ \ s@RA_State{ra_delta=delta} ->
 971   (# s, mkLoadInstr reg delta slot #)
 972
 973 getFreeRegsR :: RegM FreeRegs
 974 getFreeRegsR = RegM $ \ s@RA_State{ra_freeregs = freeregs} ->
 975   (# s, freeregs #)
 976
 977 setFreeRegsR :: FreeRegs -> RegM ()
 978 setFreeRegsR regs = RegM $ \ s ->
 979   (# s{ra_freeregs = regs}, () #)
 980
 981 getAssigR :: RegM (RegMap Loc)
 982 getAssigR = RegM $ \ s@RA_State{ra_assig = assig} ->
 983   (# s, assig #)
 984
 985 setAssigR :: RegMap Loc -> RegM ()
 986 setAssigR assig = RegM $ \ s ->
 987   (# s{ra_assig=assig}, () #)
 988
 989 getBlockAssigR :: RegM BlockAssignment
 990 getBlockAssigR = RegM $ \ s@RA_State{ra_blockassig = assig} ->
 991   (# s, assig #)
 992
 993 setBlockAssigR :: BlockAssignment -> RegM ()
 994 setBlockAssigR assig = RegM $ \ s ->
 995   (# s{ra_blockassig = assig}, () #)
 996
 997 setDeltaR :: Int -> RegM ()
 998 setDeltaR n = RegM $ \ s ->
 999   (# s{ra_delta = n}, () #)
1000
1001 getDeltaR :: RegM Int
1002 getDeltaR = RegM $ \s -> (# s, ra_delta s #)
1003
1004 getUniqueR :: RegM Unique
1005 getUniqueR = RegM $ \s ->
1006   case splitUniqSupply (ra_us s) of
1007     (us1, us2) -> (# s{ra_us = us2}, uniqFromSupply us1 #)
1008
1009 -- -----------------------------------------------------------------------------
1010 -- Utils
1011
1012 #ifdef DEBUG
1013 my_fromJust s p Nothing  = pprPanic ("fromJust: " ++ s) p
1014 my_fromJust s p (Just x) = x
1015 #else
1016 my_fromJust _ _ = fromJust
1017 #endif
1018
1019 lookItUp :: Uniquable b => String -> UniqFM a -> b -> a
1020 lookItUp str fm x = my_fromJust str (ppr (getUnique x)) (lookupUFM fm x)