feat: optional parallel kernel replay (parallel_replay)

.github/workflows/lean_action_ci.yml

@@ -54,3 +54,7 @@ jobs:
         run: |
           export PATH="$GITHUB_WORKSPACE/nanoda/target/release:$GITHUB_WORKSPACE/lean4export/.lake/build/bin:$GITHUB_WORKSPACE/landrun:$PATH"
           lean --run runtests.lean
+
+      - name: Run parallel replay kernel tests
+        working-directory: comparator
+        run: lake env lean --run tests/ParallelReplayUnit.lean

Comparator.lean

@@ -5,3 +5,4 @@ Authors: Henrik Böving
 -/
 import Comparator.Axioms
 import Comparator.Compare
+import Comparator.Replay

Comparator/Replay.lean

@@ -0,0 +1,178 @@
+/-
+Copyright (c) 2023 Kim Morrison. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Kim Morrison
+-/
+module
+
+public import Lean.AddDecl
+import all Lean.Replay
+import Lean.Util.FoldConsts
+
+/-!
+# `Lean.Kernel.Environment.replayParallel`
+
+`replayConstant`, `replayConstants` and `replay` in namespace `Lean.Kernel.Environment.Replay.Parallel` below are
+copies of the same definitions in `src/Lean/Replay.lean` (identical in Lean v4.35.0-rc3, v4.35.0-rc4 and master).
+Everything they use from there (`Context`, `State`, `isTodo`, `throwKernelException`, `addDecl` and the
+postponed constructor and recursor checks) is imported from `Lean.Replay`, with `import all` because those
+definitions are not public. The copies differ from Lean's only in:
+
+* the `.thmInfo` case of `replayConstant` calls `addThmAsync` instead of `addDecl`: the theorem is added with
+  `addDeclWithoutChecking`, and checked by the same kernel call (`addDeclCore`) against the environment as it was
+  just before it, in a separate task;
+* the reader context `Context` extends `Replay.Context` with the list of these tasks;
+* `replay` creates that list, and waits for all tasks before the postponed constructor and recursor checks,
+  failing if any task failed.
+
+So every declaration is checked by the same kernel call against the same environment as in `replay`; only the
+order in which theorems are checked differs.
+-/
+
+namespace Lean.Kernel.Environment.Replay.Parallel
+
+/-- `Replay.Context`, plus the kernel checks of theorems running in parallel (see `addThmAsync`). -/
+structure Context extends Replay.Context where
+  tasks : IO.Ref (Array (Name × Task (Except Kernel.Exception Kernel.Environment)))
+
+abbrev M := ReaderT Context <| StateRefT State IO
+
+/-- The helpers of `Lean.Replay` (`isTodo`, `addDecl`, ...) run unchanged in `M`. -/
+instance : MonadLift Replay.M M where
+  monadLift x := fun ctx => x.run ctx.toContext
+
+/--
+Add a theorem to the environment without checking it, and spawn a task that checks it with the kernel
+against the environment as it was before the theorem was added.
+-/
+def addThmAsync (info : TheoremVal) : M Unit := do
+  let kenv := (← get).env
+  let decl := Declaration.thmDecl info
+  match kenv.addDeclWithoutChecking decl with
+  | .ok env =>
+    let t := Task.spawn fun () => kenv.addDeclCore 0 0 decl (cancelTk? := none)
+    modify fun s => { s with env := env }
+    (← read).tasks.modify (·.push (info.name, t))
+  | .error ex => throwKernelException ex
+
+mutual
+/--
+Check if a `Name` still needs to be processed (i.e. is in `remaining`).
+
+If so, recursively replay any constants it refers to,
+to ensure we add declarations in the right order.
+
+The construct the `Declaration` from its stored `ConstantInfo`,
+and add it to the environment.
+-/
+partial def replayConstant (name : Name) : M Unit := do
+  if ← isTodo name then
+    let some ci := (← read).newConstants[name]? | unreachable!
+    replayConstants ci.getUsedConstantsAsSet
+    -- Check that this name is still pending: a mutual block may have taken care of it.
+    if (← get).pending.contains name then
+      try
+        match ci with
+        | .defnInfo   info =>
+          addDecl (Declaration.defnDecl   info)
+        | .thmInfo    info =>
+          -- Ignore duplicate theorems. This code is identical to that in `finalizeImport` before it
+          -- added extended duplicates support for the module system, which is not relevant for us
+          -- here as we always load all .olean information. We need this case *because* of the module
+          -- system -- as we have more data loaded than it, we might encounter duplicate private
+          -- theorems where elaboration under the module system would have only one of them in scope.
+          if let some (.thmInfo info') := (← get).env.find? ci.name then
+            if info.name == info'.name &&
+              info.type == info'.type &&
+              info.levelParams == info'.levelParams &&
+              info.all == info'.all
+            then
+              return
+          addThmAsync info
+        | .axiomInfo  info =>
+          addDecl (Declaration.axiomDecl  info)
+        | .opaqueInfo info =>
+          addDecl (Declaration.opaqueDecl info)
+        | .inductInfo info =>
+          let lparams := info.levelParams
+          let nparams := info.numParams
+          let all ← info.all.mapM fun n => do pure <| ((← read).newConstants[n]!)
+          for o in all do
+            modify fun s =>
+              { s with remaining := s.remaining.erase o.name, pending := s.pending.erase o.name }
+          let ctorInfo ← all.mapM fun ci => do
+            pure (ci, ← ci.inductiveVal!.ctors.mapM fun n => do
+              pure ((← read).newConstants[n]!))
+          -- Make sure we are really finished with the constructors.
+          for (_, ctors) in ctorInfo do
+            for ctor in ctors do
+              replayConstants ctor.getUsedConstantsAsSet
+          let types : List InductiveType := ctorInfo.map fun ⟨ci, ctors⟩ =>
+            { name := ci.name
+              type := ci.type
+              ctors := ctors.map fun ci => { name := ci.name, type := ci.type } }
+          addDecl (Declaration.inductDecl lparams nparams types false)
+        -- We postpone checking constructors,
+        -- and at the end make sure they are identical
+        -- to the constructors generated when we replay the inductives.
+        | .ctorInfo info =>
+          modify fun s => { s with postponedConstructors := s.postponedConstructors.insert info.name }
+        -- Similarly we postpone checking recursors.
+        | .recInfo info =>
+          modify fun s => { s with postponedRecursors := s.postponedRecursors.insert info.name }
+        | .quotInfo _ =>
+          -- `Quot.lift` and `Quot.ind` have types that reference `Eq`,
+          -- so we need to ensure `Eq` is replayed before adding the quotient declaration.
+          replayConstant `Eq
+          addDecl (Declaration.quotDecl)
+        modify fun s => { s with pending := s.pending.erase name }
+      catch ex =>
+        throw <| .userError s!"while replaying declaration '{name}':\n{ex}"
+
+/-- Replay a set of constants one at a time. -/
+partial def replayConstants (names : NameSet) : M Unit := do
+  for n in names do replayConstant n
+
+end
+
+/--
+"Replay" some constants into a `Kernel.Environment`, sending them to the kernel for checking.
+
+Throws a `IO.userError` if the kernel rejects a constant,
+or if there are malformed recursors or constructors for inductive types.
+-/
+public def replay (newConstants : Std.HashMap Name ConstantInfo) (env : Kernel.Environment) :
+    IO Kernel.Environment := do
+  let mut remaining : NameSet := ∅
+  for (n, ci) in newConstants.toList do
+    -- We skip unsafe constants, and also partial constants.
+    -- Later we may want to handle partial constants.
+    if !ci.isUnsafe && !ci.isPartial then
+      remaining := remaining.insert n
+  let tasks ← IO.mkRef #[]
+  let (_, s) ← StateRefT'.run (s := { env, remaining }) do
+    ReaderT.run (r := { newConstants, tasks }) do
+      for n in remaining do
+        replayConstant n
+      -- Wait for the kernel checks of all theorems; any failure fails the replay.
+      for (name, t) in ← tasks.get do
+        if let .error ex := t.get then
+          try throwKernelException ex
+          catch ex => throw <| .userError s!"while replaying declaration '{name}':\n{ex}"
+      checkPostponedConstructors
+      checkPostponedRecursors
+  return s.env
+
+end Lean.Kernel.Environment.Replay.Parallel
+
+namespace Lean.Kernel.Environment
+
+/--
+"Replay" some constants into a `Kernel.Environment`, sending them to the kernel for checking: the variant of
+`Lean.Kernel.Environment.replay` defined above.
+-/
+public def replayParallel (newConstants : Std.HashMap Name ConstantInfo) (env : Kernel.Environment) :
+    IO Kernel.Environment :=
+  Replay.Parallel.replay newConstants env
+
+end Lean.Kernel.Environment

Main.lean

@@ -21,6 +21,7 @@ structure Context where
   whichLandrun : String
   whichLean4Export : String
   externalKernels : (Std.TreeMap String (Array String))
+  parallelReplay : Bool
 
 abbrev M := ReaderT Context IO
 
@@ -219,7 +220,11 @@ def runBuiltinKernel (solution : Export.ExportedEnv) : M (Option String) := do
   let quotTargets := [`Quot.mk, `Quot.lift, `Quot.ind]
   let kernelConstMap := quotTargets.foldl (init := origConstMap) (·.erase ·)
   try
