universe issues

Stream: lean4

[Matthew Ballard]

For the umpteenth time now, specifying universe levels explicitly significantly speeds up a file and I would like to understand why.

The new example comes from taking lean#4085, which imposes some reducibility constraints when trying to insert coercions for elaborating binary operations, and seeing the impact.

For simplicity, I just cherry picked that on top of v4.8.0-rc1 and benchmarked it in #12718.

[Matthew Ballard]

This resulted in a +800B increase in instructions for one file -- AlgebraicGeometry.Restrict

[Matthew Ballard]

The cause was the attempts to synthesize CoeT instances using docs#Lean.Meta.coerceSimple?.

[Matthew Ballard]

Specifying the universe levels for Scheme makes it much faster. The failing attempts are about 10x faster

[Matthew Ballard]

Universe metavariables are creeping in from looking at the trace and slowing down things.

[Matthew Ballard]

Why is the private function Lean.Meta.preprocessLevels in Lean.Meta.SynthInstance not hooked into anything?

[Kim Morrison]

It become dead code in 99e8a98f067054ef70fcf15ad4af87d0c27ed7a5, back in 2021. Not sure beyond that.

[Kim Morrison]

Oh, and Leo has just removed it in lean4#4112.

[Matthew Ballard]

I managed to isolate one difference in the massive trace. Compare

[Meta.isDefEq] [0.000002] ✅ ↑(Prefunctor.obj.{u + 1, u + 1, u + 1, u + 1}
                          (CategoryTheory.Functor.toPrefunctor.{u, u, u + 1, u + 1}
                            (CategoryTheory.Functor.comp.{u, u, u, u + 1, u + 1, u + 1}
                              (CategoryTheory.Functor.rightOp.{u, u, u + 1, u + 1}
                                AlgebraicGeometry.Spec.toLocallyRingedSpace.{u})
                              AlgebraicGeometry.LocallyRingedSpace.Γ.{u}))
                          R) =?= ↑(Prefunctor.obj.{u + 1, u + 1, u, u + 1}
                          (CategoryTheory.Functor.toPrefunctor.{u, u, u, u + 1}
                            (AlgebraicGeometry.PresheafedSpace.presheaf.{u + 1, u, u}
                              (AlgebraicGeometry.SheafedSpace.toPresheafedSpace.{u + 1, u, u}
                                (AlgebraicGeometry.LocallyRingedSpace.toSheafedSpace.{u}
                                  (AlgebraicGeometry.Scheme.toLocallyRingedSpace.{u}
                                    (Prefunctor.obj.{u + 1, u + 1, u + 1, u + 1}
                                      (CategoryTheory.Functor.toPrefunctor.{u, u, u + 1, u + 1}
                                        AlgebraicGeometry.Scheme.Spec.{u})
                                      (Opposite.op.{u + 2} R)))))))
                          (Opposite.op.{u + 1} Top.top.{u}))

and this

[Matthew Ballard]

It starts out as

[Meta.isDefEq] [0.011830] ✅ ↑(Prefunctor.obj.{?u.170686 + 1, ?u.170686 + 1, ?u.170686 + 1,
                            ?u.170686 + 1}
                          (CategoryTheory.Functor.toPrefunctor.{?u.170686, ?u.170686, ?u.170686 + 1, ?u.170686 + 1}
                            (CategoryTheory.Functor.comp.{?u.170686, ?u.170686, ?u.170686, ?u.170686 + 1, ?u.170686 + 1,
                                ?u.170686 + 1}
                              (CategoryTheory.Functor.rightOp.{?u.170686, ?u.170686, ?u.170686 + 1, ?u.170686 + 1}
                                AlgebraicGeometry.Spec.toLocallyRingedSpace.{?u.170686})
                              AlgebraicGeometry.LocallyRingedSpace.Γ.{?u.170686}))
                          R) =?= ↑(Prefunctor.obj.{?u.170686 + 1, ?u.170686 + 1, ?u.170686, ?u.170686 + 1}
                          (CategoryTheory.Functor.toPrefunctor.{?u.170686, ?u.170686, ?u.170686, ?u.170686 + 1}
                            (AlgebraicGeometry.PresheafedSpace.presheaf.{?u.170686 + 1, ?u.170686, ?u.170686}
                              (AlgebraicGeometry.SheafedSpace.toPresheafedSpace.{?u.170686 + 1, ?u.170686, ?u.170686}
                                (AlgebraicGeometry.LocallyRingedSpace.toSheafedSpace.{?u.170686}
                                  (AlgebraicGeometry.Scheme.toLocallyRingedSpace.{?u.170686}
                                    (Prefunctor.obj.{?u.170686 + 1, ?u.170686 + 1, ?u.170686 + 1, ?u.170686 + 1}
                                      (CategoryTheory.Functor.toPrefunctor.{?u.170686, ?u.170686, ?u.170686 + 1,
                                          ?u.170686 + 1}
                                        AlgebraicGeometry.Scheme.Spec.{?u.170686})
                                      (Opposite.op.{?u.170686 + 2} R)))))))
                          (Opposite.op.{?u.170686 + 1} Top.top.{?u.170686}))

