[Git][ghc/ghc][wip/amg/tweak-co-opt] 2 commits: Less coercion optimization for non-newtype axioms

Adam Gundry (@adamgundry) gitlab at gitlab.haskell.org
Fri May 12 20:01:52 UTC 2023



Adam Gundry pushed to branch wip/amg/tweak-co-opt at Glasgow Haskell Compiler / GHC


Commits:
05334a11 by Adam Gundry at 2023-05-12T20:59:33+01:00
Less coercion optimization for non-newtype axioms

See Note [Push transitivity inside newtype axioms only] for an explanation
of the change here.  This change substantially improves the performance of
coercion optimization for programs involving transitive type family reductions.

-------------------------
Metric Decrease:
    CoOpt_Singletons
    LargeRecord
    T12227
    T12545
    T13386
    T15703
    T5030
    T8095
-------------------------

- - - - -
9bd6bf96 by Adam Gundry at 2023-05-12T21:00:04+01:00
Move checkAxInstCo to GHC.Core.Lint

A consequence of the previous change is that checkAxInstCo is no longer
called during coercion optimization, so it can be moved back where it belongs.

Also includes some edits to Note [Conflict checking with AxiomInstCo] as
suggested by @simonpj.

- - - - -


2 changed files:

- compiler/GHC/Core/Coercion/Opt.hs
- compiler/GHC/Core/Lint.hs


Changes:

=====================================
compiler/GHC/Core/Coercion/Opt.hs
=====================================
@@ -4,7 +4,6 @@
 
 module GHC.Core.Coercion.Opt
    ( optCoercion
-   , checkAxInstCo
    , OptCoercionOpts (..)
    )
 where
@@ -804,37 +803,38 @@ opt_trans_rule is co1 co2
 -- Push transitivity inside axioms
 opt_trans_rule is co1 co2
 
-  -- See Note [Why call checkAxInstCo during optimisation]
+  -- See Note [Push transitivity inside axioms] and
+  -- Note [Push transitivity inside newtype axioms only]
   -- TrPushSymAxR
   | Just (sym, con, ind, cos1) <- co1_is_axiom_maybe
+  , isNewTyCon (coAxiomTyCon con)
   , True <- sym
   , Just cos2 <- matchAxiom sym con ind co2
   , let newAxInst = AxiomInstCo con ind (opt_transList is (map mkSymCo cos2) cos1)
-  , Nothing <- checkAxInstCo newAxInst
   = fireTransRule "TrPushSymAxR" co1 co2 $ SymCo newAxInst
 
   -- TrPushAxR
   | Just (sym, con, ind, cos1) <- co1_is_axiom_maybe
+  , isNewTyCon (coAxiomTyCon con)
   , False <- sym
   , Just cos2 <- matchAxiom sym con ind co2
   , let newAxInst = AxiomInstCo con ind (opt_transList is cos1 cos2)
-  , Nothing <- checkAxInstCo newAxInst
   = fireTransRule "TrPushAxR" co1 co2 newAxInst
 
   -- TrPushSymAxL
   | Just (sym, con, ind, cos2) <- co2_is_axiom_maybe
+  , isNewTyCon (coAxiomTyCon con)
   , True <- sym
   , Just cos1 <- matchAxiom (not sym) con ind co1
   , let newAxInst = AxiomInstCo con ind (opt_transList is cos2 (map mkSymCo cos1))
-  , Nothing <- checkAxInstCo newAxInst
   = fireTransRule "TrPushSymAxL" co1 co2 $ SymCo newAxInst
 
   -- TrPushAxL
   | Just (sym, con, ind, cos2) <- co2_is_axiom_maybe
+  , isNewTyCon (coAxiomTyCon con)
   , False <- sym
   , Just cos1 <- matchAxiom (not sym) con ind co1
   , let newAxInst = AxiomInstCo con ind (opt_transList is cos1 cos2)
-  , Nothing <- checkAxInstCo newAxInst
   = fireTransRule "TrPushAxL" co1 co2 newAxInst
 
   -- TrPushAxSym/TrPushSymAx
