perf: improve join point finding (#10999)

This PR improves join point finding in the compiler through two means:
1. We now handle situations where a function `f` can only become a join
point when a function `g`
   becomes a join point as well correctly.
2. We introduce a second join point finding pass after specialisation
and before the following
simplification pass, as the specialiser might have introduced new join
point opportunities for
   the simplifier to exploit.

Notably in the code from #10995 we now correctly detect the missing join
point which required both
of these changes to be made.

Closes: #10995

Author: hargoniX

src/Lean/Compiler/LCNF/JoinPoints.lean

@@ -34,6 +34,23 @@ structure CandidateInfo where
   associated : Std.HashSet FVarId
   deriving Inhabited
 
+structure FindCtx where
+  /--
+  The current definition depth is defined by how many `fun` binders we are
+  nested in at the current point. Note that this does *not* include `jp`
+  binders.
+  -/
+  definitionDepth : Nat := 0
+  /--
+  A map from function declarations that are currently in scope to their
+  definition depth.
+  -/
+  scope : FVarIdMap Nat := {}
+  /--
+  The current function binder we are inside of if any.
+  -/
+  currentFunction : Option FVarId := none
+
 /--
 The state for the join point candidate finder.
 -/
@@ -42,38 +59,36 @@ structure FindState where
   All current join point candidates accessible by their `FVarId`.
   -/
   candidates : Std.HashMap FVarId CandidateInfo := ∅
-  /--
-  The `FVarId`s of all `fun` declarations that were declared within the
-  current `fun`.
-  -/
-  scope : Std.HashSet FVarId := ∅
 
-abbrev ReplaceCtx := Std.HashMap FVarId Name
 
-abbrev FindM := ReaderT (Option FVarId) StateRefT FindState ScopeM
+abbrev FindM := ReaderT FindCtx StateRefT FindState CompilerM
+
+abbrev ReplaceCtx := Std.HashMap FVarId Name
 abbrev ReplaceM := ReaderT ReplaceCtx CompilerM
 
 /--
 Attempt to find a join point candidate by its `FVarId`.
 -/
+@[inline]
 private def findCandidate? (fvarId : FVarId) : FindM (Option CandidateInfo) := do
   return (← get).candidates[fvarId]?
 
+/--
+Combinator for modifying the candidates in `FindM`.
+-/
+@[inline]
+private def modifyCandidates (f : Std.HashMap FVarId CandidateInfo → Std.HashMap FVarId CandidateInfo) : FindM Unit :=
+  modify (fun state => { state with candidates := f state.candidates })
+
 /--
 Erase a join point candidate as well as all the ones that depend on it
 by its `FVarId`, no error is thrown is the candidate does not exist.
 -/
 private partial def eraseCandidate (fvarId : FVarId) : FindM Unit := do
   if let some info ← findCandidate? fvarId then
-    modify (fun state => { state with candidates := state.candidates.erase fvarId })
+    modifyCandidates  fun cs => cs.erase fvarId
     info.associated.forM eraseCandidate
 
-/--
-Combinator for modifying the candidates in `FindM`.
--/
-private def modifyCandidates (f : Std.HashMap FVarId CandidateInfo → Std.HashMap FVarId CandidateInfo) : FindM Unit :=
-  modify (fun state => {state with candidates := f state.candidates })
-
 /--
 Remove all join point candidates contained in `a`.
 -/
@@ -91,16 +106,30 @@ Add a new join point candidate to the state.
 -/
 private def addCandidate (fvarId : FVarId) (arity : Nat) : FindM Unit := do
   let cinfo := { arity, associated := ∅ }
-  modifyCandidates (fun cs => cs.insert fvarId cinfo )
+  modifyCandidates fun cs => cs.insert fvarId cinfo
 
 /--
 Add a new join point dependency from `src` to `dst`.
 -/
 private def addDependency (src : FVarId) (target : FVarId) : FindM Unit := do
-  if let some targetInfo ← findCandidate? target then
-    modifyCandidates (fun cs => cs.insert target { targetInfo with associated := targetInfo.associated.insert src })
-  else
-    eraseCandidate src
+  modifyCandidates fun cs =>
+    cs.modify target fun targetInfo =>
+      { targetInfo with associated := targetInfo.associated.insert src }
+
+@[inline]
+private def withFnBody (decl : FunDecl) (x : FindM α) : FindM α :=
+    withReader (fun ctx => {
+        ctx with
+          definitionDepth := ctx.definitionDepth + 1,
+          currentFunction := some decl.fvarId }) do
+      x
+
+@[inline]
+private def withFnDefined (decl : FunDecl) (x : FindM α) : FindM α :=
+    withReader (fun ctx => {
+        ctx with
+          scope := ctx.scope.insert decl.fvarId ctx.definitionDepth }) do
+      x
 
 /--
 Find all `fun` declarations that qualify as a join point, that is:
@@ -135,7 +164,7 @@ this. This is because otherwise the calls to `myjp` in `f` and `g` would
 produce out of scope join point jumps.
 -/
 partial def find (decl : Decl) : CompilerM FindState := do
-  let (_, candidates) ← decl.value.forCodeM go |>.run none |>.run {} |>.run' {}
+  let (_, candidates) ← decl.value.forCodeM go |>.run {} |>.run {}
   return candidates
 where
   go : Code → FindM Unit
@@ -148,29 +177,30 @@ where
         if valId != decl.fvarId || args.size != candidateInfo.arity then
           eraseCandidate fvarId
         -- Out of scope join point candidate handling
-        else if let some upperCandidate ← read then
-          if !(← isInScope fvarId) then
-            addDependency fvarId upperCandidate
-      else
-        eraseCandidate fvarId
+        else
+          let currDepth := (← read).definitionDepth
+          let calleeDepth := (← read).scope.get! fvarId
+          if currDepth == calleeDepth then
+            return ()
+          else if calleeDepth + 1 == currDepth then
+            addDependency fvarId (← read).currentFunction.get!
+          else
+            eraseCandidate fvarId
     | _, _ =>
       removeCandidatesInLetValue decl.value
       go k
   | .fun decl k => do
-    withReader (fun _ => some decl.fvarId) do
-      withNewScope do
-        go decl.value
     addCandidate decl.fvarId decl.getArity
-    addToScope decl.fvarId
-    go k
+    withFnBody decl do
+      go decl.value
+    withFnDefined decl do
+      go k
   | .jp decl k => do
     go decl.value
     go k
   | .jmp _ args => args.forM removeCandidatesInArg
   | .return val => eraseCandidate val
-  | .cases c => do
-    eraseCandidate c.discr
-    c.alts.forM (·.forCodeM go)
+  | .cases c => c.alts.forM (·.forCodeM go)
   | .unreach .. => return ()
 
 /--
@@ -602,17 +632,23 @@ where
 
 end JoinPointCommonArgs
 
+def Decl.findJoinPoints? (decl : Decl) : CompilerM (Option Decl) := do
+  let findResult ← JoinPointFinder.find decl
+  trace[Compiler.findJoinPoints] "Found {findResult.candidates.size} jp candidates for {decl.name}"
+  if findResult.candidates.isEmpty then
+    return none
+  else
+    return some (← JoinPointFinder.replace decl findResult)
+
 /--
 Find all `fun` declarations in `decl` that qualify as join points then replace
 their definitions and call sites with `jp`/`jmp`.
 -/
 def Decl.findJoinPoints (decl : Decl) : CompilerM Decl := do
-  let findResult ← JoinPointFinder.find decl
-  trace[Compiler.findJoinPoints] "Found {findResult.candidates.size} jp candidates for {decl.name}"
-  JoinPointFinder.replace decl findResult
+  return (← Decl.findJoinPoints? decl).getD decl
 
-def findJoinPoints : Pass :=
-  .mkPerDeclaration `findJoinPoints Decl.findJoinPoints .base
+def findJoinPoints (occurrence : Nat := 0) : Pass :=
+  .mkPerDeclaration `findJoinPoints Decl.findJoinPoints .base (occurrence := occurrence)
 
 builtin_initialize
   registerTraceClass `Compiler.findJoinPoints (inherited := true)

src/Lean/Compiler/LCNF/Passes.lean

@@ -94,6 +94,7 @@ def builtinPassManager : PassManager := {
     -- checked without nested functions whose bodies specialization does not require access to.
     checkTemplateVisibility,
     specialize,
+    findJoinPoints (occurrence := 1),
     simp (occurrence := 2),
     cse (shouldElimFunDecls := false) (occurrence := 1),
     saveBase, -- End of base phase

tests/lean/run/more_jps.lean

@@ -0,0 +1,80 @@
+@[specialize 3 4] def List.foldrNonTR (f : α → β → β) (init : β) : (l : List α) → β
+  | []     => init
+  | a :: l => f a (foldrNonTR f init l)
+
+@[always_inline, inline]
+def List.forBreak_ {α : Type u} {m : Type w → Type x} [Monad m] (xs : List α) (s : σ) (body : α → OptionT (StateT σ (ExceptT ρ m)) PUnit) (kreturn : ρ → m γ) (kbreak : σ → m γ) : m γ :=
+  List.foldrNonTR
+    (fun a acc s => do
+      let e ← body a s
+      match e with
+      | .error r => kreturn r
+      | .ok (.some _, s) => acc s
+      | .ok (none, s) => kbreak s)
+    kbreak
+    xs
+    s
+
+/--
+trace: [Compiler.saveBase] size: 25
+    def List.foldrNonTR._at_._example.spec_0 x.1 _y.2 : Nat :=
+      jp _jp.3 x : Nat :=
+        let _x.4 := 13;
+        let x := Nat.add x _x.4;
+        let x := Nat.add x _x.4;
+        let x := Nat.add x _x.4;
+        let x := Nat.add x _x.4;
+        let x := Nat.add x _x.4;
+        let x := Nat.add x _x.4;
+        let x := Nat.add x _x.4;
+        return x;
+      cases x.1 : Nat
+      | List.nil =>
+        goto _jp.3 _y.2
+      | List.cons head.5 tail.6 =>
+        let _x.7 := 0;
+        let _x.8 := instDecidableEqNat _y.2 _x.7;
+        cases _x.8 : Nat
+        | Decidable.isFalse x.9 =>
+          let _x.10 := 10;
+          let _x.11 := Nat.decLt _x.10 _y.2;
+          cases _x.11 : Nat
+          | Decidable.isFalse x.12 =>
+            let _x.13 := Nat.add _y.2 head.5;
+            let _x.14 := List.foldrNonTR._at_._example.spec_0 tail.6 _x.13;
+            return _x.14
+          | Decidable.isTrue x.15 =>
+            goto _jp.3 _y.2
+        | Decidable.isTrue x.16 =>
+          return _y.2
+[Compiler.saveBase] size: 9
+    def _example : Nat :=
+      let x := 42;
+      let _x.1 := 1;
+      let _x.2 := 2;
+      let _x.3 := 3;
+      let _x.4 := @List.nil _;
+      let _x.5 := @List.cons _ _x.3 _x.4;
+      let _x.6 := @List.cons _ _x.2 _x.5;
+      let _x.7 := @List.cons _ _x.1 _x.6;
+      let _x.8 := List.foldrNonTR._at_._example.spec_0 _x.7 x;
+      return _x.8
+-/
+#guard_msgs in
+set_option trace.Compiler.saveBase true in
+example := Id.run do
+  let x := 42;
+  List.forBreak_ (m:=Id) (ρ := Nat) [1, 2, 3] x (fun i => do
+      let x ← get
+      if x = 0 then throw (m := ExceptT Nat Id) x   -- return
+      else if x > 10 then failure (f := OptionT _)  -- break
+      else set (x + i) >>= fun _ => pure ())        -- continue
+      pure fun x => do
+  let x := x + 13;
+  let x := x + 13;
+  let x := x + 13;
+  let x := x + 13;
+  let x := x + 13;
+  let x := x + 13;
+  let x := x + 13;
+  return x