Refactor IPE initialization

Here we refactor the representation of info table provenance information
in object code to significantly reduce its size and link-time impact.
Specifically, we deduplicate strings and represent them as 32-bit
offsets into a common string table.

In addition, we rework the registration logic to eliminate allocation
from the registration path, which is run from a static initializer where
things like allocation are technically undefined behavior (although it
did previously seem to work). For similar reasons we eliminate lock
usage from registration path, instead relying on atomic CAS.

Closes #22077.

7 files changed, 174 insertions, 67 deletions

diff --git a/compiler/GHC/Cmm/CLabel.hs b/compiler/GHC/Cmm/CLabel.hs
index 5f2b0eca59..9edccdccf5 100644
--- a/compiler/GHC/Cmm/CLabel.hs
+++ b/compiler/GHC/Cmm/CLabel.hs
@@ -302,6 +302,7 @@ data ModuleLabelKind
     | MLK_InitializerArray
     | MLK_Finalizer String
     | MLK_FinalizerArray
+    | MLK_IPEBuffer
     deriving (Eq, Ord)
 
 instance Outputable ModuleLabelKind where
@@ -309,6 +310,7 @@ instance Outputable ModuleLabelKind where
     ppr (MLK_Initializer s)  = text ("init__" ++ s)
     ppr MLK_FinalizerArray   = text "fini_arr"
     ppr (MLK_Finalizer s)    = text ("fini__" ++ s)
+    ppr MLK_IPEBuffer        = text "ipe_buf"
 
 isIdLabel :: CLabel -> Bool
 isIdLabel IdLabel{} = True
@@ -839,10 +841,10 @@ instance OutputableP Platform InfoProvEnt where
 -- Constructing Cost Center Labels
 mkCCLabel  :: CostCentre      -> CLabel
 mkCCSLabel :: CostCentreStack -> CLabel
-mkIPELabel :: InfoProvEnt -> CLabel
+mkIPELabel :: Module          -> CLabel
 mkCCLabel           cc          = CC_Label cc
 mkCCSLabel          ccs         = CCS_Label ccs
-mkIPELabel          ipe         = IPE_Label ipe
+mkIPELabel          mod         = ModuleLabel mod MLK_IPEBuffer
 
 mkRtsApFastLabel :: FastString -> CLabel
 mkRtsApFastLabel str = RtsLabel (RtsApFast (NonDetFastString str))
@@ -1020,6 +1022,7 @@ modLabelNeedsCDecl :: ModuleLabelKind -> Bool
 -- Code for finalizers and initializers are emitted in stub objects
 modLabelNeedsCDecl (MLK_Initializer _)  = True
 modLabelNeedsCDecl (MLK_Finalizer   _)  = True
+modLabelNeedsCDecl MLK_IPEBuffer        = True
 -- The finalizer and initializer arrays are emitted in the code of the module
 modLabelNeedsCDecl MLK_InitializerArray = False
 modLabelNeedsCDecl MLK_FinalizerArray   = False
@@ -1217,6 +1220,7 @@ moduleLabelKindType kind =
     MLK_InitializerArray -> DataLabel
     MLK_Finalizer _      -> CodeLabel
     MLK_FinalizerArray   -> DataLabel
+    MLK_IPEBuffer        -> DataLabel
 
 idInfoLabelType :: IdLabelInfo -> CLabelType
 idInfoLabelType info =
diff --git a/compiler/GHC/Cmm/Parser.y b/compiler/GHC/Cmm/Parser.y
index ed47fa7a7f..ae6e126b68 100644
--- a/compiler/GHC/Cmm/Parser.y
+++ b/compiler/GHC/Cmm/Parser.y
@@ -224,6 +224,7 @@ import GHC.StgToCmm.Layout     hiding (ArgRep(..))
 import GHC.StgToCmm.Ticky
 import GHC.StgToCmm.Prof
 import GHC.StgToCmm.Bind  ( emitBlackHoleCode, emitUpdateFrame )
+import GHC.StgToCmm.InfoTableProv
 
 import GHC.Cmm.Opt
 import GHC.Cmm.Graph
