Linearisation failure in Array.mapMUnsafe

[ChiCubed]

Prerequisites

Please put an X between the brackets as you perform the following steps:

• Check that your issue is not already filed:
  https://github.com/leanprover/lean4/issues
• Reduce the issue to a minimal, self-contained, reproducible test case.
  Avoid dependencies to Mathlib or Batteries.
• Test your test case against the latest nightly release, for example on
  https://live.lean-lang.org/#project=lean-nightly
  (You can also use the settings there to switch to “Lean nightly”)

Description

Array.mapMUnsafe operates on an array in-place. Under some circumstances, the generated IR fails to use the array linearly, which causes many copies and manifests as a quadratic runtime:

-- output:
-- 2^12 41.6ms
-- 2^13 160ms
-- 2^14 632ms
#eval for n in [12:15] do
  timeit s!"2^{n}" $ discard $ Array.mkArray (2^n) (IO.rand 1 6) |>.mapM id
-- (note that .mapM uses .mapMUnsafe under the hood)

To prove this is what's going on, let's take a look.

set_option trace.compiler.ir.result true
unsafe example := Array.mkArray 100000 (IO.rand 1 6) |>.mapM id

The output contains

[result]
def Array.mapMUnsafe.map._at._example._spec_1 (x_1 : usize) (x_2 : usize) (x_3 : obj) (x_4 : obj) : obj :=
    [...]
    inc x_3;
    let x_8 : obj := Array.uset ◾ x_3 x_2 x_7 ◾;
    let x_9 : obj := Array.uget ◾ x_3 x_2 ◾;
    dec x_3;
    [...]

Clearly the array is being copied. I guess this corresponds to the lines let v := and let r := in the source of mapMUnsafe, reproduced below. (But transposed for some reason?)

@[inline]
unsafe def mapMUnsafe {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → m β) (as : Array α) : m (Array β) :=
  let sz := USize.ofNat as.size
  let rec @[specialize] map (i : USize) (r : Array NonScalar) : m (Array PNonScalar.{v}) := do
    if i < sz then
     let v    := r.uget i lcProof
     -- Replace r[i] by `box(0)`.  This ensures that `v` remains unshared if possible.
     -- Note: we assume that arrays have a uniform representation irrespective
     -- of the element type, and that it is valid to store `box(0)` in any array.
     let r    := r.uset i default lcProof
     let vNew ← f (unsafeCast v)
     map (i+1) (r.uset i (unsafeCast vNew) lcProof)
    else
     pure (unsafeCast r)
  unsafeCast <| map 0 (unsafeCast as)

In the Zulip thread (linked below) fonqL observed that swapping the line let r := with the line after it prevents the copy from happening, and so we get linear runtime again. (And naturally the same result is obtained by removing that line altogether.)

Context

I first discovered this issue when I noticed that

#eval Array.mkArray 100000 (IO.rand 1 6) |>.mapM id

took an unusually long time to run.

See the associated thread on Zulip, where user fonqL reduced the problem to the form presented in this issue.

Versions

"4.11.0-nightly-2024-07-07" on live.lean-lang.org

Impact

Add 👍 to issues you consider important. If others are impacted by this issue, please ask them to add 👍 to it.