Introduce -fprof-callers flag

This introducing a new compiler flag to provide a convenient way to
introduce profiler cost-centers on all occurrences of the named
identifier.

Closes #18566.

7 files changed, 262 insertions, 3 deletions

diff --git a/compiler/GHC/Core/Lint.hs b/compiler/GHC/Core/Lint.hs
index 4dc3ec0abe..a1eae78a60 100644
--- a/compiler/GHC/Core/Lint.hs
+++ b/compiler/GHC/Core/Lint.hs
@@ -355,6 +355,7 @@ coreDumpFlag CoreTidy                 = Just Opt_D_dump_simpl
 coreDumpFlag CorePrep                 = Just Opt_D_dump_prep
 coreDumpFlag CoreOccurAnal            = Just Opt_D_dump_occur_anal
 
+coreDumpFlag CoreAddCallerCcs         = Nothing
 coreDumpFlag CoreDoPrintCore          = Nothing
 coreDumpFlag (CoreDoRuleCheck {})     = Nothing
 coreDumpFlag CoreDoNothing            = Nothing
diff --git a/compiler/GHC/Core/Opt/CallerCC.hs b/compiler/GHC/Core/Opt/CallerCC.hs
new file mode 100644
index 0000000000..1bbf96ca73
--- /dev/null
+++ b/compiler/GHC/Core/Opt/CallerCC.hs
@@ -0,0 +1,223 @@
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE DeriveDataTypeable #-}
+{-# LANGUAGE DeriveAnyClass #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE TupleSections #-}
+
+-- | Adds cost-centers to call sites selected with the @-fprof-caller=...@
+-- flag.
+module GHC.Core.Opt.CallerCC
+    ( addCallerCostCentres
+    , CallerCcFilter
+    , parseCallerCcFilter
+    ) where
+
+import Data.Bifunctor
+import Data.Word (Word8)
+import Data.Maybe
+import qualified Text.Parsec as P
+
+import Control.Applicative
+import Control.Monad.Trans.State.Strict
+import Data.Either
+import Control.Monad
+
+import GHC.Prelude
+import GHC.Utils.Outputable as Outputable
+import GHC.Driver.Session
+import GHC.Driver.Ppr
+import GHC.Types.CostCentre
+import GHC.Types.CostCentre.State
+import GHC.Types.Name hiding (varName)
+import GHC.Unit.Module.Name
+import GHC.Unit.Module.ModGuts
+import GHC.Types.SrcLoc
+import GHC.Types.Var
+import GHC.Unit.Types
+import GHC.Data.FastString
+import GHC.Core
+import GHC.Core.Opt.Monad
+import GHC.Utils.Panic
+import qualified GHC.Utils.Binary as B
+
+addCallerCostCentres :: ModGuts -> CoreM ModGuts
+addCallerCostCentres guts = do
+  dflags <- getDynFlags
+  let filters = callerCcFilters dflags
+  let env :: Env
+      env = Env
+        { thisModule = mg_module guts
+        , ccState = newCostCentreState
+        , dflags = dflags
+        , revParents = []
+        , filters = filters
+        }
+  let guts' = guts { mg_binds = doCoreProgram env (mg_binds guts)
+                   }
+  return guts'
+
+doCoreProgram :: Env -> CoreProgram -> CoreProgram
+doCoreProgram env binds = flip evalState newCostCentreState $ do
+    mapM (doBind env) binds
+
+doBind :: Env -> CoreBind -> M CoreBind
+doBind env (NonRec b rhs) = NonRec b <$> doExpr (addParent b env) rhs
+doBind env (Rec bs) = Rec <$> mapM doPair bs
+  where
+    doPair (b,rhs) = (b,) <$> doExpr (addParent b env) rhs
+
+doExpr :: Env -> CoreExpr -> M CoreExpr
+doExpr env e@(Var v)
+  | needsCallSiteCostCentre env v = do
+    let nameDoc :: SDoc
+        nameDoc = withUserStyle alwaysQualify DefaultDepth $
+          hcat (punctuate dot (map ppr (parents env))) <> parens (text "calling:" <> ppr v)
+
+        ccName :: CcName
+        ccName = mkFastString $ showSDoc (dflags env) nameDoc
+    ccIdx <- getCCIndex' ccName
+    let span = case revParents env of
+          top:_ -> nameSrcSpan $ varName top
+          _     -> noSrcSpan
+        cc = NormalCC (ExprCC ccIdx) ccName (thisModule env) span
+        tick :: Tickish Id
+        tick = ProfNote cc True True
+    pure $ Tick tick e
+  | otherwise = pure e
+doExpr _env e@(Lit _)       = pure e
+doExpr env (f `App` x)      = App <$> doExpr env f <*> doExpr env x
+doExpr env (Lam b x)        = Lam b <$> doExpr env x
+doExpr env (Let b rhs)      = Let <$> doBind env b <*> doExpr env rhs
+doExpr env (Case scrut b ty alts) =
+    Case <$> doExpr env scrut <*> pure b <*> pure ty <*> mapM doAlt alts
+  where
+    doAlt (con, bs, rhs) = (con, bs,) <$> doExpr env rhs
+doExpr env (Cast expr co)   = Cast <$> doExpr env expr <*> pure co
+doExpr env (Tick t e)       = Tick t <$> doExpr env e
+doExpr _env e@(Type _)      = pure e
+doExpr _env e@(Coercion _)  = pure e
+
+type M = State CostCentreState
+
+getCCIndex' :: FastString -> M CostCentreIndex
+getCCIndex' name = state (getCCIndex name)
+
+data Env = Env
+  { thisModule  :: Module
+  , dflags      :: DynFlags
+  , ccState     :: CostCentreState
+  , revParents  :: [Id]
+  , filters     :: [CallerCcFilter]
+  }
+
+addParent :: Id -> Env -> Env
+addParent i env = env { revParents = i : revParents env }
+
+parents :: Env -> [Id]
+parents env = reverse (revParents env)
+
+needsCallSiteCostCentre :: Env -> Id -> Bool
+needsCallSiteCostCentre env i =
+    any matches (filters env)
+  where
+    matches :: CallerCcFilter -> Bool
+    matches ccf =
+        checkModule && checkFunc
+      where
+        checkModule =
+          case ccfModuleName ccf of
+            Just modFilt
+              | Just iMod <- nameModule_maybe (varName i)
+              -> moduleName iMod == modFilt
+              | otherwise -> False
+            Nothing -> True
+        checkFunc =
+            occNameMatches (ccfFuncName ccf) (getOccName i)
+
+data NamePattern
+    = PChar Char NamePattern
+    | PWildcard NamePattern
+    | PEnd
+
+instance Outputable NamePattern where
+  ppr (PChar c rest) = char c <> ppr rest
+  ppr (PWildcard rest) = char '*' <> ppr rest
+  ppr PEnd = Outputable.empty
+
+instance B.Binary NamePattern where
+  get bh = do
+    tag <- B.get bh
+    case tag :: Word8 of
+      0 -> PChar <$> B.get bh <*> B.get bh
+      1 -> PWildcard <$> B.get bh
+      2 -> pure PEnd
+      _ -> panic "Binary(NamePattern): Invalid tag"
+  put_ bh (PChar x y) = B.put_ bh (0 :: Word8) >> B.put_ bh x >> B.put_ bh y
+  put_ bh (PWildcard x) = B.put_ bh (1 :: Word8) >> B.put_ bh x
+  put_ bh PEnd = B.put_ bh (2 :: Word8)
+
+occNameMatches :: NamePattern -> OccName -> Bool
+occNameMatches pat = go pat . occNameString
+  where
+    go :: NamePattern -> String -> Bool
+    go PEnd "" = True
+    go (PChar c rest) (d:s)
+      = d == c && go rest s
+    go (PWildcard rest) s
+      = go rest s || go (PWildcard rest) (tail s)
+    go _ _  = False
+
+type Parser = P.Parsec String ()
+
+parseNamePattern :: Parser NamePattern
+parseNamePattern = pattern
+  where
+    pattern = star <|> wildcard <|> char <|> end
+    star = PChar '*' <$ P.string "\\*" <*> pattern
+    wildcard = do
+      void $ P.char '*'
+      PWildcard <$> pattern
+    char = PChar <$> P.anyChar <*> pattern
+    end = PEnd <$ P.eof
+
+data CallerCcFilter
+    = CallerCcFilter { ccfModuleName  :: Maybe ModuleName
+                     , ccfFuncName    :: NamePattern
+                     }
+
+instance Outputable CallerCcFilter where
+  ppr ccf =
+    maybe (char '*') ppr (ccfModuleName ccf)
+    <> char '.'
+    <> ppr (ccfFuncName ccf)
+
+instance B.Binary CallerCcFilter where
+  get bh = CallerCcFilter <$> B.get bh <*> B.get bh
+  put_ bh (CallerCcFilter x y) = B.put_ bh x >> B.put_ bh y
+
+parseCallerCcFilter :: String -> Either String CallerCcFilter
+parseCallerCcFilter =
+    first show . P.parse parseCallerCcFilter' "caller-CC filter"
+
+parseCallerCcFilter' :: Parser CallerCcFilter
+parseCallerCcFilter' =
+  CallerCcFilter
+    <$> moduleFilter
+    <*  P.char '.'
+    <*> parseNamePattern
+  where
+    moduleFilter :: Parser (Maybe ModuleName)
+    moduleFilter =
+      (Just . mkModuleName <$> moduleName)
+      <|>
+      (Nothing <$ P.char '*')
+
+    moduleName :: Parser String
+    moduleName = do
+      c <- P.upper
+      cs <- some $ P.upper <|> P.lower <|> P.digit <|> P.oneOf "_"
+      rest <- optional $ P.try $ P.char '.' >> fmap ('.':) moduleName
+      return $ c : (cs ++ fromMaybe "" rest)
+
diff --git a/compiler/GHC/Core/Opt/CallerCC.hs-boot b/compiler/GHC/Core/Opt/CallerCC.hs-boot
new file mode 100644
index 0000000000..9d367916e1
--- /dev/null
+++ b/compiler/GHC/Core/Opt/CallerCC.hs-boot
@@ -0,0 +1,8 @@
+module GHC.Core.Opt.CallerCC where
+
+import GHC.Prelude
+
+-- Necessary due to import in GHC.Driver.Session.
+data CallerCcFilter
+
+parseCallerCcFilter :: String -> Either String CallerCcFilter
diff --git a/compiler/GHC/Core/Opt/Monad.hs b/compiler/GHC/Core/Opt/Monad.hs
index bdacfba90b..21b4403e94 100644
--- a/compiler/GHC/Core/Opt/Monad.hs
+++ b/compiler/GHC/Core/Opt/Monad.hs
@@ -129,6 +129,7 @@ data CoreToDo           -- These are diff core-to-core passes,
 
   | CoreTidy
   | CorePrep
