Lean 4 Tasks

Stream: new members

[Sébastien Boisgérault]

Hi everyone! :wave:

I was (naively?) expecting the following code to implement a parallel map:

def pmap {α β} (f : α → β) (inputs : List α) : List β :=
  let t_inputs : List (Task α) := inputs.map Task.pure
  let t_outputs := t_inputs.map (fun t_input => t_input.map f)
  t_outputs.map Task.get

(since my understanding was that the calls to fun t_input => t_input.map fhappen sequentially but each call is non-blocking and starts the application of f in the background).

But my experiments seem to show that a single thread is used anyway, for example when I run

def countdown (n : Nat) : Nat :=
  match n with
  | 0 => 0
  | n + 1 => countdown n

def test : IO Unit :=
  let n := 10_000_000_000
  let inputs := [n, 2 * n, 3 * n, 4 * n, 5 * n, 6 * n, 7 * n, 8 * n]
  inputs |> pmap countdown |> IO.println

def main := test

Can you pinpoint where I am wrong?

Kind regards,

Sébastien

[Mirek Olšák]

Interesting, when I run it with metaprogramming, it paralelizes. By the way it also paralellizes with the countdown but run at your own risk.

def test : Lean.Meta.MetaM (List Nat) :=
  let n := 10_000_000_000
  let inputs := [n, 2 * n, 3 * n, 4 * n, 5 * n, 6 * n, 7 * n, 8 * n]
  return pmap (fun x ↦ dbgSleep 1000 (fun _ ↦ x)) inputs
  -- return pmap countdown inputs

run_meta test

[Robin Arnez]

The code of lean_task_map_core:

extern "C" LEAN_EXPORT obj_res lean_task_map_core(obj_arg f, obj_arg t, unsigned prio,
      bool sync, bool keep_alive) {
    if (!g_task_manager || (sync && lean_to_task(t)->m_value)) {
        return lean_task_pure(apply_1(f, lean_task_get_own(t)));
    } else {
        lean_task_object * new_task = alloc_task(mk_closure_3_2(task_map_fn, f, t), sync ? LEAN_SYNC_PRIO : prio, keep_alive);
        g_task_manager->add_dep(lean_to_task(t), new_task);
        return (lean_object*)new_task;
    }
}

I suppose if no task has ever been spawned then g_task_manager is null and it never considers to spawn a task. I think the reason this is happening is that it is expected that the task t was spawned using e.g. docs#Task.spawn and then docs#Task.map would register that f should run after the original task finished on the same thread as the original task.

That is, the intended way to use the system would rather be:

def pmap {α β} (f : α → β) (inputs : List α) : List β :=
  let t_outputs := inputs.map (fun input => Task.spawn fun _ => f input)
  t_outputs.map Task.get

or maybe

def pmap {α β} (f : α → β) (inputs : List α) : List β :=
  let t_inputs := inputs.map (fun input => Task.spawn fun _ => input)
  let t_outputs := t_inputs.map (fun t_input => t_input.map f)
  t_outputs.map Task.get

[Sébastien Boisgérault]

OK, I can understand the issue with not spawning explicitly any task. But if I actually also tried a pmap that corresponds closely to your second solution:

def pmap {α β} (f : α → β) (inputs : List α) : List β :=
  -- let t_inputs : List (Task α) := inputs.map Task.pure
  let t_inputs : List (Task α)  := inputs.map fun input => Task.spawn (fun () => input)
  let t_outputs := t_inputs.map (fun t_input => t_input.map f)
  t_outputs.map Task.get

and AFAICT, I still get only 100% for the process when I run top (which I interpret as "not running in parallel", but maybe I am wrong?).

On the other hand, I get initially a CPU load with top of ~ 800% when I run for example

def main : IO Unit :=
  let n := 10_000_000_000
  let inputs := [n, 2 * n, 3 * n, 4 * n, 5 * n, 6 * n, 7 * n, 8 * n]
  let tasks := inputs.map fun input => (Task.spawn fun () => countdown input)
  let outputs := tasks.map Task.get
  IO.println outputs

Is there something else going on here?

[Sébastien Boisgérault]

Update. Well, I need to reinvestigate all this, everything is "only" à 100% CPU now including the baseline that was using 800% CPU before. I need to examine the options of top (or htop) more closely.

[Sébastien Boisgérault]

Update. I may have messed up something with the Lake project itself, the program doesn't seem to be refreshed when I edit it and redo a lake build. I will create a smaller, simpler project with only these Tasks-related test and will report back here.

Thanks everyone for your help so far! :pray:

[Sébastien Boisgérault]

OK, I am officially confused. Running the code

def countdown (n : Nat) : Nat :=
  match n with
  | 0 => 0
  | n + 1 => countdown n

def main : IO Unit := do
  let n := 10_000_000_000
  let inputs := 8 |> List.range |>.map (· + 1) |>.map (· * n)
  IO.println s!"countdown inputs: {inputs}"

  let t0 <- IO.monoMsNow
  let tasks := inputs.map fun input => (Task.spawn fun () => countdown input)
  let outputs := tasks.map Task.get
  let dt := (<- IO.monoMsNow) - t0

  IO.println s!"countdown outputs: {outputs}"
  IO.println s!"elapsed time: {dt} ms"

takes some time, then yields

countdown inputs: [10000000000, 20000000000, 30000000000, 40000000000, 50000000000, 60000000000, 70000000000, 80000000000]
countdown outputs: [0, 0, 0, 0, 0, 0, 0, 0]
elapsed time: 0 ms

List.map is not lazy here, right? And Task.get should be blockin, so why is the elapsed time 0?

Update. Ah, I guess that Lean probably feels ok with reordering the task code since it is pure? OK, let's for get about this side issue for a moment, I can put the time measure after the final IO.println.

[Sébastien Boisgérault]

Final update. OK, sorry for the noise everyone, I was wrong.

After carefully doing all the experiments again, tasks are parallelized in all cases, including when I am using Task.pure (when Task.spawn could be more "natural").

That means that

def pmap {α β} (f : α → β) (inputs : List α) : List β :=
  let t_inputs := inputs.map Task.pure
  let t_outputs := t_inputs.map (fun t_input => t_input.map f)
  t_outputs.map Task.get

works in parallel and that I can rewrite it as:

import Lake -- declare Task as a Monad

def pmap {α β} (f : α → β) (inputs : List α) : List β :=
  inputs
    |>.map (fun input : α => return input)
    |>.map (fun t_input => return f (<- t_input))
    |>.map Task.get

if I wish.

[Sébastien Boisgérault]

Git repo: https://github.com/boisgera/lean-tasks