slow declaration debugging help

Stream: mathlib4

[Kevin Buzzard]

A student came to me with some innocuous-looking code which was taking over a minute to run on their laptop. After some simplification I've got it down to about 13 seconds on my laptop:

import Mathlib.RingTheory.SimpleModule
import Mathlib.Algebra.Ring.Equiv
import Mathlib.LinearAlgebra.Matrix.IsDiag

open Matrix BigOperators

set_option trace.profiler true in
set_option maxHeartbeats 800000 in
/--
For `A`-module `M`,
`Hom(Mⁿ, Mⁿ) ≅ Mₙ(Hom(M, M))`

-/
@[simps]
def endPowEquivMatrix (A : Type*) [Ring A]
    (M : Type*) [AddCommGroup M] [Module A M] (n : ℕ):
    Module.End A (Fin n → M) ≃+* Matrix (Fin n) (Fin n) (Module.End A M) where
  toFun f i j :=
  { toFun := fun x ↦ f (Function.update 0 j x) i
    map_add' := fun x y ↦ show  f _ i = (f _ + f _) i by
      rw [← f.map_add, ← Function.update_add, add_zero]
    map_smul' := fun x y ↦ show f _ _ = (x • f _) _ by
      rw [← f.map_smul, ← Function.update_smul, smul_zero] }
  invFun M' :=
  { toFun := fun x i ↦ ∑ j : Fin n, M' i j (x j)
    map_add' := fun x y ↦ by
      simp only [map_add, ← Finset.sum_add_distrib, Pi.add_apply]
      ext
      dsimp
      exact Finset.sum_add_distrib
    map_smul' := by
      intro a x
      simp only [Pi.smul_apply, LinearMapClass.map_smul, RingHom.id_apply, Finset.smul_sum]
      ext
      dsimp
      exact Eq.symm Finset.smul_sum
  }
  left_inv f := by
    ext i x j : 3
    simp only [of_apply, LinearMap.coe_mk, AddHom.coe_mk, LinearMap.coe_comp, LinearMap.coe_single,
      Function.comp_apply, Finset.sum_apply, Function.update, eq_rec_constant, Pi.zero_apply,
      dite_eq_ite, Finset.sum_ite_eq, Finset.mem_univ, ↓reduceIte]
    rw [← Fintype.sum_apply, ← map_sum]
    congr! 1
    ext k : 1
    simp [Pi.single, Function.update]
  right_inv M := by
    dsimp
    ext i j x : 3
    simp only [Function.update, eq_rec_constant, Pi.zero_apply, dite_eq_ite, Finset.sum_apply,
      Finset.sum_ite_eq, Finset.mem_univ, ↓reduceIte, of_apply, LinearMap.coe_mk, AddHom.coe_mk]
    rw [show ∑ k : Fin n, ((M i k) (if k = j then x else 0)) =
      ∑ k : Fin n, if k = j then (M i k x) else 0
      from Finset.sum_congr rfl fun k _ ↦ by split_ifs <;> aesop]
    simp
  map_mul' f g := by -- slow defeq?
    ext i j x : 2
    simp only [of_apply, LinearMap.coe_mk, AddHom.coe_mk, mul_apply, LinearMap.coeFn_sum,
      Finset.sum_apply, LinearMap.mul_apply, AddSubmonoid.coe_finset_sum,
      Submodule.coe_toAddSubmonoid]
    rw [← Fintype.sum_apply, ← map_sum]
    congr! 1
    ext k : 1
    simp [Function.update]
  map_add' _ _ := by ext i j x : 2; simp -- slow defeq?

I'm assuming that there is some defeq abuse going on but I can't see where -- I've looked at the data-carrying fields and I'm assuming that the proofs are irrelevant. Looking at the trace the issue is that map_mul' and map_add' are taking 5 seconds each because of defeq issues. For example map_add' looks like

[step] [5.265840] ext i j x : 2; simp

and the issue is not simp, because

        [] [5.184295] ext i j x : 2 ▶
        [] [0.081321] simp

In fact the issue is not ext either, it's just defeq issues. Unfolding further one finds about four things of the form

  [] [0.562556] ✅ { toFun := fun f i j ↦ { toFun := fun x ↦ f (Function.update 0 j x) i, map_add' := ⋯, map_smul' := ⋯ },
          invFun := fun M' ↦ { toFun := fun x i ↦ ∑ j : Fin n, (M' i j) (x j), map_add' := ⋯, map_smul' := ⋯ },
          left_inv := ⋯, right_inv := ⋯ }.toFun
      (x✝¹ + x✝) =?= ?M ▼
  [assign] [0.562541] ✅ ?M := {
          toFun := fun f i j ↦ { toFun := fun x ↦ f (Function.update 0 j x) i, map_add' := ⋯, map_smul' := ⋯ },
          invFun := fun M' ↦ { toFun := fun x i ↦ ∑ j : Fin n, (M' i j) (x j), map_add' := ⋯, map_smul' := ⋯ },
          left_inv := ⋯, right_inv := ⋯ }.toFun
      (x✝¹ + x✝)

