unnecessary `getFunInfoNArgs` in `DiscrTree`

Stream: lean4 dev

[Jovan Gerbscheid]

When building the DiscrTree, for each Expr.const, Expr.proj and Expr.fvar, the function getFunInfoNArgs is called. This function computes many things about the type of this function, but the only thing that is used is the BinderInfo for each argument. I propose that instead, the BinderInfo is obtained directly form the type, which is more efficient. I saw a speedup of 10% with this change. Note that the result of getFunInfoNArgs is cached, but in the case of Expr.proj, the expression can contain metavariables, so in that case the cache is useless and the result has to be computed.

The BinderInfo is used to determine whether an argument should be ignored. It is equivalent to this:

def ignoreArg (arg : Expr) (i : Nat) (binderInfos : Array BinderInfo) : MetaM Bool := do
  if let some binderInfo := binderInfos[i]? then
    match binderInfo with
    | .instImplicit => return true
    | .implicit
    | .strictImplicit => return not (← isType arg)
    | .default => isProof arg
  else
    isProof arg

One concern could be that getFunInfoNArgs applies whnf to the function type in case it does not match Expr.forallE, which means that it can get more BinderInfo than we could get without using whnf. Note that when no BinderInfo is available, it assumes BinderInfo.default. I found that there is only one constant in Mathlib that unfolds to a Expr.forallE in this way that is not .default, which is def Presieve [Category C] (X : C) := ∀ ⦃Y⦄, Set (Y ⟶ X), which is .strictImplicit. I'm not an expert on Presieves, but I think it is reasonable to treat the argument as a .default and not to ignore its argument in the DiscrTree.

[Jovan Gerbscheid]

I also want to ask why implicit arguments are ignored when they are not types, because I can't think of many cases where an implicit argument is not a type. Could we instead treat .implicit and strictImplicit in the same way as .default?

[Kyle Miller]

When you say "speedup of 10%", what did you test?

[Kyle Miller]

Jovan Gerbscheid said:

I'm not an expert on Presieves, but I think it is reasonable to treat the argument as a .default and not to ignore its argument in the DiscrTree.

I think there would need to be a very good justification to break the assumption that the discrimination tree processes function types using their true binder infos.

[Kyle Miller]

Incidentally, there might be a bug in the current binder infos calculation, where it doesn't take into account the arguments passed to the function. It's possible for the arity of a function to be bigger than the one implied by its type due to the particular arguments passed to it. For example, id id 2 is using id with arity 2.

I've swapped out getFunInfoNArgs with a function that calculates the correct BinderInfos (and only BinderInfos, rather than the full FunInfo), and I'm going to test it out with building mathlib.

[Kyle Miller]

Jovan Gerbscheid said:

I also want to ask why implicit arguments are ignored when they are not types, because I can't think of many cases where an implicit argument is not a type.

Probably most if not all functions about Fin have an implicit Nat argument, for example.

[Jovan Gerbscheid]

Kyle Miller said:

When you say "speedup of 10%", what did you test?

I built the DiscrTree used by the rw? tactic. It went from 31.9 sec to 28.9 sec, and the heartbeats went from 282664 to 252320. The time was measured using set_option profiler true.

[Jovan Gerbscheid]

Kyle Miller said:

Jovan Gerbscheid said:

I also want to ask why implicit arguments are ignored when they are not types, because I can't think of many cases where an implicit argument is not a type.

Probably most if not all functions about Fin have an implicit Nat argument, for example.

I still don't see why we would not treat these implicit Nat arguments the same way as we treat implicit type arguments. Most lemma's about Fin are about a general n anyways, so in that case either way this argument turns into a .star.

[Jovan Gerbscheid]

Kyle Miller said:

Incidentally, there might be a bug in the current binder infos calculation, where it doesn't take into account the arguments passed to the function.

This is indeed a bit of a problem, but I wonder how the cost of the extra computation compares to value we get out of knowing every single BinderInfo, i.e. is it worth it?

[Jovan Gerbscheid]

Another suggestion I have is that besides the BinderInfo, you can also find out if an argument is an outParam, using Expr.isOutParam on the domain of the Expr.forallE. I would ignore the outParams for the same reason as the .instImplicit arguments: they are supposed to be inferable from the other arguments. In particular this would apply to addition, where the third type parameter is redundant.

[Kyle Miller]

Jovan Gerbscheid said:

I still don't see why we would not treat these implicit Nat arguments the same way as we treat implicit type arguments. Most lemma's about Fin are about a general n anyways, so in that case either way this argument turns into a .star.

I was only answering your question about cases where implicit arguments aren't types. It happens very frequently, and you probably wouldn't want the discrimination tree to do a more expensive defeq check for such arguments during lookup. Instance arguments are another example of an implicit argument that's not a type.

There are also plenty of lemmas about Fin (n + 1), etc. This is also not just a matter about the lemmas themselves that can apply, but about the definitions that the lemmas refer to. Definitions that use Fin can have implicit arguments, and maybe the the lemma has some expression involving docs#Fin.succ being passed to some other definition.

[Kyle Miller]

Jovan Gerbscheid said:

This is indeed a bit of a problem, but I wonder how the cost of the extra computation compares to value we get out of knowing every single BinderInfo, i.e. is it worth it?

It seems that you might be only accounting for the time to create a discrimination tree? If I understand it correctly, if you assume it's a default binder then that means during lookup that entails a defeq check. This means getting it wrong can make lookups more expensive. What does that do for simp across all of mathlib for example?