+  | CoreAddCallerCcs
   | CoreOccurAnal
 
 instance Outputable CoreToDo where
@@ -149,6 +150,7 @@ instance Outputable CoreToDo where
   ppr CoreDesugar              = text "Desugar (before optimization)"
   ppr CoreDesugarOpt           = text "Desugar (after optimization)"
   ppr CoreTidy                 = text "Tidy Core"
+  ppr CoreAddCallerCcs         = text "Add caller cost-centres"
   ppr CorePrep                 = text "CorePrep"
   ppr CoreOccurAnal            = text "Occurrence analysis"
   ppr CoreDoPrintCore          = text "Print core"
diff --git a/compiler/GHC/Core/Opt/Pipeline.hs b/compiler/GHC/Core/Opt/Pipeline.hs
index a090bdfe62..14fb11bccc 100644
--- a/compiler/GHC/Core/Opt/Pipeline.hs
+++ b/compiler/GHC/Core/Opt/Pipeline.hs
@@ -16,6 +16,7 @@ import GHC.Driver.Session
 import GHC.Driver.Ppr
 import GHC.Driver.Plugins ( withPlugins, installCoreToDos )
 import GHC.Driver.Env
+import GHC.Platform.Ways  ( hasWay, Way(WayProf) )
 
 import GHC.Core
 import GHC.Core.Opt.CSE  ( cseProgram )
