Concurrency in TermElabM

Stream: lean4

[Siddhartha Gadgil]

I am trying to use tasks to get concurrency within TermElabM (to run many isDefEq instances in parallel). My attempt at this does not run concurrently in the interpreter (as I see with htop) though the non-meta variant does run concurrently. My attempt:

def exprSeq : TermElabM (Array Expr) := do
  let mut arr ← Array.mk []
  for i in [0:400000] do
    let e ← ToExpr.toExpr (i % 100)
    arr := arr.push e
  return arr

def counts: TermElabM (Array Nat) := do
  let arr : Array Nat := #[1, 2, 3, 4, 5, 6]
  let cntsAux := arr.map <| fun i => Task.spawn fun _ => count (ToExpr.toExpr i)
  let cnts ←  cntsAux.mapM <| fun t => t.get
  return cnts

set_option maxHeartbeats 100000000

#eval counts

Also not successful was a variant of this, calling get on Tasks earlier so the function counts became:

def counts: TermElabM (Array Nat) := do
  let arr : Array Nat := #[1, 2, 3, 4, 5, 6]
  let cntsAux := arr.map <| fun i => Task.spawn fun _ => count (ToExpr.toExpr i)
  let cntsAux2 := cntsAux.map <| fun t => t.get
  let cnts ←  cntsAux.mapM <| fun t => t.get
  return cnts

For reference, the non-meta code that did run concurrently was:

def slowFib (id : Nat) : Nat → Nat
  | 0   => 1
  | 1   => 1
  | n+2 => dbgTrace s!"fib {id}" fun _ => slowFib id n + slowFib id (n+1)

def conc2 : Nat :=
  let arr := #[1, 3, 4, 5, 6]
  let tsks := arr.map fun n => Task.spawn fun _ => slowFib n (30 + n)
  let res := tsks.map fun t => t.get
  res.foldl (fun acc n => acc + n) 0

#eval conc2

Is there any way around this?
Thanks

[Sebastian Ullrich]

Look at the involved types: your tasks contain TermElabM _ values, i.e. mere closures. Parallelizing TermElabM is not trivial since it contains many StateRefTs, whose mutable reference should probably not be shared across threads. And depending on your use case (e.g., are there mvars involved?), you will need to figure out how to merge the threads' states in the end.