ghc/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 Outputable
 100
 101 #ifndef DEBUG
 102 import Maybe            ( fromJust )
 103 #endif
 104 import List             ( nub, partition )
 105 import Monad            ( when )
 106 import DATA_WORD
 107 import DATA_BITS
 108
 109 -- -----------------------------------------------------------------------------
 110 -- Some useful types
 111
 112 type RegSet = UniqSet Reg
 113
 114 type RegMap a = UniqFM a
 115 emptyRegMap = emptyUFM
 116
 117 type BlockMap a = UniqFM a
 118 emptyBlockMap = emptyUFM
 119
 120 -- A basic block where the isntructions are annotated with the registers
 121 -- which are no longer live in the *next* instruction in this sequence.
 122 -- (NB. if the instruction is a jump, these registers might still be live
 123 -- at the jump target(s) - you have to check the liveness at the destination
 124 -- block to find out).
 125 type AnnBasicBlock
 126         = GenBasicBlock (Instr,
 127                          [Reg],         -- registers read (only) which die
 128                          [Reg])         -- registers written which die
 129
 130 -- -----------------------------------------------------------------------------
 131 -- The free register set
 132
 133 -- This needs to be *efficient*
 134
 135 {- Here's an inefficient 'executable specification' of the FreeRegs data type:
 136 type FreeRegs = [RegNo]
 137
 138 noFreeRegs = 0
 139 releaseReg n f = if n `elem` f then f else (n : f)
 140 initFreeRegs = allocatableRegs
 141 getFreeRegs cls f = filter ( (==cls) . regClass . RealReg ) f
 142 allocateReg f r = filter (/= r) f
 143 -}
 144
 145 #if defined(powerpc_TARGET_ARCH)
 146
 147 -- The PowerPC has 32 integer and 32 floating point registers.
 148 -- This is 32bit PowerPC, so Word64 is inefficient - two Word32s are much
 149 -- better.
 150 -- Note that when getFreeRegs scans for free registers, it starts at register
 151 -- 31 and counts down. This is a hack for the PowerPC - the higher-numbered
 152 -- registers are callee-saves, while the lower regs are caller-saves, so it
 153 -- makes sense to start at the high end.
 154 -- Apart from that, the code does nothing PowerPC-specific, so feel free to
 155 -- add your favourite platform to the #if (if you have 64 registers but only
 156 -- 32-bit words).
 157
 158 data FreeRegs = FreeRegs !Word32 !Word32
 159
 160 noFreeRegs :: FreeRegs
 161 noFreeRegs = FreeRegs 0 0
 162
 163 releaseReg :: RegNo -> FreeRegs -> FreeRegs
 164 releaseReg r (FreeRegs g f)
 165     | r > 31    = FreeRegs g (f .|. (1 `shiftL` (fromIntegral r - 32)))
 166     | otherwise = FreeRegs (g .|. (1 `shiftL` fromIntegral r)) f
 167
 168 initFreeRegs :: FreeRegs
 169 initFreeRegs = foldr releaseReg noFreeRegs allocatableRegs
 170
 171 getFreeRegs :: RegClass -> FreeRegs -> [RegNo]  -- lazilly
 172 getFreeRegs cls (FreeRegs g f)
 173     | RcDouble <- cls = go f (0x80000000) 63
 174     | RcInteger <- cls = go g (0x80000000) 31
 175     where
 176         go x 0 i = []
 177         go x m i | x .&. m /= 0 = i : (go x (m `shiftR` 1) $! i-1)
 178                  | otherwise    = go x (m `shiftR` 1) $! i-1
 179
 180 allocateReg :: RegNo -> FreeRegs -> FreeRegs
 181 allocateReg r (FreeRegs g f)
 182     | r > 31    = FreeRegs g (f .&. complement (1 `shiftL` (fromIntegral r - 32)))
 183     | otherwise = FreeRegs (g .&. complement (1 `shiftL` fromIntegral r)) f
 184
 185 #else
 186
 187 -- If we have less than 32 registers, or if we have efficient 64-bit words,
 188 -- we will just use a single bitfield.
 189
 190 #if defined(alpha_TARGET_ARCH)
 191 type FreeRegs = Word64
 192 #else
 193 type FreeRegs = Word32
 194 #endif
 195
 196 noFreeRegs :: FreeRegs
 197 noFreeRegs = 0
 198
 199 releaseReg :: RegNo -> FreeRegs -> FreeRegs
 200 releaseReg n f = f .|. (1 `shiftL` n)
 201
 202 initFreeRegs :: FreeRegs
 203 initFreeRegs = foldr releaseReg noFreeRegs allocatableRegs
 204
 205 getFreeRegs :: RegClass -> FreeRegs -> [RegNo]  -- lazilly
 206 getFreeRegs cls f = go f 0
 207   where go 0 m = []
 208         go n m
 209           | n .&. 1 /= 0 && regClass (RealReg m) == cls
 210           = m : (go (n `shiftR` 1) $! (m+1))
 211           | otherwise
 212           = go (n `shiftR` 1) $! (m+1)
 213         -- ToDo: there's no point looking through all the integer registers
 214         -- in order to find a floating-point one.
 215
 216 allocateReg :: RegNo -> FreeRegs -> FreeRegs
 217 allocateReg r f = f .&. complement (1 `shiftL` fromIntegral r)
 218
 219 #endif
 220
 221 -- -----------------------------------------------------------------------------
 222 -- The free list of stack slots
 223
 224 -- This doesn't need to be so efficient.  It also doesn't really need to be
 225 -- maintained as a set, so we just use an ordinary list (lazy, because it
 226 -- contains all the possible stack slots and there are lots :-).
 227
 228 type StackSlot = Int
 229 type FreeStack = [StackSlot]
 230
 231 completelyFreeStack :: FreeStack
 232 completelyFreeStack = [0..maxSpillSlots]
 233
 234 getFreeStackSlot :: FreeStack -> (FreeStack,Int)
 235 getFreeStackSlot (slot:stack) = (stack,slot)
 236
 237 freeStackSlot :: FreeStack -> Int -> FreeStack
 238 freeStackSlot stack slot = slot:stack
 239
 240
 241 -- -----------------------------------------------------------------------------
 242 -- Top level of the register allocator
 243
 244 regAlloc :: NatCmmTop -> NatCmmTop
 245 regAlloc (CmmData sec d) = CmmData sec d
 246 regAlloc (CmmProc info lbl params [])
 247   = CmmProc info lbl params []  -- no blocks to run the regalloc on
 248 regAlloc (CmmProc info lbl params blocks@(first:rest))
 249   = -- pprTrace "Liveness" (ppr block_live) $
 250     CmmProc info lbl params (first':rest')
 251   where
 252     first_id               = blockId first
 253     sccs                   = sccBlocks blocks
 254     (ann_sccs, block_live) = computeLiveness sccs
 255     final_blocks           = linearRegAlloc block_live ann_sccs
 256     ((first':_),rest')     = partition ((== first_id) . blockId) final_blocks
 257
 258
 259 sccBlocks :: [NatBasicBlock] -> [SCC NatBasicBlock]
 260 sccBlocks blocks = stronglyConnComp graph
 261   where
 262         getOutEdges :: [Instr] -> [BlockId]
 263         getOutEdges instrs = foldr jumpDests [] instrs
 264
 265         graph = [ (block, getUnique id, map getUnique (getOutEdges instrs))
 266                 | block@(BasicBlock id instrs) <- blocks ]
 267
 268
 269 -- -----------------------------------------------------------------------------
 270 -- Computing liveness
 271
 272 computeLiveness
 273    :: [SCC NatBasicBlock]
 274    -> ([SCC AnnBasicBlock],     -- instructions annotated with list of registers
 275                                 -- which are "dead after this instruction".
 276        BlockMap RegSet)         -- blocks annontated with set of live registers
 277                                 -- on entry to the block.
 278
 279   -- NOTE: on entry, the SCCs are in "reverse" order: later blocks may transfer
 280   -- control to earlier ones only.  The SCCs returned are in the *opposite*
 281   -- order, which is exactly what we want for the next pass.
 282
 283 computeLiveness sccs
 284   = livenessSCCs emptyBlockMap [] sccs
 285   where
 286   livenessSCCs
 287          :: BlockMap RegSet
 288          -> [SCC AnnBasicBlock]         -- accum
 289          -> [SCC NatBasicBlock]
 290          -> ([SCC AnnBasicBlock], BlockMap RegSet)
 291
 292   livenessSCCs blockmap done [] = (done, blockmap)
 293   livenessSCCs blockmap done
 294         (AcyclicSCC (BasicBlock block_id instrs) : sccs) =
 295           {- pprTrace "live instrs" (ppr (getUnique block_id) $$
 296                                   vcat (map (\(instr,regs) -> docToSDoc (pprInstr instr) $$ ppr regs) instrs')) $
 297           -}
 298           livenessSCCs blockmap'
 299                 (AcyclicSCC (BasicBlock block_id instrs'):done) sccs
 300         where (live,instrs') = liveness emptyUniqSet blockmap []
 301                                         (reverse instrs)
 302               blockmap' = addToUFM blockmap block_id live
 303         -- TODO: cope with recursive blocks
 304
 305   liveness :: RegSet                    -- live regs
 306            -> BlockMap RegSet           -- live regs on entry to other BBs
 307            -> [(Instr,[Reg],[Reg])]     -- instructions (accum)
 308            -> [Instr]                   -- instructions
 309            -> (RegSet, [(Instr,[Reg],[Reg])])
 310
 311   liveness liveregs blockmap done []  = (liveregs, done)
 312   liveness liveregs blockmap done (instr:instrs)
 313         = liveness liveregs2 blockmap ((instr,r_dying,w_dying):done) instrs
 314         where
 315               RU read written = regUsage instr
 316
 317               -- registers that were written here are dead going backwards.
 318               -- registers that were read here are live going backwards.
 319               liveregs1 = (liveregs `delListFromUniqSet` written)
 320                                     `addListToUniqSet` read
 321
 322               -- union in the live regs from all the jump destinations of this
 323               -- instruction.
 324               targets = jumpDests instr [] -- where we go from here
 325               liveregs2 = unionManyUniqSets
 326                             (liveregs1 : map (lookItUp "liveness" blockmap)
 327                                                 targets)
 328
 329               -- registers that are not live beyond this point, are recorded
 330               --  as dying here.
 331               r_dying  = [ reg | reg <- read, reg `notElem` written,
 332                                  not (elementOfUniqSet reg liveregs) ]
 333
 334               w_dying = [ reg | reg <- written,
 335                                 not (elementOfUniqSet reg liveregs) ]
 336
 337 -- -----------------------------------------------------------------------------
 338 -- Linear sweep to allocate registers
 339
 340 data Loc = InReg   {-# UNPACK #-} !RegNo
 341          | InMem   {-# UNPACK #-} !Int          -- stack slot
 342          | InBoth  {-# UNPACK #-} !RegNo
 343                    {-# UNPACK #-} !Int          -- stack slot
 344   deriving (Eq, Show)
 345
 346 {-
 347 A temporary can be marked as living in both a register and memory
 348 (InBoth), for example if it was recently loaded from a spill location.
 349 This makes it cheap to spill (no save instruction required), but we
 350 have to be careful to turn this into InReg if the value in the
 351 register is changed.
 352
 353 This is also useful when a temporary is about to be clobbered.  We
 354 save it in a spill location, but mark it as InBoth because the current
 355 instruction might still want to read it.
 356 -}
 357
 358 #ifdef DEBUG
 359 instance Outputable Loc where
 360   ppr l = text (show l)
 361 #endif
 362
 363 linearRegAlloc
 364    :: BlockMap RegSet           -- live regs on entry to each basic block
 365    -> [SCC AnnBasicBlock]       -- instructions annotated with "deaths"
 366    -> [NatBasicBlock]
 367 linearRegAlloc block_live sccs = linearRA_SCCs emptyBlockMap sccs
 368   where
 369   linearRA_SCCs
 370         :: BlockAssignment
 371         -> [SCC AnnBasicBlock]
 372         -> [NatBasicBlock]
 373   linearRA_SCCs block_assig [] = []
 374   linearRA_SCCs block_assig
 375         (AcyclicSCC (BasicBlock id instrs) : sccs)
 376         = BasicBlock id instrs' : linearRA_SCCs block_assig' sccs
 377     where
 378         (block_assig',(instrs',fixups)) =
 379            case lookupUFM block_assig id of
 380                 -- no prior info about this block: assume everything is
 381                 -- free and the assignment is empty.
 382                 Nothing ->
 383                    runR block_assig initFreeRegs
 384                                 emptyRegMap completelyFreeStack $
 385                         linearRA [] [] instrs
 386                 Just (freeregs,stack,assig) ->
 387                    runR block_assig freeregs assig stack $
 388                         linearRA [] [] instrs
 389
 390   linearRA :: [Instr] -> [NatBasicBlock] -> [(Instr,[Reg],[Reg])]
 391         -> RegM ([Instr], [NatBasicBlock])
 392   linearRA instr_acc fixups [] =
 393     return (reverse instr_acc, fixups)
 394   linearRA instr_acc fixups (instr:instrs) = do
 395     (instr_acc', new_fixups) <- raInsn block_live instr_acc instr
 396     linearRA instr_acc' (new_fixups++fixups) instrs
 397
 398 -- -----------------------------------------------------------------------------
 399 -- Register allocation for a single instruction
 400
 401 type BlockAssignment = BlockMap (FreeRegs, FreeStack, RegMap Loc)
 402
 403 raInsn  :: BlockMap RegSet              -- Live temporaries at each basic block
 404         -> [Instr]                      -- new instructions (accum.)
 405         -> (Instr,[Reg],[Reg])          -- the instruction (with "deaths")
 406         -> RegM (
 407              [Instr],                   -- new instructions
 408              [NatBasicBlock]            -- extra fixup blocks
 409            )
 410
 411 raInsn block_live new_instrs (instr@(DELTA n), _, _) = do
 412     setDeltaR n
 413     return (new_instrs, [])
 414
 415 raInsn block_live new_instrs (instr, r_dying, w_dying) = do
 416     assig    <- getAssigR
 417
 418     -- If we have a reg->reg move between virtual registers, where the
 419     -- src register is not live after this instruction, and the dst
 420     -- register does not already have an assignment, then we can
 421     -- eliminate the instruction.
 422     case isRegRegMove instr of
 423         Just (src,dst)
 424                 | src `elem` r_dying,
 425                   isVirtualReg dst,
 426                   Just loc <- lookupUFM assig src,
 427                   not (dst `elemUFM` assig) -> do
 428                         setAssigR (addToUFM (delFromUFM assig src) dst loc)
 429                         return (new_instrs, [])
 430
 431         other -> genRaInsn block_live new_instrs instr r_dying w_dying
 432
 433
 434 genRaInsn block_live new_instrs instr r_dying w_dying =
 435     case regUsage instr              of { RU read written ->
 436     case partition isRealReg written of { (real_written1,virt_written) ->
 437     do
 438     let
 439         real_written = [ r | RealReg r <- real_written1 ]
 440
 441         -- we don't need to do anything with real registers that are
 442         -- only read by this instr.  (the list is typically ~2 elements,
 443         -- so using nub isn't a problem).
 444         virt_read = nub (filter isVirtualReg read)
 445     -- in
 446
 447     -- (a) save any temporaries which will be clobbered by this instruction
 448     clobber_saves <- saveClobberedTemps real_written r_dying
 449
 450     {-
 451     freeregs <- getFreeRegsR
 452     assig <- getAssigR
 453     pprTrace "raInsn" (docToSDoc (pprInstr instr) $$ ppr r_dying <+> ppr w_dying $$ ppr virt_read <+> ppr virt_written $$ text (show freeregs) $$ ppr assig) $ do
 454     -}
 455
 456     -- (b), (c) allocate real regs for all regs read by this instruction.
 457     (r_spills, r_allocd) <-
 458         allocateRegsAndSpill True{-reading-} virt_read [] [] virt_read
 459
 460     -- (d) Update block map for new destinations
 461     -- NB. do this before removing dead regs from the assignment, because
 462     -- these dead regs might in fact be live in the jump targets (they're
 463     -- only dead in the code that follows in the current basic block).
 464     (fixup_blocks, adjusted_instr)
 465         <- joinToTargets block_live [] instr (jumpDests instr [])
 466
 467     -- (e) Delete all register assignments for temps which are read
 468     --     (only) and die here.  Update the free register list.
 469     releaseRegs r_dying
 470
 471     -- (f) Mark regs which are clobbered as unallocatable
 472     clobberRegs real_written
 473
 474     -- (g) Allocate registers for temporaries *written* (only)
 475     (w_spills, w_allocd) <-
 476         allocateRegsAndSpill False{-writing-} virt_written [] [] virt_written
 477
 478     -- (h) Release registers for temps which are written here and not
 479     -- used again.
 480     releaseRegs w_dying
 481
 482     let
 483         -- (i) Patch the instruction
 484         patch_map = listToUFM   [ (t,RealReg r) |
 485                                   (t,r) <- zip virt_read r_allocd
 486                                           ++ zip virt_written w_allocd ]
 487
 488         patched_instr = patchRegs adjusted_instr patchLookup
 489         patchLookup x = case lookupUFM patch_map x of
 490                                 Nothing -> x
 491                                 Just y  -> y
 492     -- in
 493
 494     -- pprTrace "patched" (docToSDoc (pprInstr patched_instr)) $ do
 495
 496     -- (j) free up stack slots for dead spilled regs
 497     -- TODO (can't be bothered right now)
 498
 499     return (patched_instr : w_spills ++ reverse r_spills
 500                  ++ clobber_saves ++ new_instrs,
 501             fixup_blocks)
 502   }}
 503
 504 -- -----------------------------------------------------------------------------
 505 -- releaseRegs
 506
 507 releaseRegs regs = do
 508   assig <- getAssigR
 509   free <- getFreeRegsR
 510   loop assig free regs
 511  where
 512   loop assig free _ | free `seq` False = undefined
 513   loop assig free [] = do setAssigR assig; setFreeRegsR free; return ()
 514   loop assig free (RealReg r : rs) = loop assig (releaseReg r free) rs
 515   loop assig free (r:rs) =
 516      case lookupUFM assig r of
 517         Just (InBoth real _) -> loop (delFromUFM assig r) (releaseReg real free) rs
 518         Just (InReg real) -> loop (delFromUFM assig r) (releaseReg real free) rs
 519         _other            -> loop (delFromUFM assig r) free rs
 520
 521 -- -----------------------------------------------------------------------------
 522 -- Clobber real registers
 523
 524 {-
 525 For each temp in a register that is going to be clobbered:
 526   - if the temp dies after this instruction, do nothing
 527   - otherwise, put it somewhere safe (another reg if possible,
 528     otherwise spill and record InBoth in the assignment).
 529
 530 for allocateRegs on the temps *read*,
 531   - clobbered regs are allocatable.
 532
 533 for allocateRegs on the temps *written*,
 534   - clobbered regs are not allocatable.
 535 -}
 536
 537 saveClobberedTemps
 538    :: [RegNo]              -- real registers clobbered by this instruction
 539    -> [Reg]                -- registers which are no longer live after this insn
 540    -> RegM [Instr]         -- return: instructions to spill any temps that will
 541                            -- be clobbered.
 542
 543 saveClobberedTemps [] _ = return [] -- common case
 544 saveClobberedTemps clobbered dying =  do
 545   assig <- getAssigR
 546   let
 547         to_spill  = [ (temp,reg) | (temp, InReg reg) <- ufmToList assig,
 548                                    reg `elem` clobbered,
 549                                    temp `notElem` map getUnique dying  ]
 550   -- in
 551   (instrs,assig') <- clobber assig [] to_spill
 552   setAssigR assig'
 553   return instrs
 554  where
 555   clobber assig instrs [] = return (instrs,assig)
 556   clobber assig instrs ((temp,reg):rest)
 557     = do
 558         --ToDo: copy it to another register if possible
 559       (spill,slot) <- spillR (RealReg reg)
 560       clobber (addToUFM assig temp (InBoth reg slot)) (spill:instrs) rest
 561
 562 clobberRegs :: [RegNo] -> RegM ()
 563 clobberRegs [] = return () -- common case
 564 clobberRegs clobbered = do
 565   freeregs <- getFreeRegsR
 566   setFreeRegsR $! foldr allocateReg freeregs clobbered
 567   assig <- getAssigR
 568   setAssigR $! clobber assig (ufmToList assig)
 569  where
 570     -- if the temp was InReg and clobbered, then we will have
 571     -- saved it in saveClobberedTemps above.  So the only case
 572     -- we have to worry about here is InBoth.  Note that this
 573     -- also catches temps which were loaded up during allocation
 574     -- of read registers, not just those saved in saveClobberedTemps.
 575   clobber assig [] = assig
 576   clobber assig ((temp, InBoth reg slot) : rest)
 577         | reg `elem` clobbered
 578         = clobber (addToUFM assig temp (InMem slot)) rest
 579   clobber assig (entry:rest)
 580         = clobber assig rest
 581
 582 -- -----------------------------------------------------------------------------
 583 -- allocateRegsAndSpill
 584
 585 -- This function does several things:
 586 --   For each temporary referred to by this instruction,
 587 --   we allocate a real register (spilling another temporary if necessary).
 588 --   We load the temporary up from memory if necessary.
 589 --   We also update the register assignment in the process, and
 590 --   the list of free registers and free stack slots.
 591
 592 allocateRegsAndSpill
 593         :: Bool                 -- True <=> reading (load up spilled regs)
 594         -> [Reg]                -- don't push these out
 595         -> [Instr]              -- spill insns
 596         -> [RegNo]              -- real registers allocated (accum.)
 597         -> [Reg]                -- temps to allocate
 598         -> RegM ([Instr], [RegNo])
 599
 600 allocateRegsAndSpill reading keep spills alloc []
 601   = return (spills,reverse alloc)
 602
 603 allocateRegsAndSpill reading keep spills alloc (r:rs) = do
 604   assig <- getAssigR
 605   case lookupUFM assig r of
 606   -- case (1a): already in a register
 607      Just (InReg my_reg) ->
 608         allocateRegsAndSpill reading keep spills (my_reg:alloc) rs
 609
 610   -- case (1b): already in a register (and memory)
 611   -- NB1. if we're writing this register, update its assignemnt to be
 612   -- InReg, because the memory value is no longer valid.
 613   -- NB2. This is why we must process written registers here, even if they
 614   -- are also read by the same instruction.
 615      Just (InBoth my_reg mem) -> do
 616         when (not reading) (setAssigR (addToUFM assig r (InReg my_reg)))
 617         allocateRegsAndSpill reading keep spills (my_reg:alloc) rs
 618
 619   -- Not already in a register, so we need to find a free one...
 620      loc -> do
 621         freeregs <- getFreeRegsR
 622
 623         case getFreeRegs (regClass r) freeregs of
 624
 625         -- case (2): we have a free register
 626           my_reg:_ -> do
 627             spills'   <- do_load reading loc my_reg spills
 628             let new_loc
 629                  | Just (InMem slot) <- loc, reading = InBoth my_reg slot
 630                  | otherwise                         = InReg my_reg
 631             setAssigR (addToUFM assig r $! new_loc)
 632             setFreeRegsR (allocateReg my_reg freeregs)
 633             allocateRegsAndSpill reading keep spills' (my_reg:alloc) rs
 634
 635         -- case (3): we need to push something out to free up a register
 636           [] -> do
 637             let
 638               keep' = map getUnique keep
 639               candidates1 = [ (temp,reg,mem)
 640                             | (temp, InBoth reg mem) <- ufmToList assig,
 641                               temp `notElem` keep', regClass (RealReg reg) == regClass r ]
 642               candidates2 = [ (temp,reg)
 643                             | (temp, InReg reg) <- ufmToList assig,
 644                               temp `notElem` keep', regClass (RealReg reg) == regClass r  ]
 645             -- in
 646             ASSERT2(not (null candidates1 && null candidates2),
 647                     text (show freeregs) <+> ppr r <+> ppr assig) do
 648
 649             case candidates1 of
 650
 651              -- we have a temporary that is in both register and mem,
 652              -- just free up its register for use.
 653              --
 654              (temp,my_reg,slot):_ -> do
 655                 spills' <- do_load reading loc my_reg spills
 656                 let
 657                   assig1  = addToUFM assig temp (InMem slot)
 658                   assig2  = addToUFM assig1 r (InReg my_reg)
 659                 -- in
 660                 setAssigR assig2
 661                 allocateRegsAndSpill reading keep spills' (my_reg:alloc) rs
 662
 663              -- otherwise, we need to spill a temporary that currently
 664              -- resides in a register.
 665              [] -> do
 666                 let
 667                   (temp_to_push_out, my_reg) = myHead "regalloc" candidates2
 668                   -- TODO: plenty of room for optimisation in choosing which temp
 669                   -- to spill.  We just pick the first one that isn't used in
 670                   -- the current instruction for now.
 671                 -- in
 672                 (spill_insn,slot) <- spillR (RealReg my_reg)
 673                 let
 674                   assig1  = addToUFM assig temp_to_push_out (InMem slot)
 675                   assig2  = addToUFM assig1 r (InReg my_reg)
 676                 -- in
 677                 setAssigR assig2
 678                 spills' <- do_load reading loc my_reg spills
 679                 allocateRegsAndSpill reading keep (spill_insn:spills')
 680                         (my_reg:alloc) rs
 681   where
 682         -- load up a spilled temporary if we need to
 683         do_load True (Just (InMem slot)) reg spills = do
 684            insn <- loadR (RealReg reg) slot
 685            return (insn : spills)
 686         do_load _ _ _ spills =
 687            return spills
 688
 689 myHead s [] = panic s
 690 myHead s (x:xs) = x
 691
 692 -- -----------------------------------------------------------------------------
 693 -- Joining a jump instruction to its targets
 694
 695 -- The first time we encounter a jump to a particular basic block, we
 696 -- record the assignment of temporaries.  The next time we encounter a
 697 -- jump to the same block, we compare our current assignment to the
 698 -- stored one.  They might be different if spilling has occrred in one
 699 -- branch; so some fixup code will be required to match up the
 700 -- assignments.
 701
 702 joinToTargets
 703         :: BlockMap RegSet
 704         -> [NatBasicBlock]
 705         -> Instr
 706         -> [BlockId]
 707         -> RegM ([NatBasicBlock], Instr)
 708
 709 joinToTargets block_live new_blocks instr []
 710   = return (new_blocks, instr)
 711 joinToTargets block_live new_blocks instr (dest:dests) = do
 712   block_assig <- getBlockAssigR
 713   assig <- getAssigR
 714   let
 715         -- adjust the assignment to remove any registers which are not
 716         -- live on entry to the destination block.
 717         adjusted_assig = filterUFM_Directly still_live assig
 718         still_live uniq _ = uniq `elemUniqSet_Directly` live_set
 719
 720         -- and free up those registers which are now free.
 721         to_free =
 722           [ r | (reg, loc) <- ufmToList assig,
 723                 not (elemUniqSet_Directly reg live_set),
 724                 r <- regsOfLoc loc ]
 725
 726         regsOfLoc (InReg r)    = [r]
 727         regsOfLoc (InBoth r _) = [r]
 728         regsOfLoc (InMem _)    = []
 729   -- in
 730   case lookupUFM block_assig dest of
 731         -- Nothing <=> this is the first time we jumped to this
 732         -- block.
 733         Nothing -> do
 734           freeregs <- getFreeRegsR
 735           let freeregs' = foldr releaseReg freeregs to_free
 736           stack <- getStackR
 737           setBlockAssigR (addToUFM block_assig dest
 738                                 (freeregs',stack,adjusted_assig))
 739           joinToTargets block_live new_blocks instr dests
 740
 741         Just (freeregs,stack,dest_assig)
 742            | ufmToList dest_assig == ufmToList adjusted_assig
 743            -> -- ok, the assignments match
 744              joinToTargets block_live new_blocks instr dests
 745            | otherwise
 746            -> -- need fixup code
 747              panic "joinToTargets: ToDo: need fixup code"
 748   where
 749         live_set = lookItUp "joinToTargets" block_live dest
 750
 751 -- -----------------------------------------------------------------------------
 752 -- The register allocator's monad.
 753
 754 -- Here we keep all the state that the register allocator keeps track
 755 -- of as it walks the instructions in a basic block.
 756
 757 data RA_State
 758   = RA_State {
 759         ra_blockassig :: BlockAssignment,
 760                 -- The current mapping from basic blocks to
 761                 -- the register assignments at the beginning of that block.
 762         ra_freeregs   :: {-#UNPACK#-}!FreeRegs, -- free machine registers
 763         ra_assig      :: RegMap Loc,    -- assignment of temps to locations
 764         ra_delta      :: Int,           -- current stack delta
 765         ra_stack      :: FreeStack      -- free stack slots for spilling
 766   }
 767
 768 newtype RegM a = RegM { unReg :: RA_State -> (# RA_State, a #) }
 769
 770 instance Monad RegM where
 771   m >>= k   =  RegM $ \s -> case unReg m s of { (# s, a #) -> unReg (k a) s }
 772   return a  =  RegM $ \s -> (# s, a #)
 773
 774 runR :: BlockAssignment -> FreeRegs -> RegMap Loc -> FreeStack -> RegM a ->
 775   (BlockAssignment, a)
 776 runR block_assig freeregs assig stack thing =
 777   case unReg thing (RA_State{ ra_blockassig=block_assig, ra_freeregs=freeregs,
 778                         ra_assig=assig, ra_delta=0{-???-}, ra_stack=stack }) of
 779         (# RA_State{ ra_blockassig=block_assig }, returned_thing #)
 780                 -> (block_assig, returned_thing)
 781
 782 spillR :: Reg -> RegM (Instr, Int)
 783 spillR reg = RegM $ \ s@RA_State{ra_delta=delta, ra_stack=stack} ->
 784   let (stack',slot) = getFreeStackSlot stack
 785       instr  = mkSpillInstr reg delta slot
 786   in
 787   (# s{ra_stack=stack'}, (instr,slot) #)
 788
 789 loadR :: Reg -> Int -> RegM Instr
 790 loadR reg slot = RegM $ \ s@RA_State{ra_delta=delta, ra_stack=stack} ->
 791   (# s, mkLoadInstr reg delta slot #)
 792
 793 freeSlotR :: Int -> RegM ()
 794 freeSlotR slot = RegM $ \ s@RA_State{ra_stack=stack} ->
 795   (# s{ra_stack=freeStackSlot stack slot}, () #)
 796
 797 getFreeRegsR :: RegM FreeRegs
 798 getFreeRegsR = RegM $ \ s@RA_State{ra_freeregs = freeregs} ->
 799   (# s, freeregs #)
 800
 801 setFreeRegsR :: FreeRegs -> RegM ()
 802 setFreeRegsR regs = RegM $ \ s ->
 803   (# s{ra_freeregs = regs}, () #)
 804
 805 getAssigR :: RegM (RegMap Loc)
 806 getAssigR = RegM $ \ s@RA_State{ra_assig = assig} ->
 807   (# s, assig #)
 808
 809 setAssigR :: RegMap Loc -> RegM ()
 810 setAssigR assig = RegM $ \ s ->
 811   (# s{ra_assig=assig}, () #)
 812
 813 getStackR :: RegM FreeStack
 814 getStackR = RegM $ \ s@RA_State{ra_stack = stack} ->
 815   (# s, stack #)
 816
 817 setStackR :: FreeStack -> RegM ()
 818 setStackR stack = RegM $ \ s ->
 819   (# s{ra_stack=stack}, () #)
 820
 821 getBlockAssigR :: RegM BlockAssignment
 822 getBlockAssigR = RegM $ \ s@RA_State{ra_blockassig = assig} ->
 823   (# s, assig #)
 824
 825 setBlockAssigR :: BlockAssignment -> RegM ()
 826 setBlockAssigR assig = RegM $ \ s ->
 827   (# s{ra_blockassig = assig}, () #)
 828
 829 setDeltaR :: Int -> RegM ()
 830 setDeltaR n = RegM $ \ s ->
 831   (# s{ra_delta = n}, () #)
 832
 833 -- -----------------------------------------------------------------------------
 834 -- Utils
 835
 836 #ifdef DEBUG
 837 my_fromJust s p Nothing  = pprPanic ("fromJust: " ++ s) p
 838 my_fromJust s p (Just x) = x
 839 #else
 840 my_fromJust _ _ = fromJust
 841 #endif
 842
 843 lookItUp :: Uniquable b => String -> UniqFM a -> b -> a
 844 lookItUp str fm x = my_fromJust str (ppr (getUnique x)) (lookupUFM fm x)