-    kernelEnv ← kernelEnv.replay kernelConstMap
+    if (← read).parallelReplay then
+      IO.println "Using parallel replay: theorems are checked by the kernel in parallel tasks."
+      kernelEnv ← kernelEnv.replayParallel kernelConstMap
+    else
+      kernelEnv ← kernelEnv.replay kernelConstMap
     IO.println "Lean default kernel accepts the solution"
   catch e =>
     IO.println "Lean default kernel rejects the solution"
@@ -325,6 +330,7 @@ structure Config where
   permitted_axioms : Array String
   enable_nanoda? : Option Bool
   external_kernels? : Option (Std.TreeMap String (Array String))
+  parallel_replay? : Option Bool
   deriving Lean.FromJson, Lean.ToJson, Repr
 
 def M.run (x : M α) (cfg : Config) : IO α := do
@@ -361,7 +367,8 @@ def M.run (x : M α) (cfg : Config) : IO α := do
     gitLocation := gitLocation,
     whichLean4Export := whichLean4Export,
     whichLandrun := whichLandrun,
-    externalKernels := externalKernels
+    externalKernels := externalKernels,
+    parallelReplay := cfg.parallel_replay?.getD false
   }
 
 end Comparator

README.md

@@ -85,6 +85,14 @@ moves toward having an option to receive the input file as a `CLI` argument.
 
 For development purposes, comparator supports overriding `nanoda` specifically using the
 `COMPARATOR_NANODA` environment variable.
