perf: sort before elim dead branches (#11366)

This PR sorts the declarations fed into ElimDeadBranches in increasing
size. This can improve performance when we are dealing with a lot of
iterations.

The motivation for this change is as follows. Currently the algorithm
for doing one step of abstract interpretation is:
```
for decl in scc do
  interpDecl
  if summaryChanged decl then
    return true
return false
```
whenever we return true we run another step. Now suppose we are in a
situation where we have an SCC with one big decl in the front and then
`n` small ones afterwards. For each time that the small ones change
their summary, we will re-run analysis of the big one in the front.
Currently the ordering is basically at "random" based on how other
compilers inject things into the SCC. This change ensures the behavior
is consistent and at least somewhat intelligent. By putting the small
declarations first, whenever we trigger a rerun of the loop we bias
analyzing the small declarations first, thus decreasing run time.

Note that this change does not have much effect on the current pipeline
because: We usually construct the SCCs in a way such that small ones
happen to be in front anyways. However, with upcomping changes on
specialization this is about to change.

Author: hargoniX

src/Lean/Compiler/LCNF/ElimDeadBranches.lean

@@ -581,7 +581,9 @@ def inferStep : InterpM Bool := do
     withReader (fun ctx => { ctx with currFnIdx := idx }) do
       decl.params.forM fun p => updateVarAssignment p.fvarId .top
       match decl.value with
-      | .code code .. => interpCode code
+      | .code code .. =>
+        withTraceNode `Compiler.elimDeadBranches (fun _ => return m!"Analyzing {decl.name}") do
+          interpCode code
       | .extern .. => updateCurrFnSummary .top
     let newVal ← getFunVal idx
     if currentVal != newVal then
@@ -591,13 +593,14 @@ def inferStep : InterpM Bool := do
 /--
 Run `inferStep` until it reaches a fix point.
 -/
-partial def inferMain : InterpM Unit := do
+partial def inferMain (n : Nat := 0) : InterpM Unit := do
   let ctx ← read
   modify fun s => { s with assignments := ctx.decls.map fun _ => {} }
   let modified ← inferStep
   if modified then
-    inferMain
+    inferMain (n + 1)
   else
+    trace[Compiler.elimDeadBranches] m!"Termination after {n} steps"
     return ()
 
 /--
@@ -647,6 +650,11 @@ end UnreachableBranches
 
 open UnreachableBranches in
 def Decl.elimDeadBranches (decls : Array Decl) : CompilerM (Array Decl) := do
+  /-
+  We sort declarations by size here to ensure that when we restart in inferStep it will mostly be
+  small declarations that get re-analyzed.
+  -/
+  let decls := decls.qsort (fun l r => (l.size, l.name.toString).lexLt (r.size, r.name.toString))
   let mut assignments := decls.map fun _ => {}
   let initialVal i :=
     /-
@@ -659,7 +667,9 @@ def Decl.elimDeadBranches (decls : Array Decl) : CompilerM (Array Decl) := do
   let mut funVals := decls.size.fold (init := .empty) fun i _ p => p.push (initialVal i)
   let ctx := { decls }
   let mut state := { assignments, funVals }
-  (_, state) ← inferMain |>.run ctx |>.run state
+  (_, state) ←
+    withTraceNode `Compiler.elimDeadBranches (fun _ => return m!"Analyzing block: {decls.map (·.name)}")
+      inferMain |>.run ctx |>.run state
   funVals := state.funVals
   assignments := state.assignments
   modifyEnv fun e =>

tests/lean/4089.lean.expected.out

@@ -8,6 +8,9 @@
         let x_4 : obj := reuse x_5 in ctor_0[Prod.mk] x_3 x_2;
         ret x_4
 [Compiler.IR] [reset_reuse]
+    def Sigma.toProd (x_1 : ◾) (x_2 : ◾) (x_3 : obj) : obj :=
+      let x_4 : obj := Sigma.toProd._redArg x_3;
+      ret x_4
     def Sigma.toProd._redArg (x_1 : obj) : obj :=
       case x_1 : obj of
       Sigma.mk →
@@ -16,9 +19,6 @@
         let x_5 : tobj := reset[2] x_1;
         let x_4 : obj := reuse x_5 in ctor_0[Prod.mk] x_2 x_3;
         ret x_4
-    def Sigma.toProd (x_1 : ◾) (x_2 : ◾) (x_3 : obj) : obj :=
-      let x_4 : obj := Sigma.toProd._redArg x_3;
-      ret x_4
 [Compiler.IR] [reset_reuse]
     def foo (x_1 : tobj) : tobj :=
       case x_1 : tobj of

tests/lean/4240.lean.expected.out

@@ -1,4 +1,10 @@
 [Compiler.IR] [result]
+    def isSomeWithInstanceNat (x_1 : @& obj) : u8 :=
+      let x_2 : usize := 0;
+      let x_3 : tobj := Array.uget ◾ x_1 x_2 ◾;
+      let x_4 : u8 := MyOption.isSomeWithInstance._at_.isSomeWithInstanceNat.spec_0 x_3;
+      dec x_3;
+      ret x_4
     def MyOption.isSomeWithInstance._at_.isSomeWithInstanceNat.spec_0 (x_1 : @& tobj) : u8 :=
       case x_1 : tobj of
       MyOption.none →
@@ -7,19 +13,13 @@
       MyOption.some →
         let x_3 : u8 := 1;
         ret x_3
-    def isSomeWithInstanceNat (x_1 : @& obj) : u8 :=
-      let x_2 : usize := 0;
-      let x_3 : tobj := Array.uget ◾ x_1 x_2 ◾;
-      let x_4 : u8 := MyOption.isSomeWithInstance._at_.isSomeWithInstanceNat.spec_0 x_3;
-      dec x_3;
-      ret x_4
-    def MyOption.isSomeWithInstance._at_.isSomeWithInstanceNat.spec_0._boxed (x_1 : tobj) : tagged :=
-      let x_2 : u8 := MyOption.isSomeWithInstance._at_.isSomeWithInstanceNat.spec_0 x_1;
+    def isSomeWithInstanceNat._boxed (x_1 : obj) : tagged :=
+      let x_2 : u8 := isSomeWithInstanceNat x_1;
       dec x_1;
       let x_3 : tobj := box x_2;
       ret x_3
-    def isSomeWithInstanceNat._boxed (x_1 : obj) : tagged :=
-      let x_2 : u8 := isSomeWithInstanceNat x_1;
+    def MyOption.isSomeWithInstance._at_.isSomeWithInstanceNat.spec_0._boxed (x_1 : tobj) : tagged :=
+      let x_2 : u8 := MyOption.isSomeWithInstance._at_.isSomeWithInstanceNat.spec_0 x_1;
       dec x_1;
       let x_3 : tobj := box x_2;
       ret x_3

tests/lean/computedFieldsCode.lean.expected.out

@@ -27,6 +27,9 @@
       dec x_1;
       ret x_2
 [Compiler.IR] [result]
+    def Exp.ctorElim (x_1 : ◾) (x_2 : @& tobj) (x_3 : tobj) (x_4 : ◾) (x_5 : tobj) : tobj :=
+      let x_6 : tobj := Exp.ctorElim._redArg x_3 x_5;
+      ret x_6
     def Exp.ctorElim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
       case x_1 : tobj of
       Exp.var →
@@ -46,106 +49,60 @@
       default →
         dec x_1;
         ret x_2
-    def Exp.ctorElim (x_1 : ◾) (x_2 : @& tobj) (x_3 : tobj) (x_4 : ◾) (x_5 : tobj) : tobj :=
-      let x_6 : tobj := Exp.ctorElim._redArg x_3 x_5;
-      ret x_6
     def Exp.ctorElim._boxed (x_1 : tobj) (x_2 : tobj) (x_3 : tobj) (x_4 : tobj) (x_5 : tobj) : tobj :=
       let x_6 : tobj := Exp.ctorElim x_1 x_2 x_3 x_4 x_5;
       dec x_2;
       ret x_6
 [Compiler.IR] [result]
-    def Exp.var.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
-      let x_3 : tobj := Exp.ctorElim._redArg x_1 x_2;
-      ret x_3
     def Exp.var.elim (x_1 : ◾) (x_2 : tobj) (x_3 : ◾) (x_4 : tobj) : tobj :=
       let x_5 : tobj := Exp.ctorElim._redArg x_2 x_4;
       ret x_5
-[Compiler.IR] [result]
-    def Exp.app.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
+    def Exp.var.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
       let x_3 : tobj := Exp.ctorElim._redArg x_1 x_2;
       ret x_3
+[Compiler.IR] [result]
     def Exp.app.elim (x_1 : ◾) (x_2 : tobj) (x_3 : ◾) (x_4 : tobj) : tobj :=
       let x_5 : tobj := Exp.ctorElim._redArg x_2 x_4;
       ret x_5
-[Compiler.IR] [result]
-    def Exp.a1.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
+    def Exp.app.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
       let x_3 : tobj := Exp.ctorElim._redArg x_1 x_2;
       ret x_3
+[Compiler.IR] [result]
     def Exp.a1.elim (x_1 : ◾) (x_2 : tobj) (x_3 : ◾) (x_4 : tobj) : tobj :=
       let x_5 : tobj := Exp.ctorElim._redArg x_2 x_4;
       ret x_5
-[Compiler.IR] [result]
-    def Exp.a2.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
+    def Exp.a1.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
       let x_3 : tobj := Exp.ctorElim._redArg x_1 x_2;
       ret x_3
+[Compiler.IR] [result]
     def Exp.a2.elim (x_1 : ◾) (x_2 : tobj) (x_3 : ◾) (x_4 : tobj) : tobj :=
       let x_5 : tobj := Exp.ctorElim._redArg x_2 x_4;
       ret x_5
-[Compiler.IR] [result]
-    def Exp.a3.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
+    def Exp.a2.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
       let x_3 : tobj := Exp.ctorElim._redArg x_1 x_2;
       ret x_3
+[Compiler.IR] [result]
     def Exp.a3.elim (x_1 : ◾) (x_2 : tobj) (x_3 : ◾) (x_4 : tobj) : tobj :=
       let x_5 : tobj := Exp.ctorElim._redArg x_2 x_4;
       ret x_5
-[Compiler.IR] [result]
-    def Exp.a4.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
+    def Exp.a3.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
       let x_3 : tobj := Exp.ctorElim._redArg x_1 x_2;
       ret x_3
+[Compiler.IR] [result]
     def Exp.a4.elim (x_1 : ◾) (x_2 : tobj) (x_3 : ◾) (x_4 : tobj) : tobj :=
       let x_5 : tobj := Exp.ctorElim._redArg x_2 x_4;
       ret x_5
-[Compiler.IR] [result]
-    def Exp.a5.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
+    def Exp.a4.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
       let x_3 : tobj := Exp.ctorElim._redArg x_1 x_2;
       ret x_3
+[Compiler.IR] [result]
     def Exp.a5.elim (x_1 : ◾) (x_2 : tobj) (x_3 : ◾) (x_4 : tobj) : tobj :=
       let x_5 : tobj := Exp.ctorElim._redArg x_2 x_4;
       ret x_5
+    def Exp.a5.elim._redArg (x_1 : tobj) (x_2 : tobj) : tobj :=
+      let x_3 : tobj := Exp.ctorElim._redArg x_1 x_2;
+      ret x_3
 [Compiler.IR] [result]
-    def Exp.casesOn._override._redArg (x_1 : tobj) (x_2 : tobj) (x_3 : tobj) (x_4 : @& tobj) (x_5 : @& tobj) (x_6 : @& tobj) (x_7 : @& tobj) (x_8 : @& tobj) : tobj :=
-      case x_1 : tobj of
-      Exp.var._impl →
-        dec x_3;
-        let x_9 : u32 := sproj[0, 8] x_1;
-        dec x_1;
-        let x_10 : tobj := box x_9;
-        let x_11 : tobj := app x_2 x_10;
-        ret x_11
-      Exp.app._impl →
-        dec x_2;
-        let x_12 : tobj := proj[0] x_1;
-        inc x_12;
-        let x_13 : tobj := proj[1] x_1;
-        inc x_13;
-        dec x_1;
-        let x_14 : tobj := app x_3 x_12 x_13;
-        ret x_14
-      Exp.a1._impl →
-        dec x_3;
-        dec x_2;
-        inc x_4;
-        ret x_4
-      Exp.a2._impl →
-        dec x_3;
-        dec x_2;
-        inc x_5;
-        ret x_5
-      Exp.a3._impl →
-        dec x_3;
-        dec x_2;
-        inc x_6;
-        ret x_6
-      Exp.a4._impl →
-        dec x_3;
-        dec x_2;
-        inc x_7;
-        ret x_7
-      Exp.a5._impl →
-        dec x_3;
-        dec x_2;
-        inc x_8;
-        ret x_8
     def Exp.casesOn._override (x_1 : ◾) (x_2 : tobj) (x_3 : tobj) (x_4 : tobj) (x_5 : @& tobj) (x_6 : @& tobj) (x_7 : @& tobj) (x_8 : @& tobj) (x_9 : @& tobj) : tobj :=
       case x_2 : tobj of
       Exp.var._impl →
@@ -189,14 +146,49 @@
         dec x_3;
         inc x_9;
         ret x_9
-    def Exp.casesOn._override._redArg._boxed (x_1 : tobj) (x_2 : tobj) (x_3 : tobj) (x_4 : tobj) (x_5 : tobj) (x_6 : tobj) (x_7 : tobj) (x_8 : tobj) : tobj :=
-      let x_9 : tobj := Exp.casesOn._override._redArg x_1 x_2 x_3 x_4 x_5 x_6 x_7 x_8;
-      dec x_8;
-      dec x_7;
-      dec x_6;
-      dec x_5;
-      dec x_4;
-      ret x_9
+    def Exp.casesOn._override._redArg (x_1 : tobj) (x_2 : tobj) (x_3 : tobj) (x_4 : @& tobj) (x_5 : @& tobj) (x_6 : @& tobj) (x_7 : @& tobj) (x_8 : @& tobj) : tobj :=
+      case x_1 : tobj of
+      Exp.var._impl →
+        dec x_3;
+        let x_9 : u32 := sproj[0, 8] x_1;
+        dec x_1;
+        let x_10 : tobj := box x_9;
+        let x_11 : tobj := app x_2 x_10;
+        ret x_11
+      Exp.app._impl →
+        dec x_2;
+        let x_12 : tobj := proj[0] x_1;
+        inc x_12;
+        let x_13 : tobj := proj[1] x_1;
+        inc x_13;
+        dec x_1;
+        let x_14 : tobj := app x_3 x_12 x_13;
+        ret x_14
+      Exp.a1._impl →
+        dec x_3;
+        dec x_2;
+        inc x_4;
+        ret x_4
+      Exp.a2._impl →
+        dec x_3;
+        dec x_2;
+        inc x_5;
+        ret x_5
+      Exp.a3._impl →
+        dec x_3;
+        dec x_2;
+        inc x_6;
+        ret x_6
+      Exp.a4._impl →
+        dec x_3;
+        dec x_2;
+        inc x_7;
+        ret x_7
+      Exp.a5._impl →
+        dec x_3;
+        dec x_2;
+        inc x_8;
+        ret x_8
     def Exp.casesOn._override._boxed (x_1 : tobj) (x_2 : tobj) (x_3 : tobj) (x_4 : tobj) (x_5 : tobj) (x_6 : tobj) (x_7 : tobj) (x_8 : tobj) (x_9 : tobj) : tobj :=
       let x_10 : tobj := Exp.casesOn._override x_1 x_2 x_3 x_4 x_5 x_6 x_7 x_8 x_9;
       dec x_9;
@@ -205,22 +197,15 @@
       dec x_6;
       dec x_5;
       ret x_10
+    def Exp.casesOn._override._redArg._boxed (x_1 : tobj) (x_2 : tobj) (x_3 : tobj) (x_4 : tobj) (x_5 : tobj) (x_6 : tobj) (x_7 : tobj) (x_8 : tobj) : tobj :=
+      let x_9 : tobj := Exp.casesOn._override._redArg x_1 x_2 x_3 x_4 x_5 x_6 x_7 x_8;
+      dec x_8;
+      dec x_7;
+      dec x_6;
+      dec x_5;
+      dec x_4;
+      ret x_9
 [Compiler.IR] [result]
-    def Exp.var._override (x_1 : u32) : tobj :=
-      let x_2 : u64 := UInt32.toUInt64 x_1;
-      let x_3 : u64 := 1000;
-      let x_4 : u64 := UInt64.add x_2 x_3;
-      let x_5 : obj := ctor_0.0.12[Exp.var._impl];
-      sset x_5[0, 0] : u64 := x_4;
-      sset x_5[0, 8] : u32 := x_1;
-      ret x_5
-    def Exp.app._override (x_1 : tobj) (x_2 : tobj) : tobj :=
-      let x_3 : u64 := Exp.hash._override x_1;
-      let x_4 : u64 := Exp.hash._override x_2;
-      let x_5 : u64 := mixHash x_3 x_4;
-      let x_6 : obj := ctor_1.0.8[Exp.app._impl] x_1 x_2;
-      sset x_6[2, 0] : u64 := x_5;
-      ret x_6
     def Exp.a1._override : tobj :=
       let x_1 : tagged := ctor_2[Exp.a1._impl];
       ret x_1
@@ -236,6 +221,21 @@
     def Exp.a5._override : tobj :=
       let x_1 : tagged := ctor_6[Exp.a5._impl];
       ret x_1
+    def Exp.app._override (x_1 : tobj) (x_2 : tobj) : tobj :=
+      let x_3 : u64 := Exp.hash._override x_1;
+      let x_4 : u64 := Exp.hash._override x_2;
+      let x_5 : u64 := mixHash x_3 x_4;
+      let x_6 : obj := ctor_1.0.8[Exp.app._impl] x_1 x_2;
+      sset x_6[2, 0] : u64 := x_5;
+      ret x_6
+    def Exp.var._override (x_1 : u32) : tobj :=
+      let x_2 : u64 := UInt32.toUInt64 x_1;
+      let x_3 : u64 := 1000;
+      let x_4 : u64 := UInt64.add x_2 x_3;
+      let x_5 : obj := ctor_0.0.12[Exp.var._impl];
+      sset x_5[0, 0] : u64 := x_4;
+      sset x_5[0, 8] : u32 := x_1;
+      ret x_5
     def Exp.hash._override (x_1 : @& tobj) : u64 :=
       case x_1 : tobj of
       Exp.var._impl →

tests/lean/doubleReset.lean.expected.out

@@ -1,4 +1,15 @@
 [Compiler.IR] [reset_reuse]
+    def _private.Init.Data.Array.Basic.0.Array.mapMUnsafe.map._at_.applyProjectionRules.spec_0 (x_1 : ◾) (x_2 : ◾) (x_3 : ◾) (x_4 : tobj) (x_5 : usize) (x_6 : usize) (x_7 : obj) : obj :=
+      let x_8 : obj := _private.Init.Data.Array.Basic.0.Array.mapMUnsafe.map._at_.applyProjectionRules.spec_0._redArg x_4 x_5 x_6 x_7;
+      ret x_8
+    def applyProjectionRules (x_1 : ◾) (x_2 : ◾) (x_3 : ◾) (x_4 : obj) (x_5 : tobj) : obj :=
+      let x_6 : obj := applyProjectionRules._redArg x_4 x_5;
+      ret x_6
+    def applyProjectionRules._redArg (x_1 : obj) (x_2 : tobj) : obj :=
+      let x_3 : usize := Array.usize ◾ x_1;
+      let x_4 : usize := 0;
+      let x_5 : obj := _private.Init.Data.Array.Basic.0.Array.mapMUnsafe.map._at_.applyProjectionRules.spec_0._redArg x_2 x_3 x_4 x_1;
+      ret x_5
     def _private.Init.Data.Array.Basic.0.Array.mapMUnsafe.map._at_.applyProjectionRules.spec_0._redArg (x_1 : tobj) (x_2 : usize) (x_3 : usize) (x_4 : obj) : obj :=
       let x_5 : u8 := USize.decLt x_3 x_2;
       case x_5 : u8 of
@@ -24,14 +35,3 @@
             let x_16 : obj := Array.uset ◾ x_11 x_3 x_13 ◾;
             let x_17 : obj := _private.Init.Data.Array.Basic.0.Array.mapMUnsafe.map._at_.applyProjectionRules.spec_0._redArg x_1 x_2 x_15 x_16;
             ret x_17
-    def _private.Init.Data.Array.Basic.0.Array.mapMUnsafe.map._at_.applyProjectionRules.spec_0 (x_1 : ◾) (x_2 : ◾) (x_3 : ◾) (x_4 : tobj) (x_5 : usize) (x_6 : usize) (x_7 : obj) : obj :=
-      let x_8 : obj := _private.Init.Data.Array.Basic.0.Array.mapMUnsafe.map._at_.applyProjectionRules.spec_0._redArg x_4 x_5 x_6 x_7;
-      ret x_8
-    def applyProjectionRules._redArg (x_1 : obj) (x_2 : tobj) : obj :=
-      let x_3 : usize := Array.usize ◾ x_1;
-      let x_4 : usize := 0;
-      let x_5 : obj := _private.Init.Data.Array.Basic.0.Array.mapMUnsafe.map._at_.applyProjectionRules.spec_0._redArg x_2 x_3 x_4 x_1;
-      ret x_5
-    def applyProjectionRules (x_1 : ◾) (x_2 : ◾) (x_3 : ◾) (x_4 : obj) (x_5 : tobj) : obj :=
-      let x_6 : obj := applyProjectionRules._redArg x_4 x_5;
-      ret x_6

tests/lean/reduceArity.lean.expected.out

@@ -1,3 +1,7 @@
+[Compiler.result] size: 1
+    def g (α : ◾) (n : Nat) (a : lcAny) (b : lcAny) (f : lcAny → lcAny) : lcAny :=
+      let _x.1 := g._redArg n a f;
+      return _x.1
 [Compiler.result] size: 9
     def g._redArg (n : Nat) (a : lcAny) (f : lcAny → lcAny) : lcAny :=
       let zero := 0;
@@ -11,10 +15,6 @@
         let _x.2 := g._redArg n.1 a f;
         let _x.3 := f _x.2;
         return _x.3
-[Compiler.result] size: 1
-    def g (α : ◾) (n : Nat) (a : lcAny) (b : lcAny) (f : lcAny → lcAny) : lcAny :=
-      let _x.1 := g._redArg n a f;
-      return _x.1
 [Compiler.result] size: 1
     def h._closed_0 : Nat → Nat :=
       let _x.1 := double;