code gen: Improve efficiency of findPrefRealReg

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
(cherry picked from commit 9600a5fbd79c8434cf68aad078f26577da0530ad)

4 files changed, 71 insertions, 20 deletions

diff --git a/compiler/GHC/CmmToAsm/Reg/Linear.hs b/compiler/GHC/CmmToAsm/Reg/Linear.hs
index 6acd34c9be..1ab353cc34 100644
--- a/compiler/GHC/CmmToAsm/Reg/Linear.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Linear.hs
@@ -139,7 +139,6 @@ import GHC.Platform
 import Data.Maybe
 import Data.List
 import Control.Monad
-import Control.Applicative
 
 -- -----------------------------------------------------------------------------
 -- Top level of the register allocator
@@ -248,8 +247,8 @@ linearRegAlloc'
 
 linearRegAlloc' config initFreeRegs entry_ids block_live sccs
  = do   us      <- getUniqueSupplyM
-        let (_, stack, stats, blocks) =
-                runR config mapEmpty initFreeRegs emptyRegMap emptyStackMap us
+        let !(_, !stack, !stats, !blocks) =
+                runR config emptyBlockAssignment initFreeRegs emptyRegMap emptyStackMap us
                     $ linearRA_SCCs entry_ids block_live [] sccs
         return  (blocks, stats, getStackUse stack)
 
@@ -319,7 +318,7 @@ process entry_ids block_live (b@(BasicBlock id _) : blocks)
  = do
         block_assig <- getBlockAssigR
 
-        if isJust (mapLookup id block_assig)
+        if isJust (lookupBlockAssignment id block_assig)
              || id `elem` entry_ids
          then do
                 b'  <- processBlock block_live b
@@ -355,7 +354,7 @@ initBlock :: FR freeRegs
 initBlock id block_live
  = do   platform    <- getPlatform
         block_assig <- getBlockAssigR
-        case mapLookup id block_assig of
+        case lookupBlockAssignment id block_assig of
                 -- no prior info about this block: we must consider
                 -- any fixed regs to be allocated, but we can ignore
                 -- virtual regs (presumably this is part of a loop,
@@ -813,19 +812,11 @@ allocateRegsAndSpill reading keep spills alloc (r:rs)
 -- variables are likely to end up in the same registers at the
 -- end and start of the loop, avoiding redundant reg-reg moves.
 -- Note: I tried returning a list of past assignments, but that
--- turned out to barely matter but added a few tenths of
--- a percent to compile time.
+-- turned out to barely matter.
 findPrefRealReg :: VirtualReg -> RegM freeRegs (Maybe RealReg)
 findPrefRealReg vreg = do
-  bassig <- getBlockAssigR :: RegM freeRegs (BlockMap (freeRegs,RegMap Loc))
-  return $ foldr (findVirtRegAssig) Nothing bassig
-  where
-    findVirtRegAssig :: (freeRegs,RegMap Loc) -> Maybe RealReg -> Maybe RealReg
-    findVirtRegAssig assig z =
-        z <|>   case lookupUFM (toVRegMap $ snd assig) vreg of
-                        Just (InReg real_reg) -> Just real_reg
-                        Just (InBoth real_reg _) -> Just real_reg
-                        _ -> z
+  bassig <- getBlockAssigR :: RegM freeRegs (BlockAssignment freeRegs)
+  return $ lookupFirstUsed vreg bassig
 
 -- reading is redundant with reason, but we keep it around because it's
 -- convenient and it maintains the recursive structure of the allocator. -- EZY
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 8d3a46f490..3f1bea15b0 100644
--- a/compiler/GHC/CmmToAsm/Reg/Linear/JoinToTargets.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Linear/JoinToTargets.hs
@@ -97,7 +97,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
@@ -133,7 +133,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
 
diff --git a/compiler/GHC/Types/Unique/FM.hs b/compiler/GHC/Types/Unique/FM.hs
index 41f3018a05..7278587c13 100644
--- a/compiler/GHC/Types/Unique/FM.hs
+++ b/compiler/GHC/Types/Unique/FM.hs
@@ -52,6 +52,7 @@ module GHC.Types.Unique.FM (
         plusUFM_C,
         plusUFM_CD,
         plusUFM_CD2,
+        mergeUFM,
         plusMaybeUFM_C,
         plusUFMList,
         minusUFM,
@@ -62,6 +63,7 @@ module GHC.Types.Unique.FM (
         nonDetStrictFoldUFM, foldUFM, nonDetStrictFoldUFM_Directly,
         anyUFM, allUFM, seqEltsUFM,
         mapUFM, mapUFM_Directly,
+        mapMaybeUFM,
         elemUFM, elemUFM_Directly,
         filterUFM, filterUFM_Directly, partitionUFM,
         sizeUFM,
@@ -85,6 +87,7 @@ import qualified Data.IntSet as S
 import Data.Data
 import qualified Data.Semigroup as Semi
 import Data.Functor.Classes (Eq1 (..))
+import Data.Coerce
 
 -- | A finite map from @uniques@ of one type to
 -- elements in another type.
@@ -250,6 +253,20 @@ plusUFM_CD2 f (UFM xm) (UFM ym)
       (M.map (\y -> Nothing `f` Just y))
       xm ym
 
+mergeUFM
+  :: (elta -> eltb -> Maybe eltc)
+  -> (UniqFM key elta -> UniqFM key eltc)  -- map X
+  -> (UniqFM key eltb -> UniqFM key eltc) -- map Y
+  -> UniqFM key elta
+  -> UniqFM key eltb
+  -> UniqFM key eltc
+mergeUFM f g h (UFM xm) (UFM ym)
+  = UFM $ M.mergeWithKey
+      (\_ x y -> (x `f` y))
+      (coerce g)
+      (coerce h)
+      xm ym
+
 plusMaybeUFM_C :: (elt -> elt -> Maybe elt)
                -> UniqFM key elt -> UniqFM key elt -> UniqFM key elt
 plusMaybeUFM_C f (UFM xm) (UFM ym)
@@ -284,6 +301,9 @@ foldUFM k z (UFM m) = M.foldr k z m
 mapUFM :: (elt1 -> elt2) -> UniqFM key elt1 -> UniqFM key elt2
 mapUFM f (UFM m) = UFM (M.map f m)
 
+mapMaybeUFM :: (elt1 -> Maybe elt2) -> UniqFM key elt1 -> UniqFM key elt2
+mapMaybeUFM f (UFM m) = UFM (M.mapMaybe f m)
+
 mapUFM_Directly :: (Unique -> elt1 -> elt2) -> UniqFM key elt1 -> UniqFM key elt2
 mapUFM_Directly f (UFM m) = UFM (M.mapWithKey (f . getUnique) m)