@@ -44,6 +45,7 @@ import GHC.Core.Opt.CprAnal      ( cprAnalProgram )
 import GHC.Core.Opt.CallArity    ( callArityAnalProgram )
 import GHC.Core.Opt.Exitify      ( exitifyProgram )
 import GHC.Core.Opt.WorkWrap     ( wwTopBinds )
+import GHC.Core.Opt.CallerCC     ( addCallerCostCentres )
 import GHC.Core.Seq (seqBinds)
 import GHC.Core.FamInstEnv
 
@@ -155,6 +157,7 @@ getCoreToDo dflags
     pre_inline_on = gopt Opt_SimplPreInlining             dflags
     ww_on         = gopt Opt_WorkerWrapper                dflags
     static_ptrs   = xopt LangExt.StaticPointers           dflags
+    profiling     = ways dflags `hasWay` WayProf
 
     maybe_rule_check phase = runMaybe rule_check (CoreDoRuleCheck phase)
 
@@ -221,12 +224,16 @@ getCoreToDo dflags
           }
         ]
 
+    add_caller_ccs =
+        runWhen (profiling && not (null $ callerCcFilters dflags)) CoreAddCallerCcs
+
     core_todo =
      if opt_level == 0 then
        [ static_ptrs_float_outwards,
          CoreDoSimplify max_iter
              (base_mode { sm_phase = FinalPhase
                         , sm_names = ["Non-opt simplification"] })
+       , add_caller_ccs
        ]
 
      else {- opt_level >= 1 -} [
@@ -370,7 +377,9 @@ getCoreToDo dflags
         -- can become /exponentially/ more expensive. See #11731, #12996.
         runWhen (strictness || late_dmd_anal) CoreDoDemand,
 
-        maybe_rule_check FinalPhase
+        maybe_rule_check FinalPhase,
+
+        add_caller_ccs
      ]
 
     -- Remove 'CoreDoNothing' and flatten 'CoreDoPasses' for clarity.