@@ -915,30 +915,87 @@ fireTransRule _rule _co1 _co2 res
   = Just res
 
 {-
-Note [Conflict checking with AxiomInstCo]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-Consider the following type family and axiom:
+Note [Push transitivity inside axioms]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+opt_trans_rule tries to push transitivity inside axioms to deal with cases like
+the following:
+
+    newtype N a = MkN a
+
+    axN :: N a ~R# a
+
+    covar :: a ~R# b
+    co1 = axN <a> :: N a ~R# a
+    co2 = axN <b> :: N b ~R# b
+
+    co :: a ~R# b
+    co = sym co1 ; N covar ; co2
+
+When we are optimising co, we want to notice that the two axiom instantiations
+cancel out. This is implemented by rules such as TrPushSymAxR, which transforms
+    sym (axN <a>) ; N covar
+into
+    sym (axN covar)
+so that TrPushSymAx can subsequently transform
+    sym (axN covar) ; axN <b>
+into
+    covar
+which is much more compact. In some perf test cases this kind of pattern can be
+generated repeatedly during simplification, so it is very important we squash it
+to stop coercions growing exponentially.  For more details see the paper:
+
+    Evidence normalisation in System FC
+    Dimitrios Vytiniotis and Simon Peyton Jones
+    RTA'13, 2013
+    https://www.microsoft.com/en-us/research/publication/evidence-normalization-system-fc-2/
+
+
+Note [Push transitivity inside newtype axioms only]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The optimization described in Note [Push transitivity inside axioms] is possible
+for both newtype and type family axioms.  However, for type family axioms it is
+relatively common to have transitive sequences of axioms instantiations, for
+example:
+
+    data Nat = Zero | Suc Nat
+
+    type family Index (n :: Nat) (xs :: [Type]) :: Type where
+      Index Zero    (x : xs) = x
+      Index (Suc n) (x : xs) = Index n xs
+
+    axIndex :: { forall x::Type. forall xs::[Type]. Index Zero (x : xs) ~ x
+               ; forall n::Nat. forall x::Type. forall xs::[Type]. Index (Suc n) (x : xs) ~ Index n xs }
+
+    co :: Index (Suc (Suc Zero)) [a, b, c] ~ c
+    co = axIndex[1] <Suc Zero> <a> <[b, c]>
+       ; axIndex[1] <Zero> <b> <[c]>
+       ; axIndex[0] <c> <[]>
+
+Not only are there no cancellation opportunities here, but calling matchAxiom
+repeatedly down the transitive chain is very expensive.  Hence we do not attempt
+to push transitivity inside type family axioms.  See #8095, !9210 and related tickets.
+
+This is implemented by opt_trans_rule checking that the axiom is for a newtype
+constructor (i.e. not a type family).  Adding these guards substantially
+improved performance (reduced bytes allocated by more than 10%) for the tests
+CoOpt_Singletons, LargeRecord, T12227, T12545, T13386, T15703, T5030, T8095.
+
+A side benefit is that we do not risk accidentally creating an ill-typed
+coercion; see Note [Why call checkAxInstCo during optimisation].
+
+There may exist programs that previously relied on pushing transitivity inside
+type family axioms to avoid creating huge coercions, which will regress in
+compile time performance as a result of this change.  We do not currently know
+of any examples, but if any come to light we may need to reconsider this
+behaviour.
 
-type family Equal (a :: k) (b :: k) :: Bool
-type instance where
-  Equal a a = True
-  Equal a b = False
---
-Equal :: forall k::*. k -> k -> Bool
-axEqual :: { forall k::*. forall a::k. Equal k a a ~ True
-           ; forall k::*. forall a::k. forall b::k. Equal k a b ~ False }
-
-We wish to disallow (axEqual[1] <*> <Int> <Int). (Recall that the index is
-0-based, so this is the second branch of the axiom.) The problem is that, on
-the surface, it seems that (axEqual[1] <*> <Int> <Int>) :: (Equal * Int Int ~
-False) and that all is OK. But, all is not OK: we want to use the first branch
-of the axiom in this case, not the second. The problem is that the parameters
-of the first branch can unify with the supplied coercions, thus meaning that
-the first branch should be taken. See also Note [Apartness] in
-"GHC.Core.FamInstEnv".
 
 Note [Why call checkAxInstCo during optimisation]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+NB: The following is no longer relevant, because we no longer push transitivity
+into type family axioms (Note [Push transitivity inside newtype axioms only]).
+It is retained for reference in case we change this behaviour in the future.
+
 It is possible that otherwise-good-looking optimisations meet with disaster
 in the presence of axioms with multiple equations. Consider
 
@@ -1029,39 +1086,6 @@ The problem described here was first found in dependent/should_compile/dynamic-p
 
 -}
 
--- | Check to make sure that an AxInstCo is internally consistent.
--- Returns the conflicting branch, if it exists
--- See Note [Conflict checking with AxiomInstCo]
-checkAxInstCo :: Coercion -> Maybe CoAxBranch
--- defined here to avoid dependencies in GHC.Core.Coercion
--- If you edit this function, you may need to update the GHC formalism
--- See Note [GHC Formalism] in GHC.Core.Lint
-checkAxInstCo (AxiomInstCo ax ind cos)
-  = let branch       = coAxiomNthBranch ax ind
-        tvs          = coAxBranchTyVars branch
-        cvs          = coAxBranchCoVars branch
-        incomps      = coAxBranchIncomps branch
-        (tys, cotys) = splitAtList tvs (map coercionLKind cos)
-        co_args      = map stripCoercionTy cotys
-        subst        = zipTvSubst tvs tys `composeTCvSubst`
-                       zipCvSubst cvs co_args
-        target   = Type.substTys subst (coAxBranchLHS branch)
-        in_scope = mkInScopeSet $
-                   unionVarSets (map (tyCoVarsOfTypes . coAxBranchLHS) incomps)
-        flattened_target = flattenTys in_scope target in
-    check_no_conflict flattened_target incomps
-  where
-    check_no_conflict :: [Type] -> [CoAxBranch] -> Maybe CoAxBranch
-    check_no_conflict _    [] = Nothing
-    check_no_conflict flat (b at CoAxBranch { cab_lhs = lhs_incomp } : rest)
-         -- See Note [Apartness] in GHC.Core.FamInstEnv
-      | SurelyApart <- tcUnifyTysFG alwaysBindFun flat lhs_incomp
-      = check_no_conflict flat rest
-      | otherwise
-      = Just b
-checkAxInstCo _ = Nothing
-
-
 -----------
 wrapSym :: SymFlag -> Coercion -> Coercion
 wrapSym sym co | sym       = mkSymCo co


