Refactor and simplify the SRT handling

compiler/cmm/CmmBuildInfoTables.hs

@@ -13,16 +13,15 @@
 -- Todo: remove -fno-warn-warnings-deprecations
 {-# OPTIONS_GHC -fno-warn-warnings-deprecations #-}
 module CmmBuildInfoTables
-    ( CAFSet, CAFEnv, cafAnal, localCAFInfo, mkTopCAFInfo
-    , setInfoTableSRT
-    , TopSRT, emptySRT, srtToData
-    , bundleCAFs
-    ) where
+    ( CAFSet, CAFEnv, cafAnal
+    , doSRTs, TopSRT, emptySRT, srtToData )
+where
 
 #include "HsVersions.h"
 
 -- These should not be imported here!
 import StgCmmUtils
+import Hoopl
 
 import Digraph
 import qualified Prelude as P
@@ -40,13 +39,13 @@ import Name
 import Outputable
 import SMRep
 import UniqSupply
-
-import Hoopl
+import Util
 
 import Data.Map (Map)
 import qualified Data.Map as Map
 import Data.Set (Set)
 import qualified Data.Set as Set
+import Control.Monad
 
 foldSet :: (a -> b -> b) -> b -> Set a -> b
 #if __GLASGOW_HASKELL__ < 704
@@ -184,16 +183,13 @@ srtToData srt = [CmmData RelocatableReadOnlyData (Statics (lbl srt) tbl)]
 -- in the SRT. Then, if the number of CAFs is small enough to fit in a bitmap,
 -- we make sure they're all close enough to the bottom of the table that the
 -- bitmap will be able to cover all of them.
-buildSRTs :: TopSRT -> Map CLabel CAFSet -> CAFSet ->
-             UniqSM (TopSRT, Maybe CmmDecl, C_SRT)
-buildSRTs topSRT topCAFMap cafs =
-  do let liftCAF lbl z = -- get CAFs for functions without static closures
-           case Map.lookup lbl topCAFMap of Just cafs -> z `Set.union` cafs
-                                            Nothing   -> Set.insert lbl z
+buildSRTs :: TopSRT -> CAFSet -> UniqSM (TopSRT, Maybe CmmDecl, C_SRT)
+buildSRTs topSRT cafs =
+  do let
          -- For each label referring to a function f without a static closure,
          -- replace it with the CAFs that are reachable from f.
          sub_srt topSRT localCafs =
-           let cafs = Set.elems (foldSet liftCAF Set.empty localCafs)
+           let cafs = Set.elems localCafs
                mkSRT topSRT =
                  do localSRTs <- procpointSRT (lbl topSRT) (elt_map topSRT) cafs
                     return (topSRT, localSRTs)
@@ -267,15 +263,15 @@ to_SRT top_srt off len bmp
 --  keep its CAFs live.)
 -- Any procedure referring to a non-static CAF c must keep live
 -- any CAF that is reachable from c.
-localCAFInfo :: CAFEnv -> CmmDecl -> Maybe (CLabel, CAFSet)
-localCAFInfo _      (CmmData _ _) = Nothing
+localCAFInfo :: CAFEnv -> CmmDecl -> (CAFSet, Maybe CLabel)
+localCAFInfo _      (CmmData _ _) = (Set.empty, Nothing)
 localCAFInfo cafEnv (CmmProc top_info top_l (CmmGraph {g_entry=entry})) =
   case info_tbl top_info of
-    CmmInfoTable { cit_rep = rep } 
-      | not (isStaticRep rep) 
-      -> Just (toClosureLbl top_l,
-               expectJust "maybeBindCAFs" $ mapLookup entry cafEnv)
-    _ -> Nothing
+    CmmInfoTable { cit_rep = rep } | not (isStaticRep rep)
+      -> (cafs, Just (toClosureLbl top_l))
+    _other -> (cafs, Nothing)
+  where
+    cafs = expectJust "maybeBindCAFs" $ mapLookup entry cafEnv
 
 -- Once we have the local CAF sets for some (possibly) mutually
 -- recursive functions, we can create an environment mapping
@@ -288,54 +284,77 @@ localCAFInfo cafEnv (CmmProc top_info top_l (CmmGraph {g_entry=entry})) =
 -- the environment with every reference to f replaced by its set of CAFs.
 -- To do this replacement efficiently, we gather strongly connected
 -- components, then we sort the components in topological order.
-mkTopCAFInfo :: [(CLabel, CAFSet)] -> Map CLabel CAFSet
+mkTopCAFInfo :: [(CAFSet, Maybe CLabel)] -> Map CLabel CAFSet
 mkTopCAFInfo localCAFs = foldl addToTop Map.empty g
