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