Run tactics in parallel (asynchronously)

Stream: lean4

[Son Ho]

Is there a way of running tactics in parallel? For instance, given a list of MVarIds and a list of tactics to solve those subgoals, I would like to run the tactics on those subgoals in an asynchronous manner. As a side note: I actually tried implementing something which goes in that direction, but it doesn't work extremely well (in particular I have random failures on some of the subgoals because the asynchronous computations seem to mess up with the heartbeats count - I can share the code if people are interested to look at it).

[Sebastian Ullrich]

For a robust solution you will want to build on Lean.Core.wrapAsync, extending that to TacticM. Depending on how you want to handle messages created by the async tactics, you might also want to consider the AsSnapshot variants.

[Sebastian Ullrich]

Needless to say, I am very interested in the outcome, so please let me know if you get stuck :)

[Son Ho]

Thanks! I had attempted to mimic what Lean.Core.wrapAsync does, but building my tactics directly on top of this function solves the issue with the heart beats (I don't know why). Here is what I came up with :)

import Lean

open Lean Elab Term Meta Tactic

/-!
# Asynchronous execution
-/
/--
One issue with tactics like `omega` is that they wrap proofs in auxiliary theorems.
As a temporary workaround I'm using this helper to inline the proof terms.
-/
def inlineFreshProofs (env0 : Environment) (e : Expr) : MetaM Expr := do
  let rec inline (e : Expr) : MetaM Expr := do
    match e with
    | .bvar _ | .fvar _ | .mvar _ | .sort _ | .lit _ => pure e
    | .const declName us =>
      if env0.constants.contains declName then pure e
      else
        -- We need to inline this constant
        let const ← pure ((← getEnv).constants.find! declName)
        let some body := const.value? (allowOpaque := true)
          | throwError "Could not inline constant: {e}"
        -- Replace the levels in the body
        let levels := const.levelParams
        let levels := Std.HashMap.ofList (List.zip (List.map Level.param levels) us)
        let body := body.replaceLevel (Std.HashMap.get? levels)
        pure body
    | .app fn arg => pure (.app (← inline fn) (← inline arg))
    | .lam name ty body info =>
      pure (.lam name (← inline ty) (← inline body) info)
    | .forallE name ty body info =>
      pure (.forallE name (← inline ty) (← inline body) info)
    | .letE name ty value body nonDep =>
      pure (.letE name (← inline ty) (← inline value) (← inline body) nonDep)
    | .mdata data e => pure (.mdata data (← inline e))
    | .proj name idx struct =>
      pure (.proj name idx (← inline struct))
  inline (← instantiateMVars e)

def wrapTactic {α : Type} (tactic : α → TacticM Unit) (cancelTk? : Option IO.CancelToken)  :
  TacticM (α → BaseIO (Except Exception (Option Expr))) := do
  let env0 ← getEnv
  -- The code below is adapted from `Lean.Elab.Tactic.tacticToDischarge` --
  let ref ← IO.mkRef (← getThe Term.State)
  let ctx ← readThe Term.Context
  let mvar ← mkFreshExprSyntheticOpaqueMVar (← getMainTarget) `simp.discharger
  let s ← ref.get
  let runTac? (x : α) : TermElabM (Option Expr) :=
    try
      --Term.withoutModifyingElabMetaStateWithInfo do
        let ngoals ← --Term.withSynthesize (postpone := .no) do
          Tactic.run mvar.mvarId! (tactic x)
        if ngoals.isEmpty then
          let result ← instantiateMVars mvar
          if result.hasExprMVar then
            return none
          else
            let result ← inlineFreshProofs env0 result
            return some result
        else return none
    catch _ =>
      return none
  --------------------------------------------------------------------------
  let metaCtx ← readThe Meta.Context
  let metaState ← getThe Meta.State
  let tac (x : α) := do
    MetaM.run' (ctx := metaCtx) (s := metaState)
    <| Term.TermElabM.run' (runTac? x) ctx s
  let tac ← Lean.Core.wrapAsync tac cancelTk?
  pure (fun x => (tac x).toBaseIO)

/- Run `tac` on the current goals in parallel -/
def allGoalsAsync' (tac : TacticM Unit) : TacticM Unit := do
  let mvarIds ← getGoals
  let mut tasks := #[]
  -- Create tasks
  for mvarId in mvarIds do
    unless (← mvarId.isAssigned) do
      setGoals [mvarId]
      -- TODO: how to check for task cancellation?
      let tac ← pure ((← wrapTactic (fun () => tac) none) ())
      tasks := tasks.push (← (BaseIO.asTask tac))
  -- Wait for the tasks
  let mut unsolved := #[]
  for (i, (mvarId, task)) in (List.zip mvarIds tasks.toList).mapIdx (fun id x => (id, x)) do
    match task.get with
    | .error e =>
      logInfo m!"Task {i}: exception: {e.toMessageData}"
      unsolved := unsolved.push mvarId
    | .ok x =>
      -- TODO: how to propagate the information from the `Core.State`? (in particular, I
      -- want to preserve the log)
      match x with
      | .none =>
        logInfo m!"Task {i}: tactic failed"
        unsolved := unsolved.push mvarId
      | .some x =>
        /- Successfully generated a proof! Assign the meta-variable -/
        logInfo m!"Task {i}: solved!"
        /- TODO: would it be interesting to introduce theorems for all the proofs
           of the subgoals, so that they get type-checked in parallel? -/
        mvarId.assign x
  setGoals unsolved.toList


/-!
# Test
-/

/-!
## Small helpers, unrelated to asynchronous executions
-/

-- Copy/pasted from Mathlib to make this file more self-containted
def allGoals (tac : TacticM Unit) : TacticM Unit := do
  let mvarIds ← getGoals
  let mut mvarIdsNew := #[]
  for mvarId in mvarIds do
    unless (← mvarId.isAssigned) do
      setGoals [mvarId]
      try
        tac
        mvarIdsNew := mvarIdsNew ++ (← getUnsolvedGoals)
      catch ex =>
        if (← read).recover then
          logException ex
          mvarIdsNew := mvarIdsNew.push mvarId
        else
          throw ex
  setGoals mvarIdsNew.toList

-- Repeatedly apply a tactic
partial def repeatTac (tac : TacticM Unit) : TacticM Unit := do
  try
    tac
    allGoals (focus (repeatTac tac))
  -- TODO: does this restore the state?
  catch _ => pure ()

def optSplitConj (e : Expr) : Expr × Option Expr :=
  e.consumeMData.withApp fun f args =>
  if f.isConstOf ``And ∧ args.size = 2 then (args[0]!, some (args[1]!))
  else (e, none)

-- Split the goal if it is a conjunction
def splitConjTarget : TacticM Unit := do
  withMainContext do
  let g ← getMainTarget
  -- The tactic was initially implemened with `_root_.Lean.MVarId.apply`
  -- but it tended to mess the goal by unfolding terms, even when it failed
  let (l, r) := optSplitConj g
  match r with
  | none => do throwError "The goal is not a conjunction"
  | some r => do
    let lmvar ← mkFreshExprSyntheticOpaqueMVar l
    let rmvar ← mkFreshExprSyntheticOpaqueMVar r
    let and_intro ← mkAppM ``And.intro #[lmvar, rmvar]
    let g ← getMainGoal
    g.assign and_intro
    let goals ← getUnsolvedGoals
    setGoals (lmvar.mvarId! :: rmvar.mvarId! :: goals)

/-!
## The test
-/

/- Solve the goal by splitting the conjunctions -/
def trySync : TacticM Unit := do
  repeatTac splitConjTarget
  allGoals (do
    --IO.sleep 100
    Tactic.Omega.omegaTactic {})

syntax "sync_solve" : tactic
elab "sync_solve" : tactic => do trySync

def tryAsync : TacticM Unit := do
  repeatTac splitConjTarget
  allGoalsAsync' (do
    --IO.sleep 100
    Tactic.Omega.omegaTactic {})

syntax "async_solve" : tactic
elab "async_solve" : tactic => do tryAsync

abbrev goal (x y : Nat) : Prop :=
  x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧
  x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧
  x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧
  x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧
  x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧
  x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧
  x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧
  x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧
  x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧
  x + y < 30 ∧ x + y < 30 ∧ x + y < 30 ∧ x + y < 30

set_option trace.profiler true in
def test (x y : Nat) (h : x < 10) (h : y < 20) : goal x y
  := by
  /- The async execution is actually slower than the synchronous one (I get
     ~0.30s with a quite high variance versus 0.10s), but when using a more
     expensive tactic on a more difficult goal (I'm not doing this here for
     the purpose of keeping the file self-contained) I do have a speedup
     (I need to do more benchmarks).
  -/
  unfold goal
  async_solve -- compare with `sync_solve`

[Son Ho]

I do have a few questions though:

• as you can see from the code, I have an issue with the fact that tactics like omega introduce auxiliary theorems. What would be the best way of handling this? I'm not sure about how to retrieve the list of fresh theorems, and how to add them to the main environment. (I just need pointers)
• following what is being done in MutualDef.lean, I believe I should use promises to turn wrapAsyncAsSnapshot into something which outputs arbitrary values. Also, in order to preserve the log all I need to do is to give the snapshot tree to Core.logSnapshotTask. Is this correct? I managed to make this work
• finally, I would like to introduce auxiliary theorems for the proofs generated by the tactics. Do you have some pointers to do this in an asynchronous manner (this may be redundant with my first question)?

[𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒]

I'm not sure about how to retrieve the list of fresh theorems,

I am guessing you mean the list of theorems introduced by a tactic that you ran via Tactic.run or a similar mechanism? You could look at the implementation of whatsnew for that. If you were doing everything synchronously, then I think you wouldn't have to do this because omega would just modify the Environment and the TacticM state would pass the new Environment through. With async elaboration, you may have to commit the async environment branch containing all those new theorems.

I would like to introduce auxiliary theorems for the proofs generated by the tactics. Do you have some pointers to do this in an asynchronous manner

Afaict the mechanism for adding new constants to the environment asynchronously is docs#Lean.Environment.addConstAsync. The AddConstAsyncResult structure contains two Environments: the 'main' one which will ultimately contain all the constants you add, and an 'async branch' which is what the other thread should add its constants into. Then you can you use commitConst and commitCheckEnv (I think the difference between these two is that the latter commits a kernel-checked environment, and the former just the (possibly malformed) Lean-userspace-environment) to send results from async construction back to the main thread. addConstAsync can only add one constant, it seems like, but you can nest calls to it to commit multiple constants.

As far as adding new lemmas itself is concerned (sync or async), I would look at as_aux_lemma. I started writing an async version of this tactic (in the sense that it would commit the lemma's signature to the main thread immediately, but then continue checking the proof asynchronously) but ran into issues with interactive feedback as in this thread.

[𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒]

One think I'm not sure how to deal with is that addConstAsync expects you to give it the name of the async constant before you launch the async task. But if you want to asyncly add a bunch of theorems dumped by a tactic, you will not know their names (and possibly their kinds) a priori. Then maybe you have to do write some more complex, manual Promise-based code.

This is speculative, so I welcome corrections from the core devs.

[Sebastian Ullrich]

The given name is used to collect any asynchronously added constants, you don't need to know those a priori. You can use an auxiliary decl for addConstAsync but if we're already in an asynchronous context, it might be possible to extend the API so that this intermediary constant is not needed.

[Son Ho]

𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒 said:

You could look at the implementation of whatsnew for that

Yes, that's what I need to do, but I hoped there would be a way of achieving it without performing a linear scan of the environment? This seems to be what I need, thanks! If I understand correctly, this function ignores the imported definitions, it only goes through the ones local to the file (because of this line)?

[Son Ho]

𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒 said:

If you were doing everything synchronously, then I think you wouldn't have to do this because omega would just modify the Environment and the TacticM state would pass the new Environment through. With async elaboration, you may have to commit the async environment branch containing all those new theorems.

Yes indeed, the problem only arises in an asynchronous setting.

[Son Ho]

𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒 said:

As far as adding new lemmas itself is concerned (sync or async), I would look at as_aux_lemma. I started writing an async version of this tactic (in the sense that it would commit the lemma's signature to the main thread immediately, but then continue checking the proof asynchronously) but ran into issues with interactive feedback as in this thread.

Thanks for the pointer! I actually ended up using mkAuxTheorem which I believe is asynchronous (it calls this function under the hood)?

[𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒]

this function ignores the imported definitions, it only goes through the ones local to the file

Yeah; that's fine for you, right? You can only import definitions at the top of a module in the import block anyway, not after that.

[𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒]

Son Ho said:

Thanks for the pointer! I actually ended up using mkAuxTheorem which I believe is asynchronous (it calls this function under the hood)?

I don't think mkAuxTheorem on its own is async, in the sense that if you just call it on the main thread it would run there and block. It would be async if you fork off an async environment branch, set that to be your MetaM env in a new task, and call mkAuxTheorem then.

[Son Ho]

𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒 said:

this function ignores the imported definitions, it only goes through the ones local to the file

Yeah; that's fine for you, right? You can only import definitions at the top of a module in the import block anyway, not after that.

Yes, it's exactly what I need: I was afraid that it would go through the whole environment :)

[Son Ho]

𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒 said:

I don't think mkAuxTheorem on its own is async

Ok, I'll use addConstAsync then