Speed up Fin Vectors used over and over again

Stream: general

[Niels Voss]

In the following:

import Mathlib

def growing_list {α : Type*} (g : List α → α) : ℕ → List α
| 0 => []
| n + 1 => let l := growing_list g n; List.concat l (g l)

-- A list of 8 elements where each element is one more than the sum of the previous elements
#eval growing_list (fun (l : List ℕ) => List.sum l + 1) 30 -- fast
-- [1, 2, 4, 8, 16, 32, 64, 128, ...]

def growing_fin_vec {α : Type*} (g : (n : ℕ) → (Fin n → α) → α) : (n : ℕ) → Fin n → α
| 0 => fun _ => g 0 isEmptyElim
| n + 1 => let f := growing_fin_vec g n; Fin.snoc f (g n f)

#eval growing_fin_vec (fun n (f : Fin n → ℕ) => Finset.univ.sum f + 1) 10 -- Takes several minutes
-- ![1, 2, 4, 8, 16, 32, 64, 128, 256, 512]

growing_list evaluates a lot faster than growing_fin_vec. I suspect that this is because the Fin vectors in growing_fin_vec constantly have to recompute the same values, leading to high time complexity very fast. Is there a way to "freeze" a Fin vector to speed up performance?

(For relevance, growing_fin_vec takes in a function that, given a vector, returns the next element to be added to it, and uses that to build the actual vector of a given size. This is supposed to help define permutations on Fin n for the purposes of helping prove the Levy-Steinitz theorem. I don't know if there is an equivalent function in Mathlib. Performance for this isn't really that important because it's mostly going to be noncomputable anyways, but I'm still curious.)

[Anne Baanen]

Funnily enough I tried this last Thursday:

import Mathlib

variable {n : ℕ}

/- Evaluate all values of `f` once and then construct a vector from it. Doesn't actually seem to speed up evaluation... :/ -/
def unthunk_vector {α : Type*} (f : Fin n → α) : Fin n → α :=
  let f' := List.ofFn f
  fun i => f'.get (i.cast (by simp [f']))

@[simp] lemma unthunk_vector_apply {α : Type*} (f : Fin n → α) (i : Fin n) :
    unthunk_vector f i = f i :=
  by simp [unthunk_vector]

/- Evaluate all values of `f` once and then construct a vector from it. Doesn't actually seem to speed up evaluation... :/ -/
def unthunk_vector₂ {α : Type*} (f : Fin n → Fin n → α) : Fin n → Fin n → α :=
  let f' := List.ofFn (fun i => unthunk_vector (f i))
  fun i => f'.get (i.cast (by simp [f']))

@[simp] lemma unthunk_vector₂_apply {α : Type*} (f : Fin n → Fin n → α) (i j : Fin n) :
    unthunk_vector₂ f i j = f i j :=
  by simp [unthunk_vector₂]

/- Evaluate all values of `f` once and then construct a vector from it. Doesn't actually seem to speed up evaluation... :/ -/
def unthunk_vector₃ {α : Type*} (f : Fin n → Fin n → Fin n → α) : Fin n → Fin n → Fin n → α :=
  let f' := List.ofFn (fun i => unthunk_vector₂ (f i))
  fun i => f'.get (i.cast (by simp [f']))

[Anne Baanen]

But it only seems to slow down things:

def T' := ![![![(1 : ZMod 37), 0, 0, 0, 0, 0, 0, 0], ![0, 1, 0, 0, 0, 0, 0, 0], ![0, 0, 1, 0, 0, 0, 0, 0], ![0, 0, 0, 1, 0, 0, 0, 0], ![0, 0, 0, 0, 1, 0, 0, 0], ![0, 0, 0, 0, 0, 1, 0, 0], ![0, 0, 0, 0, 0, 0, 1, 0], ![0, 0, 0, 0, 0, 0, 0, 1]],
      ![![0, 1, 0, 0, 0, 0, 0, 0], ![0, 0, 1, 0, 0, 0, 0, 0], ![0, 0, 0, 1, 0, 0, 0, 0], ![0, 0, 0, 0, 1, 0, 0, 0], ![0, 0, 0, 0, 0, 1, 0, 0], ![0, 0, 0, 0, 0, 0, 1, 0], ![0, -1, 0, -4, 0, -1, 0, 7], ![-49, -14, 1, 14, -1, -14, 1, 14]],
      ![![0, 0, 1, 0, 0, 0, 0, 0], ![0, 0, 0, 1, 0, 0, 0, 0], ![0, 0, 0, 0, 1, 0, 0, 0], ![0, 0, 0, 0, 0, 1, 0, 0], ![0, 0, 0, 0, 0, 0, 1, 0], ![0, -1, 0, -4, 0, -1, 0, 7], ![-343, -98, 6, 98, -11, -98, 6, 98], ![-686, -246, 0, 193, 0, -198, 0, 203]],
      ![![0, 0, 0, 1, 0, 0, 0, 0], ![0, 0, 0, 0, 1, 0, 0, 0], ![0, 0, 0, 0, 0, 1, 0, 0], ![0, 0, 0, 0, 0, 0, 1, 0], ![0, -1, 0, -4, 0, -1, 0, 7], ![-343, -98, 6, 98, -11, -98, 6, 98], ![-4802, -1721, 0, 1354, 0, -1389, 0, 1414], ![-9947, -3528, -43, 2842, -10, -2842, 5, 2842]],
      ![![0, 0, 0, 0, 1, 0, 0, 0], ![0, 0, 0, 0, 0, 1, 0, 0], ![0, 0, 0, 0, 0, 0, 1, 0], ![0, -1, 0, -4, 0, -1, 0, 7], ![-343, -98, 6, 98, -11, -98, 6, 98], ![-4802, -1721, 0, 1354, 0, -1389, 0, 1414], ![-69286, -24598, -307, 19796, -60, -19796, 25, 19796], ![-139258, -49740, -686, 39725, 0, -39803, 0, 39823]],
      ![![0, 0, 0, 0, 0, 1, 0, 0], ![0, 0, 0, 0, 0, 0, 1, 0], ![0, -1, 0, -4, 0, -1, 0, 7], ![-343, -98, 6, 98, -11, -98, 6, 98], ![-4802, -1721, 0, 1354, 0, -1389, 0, 1414], ![-69286, -24598, -307, 19796, -60, -19796, 25, 19796], ![-970004, -346455, -4802, 276737, 0, -277229, 0, 277319], ![-1951327, -696780, -9917, 556836, -98, -557522, 20, 557522]],
      ![![0, 0, 0, 0, 0, 0, 1, 0], ![0, -1, 0, -4, 0, -1, 0, 7], ![-343, -98, 6, 98, -11, -98, 6, 98], ![-4802, -1721, 0, 1354, 0, -1389, 0, 1414], ![-69286, -24598, -307, 19796, -60, -19796, 25, 19796], ![-970004, -346455, -4802, 276737, 0, -277229, 0, 277319], ![-13588631, -4852470, -69136, 3877664, -582, -3882466, 90, 3882466], ![-27318578, -9756655, -139258, 7795311, -686, -7805426, 0, 7805448]],
      ![![0, 0, 0, 0, 0, 0, 0, 1], ![-49, -14, 1, 14, -1, -14, 1, 14], ![-686, -246, 0, 193, 0, -198, 0, 203], ![-9947, -3528, -43, 2842, -10, -2842, 5, 2842], ![-139258, -49740, -686, 39725, 0, -39803, 0, 39823], ![-1951327, -696780, -9917, 556836, -98, -557522, 20, 557522], ![-27318578, -9756655, -139258, 7795311, -686, -7805426, 0, 7805448], ![-54922588, -19615108, -280185, 15672176, -1468, -15692266, 9, 15692168]]]

set_option trace.profiler true
#eval fun i j k => T' i j k -- [Elab.command] [0.071371s]
#eval let T := (unthunk_vector₃ T'); fun i j k => T i j k -- [Elab.command] [0.455452s]

[Eric Wieser]

I get a speedup by using Array instead of List

[Joachim Breitner]

I would be very curious if https://github.com/leanprover/lean4/pull/4010 helps with evaluating these vectors. Although I wouldn’t expect it to help with asymptotically bad behavior.

Ah, nevermind, that change wouldn’t help with #eval, only potentially during type checking.