code gen: Improve efficiency of findPrefRealRegwip/linear-reg-alloc

Old strategy: For each variable linearly scan through all the blocks and
check to see if the variable is any of the block register mappings. This
is very slow when you have a lot of blocks.

New strategy: Maintain a map from virtual registers to the first real
register the virtual register was assigned to. Consult this map in
findPrefRealReg.

The map is updated when the register mapping is updated and is hidden
behind the BlockAssigment abstraction.

On the mmark package this reduces compilation time from about 44s to
32s.

Ticket: #19471

2 files changed, 44 insertions, 4 deletions

diff --git a/compiler/GHC/CmmToAsm/Reg/Linear/Base.hs b/compiler/GHC/CmmToAsm/Reg/Linear/Base.hs
index 33a15fd7b8..180926d0bf 100644
--- a/compiler/GHC/CmmToAsm/Reg/Linear/Base.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Linear/Base.hs
@@ -1,8 +1,13 @@
+{-# LANGUAGE RecordWildCards #-}
 
 -- | Put common type definitions here to break recursive module dependencies.
 
 module GHC.CmmToAsm.Reg.Linear.Base (
         BlockAssignment,
+        lookupBlockAssignment,
+        lookupFirstUsed,
+        emptyBlockAssignment,
+        updateBlockAssignment,
 
         Loc(..),
         regsOfLoc,
@@ -29,6 +34,8 @@ import GHC.Types.Unique
 import GHC.Types.Unique.FM
 import GHC.Types.Unique.Supply
 import GHC.Cmm.BlockId
+import GHC.Cmm.Dataflow.Collections
+import GHC.CmmToAsm.Reg.Utils
 
 data ReadingOrWriting = Reading | Writing deriving (Eq,Ord)
 
@@ -37,8 +44,41 @@ data ReadingOrWriting = Reading | Writing deriving (Eq,Ord)
 --      target a particular label. We have to insert fixup code to make
 --      the register assignments from the different sources match up.
 --
-type BlockAssignment freeRegs
-        = BlockMap (freeRegs, RegMap Loc)
+data BlockAssignment freeRegs
+        = BlockAssignment { blockMap :: !(BlockMap (freeRegs, RegMap Loc))
+                          , firstUsed :: !(UniqFM VirtualReg RealReg) }
+
+-- | Find the register mapping for a specific BlockId.
+lookupBlockAssignment :: BlockId -> BlockAssignment freeRegs -> Maybe (freeRegs, RegMap Loc)
+lookupBlockAssignment bid ba = mapLookup bid (blockMap ba)
+
+-- | Lookup which register a virtual register was first assigned to.
+lookupFirstUsed :: VirtualReg -> BlockAssignment freeRegs -> Maybe RealReg
+lookupFirstUsed vr ba = lookupUFM (firstUsed ba) vr
+
+-- | An initial empty 'BlockAssignment'
+emptyBlockAssignment :: BlockAssignment freeRegs
+emptyBlockAssignment = BlockAssignment mapEmpty mempty
+
+-- | Add new register mappings for a specific block.
+updateBlockAssignment :: BlockId
+  -> (freeRegs, RegMap Loc)
+  -> BlockAssignment freeRegs
+  -> BlockAssignment freeRegs
+updateBlockAssignment dest (freeRegs, regMap) (BlockAssignment {..}) =
+  BlockAssignment (mapInsert dest (freeRegs, regMap) blockMap)
+                  (mergeUFM combWithExisting id (mapMaybeUFM fromLoc) (firstUsed) (toVRegMap regMap))
+  where
+    -- The blocks are processed in dependency order, so if there's already an
+    -- entry in the map then keep that assignment rather than writing the new
+    -- assignment.
+    combWithExisting :: RealReg -> Loc -> Maybe RealReg
+    combWithExisting old_reg _ = Just $ old_reg
+
+    fromLoc :: Loc -> Maybe RealReg
+    fromLoc (InReg rr) = Just rr
+    fromLoc (InBoth rr _) = Just rr
+    fromLoc _ = Nothing
 
 
 -- | Where a vreg is currently stored
diff --git a/compiler/GHC/CmmToAsm/Reg/Linear/JoinToTargets.hs b/compiler/GHC/CmmToAsm/Reg/Linear/JoinToTargets.hs
index cbdf5d030b..ab63e18bbd 100644
--- a/compiler/GHC/CmmToAsm/Reg/Linear/JoinToTargets.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Linear/JoinToTargets.hs
@@ -100,7 +100,7 @@ joinToTargets' block_live new_blocks block_id instr (dest:dests)
                         , not (elemUniqSet_Directly reg live_set)
                         , r          <- regsOfLoc loc ]
 
-        case mapLookup dest block_assig of
+        case lookupBlockAssignment  dest block_assig of
          Nothing
           -> joinToTargets_first
                         block_live new_blocks block_id instr dest dests
@@ -136,7 +136,7 @@ joinToTargets_first block_live new_blocks block_id instr dest dests
         let freeregs' = foldl' (flip $ frReleaseReg platform) freeregs to_free
 
         -- remember the current assignment on entry to this block.
-        setBlockAssigR (mapInsert dest (freeregs', src_assig) block_assig)
+        setBlockAssigR (updateBlockAssignment dest (freeregs', src_assig) block_assig)
 
         joinToTargets' block_live new_blocks block_id instr dests