@@ -1517,9 +1518,12 @@ parseCmmFile cmmpConfig this_mod home_unit filename = do
         let fcode = do
               ((), cmm) <- getCmm $ unEC code "global" (initEnv (pdProfile pdConfig)) [] >> return ()
               -- See Note [Mapping Info Tables to Source Positions] (IPE Maps)
-              let used_info = map (cmmInfoTableToInfoProvEnt this_mod)
-                                              (mapMaybe topInfoTable cmm)
-              ((), cmm2) <- getCmm $ mapM_ emitInfoTableProv used_info
+              let used_info
+                    | do_ipe    = map (cmmInfoTableToInfoProvEnt this_mod) (mapMaybe topInfoTable cmm)
+                    | otherwise = []
+                    where
+                      do_ipe = stgToCmmInfoTableMap $ cmmpStgToCmmConfig cmmpConfig
+              ((), cmm2) <- getCmm $ emitIpeBufferListNode this_mod used_info
               return (cmm ++ cmm2, used_info)
             (cmm, _) = runC (cmmpStgToCmmConfig cmmpConfig) fstate st fcode
             (warnings,errors) = getPsMessages pst
diff --git a/compiler/GHC/Driver/CodeOutput.hs b/compiler/GHC/Driver/CodeOutput.hs
index e07c0af91f..09315c4f05 100644
--- a/compiler/GHC/Driver/CodeOutput.hs
+++ b/compiler/GHC/Driver/CodeOutput.hs
@@ -369,24 +369,17 @@ ipInitCode
   :: Bool            -- is Opt_InfoTableMap enabled or not
   -> Platform
   -> Module
-  -> [InfoProvEnt]
   -> CStub
-ipInitCode do_info_table platform this_mod ents
+ipInitCode do_info_table platform this_mod
   | not do_info_table = mempty
-  | otherwise = initializerCStub platform fn_nm decls body
+  | otherwise = initializerCStub platform fn_nm ipe_buffer_decl body
  where
    fn_nm = mkInitializerStubLabel this_mod "ip_init"
-   decls = vcat
-        $  map emit_ipe_decl ents
-        ++ [emit_ipe_list ents]
-   body = text "registerInfoProvList" <> parens local_ipe_list_label <> semi
-   emit_ipe_decl ipe =
-       text "extern InfoProvEnt" <+> ipe_lbl <> text "[];"
-     where ipe_lbl = pprCLabel platform CStyle (mkIPELabel ipe)
-   local_ipe_list_label = text "local_ipe_" <> ppr this_mod
-   emit_ipe_list ipes =
-      text "static InfoProvEnt *" <> local_ipe_list_label <> text "[] ="
-      <+> braces (vcat $ [ pprCLabel platform CStyle (mkIPELabel ipe) <> comma
-                         | ipe <- ipes
-                         ] ++ [text "NULL"])
-      <> semi
+
+   body = text "registerInfoProvList" <> parens (text "&" <> ipe_buffer_label) <> semi
+
+   ipe_buffer_label = pprCLabel platform CStyle (mkIPELabel this_mod)
+
+   ipe_buffer_decl =
+       text "extern IpeBufferListNode" <+> ipe_buffer_label <> text ";"
+
diff --git a/compiler/GHC/Driver/Main.hs b/compiler/GHC/Driver/Main.hs
index 175a78962e..5f22840395 100644
--- a/compiler/GHC/Driver/Main.hs
+++ b/compiler/GHC/Driver/Main.hs
@@ -1970,7 +1970,7 @@ hscCompileCmmFile hsc_env original_filename filename output_filename = runHsc hs
         mod_name = mkModuleName $ "Cmm$" ++ original_filename
         cmm_mod = mkHomeModule home_unit mod_name
         cmmpConfig = initCmmParserConfig dflags
-    (cmm, ents) <- ioMsgMaybe
+    (cmm, ipe_ents) <- ioMsgMaybe
                $ do
                   (warns,errs,cmm) <- withTiming logger (text "ParseCmm"<+>brackets (text filename)) (\_ -> ())
                                        $ parseCmmFile cmmpConfig cmm_mod home_unit filename
