[Git][ghc/ghc][wip/T23083] CorePrep: Eta expand arguments (#23083)

Sebastian Graf (@sgraf812) gitlab at gitlab.haskell.org
Wed Apr 26 15:34:56 UTC 2023



Sebastian Graf pushed to branch wip/T23083 at Glasgow Haskell Compiler / GHC


Commits:
e721cc6f by Sebastian Graf at 2023-04-26T17:34:47+02:00
CorePrep: Eta expand arguments (#23083)

Previously, we'd only eta expand let bindings and lambdas,
now we'll also eta expand arguments such as in T23083:
```hs
g f h = f (h `seq` (h $))
```
Unless `-fpedantic-bottoms` is set, we'll now transform to
```hs
g f h = f (\eta -> h eta)
```
in CorePrep.

See the new `Note [Eta expansion of arguments in CorePrep]` for the details.

We only do this optimisation with -O2 because we saw 2-3% ghc/alloc regressions
in T4801 and T5321FD.

Fixes #23083.

- - - - -


10 changed files:

- compiler/GHC/Core/Utils.hs
- compiler/GHC/CoreToStg.hs
- compiler/GHC/CoreToStg/Prep.hs
- compiler/GHC/Driver/Config/CoreToStg/Prep.hs
- compiler/GHC/Driver/Flags.hs
- compiler/GHC/Driver/Session.hs
- docs/users_guide/using-optimisation.rst
- + testsuite/tests/simplCore/should_compile/T23083.hs
- + testsuite/tests/simplCore/should_compile/T23083.stderr
- testsuite/tests/simplCore/should_compile/all.T


Changes:

=====================================
compiler/GHC/Core/Utils.hs
=====================================
@@ -1047,16 +1047,16 @@ it off at source.
 -}
 
 {-# INLINE trivial_expr_fold #-}
-trivial_expr_fold :: (Id -> r) -> (Literal -> r) -> r -> r -> CoreExpr -> r
+trivial_expr_fold :: (Id -> r) -> (Literal -> r) -> r -> (CoreExpr -> r) -> CoreExpr -> r
 -- ^ The worker function for Note [exprIsTrivial] and Note [getIdFromTrivialExpr]
 -- This is meant to have the code of both functions in one place and make it
 -- easy to derive custom predicates.
 --
 -- (trivial_expr_fold k_id k_triv k_not_triv e)
--- * returns (k_id x) if `e` is a variable `x` (with trivial wrapping)
+-- * returns (k_id x) if `e` is a variable `x` (with trivial wrapping W such that e = W[x])
 -- * returns (k_lit x) if `e` is a trivial literal `l` (with trivial wrapping)
 -- * returns k_triv if `e` is a literal, type, or coercion (with trivial wrapping)
--- * returns k_not_triv otherwise
+-- * returns (k_not_triv e') if e' is not trivial (with trivial wrapping W such that e = W[e'])
 --
 -- where "trivial wrapping" is
 -- * Type application or abstraction
@@ -1073,10 +1073,10 @@ trivial_expr_fold k_id k_lit k_triv k_not_triv = go
     go (Tick t e) | not (tickishIsCode t) = go e              -- See Note [Tick trivial]
     go (Cast e _)                         = go e
     go (Case e _ _ [])                    = go e              -- See Note [Empty case is trivial]
-    go _                                  = k_not_triv
+    go e                                  = k_not_triv e
 
 exprIsTrivial :: CoreExpr -> Bool
-exprIsTrivial e = trivial_expr_fold (const True) (const True) True False e
+exprIsTrivial e = trivial_expr_fold (const True) (const True) True (const False) e
 
 {-
 Note [getIdFromTrivialExpr]
@@ -1097,12 +1097,12 @@ T12076lit for an example where this matters.
 
 getIdFromTrivialExpr :: HasDebugCallStack => CoreExpr -> Id
 -- See Note [getIdFromTrivialExpr]
-getIdFromTrivialExpr e = trivial_expr_fold id (const panic) panic panic e
+getIdFromTrivialExpr e = trivial_expr_fold id (const panic) panic (const panic) e
   where
     panic = pprPanic "getIdFromTrivialExpr" (ppr e)
 
 getIdFromTrivialExpr_maybe :: CoreExpr -> Maybe Id
-getIdFromTrivialExpr_maybe e = trivial_expr_fold Just (const Nothing) Nothing Nothing e
+getIdFromTrivialExpr_maybe e = trivial_expr_fold Just (const Nothing) Nothing (const Nothing) e
 
 {- *********************************************************************
 *                                                                      *


=====================================
compiler/GHC/CoreToStg.hs
=====================================
@@ -552,9 +552,11 @@ coreToStgApp f args ticks = do
 -- ---------------------------------------------------------------------------
 
 getStgArgFromTrivialArg :: HasDebugCallStack => CoreArg -> StgArg
-getStgArgFromTrivialArg e = trivial_expr_fold StgVarArg StgLitArg panic panic e
+getStgArgFromTrivialArg e = trivial_expr_fold StgVarArg StgLitArg panic try_string e
   where
     panic = pprPanic "getStgArgFromTrivialArg" (ppr e)
+    try_string (Lit l at LitString{}) = StgLitArg l -- string literals are not considered trivial, but atomic
+    try_string _                   = panic
 
 coreToStgArgs :: [CoreArg] -> CtsM ([StgArg], [StgTickish])
 coreToStgArgs []


=====================================
compiler/GHC/CoreToStg/Prep.hs
=====================================
@@ -1445,12 +1445,33 @@ cpeArg env dmd arg
 
        ; if okCpeArg arg2
          then do { v <- newVar arg_ty
-                 ; let arg3      = cpeEtaExpand (exprArity arg2) arg2
+                 -- See Note [Eta expansion of arguments in CorePrep]
+                 ; let arity | Just ao <- cp_arityOpts (cpe_config env) -- Just <=> -O2
+                             , not (is_join_head arg2)
+                             -- See Note [Eta expansion for join points]
+                             -- Eta expanding the join point would
+                             -- introduce crap that we can't generate
+                             -- code for
+                             = case exprEtaExpandArity ao arg2 of
+                                 Nothing -> 0
+                                 Just at -> arityTypeArity at
+                             | otherwise
+                             = exprArity arg2 -- this is cheap enough for -O0 and -O1
+                       arg3 = cpeEtaExpand arity arg2
                        arg_float = mkFloat env dmd is_unlifted v arg3
                  ; return (addFloat floats2 arg_float, varToCoreExpr v) }
          else return (floats2, arg2)
        }
 
+is_join_head :: CoreExpr -> Bool
+-- ^ Identify the cases where our mishandling described in
+-- Note [Eta expansion for join points] would generate crap
+is_join_head (Let bs e)        = isJoinBind bs || is_join_head e
+is_join_head (Cast e _)        = is_join_head e
+is_join_head (Tick _ e)        = is_join_head e
+is_join_head (Case _ _ _ alts) = any is_join_head (rhssOfAlts alts)
+is_join_head _                 = False
+
 {-
 Note [Floating unlifted arguments]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1568,6 +1589,44 @@ and now we do NOT want eta expansion to give
 Instead GHC.Core.Opt.Arity.etaExpand gives
                 f = /\a -> \y -> let s = h 3 in g s y
 
+Note [Eta expansion of arguments in CorePrep]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Suppose `g = \x y. blah` and consider the expression `f (g x)`; we ANFise to
+
+  let t = g x
+  in f t
+
+We really don't want that `t` to be a thunk! That just wastes runtime, updating
+a thunk with a PAP etc. The code generator could in principle allocate a PAP,
+but in fact it does not know how to do that -- it's easier just to eta-expand:
+
+  let t = \y. g x y
+  in f t
+
+To what arity should we eta-expand the argument? `cpeArg` uses two strategies,
+governed by the presence of `-fdo-clever-arg-eta-expansion` (implied by -O):
+
+  1. Cheap, with -O0: just use `exprArity`.
+  2. More clever but expensive, with -O1 -O2: use `exprEtaExpandArity`,
+     same function the Simplifier uses to eta expand RHSs and lambda bodies.
+
+The only reason for using (1) rather than (2) is to keep compile times down.
+Using (2) in -O0 bumped up compiler allocations by 2-3% in tests T4801 and
+T5321*. However, Plan (2) catches cases that (1) misses.
+For example (#23083, assuming -fno-pedantic-bottoms):
+
+  let t = case z of __DEFAULT -> g x
+  in f t
+
+to
+
+  let t = \y -> case z of __DEFAULT -> g x y
+  in f t
+
+Note that there is a missed opportunity in eta expanding `t` earlier, in the
+Simplifier: It would allow us to inline `g`, potentially enabling further
+simplification. But then we could have inlined `g` into the PAP to begin with,
+and that is discussed in #23150; hence we needn't worry about that in CorePrep.
 -}
 
 cpeEtaExpand :: Arity -> CpeRhs -> CpeRhs
@@ -1931,6 +1990,11 @@ data CorePrepConfig = CorePrepConfig
   , cp_convertNumLit           :: !(LitNumType -> Integer -> Maybe CoreExpr)
   -- ^ Convert some numeric literals (Integer, Natural) into their final
   -- Core form.
+
+  , cp_arityOpts               :: !(Maybe ArityOpts)
+  -- ^ Configuration for arity analysis ('exprEtaExpandArity').
+  -- See Note [Eta expansion of arguments in CorePrep]
+  -- When 'Nothing' (e.g., -O0, -O1), use the cheaper 'exprArity' instead
   }
 
 data CorePrepEnv
@@ -1941,6 +2005,7 @@ data CorePrepEnv
         -- enabled we instead produce an 'error' expression to catch
         -- the case where a function we think should bottom
         -- unexpectedly returns.
+
         , cpe_env             :: IdEnv CoreExpr   -- Clone local Ids
         -- ^ This environment is used for three operations:
         --


=====================================
compiler/GHC/Driver/Config/CoreToStg/Prep.hs
=====================================
@@ -9,6 +9,7 @@ import GHC.Core.Opt.Pipeline.Types ( CoreToDo(..) )
 import GHC.Driver.Env
 import GHC.Driver.Session
 import GHC.Driver.Config.Core.Lint
+import GHC.Driver.Config.Core.Opt.Arity
 import GHC.Tc.Utils.Env
 import GHC.Types.Var
 import GHC.Utils.Outputable ( alwaysQualify )
@@ -17,14 +18,18 @@ import GHC.CoreToStg.Prep
 
 initCorePrepConfig :: HscEnv -> IO CorePrepConfig
 initCorePrepConfig hsc_env = do
+   let dflags = hsc_dflags hsc_env
    convertNumLit <- do
-     let platform = targetPlatform $ hsc_dflags hsc_env
+     let platform = targetPlatform dflags
          home_unit = hsc_home_unit hsc_env
          lookup_global = lookupGlobal hsc_env
      mkConvertNumLiteral platform home_unit lookup_global
    return $ CorePrepConfig
-      { cp_catchNonexhaustiveCases = gopt Opt_CatchNonexhaustiveCases $ hsc_dflags hsc_env
+      { cp_catchNonexhaustiveCases = gopt Opt_CatchNonexhaustiveCases dflags
       , cp_convertNumLit = convertNumLit
+      , cp_arityOpts = if gopt Opt_DoCleverArgEtaExpansion dflags
+                       then Just (initArityOpts dflags)
+                       else Nothing
       }
 
 initCorePrepPgmConfig :: DynFlags -> [Var] -> CorePrepPgmConfig


=====================================
compiler/GHC/Driver/Flags.hs
=====================================
@@ -262,6 +262,7 @@ data GeneralFlag
    | Opt_SpecConstr
    | Opt_SpecConstrKeen
    | Opt_DoLambdaEtaExpansion
+   | Opt_DoCleverArgEtaExpansion        -- More sophisticated eta expansion of arguments in CorePrep
    | Opt_IgnoreAsserts
    | Opt_DoEtaReduction
    | Opt_CaseMerge


=====================================
compiler/GHC/Driver/Session.hs
=====================================
@@ -3467,6 +3467,7 @@ fFlagsDeps = [
       Opt_DmdTxDictSel "effect is now unconditionally enabled",
   flagSpec "do-eta-reduction"                 Opt_DoEtaReduction,
   flagSpec "do-lambda-eta-expansion"          Opt_DoLambdaEtaExpansion,
+  flagSpec "do-clever-arg-eta-expansion"      Opt_DoCleverArgEtaExpansion,
   flagSpec "eager-blackholing"                Opt_EagerBlackHoling,
   flagSpec "embed-manifest"                   Opt_EmbedManifest,
   flagSpec "enable-rewrite-rules"             Opt_EnableRewriteRules,
@@ -4059,6 +4060,7 @@ optLevelFlags :: [([Int], GeneralFlag)]
 -- Default settings of flags, before any command-line overrides
 optLevelFlags -- see Note [Documenting optimisation flags]
   = [ ([0,1,2], Opt_DoLambdaEtaExpansion)
+    , ([1,2],   Opt_DoCleverArgEtaExpansion)
     , ([0,1,2], Opt_DoEtaReduction)       -- See Note [Eta-reduction in -O0]
     , ([0,1,2], Opt_LlvmTBAA)
     , ([0,1,2], Opt_ProfManualCcs )


=====================================
docs/users_guide/using-optimisation.rst
=====================================
@@ -467,6 +467,17 @@ by saying ``-fno-wombat``.
 
     Eta-expand let-bindings to increase their arity.
 
+.. ghc-flag:: -fdo-clever-arg-eta-expansion
+    :shortdesc: Enable sophisticated argument eta-expansion. Implied by :ghc-flag:`-O2`.
+    :type: dynamic
+    :reverse: -fno-do-clever-arg-eta-expansion
+    :category:
+
+    :default: off
+
+    Eta-expand arguments to increase their arity to avoid allocating unnecessary
+    thunks for them.
+
 .. ghc-flag:: -feager-blackholing
     :shortdesc: Turn on :ref:`eager blackholing <parallel-compile-options>`
     :type: dynamic


=====================================
testsuite/tests/simplCore/should_compile/T23083.hs
=====================================
@@ -0,0 +1,10 @@
+module T23083 where
+
+-- Just ($), but NOINLINE so that we don't inline it eagerly, subverting the
+-- test case
+($$) :: (a -> b) -> a -> b
+($$) f x = f x
+{-# NOINLINE ($$) #-}
+
+g :: ((Integer -> Integer) -> Integer) -> (Integer -> Integer) -> Integer
+g f h = f (h `seq` (h $$))


=====================================
testsuite/tests/simplCore/should_compile/T23083.stderr
=====================================
@@ -0,0 +1,47 @@
+
+==================== CorePrep ====================
+Result size of CorePrep = {terms: 34, types: 34, coercions: 0, joins: 0/1}
+
+-- RHS size: {terms: 6, types: 5, coercions: 0, joins: 0/0}
+(T23083.$$) [InlPrag=NOINLINE] :: forall a b. (a -> b) -> a -> b
+[GblId, Arity=2, Str=<1C(1,L)><L>, Unf=OtherCon []]
+(T23083.$$) = \ (@a) (@b) (f [Occ=Once1!] :: a -> b) (x [Occ=Once1] :: a) -> f x
+
+-- RHS size: {terms: 12, types: 12, coercions: 0, joins: 0/1}
+T23083.g :: ((GHC.Num.Integer.Integer -> GHC.Num.Integer.Integer) -> GHC.Num.Integer.Integer) -> (GHC.Num.Integer.Integer -> GHC.Num.Integer.Integer) -> GHC.Num.Integer.Integer
+[GblId, Arity=2, Str=<1C(1,L)><ML>, Unf=OtherCon []]
+T23083.g
+  = \ (f [Occ=Once1!] :: (GHC.Num.Integer.Integer -> GHC.Num.Integer.Integer) -> GHC.Num.Integer.Integer) (h [Occ=OnceL1] :: GHC.Num.Integer.Integer -> GHC.Num.Integer.Integer) ->
+      let {
+        sat [Occ=Once1] :: GHC.Num.Integer.Integer -> GHC.Num.Integer.Integer
+        [LclId]
+        sat = \ (eta [Occ=Once1] :: GHC.Num.Integer.Integer) -> case h of h1 [Occ=Once1] { __DEFAULT -> T23083.$$ @GHC.Num.Integer.Integer @GHC.Num.Integer.Integer h1 eta } } in
+      f sat
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+T23083.$trModule4 :: GHC.Prim.Addr#
+[GblId, Unf=OtherCon []]
+T23083.$trModule4 = "main"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+T23083.$trModule3 :: GHC.Types.TrName
+[GblId, Unf=OtherCon []]
+T23083.$trModule3 = GHC.Types.TrNameS T23083.$trModule4
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+T23083.$trModule2 :: GHC.Prim.Addr#
+[GblId, Unf=OtherCon []]
+T23083.$trModule2 = "T23083"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+T23083.$trModule1 :: GHC.Types.TrName
+[GblId, Unf=OtherCon []]
+T23083.$trModule1 = GHC.Types.TrNameS T23083.$trModule2
+
+-- RHS size: {terms: 3, types: 0, coercions: 0, joins: 0/0}
+T23083.$trModule :: GHC.Types.Module
+[GblId, Unf=OtherCon []]
+T23083.$trModule = GHC.Types.Module T23083.$trModule3 T23083.$trModule1
+
+
+


=====================================
testsuite/tests/simplCore/should_compile/all.T
=====================================
@@ -477,3 +477,4 @@ test('T23012', normal, compile, ['-O'])
 test('RewriteHigherOrderPatterns', normal, compile, ['-O -ddump-rule-rewrites -dsuppress-all -dsuppress-uniques'])
 test('T23024', normal, multimod_compile, ['T23024', '-O -v0'])
 test('T23026', normal, compile, ['-O'])
+test('T23083', [ grep_errmsg(r'eta.+::.+Integer') ], compile, ['-O -ddump-prep -dsuppress-uniques -dppr-cols=99999'])



View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/e721cc6f460f3c44e2ff3c846e8ac60e4a980e01

-- 
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/e721cc6f460f3c44e2ff3c846e8ac60e4a980e01
You're receiving this email because of your account on gitlab.haskell.org.


-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.haskell.org/pipermail/ghc-commits/attachments/20230426/9f208f1d/attachment-0001.html>


More information about the ghc-commits mailing list