@@ -509,6 +518,9 @@ doCorePass CoreDoSpecialising        = {-# SCC "Specialise" #-}
 doCorePass CoreDoSpecConstr          = {-# SCC "SpecConstr" #-}
                                        specConstrProgram
 
+doCorePass CoreAddCallerCcs          = {-# SCC "AddCallerCcs" #-}
+                                       addCallerCostCentres
+
 doCorePass CoreDoPrintCore              = observe   printCore
 doCorePass (CoreDoRuleCheck phase pat)  = ruleCheckPass phase pat
 doCorePass CoreDoNothing                = return
diff --git a/compiler/GHC/Driver/Session.hs b/compiler/GHC/Driver/Session.hs
index 687c6e1598..2f26276eb7 100644
--- a/compiler/GHC/Driver/Session.hs
+++ b/compiler/GHC/Driver/Session.hs
@@ -263,6 +263,7 @@ import GHC.Utils.Fingerprint
 import GHC.Utils.Outputable
 import GHC.Settings
 import GHC.CmmToAsm.CFG.Weight
+import {-# SOURCE #-} GHC.Core.Opt.CallerCC
 
 import GHC.Types.Error
 import {-# SOURCE #-} GHC.Utils.Error
@@ -699,6 +700,7 @@ data DynFlags = DynFlags {
 
   -- | what kind of {-# SCC #-} to add automatically
   profAuto              :: ProfAuto,
+  callerCcFilters       :: [CallerCcFilter],
 
   interactivePrint      :: Maybe String,
 
@@ -1313,6 +1315,7 @@ defaultDynFlags mySettings llvmConfig =
         canUseColor = False,
         colScheme = Col.defaultScheme,
         profAuto = NoProfAuto,
+        callerCcFilters = [],
         interactivePrint = Nothing,
         nextWrapperNum = panic "defaultDynFlags: No nextWrapperNum",
         sseVersion = Nothing,
@@ -2947,6 +2950,10 @@ dynamic_flags_deps = [
   , make_ord_flag defGhcFlag "fno-prof-auto"
       (noArg (\d -> d { profAuto = NoProfAuto } ))
 
+        -- Caller-CC
+  , make_ord_flag defGhcFlag "fprof-callers"
+         (HasArg setCallerCcFilters)
+
         ------ Compiler flags -----------------------------------------------
 
   , make_ord_flag defGhcFlag "fasm"             (NoArg (setObjBackend NCG))
@@ -4548,6 +4555,12 @@ checkOptLevel n dflags
    | otherwise
      = Right dflags
 
+setCallerCcFilters :: String -> DynP ()
+setCallerCcFilters arg =
+  case parseCallerCcFilter arg of
+    Right filt -> upd $ \d -> d { callerCcFilters = filt : callerCcFilters d }
+    Left err -> addErr err
+
 setMainIs :: String -> DynP ()
 setMainIs arg
   | not (null main_fn) && isLower (head main_fn)
diff --git a/compiler/GHC/Iface/Recomp/Flags.hs b/compiler/GHC/Iface/Recomp/Flags.hs
index e73c061018..1c52f4e326 100644
--- a/compiler/GHC/Iface/Recomp/Flags.hs
+++ b/compiler/GHC/Iface/Recomp/Flags.hs
@@ -17,7 +17,7 @@ import GHC.Types.Name
 import GHC.Types.SafeHaskell
 import GHC.Utils.Fingerprint
 import GHC.Iface.Recomp.Binary
--- import GHC.Utils.Outputable
+import GHC.Core.Opt.CallerCC () -- for Binary instances
 
 import GHC.Data.EnumSet as EnumSet
 import System.FilePath (normalise)
@@ -61,7 +61,7 @@ fingerprintDynFlags dflags@DynFlags{..} this_mod nameio =
         ticky =
           map (`gopt` dflags) [Opt_Ticky, Opt_Ticky_Allocd, Opt_Ticky_LNE, Opt_Ticky_Dyn_Thunk]
 
-        flags = ((mainis, safeHs, lang, cpp), (paths, prof, ticky, debugLevel))
+        flags = ((mainis, safeHs, lang, cpp), (paths, prof, ticky, debugLevel, callerCcFilters))
 
     in -- pprTrace "flags" (ppr flags) $
        computeFingerprint nameio flags