=====================================
compiler/GHC/Core/Lint.hs
=====================================
@@ -56,7 +56,6 @@ import GHC.Core.TyCon as TyCon
 import GHC.Core.Coercion.Axiom
 import GHC.Core.FamInstEnv( compatibleBranches )
 import GHC.Core.Unify
-import GHC.Core.Coercion.Opt ( checkAxInstCo )
 import GHC.Core.Opt.Arity    ( typeArity, exprIsDeadEnd )
 
 import GHC.Core.Opt.Monad
@@ -2531,6 +2530,70 @@ lintCoercion (HoleCo h)
   = do { addErrL $ text "Unfilled coercion hole:" <+> ppr h
        ; lintCoercion (CoVarCo (coHoleCoVar h)) }
 
+
+{-
+Note [Conflict checking with AxiomInstCo]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider the following type family and axiom:
+
+type family Equal (a :: k) (b :: k) :: Bool
+type instance where
+  Equal a a = True
+  Equal a b = False
+--
+Equal :: forall k::*. k -> k -> Bool
+axEqual :: { forall k::*. forall a::k. Equal k a a ~ True
+           ; forall k::*. forall a::k. forall b::k. Equal k a b ~ False }
+
+The coercion (axEqual[1] <*> <Int> <Int) is ill-typed, and Lint should reject it.
+(Recall that the index is 0-based, so this is the second branch of the axiom.)
+The problem is that, on the surface, it seems that
+
+  (axEqual[1] <*> <Int> <Int>) :: (Equal * Int Int ~ False)
+
+and that all is OK. But, all is not OK: we want to use the first branch of the
+axiom in this case, not the second. The problem is that the parameters of the
+first branch can unify with the supplied coercions, thus meaning that the first
+branch should be taken. See also Note [Apartness] in "GHC.Core.FamInstEnv".
+
+For more details, see the section "Branched axiom conflict checking" in
+docs/core-spec, which defines the corresponding no_conflict function used by the
+Co_AxiomInstCo rule in the section "Coercion typing".
+-}
+
+-- | Check to make sure that an AxInstCo is internally consistent.
+-- Returns the conflicting branch, if it exists
+-- See Note [Conflict checking with AxiomInstCo]
+checkAxInstCo :: Coercion -> Maybe CoAxBranch
+-- defined here to avoid dependencies in GHC.Core.Coercion
+-- If you edit this function, you may need to update the GHC formalism
+-- See Note [GHC Formalism] in GHC.Core.Lint
+checkAxInstCo (AxiomInstCo ax ind cos)
+  = let branch       = coAxiomNthBranch ax ind
+        tvs          = coAxBranchTyVars branch
+        cvs          = coAxBranchCoVars branch
+        incomps      = coAxBranchIncomps branch
+        (tys, cotys) = splitAtList tvs (map coercionLKind cos)
+        co_args      = map stripCoercionTy cotys
+        subst        = zipTvSubst tvs tys `composeTCvSubst`
+                       zipCvSubst cvs co_args
+        target   = Type.substTys subst (coAxBranchLHS branch)
+        in_scope = mkInScopeSet $
+                   unionVarSets (map (tyCoVarsOfTypes . coAxBranchLHS) incomps)
+        flattened_target = flattenTys in_scope target in
+    check_no_conflict flattened_target incomps
+  where
+    check_no_conflict :: [Type] -> [CoAxBranch] -> Maybe CoAxBranch
+    check_no_conflict _    [] = Nothing
+    check_no_conflict flat (b at CoAxBranch { cab_lhs = lhs_incomp } : rest)
+         -- See Note [Apartness] in GHC.Core.FamInstEnv
+      | SurelyApart <- tcUnifyTysFG alwaysBindFun flat lhs_incomp
+      = check_no_conflict flat rest
+      | otherwise
+      = Just b
+checkAxInstCo _ = Nothing
+
+
 {-
 ************************************************************************
 *                                                                      *



View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/compare/6846ec0e2d21011e6119414a5002aab1395cca77...9bd6bf961369aab708e73fe45007160a50137f51

-- 
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/compare/6846ec0e2d21011e6119414a5002aab1395cca77...9bd6bf961369aab708e73fe45007160a50137f51
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/20230512/f9084782/attachment-0001.html>


More information about the ghc-commits mailing list