Parallel Execution of `List.map`

Stream: Is there code for X?

[Asei Inoue]

How can I execute List.map in parallel?

[Asei Inoue]

-- Define the syntax for a tactic named `my_exact?` (make `my_exact?` recognized as a syntax)

syntax (name := myExact?) "my_exact?" : tactic

open Lean Elab Tactic in

-- Define the implementation of the `my_exact?` tactic
@[tactic myExact?]
def evalMyExact? : Tactic := fun _stx => do
  -- Retrieve the current environment (which stores theorems and other constants)
  let env ← getEnv

  -- Get the constants in the environment and filter them under the following conditions:
  -- 1. They are not marked as unsafe
  -- 2. Their kind is either `axiom` or `thm` (theorem)
  let theorems : List Name := SMap.toList (Environment.constants env)
    |>.filter (fun (_name, info) => ! ConstantInfo.isUnsafe info)
    |>.filterMap (fun (name, _info) => do
        let kind ← getOriginalConstKind? env name
        match kind with
        | .axiom | .thm => name
        | _ => none
      )

  -- Try each of the theorems that meet the conditions
  for name in theorems do
    try
      -- Convert the name into a syntax node
      let nameStx := mkIdent name

      -- Expand and execute `apply name <;> assumption` as a tactic expression
      evalTactic <| ← `(tactic| apply $nameStx <;> assumption)

      -- If successful, log a message and exit the tactic
      logInfo m!"Applied {name} successfully."
      return

    catch _ =>
      -- On failure, continue to the next theorem
      continue

  -- If no theorem could be applied, fail as a tactic
  failure

set_option maxHeartbeats 500000 in

-- Example usage
example (x y : Nat) (h : x = y) : y = x := by
  my_exact?

I wrote a tactic that imitates exact? as shown above, but it is extremely slow, which is a problem.
I want to speed it up by using parallel execution.

[Damiano Testa]

This is not an answer to your question, but I think that exact? is relatively fast because it preprocesses lemmas, so that it does not even try to apply the ones that are clearly not applicable. For that, I think that it uses a discrimination tree. I suspect that using some quicker test to pre-filter the lemmas that you apply has a larger impact on performance than just dividing by the number of cores.

[Markus Himmel]

Using the new Async monad in Lean 4.22 (it's still Internal in Lean 4.22, so use at your own risk. It will be officially released later this year):

import Std.Internal.Async.Basic

open Std.Internal.IO.Async

def expensiveOperation (n : Nat) : Async Nat := do
  IO.sleep 1000
  return 2 * n

def manyInParallel (n : Nat) : Async Nat := do
  let tasks := (Array.range n).map expensiveOperation
  let results ← concurrentlyAll tasks
  return results.sum

-- Only runs about a second
#eval (manyInParallel 10).wait

[Niklas Halonen]

Thanks @Markus Himmel , this is really cool!

To more or less answer the question, I have devised the following definition for the Async map function for Array:

import Std.Internal.Async.Basic

open Std.Internal.IO.Async

def process (x : Nat) : Async Nat := do
  -- IO.println x
  IO.sleep 1000
  return 2 * x

def concMap (f : α → (Async β)) (xs : Array α) : Async (Array β) := concurrentlyAll (xs.map f)

#eval (concMap process (Array.range 10)).wait

List version can be gotten by pre and post-composition with Array.toList and List.toArray.

Also I have a question of my own: Why does the code with IO.println get stuck at the #eval?

In the process of experimenting with this I noticed that the Lean server at live.lean-lang.org keeps crashing after almost any code change (usually when there's a type error). I wonder if this is a known bug...

[Joe Hendrix]

It's been a while since I worked on it, but exact? is using a few things to run quickly.

It uses a discrimination tree so that one can quickly look up a pattern that could potentially match using apply in the tree. Furthermore, the tree data type isn't the standard Lean discrimination tree used by simp, but is a lazy discrimination tree that is tuned to allow patterns to be inserted at the root, but not fully expanded until needed in a lookup.

For concurrency, multiple lazy discrimination trees on subsets of available definitions are constructed in parallel and then merged.

Concurrency is likely fine as long as you have a relatively small number of theorems. The laziness tricks become important if you want to do things like quickly search all theorems in Mathlib.

[Markus Himmel]

Niklas Halonen said:

Also I have a question of my own: Why does the code with IO.println get stuck at the #eval?

In the process of experimenting with this I noticed that the Lean server at live.lean-lang.org keeps crashing after almost any code change (usually when there's a type error). I wonder if this is a known bug...

Yes, this is a known issue: lean4#738.