@@ -1996,10 +1996,11 @@ hscCompileCmmFile hsc_env original_filename filename output_filename = runHsc hs
           Nothing -> cmmToRawCmm logger profile (Stream.yield cmmgroup)
           Just h  -> h           dflags Nothing (Stream.yield cmmgroup)
 
-        let foreign_stubs _ =
-              let ip_init   = ipInitCode do_info_table platform cmm_mod ents
-              in NoStubs `appendStubC` ip_init
-
+        let foreign_stubs _
+              | not $ null ipe_ents =
+                  let ip_init = ipInitCode do_info_table platform cmm_mod
+                  in NoStubs `appendStubC` ip_init
+              | otherwise     = NoStubs
         (_output_filename, (_stub_h_exists, stub_c_exists), _foreign_fps, _caf_infos)
           <- codeOutput logger tmpfs llvm_config dflags (hsc_units hsc_env) cmm_mod output_filename no_loc foreign_stubs [] S.empty
              rawCmms
diff --git a/compiler/GHC/StgToCmm/InfoTableProv.hs b/compiler/GHC/StgToCmm/InfoTableProv.hs
new file mode 100644
index 0000000000..f0bc0f12c6
--- /dev/null
+++ b/compiler/GHC/StgToCmm/InfoTableProv.hs
@@ -0,0 +1,131 @@
+module GHC.StgToCmm.InfoTableProv (emitIpeBufferListNode) where
+
+import GHC.Prelude
+import GHC.Platform
+import GHC.Unit.Module
+import GHC.Utils.Outputable
+
+import GHC.Cmm.CLabel
+import GHC.Cmm.Expr
+import GHC.Cmm.Utils
+import GHC.StgToCmm.Config
+import GHC.StgToCmm.Lit (newByteStringCLit)
+import GHC.StgToCmm.Monad
+import GHC.StgToCmm.Utils
+
+import GHC.Data.ShortText (ShortText)
+import qualified GHC.Data.ShortText as ST
+
+import Data.Bifunctor (first)
+import qualified Data.Map.Strict as M
+import Control.Monad.Trans.State.Strict
+import qualified Data.ByteString as BS
+import qualified Data.ByteString.Builder as BSB
+import qualified Data.ByteString.Lazy as BSL
+
+emitIpeBufferListNode :: Module
+                      -> [InfoProvEnt]
+                      -> FCode ()
+emitIpeBufferListNode _ [] = return ()
+emitIpeBufferListNode this_mod ents = do
+    cfg <- getStgToCmmConfig
+    let ctx      = stgToCmmContext  cfg
+        platform = stgToCmmPlatform cfg
+
+    let (cg_ipes, strtab) = flip runState emptyStringTable $ do
+            module_name <- lookupStringTable $ ST.pack $ renderWithContext ctx (ppr this_mod)
+            mapM (toCgIPE platform ctx module_name) ents
+
+    let -- Emit the fields of an IpeBufferEntry struct.
+        toIpeBufferEntry :: CgInfoProvEnt -> [CmmLit]
+        toIpeBufferEntry cg_ipe =
+            [ CmmLabel (ipeInfoTablePtr cg_ipe)
+            , strtab_offset (ipeTableName cg_ipe)
+            , strtab_offset (ipeClosureDesc cg_ipe)
+            , strtab_offset (ipeTypeDesc cg_ipe)
+            , strtab_offset (ipeLabel cg_ipe)
+            , strtab_offset (ipeModuleName cg_ipe)
+            , strtab_offset (ipeSrcLoc cg_ipe)
+            ]
+
+        int n = mkIntCLit platform n
+        int32 n = CmmInt n W32
+        strtab_offset (StrTabOffset n) = int32 (fromIntegral n)
+
+    strings <- newByteStringCLit (getStringTableStrings strtab)
+    let lits = [ zeroCLit platform     -- 'next' field
+               , strings               -- 'strings' field
+               , int $ length cg_ipes  -- 'count' field
+               ] ++ concatMap toIpeBufferEntry cg_ipes
+    emitDataLits (mkIPELabel this_mod) lits
+
+toCgIPE :: Platform -> SDocContext -> StrTabOffset -> InfoProvEnt -> State StringTable CgInfoProvEnt
+toCgIPE platform ctx module_name ipe = do
+    table_name <- lookupStringTable $ ST.pack $ renderWithContext ctx (pprCLabel platform CStyle (infoTablePtr ipe))
+    closure_desc <- lookupStringTable $ ST.pack $ show (infoProvEntClosureType ipe)
+    type_desc <- lookupStringTable $ ST.pack $ infoTableType ipe
+    let (src_loc_str, label_str) = maybe ("", "") (first (renderWithContext ctx . ppr)) (infoTableProv ipe)
+    label <- lookupStringTable $ ST.pack label_str
+    src_loc <- lookupStringTable $ ST.pack src_loc_str
+    return $ CgInfoProvEnt { ipeInfoTablePtr = infoTablePtr ipe
+                           , ipeTableName = table_name
+                           , ipeClosureDesc = closure_desc
+                           , ipeTypeDesc = type_desc
+                           , ipeLabel = label
+                           , ipeModuleName = module_name
+                           , ipeSrcLoc = src_loc
+                           }
+
+data CgInfoProvEnt = CgInfoProvEnt
+                               { ipeInfoTablePtr :: !CLabel
+                               , ipeTableName :: !StrTabOffset
+                               , ipeClosureDesc :: !StrTabOffset
+                               , ipeTypeDesc :: !StrTabOffset
+                               , ipeLabel :: !StrTabOffset
+                               , ipeModuleName :: !StrTabOffset
+                               , ipeSrcLoc :: !StrTabOffset
+                               }
+
+data StringTable = StringTable { stStrings :: DList ShortText
+                               , stLength :: !Int
+                               , stLookup :: !(M.Map ShortText StrTabOffset)
+                               }
+
+newtype StrTabOffset = StrTabOffset Int
+
+emptyStringTable :: StringTable
+emptyStringTable =
+    StringTable { stStrings = emptyDList
+                , stLength = 0
+                , stLookup = M.empty
+                }
+
+getStringTableStrings :: StringTable -> BS.ByteString
+getStringTableStrings st =
+    BSL.toStrict $ BSB.toLazyByteString
+    $ foldMap f $ dlistToList (stStrings st)
+  where
+    f x = BSB.shortByteString (ST.contents x) `mappend` BSB.word8 0
+
+lookupStringTable :: ShortText -> State StringTable StrTabOffset
+lookupStringTable str = state $ \st ->
+    case M.lookup str (stLookup st) of
+      Just off -> (off, st)
+      Nothing ->
+          let !st' = st { stStrings = stStrings st `snoc` str
+                        , stLength  = stLength st + ST.byteLength str + 1
+                        , stLookup  = M.insert str res (stLookup st)
+                        }
+              res = StrTabOffset (stLength st)
+          in (res, st')
+
+newtype DList a = DList ([a] -> [a])
+
+emptyDList :: DList a
+emptyDList = DList id
+
+snoc :: DList a -> a -> DList a
+snoc (DList f) x = DList (f . (x:))
+
+dlistToList :: DList a -> [a]
+dlistToList (DList f) = f []
diff --git a/compiler/GHC/StgToCmm/Prof.hs b/compiler/GHC/StgToCmm/Prof.hs
index 8c8f89dbe9..e07ee0a272 100644
--- a/compiler/GHC/StgToCmm/Prof.hs
+++ b/compiler/GHC/StgToCmm/Prof.hs
@@ -11,7 +11,7 @@ module GHC.StgToCmm.Prof (
         mkCCostCentre, mkCCostCentreStack,
 
         -- infoTablePRov
-        initInfoTableProv, emitInfoTableProv,
+        initInfoTableProv,
 
         -- Cost-centre Profiling
         dynProfHdr, profDynAlloc, profAlloc, staticProfHdr, initUpdFrameProf,
@@ -32,6 +32,7 @@ import GHC.Platform
 import GHC.Platform.Profile
 import GHC.StgToCmm.Closure
 import GHC.StgToCmm.Config
+import GHC.StgToCmm.InfoTableProv
 import GHC.StgToCmm.Utils
 import GHC.StgToCmm.Monad
 import GHC.StgToCmm.Lit
@@ -55,7 +56,6 @@ import GHC.Utils.Encoding
 
 import Control.Monad
 import Data.Char       (ord)
-import Data.Bifunctor  (first)
 import GHC.Utils.Monad (whenM)
 
 -----------------------------------------------------------------------------
@@ -274,9 +274,8 @@ sizeof_ccs_words platform
   where
    (ws,ms) = pc_SIZEOF_CostCentreStack (platformConstants platform) `divMod` platformWordSizeInBytes platform
 
-
+-- | Emit info-table provenance declarations
 initInfoTableProv ::  [CmmInfoTable] -> InfoTableProvMap -> FCode CStub
--- Emit the declarations
 initInfoTableProv infos itmap
   = do
        cfg <- getStgToCmmConfig
@@ -284,42 +283,16 @@ initInfoTableProv infos itmap
            info_table = stgToCmmInfoTableMap cfg
            platform   = stgToCmmPlatform     cfg
            this_mod   = stgToCmmThisModule   cfg
-       -- Output the actual IPE data
-       mapM_ emitInfoTableProv ents
-       -- Create the C stub which initialises the IPE map
-       return (ipInitCode info_table platform this_mod ents)
-
---- Info Table Prov stuff
-emitInfoTableProv :: InfoProvEnt  -> FCode ()
-emitInfoTableProv ip = do
-  { cfg <- getStgToCmmConfig
-  ; let mod      = infoProvModule ip
-        ctx      = stgToCmmContext  cfg
-        platform = stgToCmmPlatform cfg
-  ; let (src, label) = maybe ("", "") (first (renderWithContext ctx . ppr)) (infoTableProv ip)
-        mk_string    = newByteStringCLit . utf8EncodeByteString
-  ; label <- mk_string label
-  ; modl  <- newByteStringCLit (bytesFS $ moduleNameFS
-                                        $ moduleName mod)
-
-  ; ty_string  <- mk_string (infoTableType ip)
-  ; loc        <- mk_string src
-  ; table_name <- mk_string (renderWithContext ctx
-                             (pprCLabel platform CStyle (infoTablePtr ip)))
-  ; closure_type <- mk_string (renderWithContext ctx
-                               (text $ show $ infoProvEntClosureType ip))
-  ; let
-     lits = [ CmmLabel (infoTablePtr ip), -- Info table pointer
-              table_name,     -- char *table_name
-              closure_type,   -- char *closure_desc -- Filled in from the InfoTable
-              ty_string,      -- char *ty_string
-              label,          -- char *label,
-              modl,           -- char *module,
-              loc,            -- char *srcloc,
-              zero platform   -- struct _InfoProvEnt *link
-            ]
-  ; emitDataLits (mkIPELabel ip) lits
-  }
+
+       case ents of
+         [] -> return mempty
+         _  -> do
+           -- Emit IPE buffer
+           emitIpeBufferListNode this_mod ents
+
+           -- Create the C stub which initialises the IPE map
+           return (ipInitCode info_table platform this_mod)
+
 -- ---------------------------------------------------------------------------
 -- Set the current cost centre stack
 
diff --git a/compiler/ghc.cabal.in b/compiler/ghc.cabal.in
index 7ae5a9b48a..39e5e5f1ea 100644
--- a/compiler/ghc.cabal.in
+++ b/compiler/ghc.cabal.in
@@ -616,6 +616,7 @@ Library
         GHC.StgToCmm.Foreign
         GHC.StgToCmm.Heap
         GHC.StgToCmm.Hpc
+        GHC.StgToCmm.InfoTableProv
         GHC.StgToCmm.Layout
         GHC.StgToCmm.Lit
         GHC.StgToCmm.Monad