-  where addToTop env (AcyclicSCC (l, cafset)) =
+  where
+        addToTop env (AcyclicSCC (l, cafset)) =
           Map.insert l (flatten env cafset) env
         addToTop env (CyclicSCC nodes) =
           let (lbls, cafsets) = unzip nodes
               cafset  = foldr Set.delete (foldl Set.union Set.empty cafsets) lbls
           in foldl (\env l -> Map.insert l (flatten env cafset) env) env lbls
-        flatten env cafset = foldSet (lookup env) Set.empty cafset
-        lookup env caf cafset' =
-          case Map.lookup caf env of Just cafs -> foldSet add cafset' cafs
-                                     Nothing -> add caf cafset'
-        add caf cafset' = Set.insert caf cafset'
+
         g = stronglyConnCompFromEdgedVertices
-              (map (\n@(l, cafs) -> (n, l, Set.elems cafs)) localCAFs)
-
--- Bundle the CAFs used at a procpoint.
-bundleCAFs :: CAFEnv -> CmmDecl -> (CAFSet, CmmDecl)
-bundleCAFs cafEnv t@(CmmProc _ _ (CmmGraph {g_entry=entry})) =
-  (expectJust "bundleCAFs" (mapLookup entry cafEnv), t)
-bundleCAFs _ t = (Set.empty, t)
-
--- Construct the SRTs for the given procedure.
-setInfoTableSRT :: Map CLabel CAFSet -> TopSRT -> (CAFSet, CmmDecl) ->
-                   UniqSM (TopSRT, [CmmDecl])
-setInfoTableSRT topCAFMap topSRT (cafs, t) =
-  setSRT cafs topCAFMap topSRT t
-
-setSRT :: CAFSet -> Map CLabel CAFSet -> TopSRT ->
-          CmmDecl -> UniqSM (TopSRT, [CmmDecl])
-setSRT cafs topCAFMap topSRT t =
-  do (topSRT, cafTable, srt) <- buildSRTs topSRT topCAFMap cafs
-     let t' = updInfo id (const srt) t
-     case cafTable of
-       Just tbl -> return (topSRT, [t', tbl])
-       Nothing  -> return (topSRT, [t'])
-
-type StackLayout = Liveness
-
-updInfo :: (StackLayout -> StackLayout) -> (C_SRT -> C_SRT) -> CmmDecl -> CmmDecl
-updInfo toVars toSrt (CmmProc top_info top_l g) =
-  CmmProc (top_info {info_tbl=updInfoTbl toVars toSrt (info_tbl top_info)}) top_l g
-updInfo _ _ t = t
-
-updInfoTbl :: (StackLayout -> StackLayout) -> (C_SRT -> C_SRT) -> CmmInfoTable -> CmmInfoTable
-updInfoTbl toVars toSrt info_tbl@(CmmInfoTable {})
-  = info_tbl { cit_srt = toSrt (cit_srt info_tbl)
-             , cit_rep = case cit_rep info_tbl of
-                           StackRep ls -> StackRep (toVars ls)
-                           other       -> other }
-updInfoTbl _ _ t@CmmNonInfoTable = t
+              [ ((l,cafs), l, Set.elems cafs) | (cafs, Just l) <- localCAFs ]
+
+flatten :: Map CLabel CAFSet -> CAFSet -> CAFSet
+flatten env cafset = foldSet (lookup env) Set.empty cafset
+  where
+      lookup env caf cafset' =
+          case Map.lookup caf env of
+             Just cafs -> foldSet Set.insert cafset' cafs
+             Nothing   -> Set.insert caf cafset'
+
+bundle :: Map CLabel CAFSet
+       -> (CAFEnv, CmmDecl)
+       -> (CAFSet, Maybe CLabel)
+       -> (CAFSet, CmmDecl)
+bundle flatmap (_, decl) (cafs, Nothing)
+  = (flatten flatmap cafs, decl)
+bundle flatmap (_, decl) (_, Just l)
+  = (expectJust "bundle" $ Map.lookup l flatmap, decl)
+
+flattenCAFSets :: [(CAFEnv, [CmmDecl])] -> [(CAFSet, CmmDecl)]
+flattenCAFSets cpsdecls = zipWith (bundle flatmap) zipped localCAFs
+   where
+     zipped    = [(e,d) | (e,ds) <- cpsdecls, d <- ds ]
+     localCAFs = unzipWith localCAFInfo zipped
+     flatmap   = mkTopCAFInfo localCAFs -- transitive closure of localCAFs
+
+doSRTs :: TopSRT
+       -> [(CAFEnv, [CmmDecl])]
+       -> IO (TopSRT, [CmmDecl])
+
+doSRTs topSRT tops
+  = do
+     let caf_decls = flattenCAFSets tops
+     us <- mkSplitUniqSupply 'u'
+     let (topSRT', gs') = initUs_ us $ foldM setSRT (topSRT, []) caf_decls
+     return (topSRT', reverse gs' {- Note [reverse gs] -})
+  where
+    setSRT (topSRT, rst) (cafs, decl@(CmmProc{})) = do
+       (topSRT, cafTable, srt) <- buildSRTs topSRT cafs
+       let decl' = updInfo (const srt) decl
+       case cafTable of
+         Just tbl -> return (topSRT, decl': tbl : rst)
+         Nothing  -> return (topSRT, decl' : rst)
+    setSRT (topSRT, rst) (_, decl) =
+      return (topSRT, decl : rst)
+
+{- Note [reverse gs]
+
+   It is important to keep the code blocks in the same order,
+   otherwise binary sizes get slightly bigger.  I'm not completely
+   sure why this is, perhaps the assembler generates bigger jump
+   instructions for forward refs.  --SDM
+-}
+
+updInfo :: (C_SRT -> C_SRT) -> CmmDecl -> CmmDecl
+updInfo toSrt (CmmProc top_info top_l g) =
+  CmmProc (top_info {info_tbl = updInfoTbl toSrt (info_tbl top_info)}) top_l g
+updInfo _ t = t
+
+updInfoTbl :: (C_SRT -> C_SRT) -> CmmInfoTable -> CmmInfoTable
+updInfoTbl toSrt info_tbl@(CmmInfoTable {})
+  = info_tbl { cit_srt = toSrt (cit_srt info_tbl) }
+updInfoTbl _ t@CmmNonInfoTable = t

compiler/cmm/CmmPipeline.hs

@@ -9,7 +9,6 @@ module CmmPipeline (
   cmmPipeline
 ) where
 
-import CLabel
 import Cmm
 import CmmLint
 import CmmBuildInfoTables
@@ -18,76 +17,41 @@ import CmmProcPoint
 import CmmContFlowOpt
 import CmmLayoutStack
 import CmmSink
+import Hoopl
 
 import UniqSupply
 import DynFlags
 import ErrUtils
 import HscTypes
-import Data.Maybe
 import Control.Monad
 import Outputable
 
-import qualified Data.Set as Set
-import Data.Map (Map)
-
 -----------------------------------------------------------------------------
 -- | Top level driver for C-- pipeline
 -----------------------------------------------------------------------------
--- There are two complications here:
--- 1. We need to compile the procedures in two stages because we need
---    an analysis of the procedures to tell us what CAFs they use.
---    The first stage returns a map from procedure labels to CAFs,
---    along with a closure that will compute SRTs and attach them to
---    the compiled procedures.
---    The second stage is to combine the CAF information into a top-level
---    CAF environment mapping non-static closures to the CAFs they keep live,
---    then pass that environment to the closures returned in the first
---    stage of compilation.
--- 2. We need to thread the module's SRT around when the SRT tables
---    are computed for each procedure.
---    The SRT needs to be threaded because it is grown lazily.
--- 3. We run control flow optimizations twice, once before any pipeline
---    work is done, and once again at the very end on all of the
---    resulting C-- blocks.  EZY: It's unclear whether or not whether
---    we actually need to do the initial pass.
+
 cmmPipeline  :: HscEnv -- Compilation env including
                        -- dynamic flags: -dcmm-lint -ddump-cps-cmm
              -> TopSRT     -- SRT table and accumulating list of compiled procs
              -> CmmGroup             -- Input C-- with Procedures
              -> IO (TopSRT, CmmGroup) -- Output CPS transformed C--
+
 cmmPipeline hsc_env topSRT prog =
   do let dflags = hsc_dflags hsc_env
-     --
-     showPass dflags "CPSZ"
-
-     (cafEnvs, tops) <- {-# SCC "tops" #-} liftM unzip $ mapM (cpsTop hsc_env) prog
-     -- tops :: [[(CmmDecl,CAFSet]]  (one list per group)
 
-     let topCAFEnv = {-# SCC "topCAFEnv" #-} mkTopCAFInfo (concat cafEnvs)
-
-     -- folding over the groups
-     (topSRT, tops) <- {-# SCC "toTops" #-} foldM (toTops topCAFEnv) (topSRT, []) tops
+     showPass dflags "CPSZ"
 
-     let cmms :: CmmGroup
-         cmms = reverse (concat tops)
+     tops <- {-# SCC "tops" #-} mapM (cpsTop hsc_env) prog
 
+     (topSRT, cmms) <- {-# SCC "toTops" #-} doSRTs topSRT tops
      dumpIfSet_dyn dflags Opt_D_dump_cps_cmm "Post CPS Cmm" (ppr cmms)
 
      return (topSRT, cmms)
 
-{- [Note global fuel]
-~~~~~~~~~~~~~~~~~~~~~
-The identity and the last pass are stored in
-mutable reference cells in an 'HscEnv' and are
-global to one compiler session.
--}
 
--- EZY: It might be helpful to have an easy way of dumping the "pre"
--- input for any given phase, besides just turning it all on with
--- -ddump-cmmz
 
-cpsTop :: HscEnv -> CmmDecl -> IO ([(CLabel, CAFSet)], [(CAFSet, CmmDecl)])
-cpsTop _ p@(CmmData {}) = return ([], [(Set.empty, p)])
+cpsTop :: HscEnv -> CmmDecl -> IO (CAFEnv, [CmmDecl])
+cpsTop _ p@(CmmData {}) = return (mapEmpty, [p])
 cpsTop hsc_env (CmmProc h@(TopInfo {stack_info=StackInfo {arg_space=entry_off}}) l g) =
     do
        ----------- Control-flow optimisations ---------------
@@ -132,31 +96,21 @@ cpsTop hsc_env (CmmProc h@(TopInfo {stack_info=StackInfo {arg_space=entry_off}})
              splitAtProcPoints l callPPs procPoints procPointMap (CmmProc h l g)
        dumps Opt_D_dump_cmmz_split "Post splitting" gs
 
-       ------------- More CAFs ------------------------------
+       ------------- CAF analysis ------------------------------
        let cafEnv = {-# SCC "cafAnal" #-} cafAnal g
-       let localCAFs = {-# SCC "localCAFs" #-} catMaybes $ map (localCAFInfo cafEnv) gs
-       mbpprTrace "localCAFs" (ppr localCAFs) $ return ()
 
-       -- NO MORE GRAPH TRANSFORMATION AFTER HERE -- JUST MAKING INFOTABLES
+       ------------- Populate info tables with stack info ------
        gs <- {-# SCC "setInfoTableStackMap" #-}
              return $ map (setInfoTableStackMap stackmaps) gs
        dumps Opt_D_dump_cmmz_info "after setInfoTableStackMap" gs
 
-       ----------- Control-flow optimisations ---------------
+       ----------- Control-flow optimisations -----------------
        gs <- {-# SCC "cmmCfgOpts(2)" #-} return $ map cmmCfgOptsProc gs
        dumps Opt_D_dump_cmmz_cfg "Post control-flow optimsations" gs
 
-       gs <- {-# SCC "bundleCAFs" #-} return $ map (bundleCAFs cafEnv) gs
-       dumps Opt_D_dump_cmmz_cafs "after bundleCAFs" gs
-
-       return (localCAFs, gs)
-
-              -- gs        :: [ (CAFSet, CmmDecl) ]
-              -- localCAFs :: [ (CLabel, CAFSet) ] -- statics filtered out(?)
+       return (cafEnv, gs)
 
   where dflags = hsc_dflags hsc_env
-        mbpprTrace x y z | dopt Opt_D_dump_cmmz dflags = pprTrace x y z
-                         | otherwise = z
         dump = dumpGraph dflags
 
         dumps flag name
@@ -188,14 +142,3 @@ dumpWith dflags flag txt g = do
    when (not (dopt flag dflags)) $
       dumpIfSet_dyn dflags Opt_D_dump_cmmz txt (ppr g)
 
--- This probably belongs in CmmBuildInfoTables?
--- We're just finishing the job here: once we know what CAFs are defined
--- in non-static closures, we can build the SRTs.
-toTops :: Map CLabel CAFSet -> (TopSRT, [[CmmDecl]])
-       -> [(CAFSet, CmmDecl)] -> IO (TopSRT, [[CmmDecl]])
-toTops topCAFEnv (topSRT, tops) gs =
-  do let setSRT (topSRT, rst) g =
-           do (topSRT, gs) <- setInfoTableSRT topCAFEnv topSRT g
-              return (topSRT, gs : rst)
-     (topSRT, gs') <- runUniqSM $ foldM setSRT (topSRT, []) gs
-     return (topSRT, concat gs' : tops)