+## Parallel Kernel Replay
+For very large solutions the replay through the Lean kernel dominates the run time, because it checks
+every declaration one at a time. Setting `"parallel_replay": true` in the configuration makes comparator
+check theorems in parallel tasks on Lean's task pool (`LEAN_NUM_THREADS` sets the number of workers).
+Definitions, inductive types and all other declarations are still checked in order on one thread. Every
+theorem is checked by the same kernel call, against the same environment, as in the default sequential
+replay, and comparator waits for all theorem checks before accepting the solution. The default is `false`.
+
 ## Definition Holes
 Sometimes challenges want to leave open definitions for solutions to fill in. This can range from
 simple things like filling in a `Prop` valued definition to resolve whether a conjecture is true or

tests/ParallelReplayUnit.lean

@@ -0,0 +1,70 @@
+import Lean
+import Comparator.Replay
+
+/-!
+Kernel-level tests of `Lean.Kernel.Environment.replayParallel`. Every case must be accepted or rejected
+exactly as by `Lean.Kernel.Environment.replay`; the rejected cases include proofs that the elaborator would
+never produce, so they are built directly as `ConstantInfo`s.
+
+Run with `lake env lean --run tests/ParallelReplayUnit.lean` (after `lake build`).
+-/
+
+open Lean
+
+/-- `inductive TestTrue : Prop | intro` and `inductive TestFalse : Prop` (no constructors). -/
+def propInductives : List ConstantInfo :=
+  let trueInd : InductiveVal := {
+    name := `TestTrue, levelParams := [], type := .sort .zero,
+    numParams := 0, numIndices := 0, all := [`TestTrue], ctors := [`TestTrue.intro],
+    numNested := 0, isRec := false, isUnsafe := false, isReflexive := false }
+  let trueCtor : ConstructorVal := {
+    name := `TestTrue.intro, levelParams := [], type := .const `TestTrue [],
+    induct := `TestTrue, cidx := 0, numParams := 0, numFields := 0, isUnsafe := false }
+  let falseInd : InductiveVal := {
+    name := `TestFalse, levelParams := [], type := .sort .zero,
+    numParams := 0, numIndices := 0, all := [`TestFalse], ctors := [],
+    numNested := 0, isRec := false, isUnsafe := false, isReflexive := false }
+  [.inductInfo trueInd, .ctorInfo trueCtor, .inductInfo falseInd]
+
+def thm (name : Name) (type value : Expr) : ConstantInfo :=
+  .thmInfo { name, levelParams := [], type, value, all := [name] }
+
+/-- Replay `propInductives ++ extra` with both `replay` and `replayParallel`; return whether both
+gave the expected verdict. -/
+def runBoth (label : String) (expectAccept : Bool) (extra : List ConstantInfo) : IO Bool := do
+  let constMap := (propInductives ++ extra).foldl (init := {}) fun m ci => m.insert ci.name ci
+  let kenv := (← mkEmptyEnvironment).toKernelEnv
+  let seqOk ← try discard <| kenv.replay constMap; pure true catch _ => pure false
+  let parOk ← try discard <| kenv.replayParallel constMap; pure true catch _ => pure false
+  let ok := seqOk == expectAccept && parOk == expectAccept
+  IO.println s!"{if ok then "ok  " else "FAIL"} {label}: expected accept={expectAccept}, replay={seqOk}, replayParallel={parOk}"
+  return ok
+
+def main : IO UInt32 := do
+  let tru := Expr.const `TestTrue []
+  let intro := Expr.const `TestTrue.intro []
+  let fls := Expr.const `TestFalse []
+  let results := #[
+    -- a valid chain: `lemA : TestTrue := intro`, `thmB : TestTrue := lemA`
+    ← runBoth "valid chain of theorems" true
+      [thm `lemA tru intro, thm `thmB tru (.const `lemA [])],
+    -- a theorem whose proof has the wrong type (`TestTrue` is a type, not a proof of `TestTrue`)
+    ← runBoth "ill-typed proof of the last theorem" false
+      [thm `lemA tru intro, thm `thmB tru tru],
+    -- an ill-typed lemma used by a theorem whose own proof is fine given the lemma's statement
+    ← runBoth "ill-typed proof of a lemma used later" false
+      [thm `lemBad fls intro, thm `thmUsesBad fls (.const `lemBad [])],
+    -- two theorems that prove `TestFalse` from each other
+    ← runBoth "cyclic theorems" false
+      [thm `cycleA fls (.const `cycleB []), thm `cycleB fls (.const `cycleA [])],
+    -- an ill-typed definition (checked on the main thread, as in `replay`)
+    ← runBoth "ill-typed definition" false
+      [.defnInfo { name := `badDef, levelParams := [], type := fls, value := intro,
+                   hints := .abbrev, safety := .safe, all := [`badDef] }]
+  ]
+  if results.all id then
+    IO.println "All parallel replay kernel tests passed."
+    return 0
+  else
+    IO.println "Some parallel replay kernel tests FAILED."
+    return 1

tests/projects/parallel_match/Challenge.lean

@@ -0,0 +1,2 @@
+theorem comm (n m : Nat) : n + m = m + n := by
+  sorry

tests/projects/parallel_match/Solution.lean

@@ -0,0 +1,5 @@
+theorem add_zero_right (n : Nat) : n + 0 = n := rfl
+
+theorem comm (n m : Nat) : n + m = m + n := by
+  have _ := add_zero_right n
+  grind

tests/projects/parallel_match/config.json

@@ -0,0 +1,8 @@
+{
+    "challenge_module": "Challenge",
+    "solution_module": "Solution",
+    "theorem_names": ["comm"],
+    "permitted_axioms": ["propext", "Quot.sound", "Classical.choice"],
+    "enable_nanoda": false,
+    "parallel_replay": true
+}

tests/projects/parallel_match/test.json

@@ -0,0 +1,3 @@
+{
+        "exit_code": 0
+}

tests/projects/parallel_olean_issue/Challenge.lean

@@ -0,0 +1 @@
+theorem boom : False := sorry

tests/projects/parallel_olean_issue/Solution.lean

@@ -0,0 +1,22 @@
+import Lean
+
+open Lean
+
+unsafe def badNatUnsafe : Nat := unsafeCast (-4294967296 : Int)
+
+@[implemented_by badNatUnsafe] opaque badNatVal : Nat
+
+run_elab
+  addDecl <| .defnDecl {
+    name        := .str .anonymous "badNat"
+    levelParams := []
+    type        := .const ``Nat []
+    value       := .lit <| .natVal badNatVal
+    hints       := .opaque
+    safety      := .safe
+  }
+
+theorem boom : False := by
+  have truly_marvelous_0 : ¬badNat ≤ 9223372036854775807 := by decide
+  have truly_marvelous_1 : ¬9223372036854775807 < badNat := by decide
+  omega

tests/projects/parallel_olean_issue/config.json

@@ -0,0 +1,8 @@
+{
+    "challenge_module": "Challenge",
+    "solution_module": "Solution",
+    "theorem_names": ["boom"],
+    "permitted_axioms": ["propext", "Quot.sound", "Classical.choice"],
+    "enable_nanoda": false,
+    "parallel_replay": true
+}

tests/projects/parallel_olean_issue/test.json

@@ -0,0 +1,3 @@
+{
+        "exit_code": 1
+}