GitLab

Simon Jakobi pushed to branch wip/sjakobi/T27507 at Glasgow Haskell Compiler / GHC

Commits:

24bccaee

by Simon Jakobi at 2026-07-17T13:04:46+02:00

Infer EPT signatures for local functions

Extend parameter tag inference to local let-bound functions, with conservative handling for partial applications and escaping uses. See Note [EPT signatures for local function arguments].

Add regression coverage for #27507.

6 changed files:

compiler/GHC/Stg/EnforceEpt.hs
compiler/GHC/Stg/EnforceEpt/Types.hs
testsuite/tests/simplStg/should_run/T27507a.hs
+ testsuite/tests/simplStg/should_run/T27507b.hs
+ testsuite/tests/simplStg/should_run/T27507b.stdout
testsuite/tests/simplStg/should_run/all.T

Changes:

compiler/GHC/Stg/EnforceEpt.hs

@@ -452,65 +452,12 @@ inferTagExpr _ (StgOpApp op args ty)
  inferTagExpr env (StgLet ext bind body)
    = (info, StgLet ext bind' body')
    where
 -    (env', bind') = inferTagBind env bind
 -    (info, body') = inferTagExpr env' body
 +    (info, bind', body') = inferTagLet env (localFunsOfBind env False bind) bind body
  inferTagExpr env (StgLetNoEscape ext bind body)
 -  | all (isJust . lookupJoinArgInfo env . fst) root_joins
 -  = (plain_info, StgLetNoEscape ext plain_bind plain_body)
 -  | otherwise
 -  = go initial_arg_infos
 +  = (info, StgLetNoEscape ext bind' body')
    where
 -    (plain_body_env, plain_bind) = inferTagBind env bind
 -    (plain_info, plain_body) = inferTagExpr plain_body_env body
+-
 -    joins = root_joins ++ joinsInBind bind ++ joinsInExpr body
 -    join_ids = map fst joins
 -    initial_arg_infos = [replicate arity TagEPT | (_, arity) <- joins]
+-
 -    go arg_infos
 -      | arg_infos == new_arg_infos
 -      = (info, StgLetNoEscape ext bind' body')
 -      | otherwise
 -      = go new_arg_infos
 -      where
 -        join_env = extendJoinArgEnv env (zipEqual join_ids arg_infos)
 -        (body_env, bind') = inferTagBind join_env bind
 -        (info, body') = inferTagExpr body_env body
 -        calls = collectJoinCalls (makeTagged join_env) join_ids bind' body'
 -        new_arg_infos =
 -          [ maybe initial_infos (combineArgs initial_infos) (lookupVarEnv calls join_id)
 -          | (join_id, initial_infos) <- zipEqual join_ids initial_arg_infos ]
+-
 -    combineArgs = zipWithEqual combineAltInfo
+-
 -    root_joins = joinsOf bind
+-
 -    joinsOf (StgNonRec bndr rhs) = [joinOf bndr rhs]
 -    joinsOf (StgRec pairs) = [joinOf bndr rhs | (bndr, rhs) <- pairs]
+-
 -    joinOf bndr (StgRhsClosure _ _ _ bndrs _ _)
 -      = (getBinderId env bndr, length bndrs)
 -    joinOf bndr (StgRhsCon {})
 -      = (getBinderId env bndr, idArity (getBinderId env bndr))
+-
 -    joinsInBind (StgNonRec _ rhs) = joinsInRhs rhs
 -    joinsInBind (StgRec pairs) = concatMap (joinsInRhs . snd) pairs
+-
 -    joinsInRhs (StgRhsClosure _ _ _ _ rhs _) = joinsInExpr rhs
 -    joinsInRhs (StgRhsCon {}) = []
+-
 -    joinsInExpr (StgApp {}) = []
 -    joinsInExpr (StgConApp {}) = []
 -    joinsInExpr (StgLit {}) = []
 -    joinsInExpr (StgTick _ expr) = joinsInExpr expr
 -    joinsInExpr (StgOpApp {}) = []
 -    joinsInExpr (StgLet _ let_bind expr)
 -      = joinsInBind let_bind ++ joinsInExpr expr
 -    joinsInExpr (StgLetNoEscape _ let_bind expr)
 -      = joinsOf let_bind ++ joinsInBind let_bind ++ joinsInExpr expr
 -    joinsInExpr (StgCase scrut _ _ alts)
 -      = joinsInExpr scrut ++ concatMap (joinsInExpr . alt_rhs) alts
 +    (info, bind', body') = inferTagLet env (localFunsOfBind env True bind) bind body
  inferTagExpr in_env (StgCase scrut bndr ty alts)
    -- Unboxed tuples get their info from the expression we scrutinise if any
@@ -564,28 +511,108 @@ inferTagExpr in_env (StgCase scrut bndr ty alts)
      (scrut_info, scrut') = inferTagExpr in_env scrut
      bndr' = (getBinderId in_env bndr, TagVal TagEPT)
 --- See Note [EPT signatures for join point arguments].
 -collectJoinCalls
 +localFunsOfBind
 +  :: TagEnv p -> Bool -> GenStgBinding p -> [(Id, Int)]
 +localFunsOfBind env is_join (StgNonRec bndr rhs)
 +  = localFunOf env is_join bndr rhs
 +localFunsOfBind env is_join (StgRec pairs)
 +  = concatMap (uncurry (localFunOf env is_join)) pairs
++
 +localFunOf
 +  :: TagEnv p -> Bool -> BinderP p -> GenStgRhs p -> [(Id, Int)]
 +localFunOf env is_join bndr (StgRhsClosure _ _ _ bndrs _ _)
 +  | is_join || notNull bndrs
 +  = [(getBinderId env bndr, length bndrs)]
 +localFunOf _ _ _ (StgRhsClosure {}) = []
 +localFunOf env is_join bndr (StgRhsCon {})
 +  | is_join
 +  = [(getBinderId env bndr, idArity (getBinderId env bndr))]
 +localFunOf _ _ _ (StgRhsCon {}) = []
++
 +-- See Note [EPT signatures for local function arguments].
 +inferTagLet
 +  :: forall p. (OutputableInferPass p, InferExtEq p)
 +  => TagEnv p
 +  -> [(Id, Int)]
 +  -> GenStgBinding p
 +  -> GenStgExpr p
 +  -> (TagInfo, InferStgBinding, InferStgExpr)
 +inferTagLet env root_funs bind body
 +  | null root_funs
 +  = (plain_info, plain_bind, plain_body)
 +  | all (isJust . lookupFunArgInfo env . fst) root_funs
 +  = (plain_info, plain_bind, plain_body)
 +  | otherwise
 +  = go initial_arg_infos
 +  where
 +    (plain_body_env, plain_bind) = inferTagBind env bind
 +    (plain_info, plain_body) = inferTagExpr plain_body_env body
++
 +    funs = root_funs ++ funsInBind bind ++ funsInExpr body
 +    fun_ids = map fst funs
 +    initial_arg_infos = [replicate arity TagEPT | (_, arity) <- funs]
++
 +    go arg_infos
 +      | arg_infos == new_arg_infos
 +      = (info, bind', body')
 +      | otherwise
 +      = go new_arg_infos
 +      where
 +        fun_env = extendFunArgEnv env (zipEqual fun_ids arg_infos)
 +        (body_env, bind') = inferTagBind fun_env bind
 +        (info, body') = inferTagExpr body_env body
 +        calls = collectFunCalls (makeTagged fun_env) funs bind' body'
 +        new_arg_infos =
 +          [ maybe initial_infos (combineArgs initial_infos) (lookupVarEnv calls fun_id)
 +          | (fun_id, initial_infos) <- zipEqual fun_ids initial_arg_infos ]
++
 +    combineArgs = zipWithEqual combineAltInfo
++
 +    funsInBind (StgNonRec _ rhs) = funsInRhs rhs
 +    funsInBind (StgRec pairs) = concatMap (funsInRhs . snd) pairs
++
 +    funsInRhs (StgRhsClosure _ _ _ _ rhs _) = funsInExpr rhs
 +    funsInRhs (StgRhsCon {}) = []
++
 +    funsInExpr (StgApp {}) = []
 +    funsInExpr (StgConApp {}) = []
 +    funsInExpr (StgLit {}) = []
 +    funsInExpr (StgTick _ expr) = funsInExpr expr
 +    funsInExpr (StgOpApp {}) = []
 +    funsInExpr (StgLet _ let_bind expr)
 +      = localFunsOfBind env False let_bind ++ funsInBind let_bind ++ funsInExpr expr
 +    funsInExpr (StgLetNoEscape _ let_bind expr)
 +      = localFunsOfBind env True let_bind ++ funsInBind let_bind ++ funsInExpr expr
 +    funsInExpr (StgCase scrut _ _ alts)
 +      = funsInExpr scrut ++ concatMap (funsInExpr . alt_rhs) alts
++
 +-- See Note [EPT signatures for local function arguments].
 +collectFunCalls
    :: TagEnv 'InferTaggedBinders
 -  -> [Id]
 +  -> [(Id, Int)]
    -> InferStgBinding
    -> InferStgExpr
    -> IdEnv [TagInfo]
 -collectJoinCalls in_env join_ids bind body
 +collectFunCalls in_env funs bind body
    = bind_calls `plusCalls` collectExpr body_env body
    where
 -    join_env = mkVarEnv [(join_id, ()) | join_id <- join_ids]
 +    fun_env = mkVarEnv funs
      (bind_calls, body_env) = collectBind in_env bind
      collectExpr env (StgApp fun args)
 -      | elemVarEnv fun join_env
 -      = unitVarEnv fun (map (lookupInfo env) args)
 -      | otherwise
 -      = emptyVarEnv
 -    collectExpr _ (StgConApp {}) = emptyVarEnv
 +      = direct_call `plusCalls` collectArgs args
 +      where
 +        direct_call
 +          | Just arity <- lookupVarEnv fun_env fun
 +          = unitVarEnv fun
 +              (take arity (map (lookupInfo env) args) ++
 +               replicate (arity - length args) TagDunno)
 +          | otherwise
 +          = emptyVarEnv
 +    collectExpr _ (StgConApp _ _ args _) = collectArgs args
      collectExpr _ (StgLit {}) = emptyVarEnv
      collectExpr env (StgTick _ expr) = collectExpr env expr
 -    collectExpr _ (StgOpApp {}) = emptyVarEnv
 +    collectExpr _ (StgOpApp _ args _) = collectArgs args
      collectExpr env (StgLet _ let_bind expr)
        = let_calls `plusCalls` collectExpr let_env expr
        where
@@ -610,39 +637,57 @@ collectJoinCalls in_env join_ids bind body
      collectRhs env (StgRhsClosure _ _ _ bndrs rhs _)
        = collectExpr (extendSigEnv env bndrs) rhs
 -    collectRhs _ (StgRhsCon {}) = emptyVarEnv
 +    collectRhs _ (StgRhsCon _ _ _ _ args _) = collectArgs args
++
 +    collectArgs = plusCallList . map collectArg
++
 +    collectArg (StgVarArg var)
 +      | Just arity <- lookupVarEnv fun_env var
 +      = unitVarEnv var (replicate arity TagDunno)
 +    collectArg _ = emptyVarEnv
      plusCalls = plusVarEnv_C (zipWithEqual combineAltInfo)
      plusCallList = foldr plusCalls emptyVarEnv
 -{- Note [EPT signatures for join point arguments]
 -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 -Join points are local, non-escaping functions, and every occurrence is a
 -saturated tail call.  Consequently we can infer EPT information for their
 -arguments by treating them like SSA block parameters: an argument is EPT when
 -the corresponding actual argument is EPT at every jump to the join point.
+-
 -At each outermost StgLetNoEscape we optimistically start every argument of that
 -join group and all nested join groups at TagEPT, infer tags for the binding and
 -its body, and collect the argument tags at all calls to those join points.
 +{- Note [EPT signatures for local function arguments]
 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 +For local functions we can inspect every use and infer that a parameter is EPT
 +when the corresponding actual argument is EPT at every call.  Join points are
 +the particularly simple case: they never escape and every occurrence is a
 +saturated tail call, so their parameters behave like SSA block parameters.
++
 +Ordinary let-bound functions can escape or be applied to too few arguments.
 +We therefore track their parameters independently.  A direct call contributes
 +the tag of each supplied argument; an undersaturated call contributes TagDunno
 +for the missing suffix.  If the function occurs as a value (for example as an
 +argument or constructor field), every parameter becomes TagDunno because a
 +later call is outside the analysis.  Arguments beyond the function's arity do
 +not describe its parameters and are ignored.  We still inspect every argument
 +for escaping occurrences of other tracked functions.
++
 +At each outermost local function binding we optimistically start every argument
 +of that function group and all nested function and join groups at TagEPT, infer
 +tags for the binding and its body, and collect the argument tags at all uses.
  Combining those call-site tags gives the next, no more optimistic approximation
 -for the join arguments.  Iterating reaches the greatest fixed point, including
 -for recursive and mutually recursive joins.  Using the greatest fixed point is
 -important for recursive calls such as @jump j x@, where @x@ is itself a
 -parameter of @j@: such a back edge preserves EPTness rather than providing an
 -independent reason to reject it.
+-
 -Solving all nested join groups simultaneously is important for compile-time
 -performance.  Independently solving a nested join group on every iteration of
 -each enclosing group causes exponential re-analysis in deeply nested code.
 -Nested StgLetNoEscapes therefore merely use the argument information installed
 -by the outer solver.
 +for the function arguments.  Iterating reaches the greatest fixed point,
 +including for recursive and mutually recursive functions.  Using the greatest
 +fixed point is important for recursive calls such as @jump j x@, where @x@ is
 +itself a parameter of @j@: such a back edge preserves EPTness rather than
 +providing an independent reason to reject it.
++
 +Solving all nested function and join groups simultaneously is important for
 +compile-time performance.  Independently solving a nested group on every
 +iteration of each enclosing group causes exponential re-analysis in deeply
 +nested code.
 +Nested bindings therefore merely use the argument information installed by the
 +outer solver.  A dead local function has no call-site evidence against the
 +optimistic approximation, which is sound because it cannot be entered.
  This is tag inference, not strictness inference.  It does not evaluate an
  argument or change the calling convention.  It merely records that every path
 -which enters the join point already supplies an EPT value.  Existing CBV marks
 -remain authoritative in inferTagRhs, since the rewriter enforces their EPT
 -precondition at call sites.
 +which enters the local function already supplies an EPT value.  Existing CBV
 +marks remain authoritative in inferTagRhs, since the rewriter enforces their
 +EPT precondition at call sites.
  -}
  -- Compute binder sigs based on the constructors strict fields.
@@ -787,7 +832,7 @@ inferTagRhs bnd_id in_env (StgRhsClosure ext cc upd bndrs body typ)
        = repeat NotMarkedCbv
      inferred_arg_infos =
 -      fromMaybe (repeat TagDunno) (lookupJoinArgInfo in_env bnd_id)
 +      fromMaybe (repeat TagDunno) (lookupFunArgInfo in_env bnd_id)
      env' = extendSigEnv in_env out_bndrs
      (info, body') = inferTagExpr env' body

compiler/GHC/Stg/EnforceEpt/Types.hs

@@ -54,9 +54,9 @@ combineAltInfo (TagTuple {})   TagEPT          = TagDunno
  combineAltInfo (TagTuple is1)  (TagTuple is2)  = TagTuple (zipWithEqual combineAltInfo is1 is2)
  type TagSigEnv = IdEnv TagSig
 -type JoinArgEnv = IdEnv [TagInfo]
 +type FunArgEnv = IdEnv [TagInfo]
  data TagEnv p = TE { te_env :: TagSigEnv
 -                   , te_join_args :: JoinArgEnv
 +                   , te_fun_args :: FunArgEnv
                     , te_get :: BinderP p -> Id
                     , te_bytecode :: !Bool
+                    }
@@ -73,7 +73,7 @@ getBinderId = te_get
  initEnv :: Bool -> TagEnv 'CodeGen
  initEnv for_bytecode = TE { te_env = emptyVarEnv
 -             , te_join_args = emptyVarEnv
 +             , te_fun_args = emptyVarEnv
               , te_get = \x -> x
               , te_bytecode = for_bytecode }
@@ -81,7 +81,7 @@ initEnv for_bytecode = TE { te_env = emptyVarEnv
  -- with no other changes.
  makeTagged :: TagEnv p -> TagEnv 'InferTaggedBinders
  makeTagged env = TE { te_env = te_env env
 -                    , te_join_args = te_join_args env
 +                    , te_fun_args = te_fun_args env
                      , te_get = fst
                      , te_bytecode = te_bytecode env }
@@ -159,9 +159,9 @@ extendSigEnv :: TagEnv p -> [(Id,TagSig)] -> TagEnv p
  extendSigEnv env@(TE { te_env = sig_env }) bndrs
    = env { te_env = extendVarEnvList sig_env bndrs }
 -lookupJoinArgInfo :: TagEnv p -> Id -> Maybe [TagInfo]
 -lookupJoinArgInfo env join_id = lookupVarEnv (te_join_args env) join_id
 +lookupFunArgInfo :: TagEnv p -> Id -> Maybe [TagInfo]
 +lookupFunArgInfo env fun_id = lookupVarEnv (te_fun_args env) fun_id
 -extendJoinArgEnv :: TagEnv p -> [(Id, [TagInfo])] -> TagEnv p
 -extendJoinArgEnv env@(TE { te_join_args = join_env }) joins
 -  = env { te_join_args = extendVarEnvList join_env joins }
 +extendFunArgEnv :: TagEnv p -> [(Id, [TagInfo])] -> TagEnv p
 +extendFunArgEnv env@(TE { te_fun_args = fun_env }) funs
 +  = env { te_fun_args = extendVarEnvList fun_env funs }

testsuite/tests/simplStg/should_run/T27507a.hs

@@ -12,10 +12,7 @@ data Box a = Box !a
  foo :: Bool -> Int -> (Int, Box Int)
  foo b !x = case x of
 -  x' -> let -- Keep fun_exit monomorphic so that it can become a join point.
 -            -- Without this signature its polymorphic result type prevents contification.
 -            fun_exit :: Int -> Int -> (Int, Box Int)
 -            fun_exit x_f i = (i, Box x_f)
 +  x' -> let fun_exit x_f i = (i, Box x_f)
              {-# NOINLINE fun_exit #-}
          in if b then fun_exit x' 0 else fun_exit x' 1
  {-# NOINLINE foo #-}

testsuite/tests/simplStg/should_run/T27507b.hs

 +{-# LANGUAGE BangPatterns #-}
++
 +-- Keep the local functions local and their boxed results visible.
 +{-# OPTIONS_GHC -fno-cpr-anal -fno-full-laziness #-}
++
 +module Main where
++
 +import GHC.Exts (noinline)
 +import GHC.Exts.Heap (GenClosure(..), getClosureData)
 +import System.Exit (exitFailure)
++
 +data Box a = Box !a
++
 +recursive :: Int -> Int -> (Int, Box Int)
 +recursive !x n =
 +  let go y i
 +        | i == 0 = (i, Box y)
 +        | otherwise = go y (i - 1)
 +      {-# NOINLINE go #-}
 +  in go x n
 +{-# NOINLINE recursive #-}
++
 +mixed :: Bool -> Int -> Int -> (Box Int, Box Int)
 +mixed b !x z =
 +  let fun a c = (Box a, Box c)
 +      {-# NOINLINE fun #-}
 +  in if b then fun x x else fun x z
 +{-# NOINLINE mixed #-}
++
 +apply :: (a -> b) -> a -> b
 +apply fun x = fun x
 +{-# NOINLINE apply #-}
++
 +partial :: Int -> Int -> (Box Int, Box Int)
 +partial !x z =
 +  let fun a c = (Box a, Box c)
 +      {-# NOINLINE fun #-}
 +      pap = fun x
 +  in apply pap z
 +{-# NOINLINE partial #-}
++
 +escaping :: Int -> (Int, Box Int)
 +escaping !x =
 +  let fun a = (0, Box a)
 +      {-# NOINLINE fun #-}
 +  in case noinline Just fun of
 +       Just escaped -> escaped x
 +       Nothing -> (0, Box x)
 +{-# NOINLINE escaping #-}
++
 +checkConstr :: String -> a -> IO ()
 +checkConstr label value = do
 +  closure <- getClosureData value
 +  case closure of
 +    ConstrClosure{} -> pure ()
 +    _ -> putStrLn ("FAIL: " ++ label ++ " was not a constructor") >> exitFailure
++
 +checkThunk :: String -> a -> IO ()
 +checkThunk label value = do
 +  closure <- getClosureData value
 +  case closure of
 +    ThunkClosure{} -> pure ()
 +    _ -> putStrLn ("FAIL: " ++ label ++ " was not a thunk") >> exitFailure
++
 +main :: IO ()
 +main = do
 +  case recursive 42 1 of
 +    (_, value) -> checkConstr "recursive local function" value
++
 +  case mixed False 42 undefined of
 +    (known, unknown) -> do
 +      checkConstr "mixed known parameter" known
 +      checkThunk "mixed unknown parameter" unknown
++
 +  case partial 42 undefined of
 +    (known, unknown) -> do
 +      checkConstr "partial known prefix" known
 +      checkThunk "partial unknown suffix" unknown
++
 +  case escaping 42 of
 +    (_, value) -> checkThunk "escaping local function" value
 +  putStrLn "OK"

testsuite/tests/simplStg/should_run/T27507b.stdout

+OK

testsuite/tests/simplStg/should_run/all.T

@@ -25,3 +25,4 @@ test('unpack_enum', normal, compile_and_run, [''])
  test('T27005a', normal, compile_and_run, [''])
  test('T27507', normal, compile_and_run, [''])
  test('T27507a', normal, compile_and_run, [''])
 +test('T27507b', normal, compile_and_run, [''])