and then finishes 2000+ lines later with

[Meta.isLevelDefEq] [0.000000] ✅ ?u.170686 =?= ?u.170686

[Matthew Ballard]

There are 33 such rfl checks in the trace

[Matthew Ballard]

I am guessing unification is going to a cache for the first one since I cannot tell they are defeq from looking at them and it dispatches it immediately.

[Matthew Ballard]

Moderate improvement if you allow caching of level metavariables

  Benchmark                                                    Metric         Change
  ==================================================================================
+ ~Mathlib.Algebra.Category.ModuleCat.Basic                    instructions   -34.7%
+ ~Mathlib.Algebra.Category.ModuleCat.ChangeOfRings            instructions    -7.6%
+ ~Mathlib.Algebra.Homology.Bifunctor                          instructions    -9.9%
+ ~Mathlib.Algebra.Homology.HomotopyCategory.Pretriangulated   instructions    -5.1%
+ ~Mathlib.AlgebraicGeometry.GammaSpecAdjunction               instructions   -19.9%
+ ~Mathlib.CategoryTheory.Closed.FunctorCategory               instructions   -33.9%
+ ~Mathlib.CategoryTheory.GradedObject.Trifunctor              instructions    -6.5%
+ ~Mathlib.CategoryTheory.Idempotents.FunctorExtension         instructions    -6.9%
+ ~Mathlib.CategoryTheory.Triangulated.Functor                 instructions   -15.3%
+ ~Mathlib.CategoryTheory.Triangulated.Opposite                instructions   -17.7%
+ ~Mathlib.Geometry.RingedSpace.OpenImmersion                  instructions    -8.3%
+ ~Mathlib.Geometry.RingedSpace.PresheafedSpace                instructions   -22.0%

-600B instructions overall

[Matthew Ballard]

#12807 for reference

[Matthew Ballard]

Well, at least it worked for the naughty declaration that generated the trace. docs#AlgebraicGeometry.SpecΓIdentity no longer needs the flag backward.isDefEq.lazyWhnfCore false set and doesn't need a heartbeat bump (previously 800000)

[Matthew Ballard]

private def getDefEqCacheKind (t s : Expr) : MetaM DefEqCacheKind := do
-  if t.hasMVar || s.hasMVar || (← read).canUnfold?.isSome then
+  if t.data.hasExprMVar || s.data.hasExprMVar || (← read).canUnfold?.isSome then
     return .transient
   else
     return .permanent

actual change for reference

[Kevin Buzzard]

So we are now about one week, two weeks in and still nobody has a clue about why we get these crazy speedups just by randomly adding explicit universe annotations?

[Matthew Ballard]

The results are cached with the explicit universes but not with the universe metavariables. So when Lean sees it again it can use the fact in the cache in the former situation whereas in the latter it has to redo the check.

[Kyle Miller]

(Any reason you wrote t.data.hasExprMVar instead of t.hasExprMVar? I'm guessing you unfolded the definition of t.hasMVar and then deleted the level mvar case?)

[Matthew Ballard]

Yes. I just wanted to do what existed except for the level mvar check.

[Kim Morrison]

I'm sceptical that it's going to be possible to cache more here, but not particularly sceptical. :-)

Actually, mostly in these examples I'm surprised that the universe metavariables don't get assigned way earlier. Perhaps we're missing some opportunity to assign them, and fixing that would then remove the desire for caching with metavariables present?