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