and my understanding is that I shouldn't expect to be able to unravel [assign] any further. What is surprising me is that this is an assignment of a metavariable ?M : Matrix (Fin n) (Fin n) (Module.End A M) to something which, it seems to me, easily unravels to such a term, so I'm surprised that things are consistently taking nearly a second (and there are about 8 in total of them, which makes things add up). Here x✝¹ : Module.End A (Fin n → M) and the same for x✝, and Module.End is an abbrev for a type of linear maps.

set_option diagnostics true gives

[reduction] unfolded declarations (max: 648, num: 24): ▶
[reduction] unfolded instances (max: 2476, num: 88): ▶
[reduction] unfolded reducible declarations (max: 4719, num: 18): ▶
[type_class] used instances (max: 82, num: 108): ▶
[def_eq] heuristic for solving `f a =?= f b` (max: 120, num: 9): ▶
[kernel] unfolded declarations (max: 1050, num: 60): ▶
use `set_option diagnostics.threshold <num>` to control threshold for reporting counters

but I'm not yet au fait with these numbers -- are they huge?

[Kim Morrison]

Generally the numbers themselves on set_option diagnostics true don't tell you much besides the scale of the proof. You need to click the ▶ to get the declaration level counts, and see if anything "looks suspicious".

[Kim Morrison]

Have you just tried breaking out the fields into preliminary declarations?

[Kevin Buzzard]

Oh nice idea! Pulling out toFun and invFun just magically makes it quick! Thanks!

import Mathlib.RingTheory.SimpleModule
import Mathlib.Algebra.Ring.Equiv
import Mathlib.LinearAlgebra.Matrix.IsDiag

open Matrix BigOperators

namespace endPowEquivMatrix

variable {A : Type*} [Ring A]
    {M : Type*} [AddCommGroup M] [Module A M] {n : ℕ}

abbrev toFun : Module.End A (Fin n → M) → Matrix (Fin n) (Fin n) (Module.End A M) :=
  fun f i j ↦
  { toFun := fun x ↦ f (Function.update 0 j x) i
    map_add' := fun x y ↦ show  f _ i = (f _ + f _) i by
      rw [← f.map_add, ← Function.update_add, add_zero]
    map_smul' := fun x y ↦ show f _ _ = (x • f _) _ by
      rw [← f.map_smul, ← Function.update_smul, smul_zero] }

abbrev invFun : Matrix (Fin n) (Fin n) (Module.End A M) → Module.End A (Fin n → M) :=
  fun M' ↦
  { toFun := fun x i ↦ ∑ j : Fin n, M' i j (x j)
    map_add' := fun x y ↦ by
      simp only [map_add, ← Finset.sum_add_distrib, Pi.add_apply]
      ext
      dsimp
      exact Finset.sum_add_distrib
    map_smul' := by
      intro a x
      simp only [Pi.smul_apply, LinearMapClass.map_smul, RingHom.id_apply, Finset.smul_sum]
      ext
      dsimp
      exact Eq.symm Finset.smul_sum
  }

end endPowEquivMatrix

-- count_heartbeats in -- 17331
def endPowEquivMatrix (A : Type*) [Ring A]
    (M : Type*) [AddCommGroup M] [Module A M] (n : ℕ):
    Module.End A (Fin n → M) ≃+* Matrix (Fin n) (Fin n) (Module.End A M) where
  toFun := endPowEquivMatrix.toFun
  invFun := endPowEquivMatrix.invFun
  left_inv f := by
    ext i x j : 3
    simp only [of_apply, LinearMap.coe_mk, AddHom.coe_mk, LinearMap.coe_comp, LinearMap.coe_single,
      Function.comp_apply, Finset.sum_apply, Function.update, eq_rec_constant, Pi.zero_apply,
      dite_eq_ite, Finset.sum_ite_eq, Finset.mem_univ, ↓reduceIte]
    rw [← Fintype.sum_apply, ← map_sum]
    congr! 1
    ext k : 1
    simp [Pi.single, Function.update]
  right_inv M := by
    ext i j x : 3
    simp only [Function.update, eq_rec_constant, Pi.zero_apply, dite_eq_ite, Finset.sum_apply,
      Finset.sum_ite_eq, Finset.mem_univ, ↓reduceIte, of_apply, LinearMap.coe_mk, AddHom.coe_mk]
    rw [show ∑ k : Fin n, ((M i k) (if k = j then x else 0)) =
      ∑ k : Fin n, if k = j then (M i k x) else 0
      from Finset.sum_congr rfl fun k _ ↦ by split_ifs <;> aesop]
    simp
  map_mul' f g := by
    ext i j x : 2
    simp only [of_apply, LinearMap.coe_mk, AddHom.coe_mk, mul_apply, LinearMap.coeFn_sum,
      Finset.sum_apply, LinearMap.mul_apply, AddSubmonoid.coe_finset_sum,
      Submodule.coe_toAddSubmonoid]
    rw [← Fintype.sum_apply, ← map_sum]
    congr! 1
    ext k : 1
    simp [Function.update]
  map_add' _ _ := by ext i j x : 2; simp

@Edison Xie if you define the functions outside the equiv it's now quick for a reason I don't understand (this is the same code as above, just rearranged a bit).

[Matthew Ballard]

ext is probably prompting unification to unfold the inlined definitions.

[Edison Xie]

Oh nice thanks everyone!