The tradeoff is whether you make discrimination trees more permissive in what they match (perhaps making them faster to construct?) at the expense of having potentially a lot more theorems to test using more expensive defeq checks.

[Kyle Miller]

Jovan Gerbscheid said:

Another suggestion I have is that besides the BinderInfo, you can also find out if an argument is an outParam

That seems sensible. It might be interesting to measure the effect.

[Jovan Gerbscheid]

Actually, I suggested that in the .implicit and .instImplicit cases we would index the argument when it is not a type (currently it only does so when it is a type), just like for .default. This would mean that the DiscrTree is less permissive in what it matches. (So thereby less defeq checks)

[Jovan Gerbscheid]

Or, maybe keep those the same as they are now, but when you don't know the BinderInfo, default to the .default, which makes the argument more likely to be indexed in the tree.

[Mario Carneiro]

Jovan Gerbscheid said:

Kyle Miller said:

When you say "speedup of 10%", what did you test?

I built the DiscrTree used by the rw? tactic. It went from 31.9 sec to 28.9 sec, and the heartbeats went from 282664 to 252320. The time was measured using set_option profiler true.

Kyle Miller said:

It seems that you might be only accounting for the time to create a discrimination tree? If I understand it correctly, if you assume it's a default binder then that means during lookup that entails a defeq check. This means getting it wrong can make lookups more expensive. What does that do for simp across all of mathlib for example?

This happened before with the DiscrTree PR in std. You reported a perf win regarding construction of discrimination trees, but discrimination trees are a kind of lookup table, they are optimized for lookup not for construction, so you really need to be including performance testing of lookups, with construction being only a secondary metric.

[Mario Carneiro]

the fastest discrimination tree to construct is of course the one that maps everything to .star :wink:

[Jovan Gerbscheid]

Is there a standard test for DiscrTree lookup performance? Because I suppose that the performance gain/loss for lookup is more in what the tactics do with the obtained result than in the lookup itself. Testing it with building Mathlib would be ideal, but I don't think I can feasibly do that.

[Jovan Gerbscheid]

Kyle Miller said:

I've swapped out getFunInfoNArgs with a function that calculates the correct BinderInfos

I've tried the same, and I now have a version that builds around 5% faster than the standard DiscrTree. It uses some nice optimization inspired by inferAppType in Lean.Meta.InferType, so I though I'd share that part:

def argsIgnore (fn : Expr) (args : Array Expr) : MetaM (Array Bool) := do
  let mut fnType ← inferType fn
  let mut result := Array.mkEmpty args.size
  let mut j := 0
  for i in [:args.size] do
    let arg := args[i]!
    unless fnType matches .forallE .. do
      fnType ← whnfD (fnType.instantiateRevRange j i args)
      j := i
    let (.forallE _ d b binderInfo) := fnType | throwError m! "expected function type {indentExpr fnType}"
    let ignore ← (do
      if d.isOutParam then
        return true
      else
        match binderInfo with
        | .instImplicit => return true
        | .default => isProof arg
        | _ => return !(← isType arg))
    fnType := b
    result := result.push ignore
  return result

[Kim Morrison]

Jovan Gerbscheid said:

Testing it with building Mathlib would be ideal, but I don't think I can feasibly do that.

Why do you say that? Just make a branch and a draft PR, and run !bench.

[Jovan Gerbscheid]

I've never used !bench before. Can it also benchmark mathlib with changes to core lean? And should that be with a PR to core or to mathlib?

[Joachim Breitner]

When you create a PR against core, it will make a “release” for just your PR.

Also, if your PR branches off the tag nightly-with-mathlib, then the lean CI will automatically create a branch in the mathlib repo that tests your changes. But you can create such a branch in mathlib youself (which is what Scott wants to say here): A normal mathlib (draft) PR with lean-toolchain changed to your PR’s release. Then you can comment !bench there and some :cloud: :computer: will benchmark your changes.

It can be a bit tricky to make sure the bench comparison really only captures your changes. Probably best to rebase your core PR onto the particular release tag that mathlib is using, so that only your changes are included.

[Kyle Miller]

By the way, I made a draft PR at lean4#3113 with conservative changes (no changes in functionality, just use an optimized binderinfo collector), but CI didn't go through and I forgot about it over the rest of the holidays.

[Jovan Gerbscheid]

So I would need to make a PR in lean branching off of leanprover/lean4:v4.5.0-rc1, and a PR in mathlib, setting the toolchain to the release corresponding to the lean PR, and then comment !bench in the mathlib PR?

[Jovan Gerbscheid]

But it seems @Kyle Miller didn't make a separate PR to mathlib for this, or just not yet?

[Kyle Miller]

@Jovan Gerbscheid I just rebased that Lean branch on top of a nightly branch (nightly-with-mathlib), and hopefully that will cause the CI to automatically create that mathlib branch

[Joachim Breitner]

If it doesn’t ping me and I can have a look why not. It should be working again these days :)

[Kyle Miller]

It built, and I got it to benchmark, but there are no benchmark results to compare against. https://github.com/leanprover-community/mathlib4/pull/9460

Could someone who has privileges to do so get f038bf8f141f7790d7c4e5903a7bfad571985094 to be benchmarked? That's the parent commit, which seems reasonable enough to compare against.

[Joachim Breitner]

Ah, yes, I keep running into this. What I do in that case is create a dummy branch from that commit, create a draft PR (against master or whater), run !bench there, and manually find the two runs on the speedcenter and compare.

[Joachim Breitner]

This is very tedious, I agree. I wonder if we should automatically bench all nightly-testing commits? Would that have helped here? I assume so.