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