Code Gen: Use more efficient block merging algorithm

The previous algorithm scaled poorly when there was a large number of
blocks and edges.

The algorithm links together block chains which have edges between them
in the CFG. The new algorithm uses a union find data structure in order
to efficiently merge together blocks and calculate which block chain
each block id belonds to.

I copied the UnionFind data structure which already existed in Cabal
into the GHC library rathert than reimplement it myself.

This change results in a very significant reduction in allocations when
compiling the mmark package.

Ticket: #19471

1 files changed, 17 insertions, 27 deletions

diff --git a/compiler/GHC/CmmToAsm/BlockLayout.hs b/compiler/GHC/CmmToAsm/BlockLayout.hs
index 6e2e7e2189..70e131c717 100644
--- a/compiler/GHC/CmmToAsm/BlockLayout.hs
+++ b/compiler/GHC/CmmToAsm/BlockLayout.hs
@@ -41,12 +41,13 @@ import GHC.Utils.Panic
 import GHC.Utils.Panic.Plain
 import GHC.Utils.Misc
 
-import Data.List (sortOn, sortBy)
+import Data.List (sortOn, sortBy, nub)
 import Data.Foldable (toList)
 import qualified Data.Set as Set
 import Data.STRef
 import Control.Monad.ST.Strict
-import Control.Monad (foldM)
+import Control.Monad (foldM, unless)
+import GHC.Data.UnionFind
 
 {-
   Note [CFG based code layout]
@@ -480,10 +481,9 @@ combineNeighbourhood edges chains
 mergeChains :: [CfgEdge] -> [BlockChain]
             -> (BlockChain)
 mergeChains edges chains
-    = -- pprTrace "combine" (ppr edges) $
-      runST $ do
+    = runST $ do
         let addChain m0 chain = do
-                ref <- newSTRef chain
+                ref <- fresh chain
                 return $ chainFoldl (\m' b -> mapInsert b ref m') m0 chain
         chainMap' <- foldM (\m0 c -> addChain m0 c) mapEmpty chains
         merge edges chainMap'
@@ -491,35 +491,25 @@ mergeChains edges chains
         -- We keep a map from ALL blocks to their respective chain (sigh)
         -- This is required since when looking at an edge we need to find
         -- the associated chains quickly.
-        -- We use a map of STRefs, maintaining a invariant of one STRef per chain.
-        -- When merging chains we can update the
-        -- STRef of one chain once (instead of writing to the map for each block).
-        -- We then overwrite the STRefs for the other chain so there is again only
-        -- a single STRef for the combined chain.
-        -- The difference in terms of allocations saved is ~0.2% with -O so actually
-        -- significant compared to using a regular map.
-
-        merge :: forall s. [CfgEdge] -> LabelMap (STRef s BlockChain) -> ST s BlockChain
+        -- We use a union-find data structure to do this efficiently.
+
+        merge :: forall s. [CfgEdge] -> LabelMap (Point s BlockChain) -> ST s BlockChain
         merge [] chains = do
-            chains' <- ordNub <$> (mapM readSTRef $ mapElems chains) :: ST s [BlockChain]
+            chains' <- mapM find =<< (nub <$> (mapM repr $ mapElems chains)) :: ST s [BlockChain]
             return $ foldl' chainConcat (head chains') (tail chains')
         merge ((CfgEdge from to _):edges) chains
         --   | pprTrace "merge" (ppr (from,to) <> ppr chains) False
         --   = undefined
-          | cFrom == cTo
-          = merge edges chains
-          | otherwise
           = do
-            chains' <- mergeComb cFrom cTo
-            merge edges chains'
+            same <- equivalent cFrom cTo
+            unless same $ do
+              cRight <- find cTo
+              cLeft <- find cFrom
+              new_point <- fresh (chainConcat cLeft cRight)
+              union cTo new_point
+              union cFrom new_point
+            merge edges chains
           where
-            mergeComb :: STRef s BlockChain -> STRef s BlockChain -> ST s (LabelMap (STRef s BlockChain))
-            mergeComb refFrom refTo = do
-                cRight <- readSTRef refTo
-                chain <- pure chainConcat <*> readSTRef refFrom <*> pure cRight
-                writeSTRef refFrom chain
-                return $ chainFoldl (\m b -> mapInsert b refFrom m) chains cRight
-
             cFrom = expectJust "mergeChains:chainMap:from" $ mapLookup from chains
             cTo = expectJust "mergeChains:chainMap:to"   $ mapLookup to   chains