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