Iterator efficiency is very sensitive to inlining

Reproduction

prelude
import Init.NotationExtra

section

variable {α : Sort _} {β : α → Sort _} {γ : (a : α) → β a → Sort _} {C₂ : (a : α) → β a → Sort _}

@[specialize]
partial def extrinsicFix₂ [∀ a b, Nonempty (C₂ a b)]
    (F : (a : α) → (b : β a) → ((a' : α) → (b' : β a') → C₂ a' b') → C₂ a b)
    (a : α) (b : β a) :
    C₂ a b :=
  F a b (fun a' b' => extrinsicFix₂ F a' b')

end

inductive IterStep (α β) where
  | yield : (it : α) → (out : β) → IterStep α β
  | skip : (it : α) → IterStep α β
  | done : IterStep α β

abbrev PlausibleIterStep (IsPlausibleStep : IterStep α β → Prop) := Subtype IsPlausibleStep

class Iterator (α : Type) (β : outParam (Type)) where
  step : α → PlausibleIterStep (fun _ : IterStep α β => True)

@[inline]
def fold {α : Type} {β : Type} {γ : Type} [Iterator α β]
    (f : γ → β → γ) (init : γ) (it : α) : γ :=
  haveI : Nonempty γ := sorry
  extrinsicFix₂ (C₂ := fun _ _ => γ)
    (fun it acc recur =>
      match Iterator.step it with
      | ⟨.yield it' out, h⟩ =>
        recur it' (f acc out)
      | ⟨.skip it', h⟩ => recur it' acc
      | ⟨.done, _⟩ => acc) it init

structure FilterMap (α : Type) {β γ : Type}
    (f : β → Subtype (fun _ : Option γ => True)) where
  inner : α

instance FilterMap.instIterator {α β γ : Type}
    {f : β → Subtype (fun _ : Option γ => True)}
    [Iterator α β] :
    Iterator (FilterMap α f) γ where
  step it :=
    match Iterator.step it.inner with
      | ⟨.yield it' out, h⟩ =>
        match f out with
        | ⟨none, h'⟩ => ⟨.skip ⟨it'⟩, sorry⟩
        | ⟨some out', h'⟩ => ⟨.yield ⟨it'⟩ out', sorry⟩
      | ⟨.skip it', h⟩ => sorry
      | ⟨.done, h⟩ => ⟨.done, sorry⟩

@[inline]
def filterMap {α : Type} {β : Type} {γ : Type} [Iterator α β]
    (f : β → Option γ) (it : α) :
    FilterMap α (fun b => pure (f := Id) ⟨(f b), sorry⟩) :=
  ⟨it⟩

structure ArrayIterator (α : Type) where
  array : Array α
  pos : Nat

def _root_.Array.iter {α : Type} (l : Array α) :
    ArrayIterator α :=
  ⟨l, 0⟩

def PlausibleIterStep.yield {IsPlausibleStep : IterStep α β → Prop}
    (it' : α) (out : β) (h : IsPlausibleStep (.yield it' out)) :
    PlausibleIterStep IsPlausibleStep :=
  ⟨.yield it' out, h⟩

@[inline]
def PlausibleIterStep.yieldI {IsPlausibleStep : IterStep α β → Prop}
    (it' : α) (out : β) (h : IsPlausibleStep (.yield it' out)) :
    PlausibleIterStep IsPlausibleStep :=
  ⟨.yield it' out, h⟩

namespace Def

local instance ArrayIterator.instIterator {α : Type} : Iterator (ArrayIterator α) α where
  step it := if h : it.pos < it.array.size then
      PlausibleIterStep.yield ⟨it.array, it.pos + 1⟩ it.array[it.pos] sorry
    else
      ⟨.done, sorry⟩

set_option trace.Compiler true in
def f (xs : Array Nat) : Nat :=
  filterMap (fun x => Option.guard (· % 2 = 0) (3 * x)) (Array.iter xs) |> fold (init := 0) (· + ·)

end Def

namespace Inline

local instance ArrayIterator.instIterator {α : Type} : Iterator (ArrayIterator α) α where
  step it := if h : it.pos < it.array.size then
      PlausibleIterStep.yieldI ⟨it.array, it.pos + 1⟩ it.array[it.pos] sorry
    else
      ⟨.done, sorry⟩

set_option trace.Compiler true in
def f (xs : Array Nat) : Nat :=
  filterMap (fun x => Option.guard (· % 2 = 0) (3 * x)) (Array.iter xs) |> fold (init := 0) (· + ·)

end Inline

The difference between Def.f and Inline.f lies in whether their iterator instance uses the (not inlined) .yield or the (inlined) yieldI.

The IR for Def.f is less efficient because the JPCases optimization pass does not inline a join point that would be desirable to inline.

Context

This kind of behavior caused a 24% slowdown when I attempted to convert PlausibleIterStep.yield/step/don into abbrevs, which implicitly prompted the compiler to more eagerly inline them. (The always were inlined even before, but at a later stage.)

Why this happens

Compiler Optimization Analysis: filterMap Iterator Performance

Def.f is the efficient variant (it uses PlausibleIterStep.yield, which has no inline annotation) and Inline.f is the inefficient one (uses PlausibleIterStep.yieldI, which has an inline annotation).

High-level overview

We sum over an iterator of the form array.iter.filterMap f.

1. The inner iterator returns a subtype of IterStep, which optionally provides the next element in the array and the successor iterator.
2. Pattern matching is performed on this subtype of IterStep. If it contains a yield, f is called on the output value (next element in the array).
3. f itself returns a subtype that contains an Option.
4. If f returns some x, x is added to the accumulator and the loop returns to the beginning to continue with the successor iterator.

The bad case

In the following IR, PlausibleIterStep.yieldI has already been inlined:

| Decidable.isTrue h =>
    -- Build the arguments for `PlausibleIterStep.yieldI`
    let _x.33 := 1;
    let _x.34 := Nat.add _x.29 _x.33;
    let _x.35 := @ArrayIterator.mk _ _x.30 _x.34;
    let _x.36 := @Array.getInternal _ _x.30 _x.29 ◾;
    -- Normally we would call `PlausibleIterStep.yieldI` here
    -- and pattern match directly on it, but the call has already been inlined
    -- and the match block simplified away.
    let _x.37 := _f.1 _x.36; -- the function passed to `filterMap`
    cases _x.37 : Nat
    | Subtype.mk ...

After inlining, we see that we could actually compute x.33-35 later. floatLetIn therefore moves them into the match arm, which normally saves work, but here is somewhat unnecessary since there is only one constructor (Subtype.mk). Later in mono, the IR looks like this:

  | Bool.true =>                                                                                                                                                                                                                                    
    let _x.19 := Array.getInternal ...;
    let _x.20 := _f.2 _x.19; -- the function passed to `filterMap`
    -- When inlining `_f.2`, everything following becomes the continuation.
    -- `PlausibleIterStep.yieldI` has already been inlined.
    -- These `let`s are the returned components of `IterStep.yield (@ArrayIterator.mk ..) _`.
    -- `floatLetIn` has moved them after `_f.2` into the `Subtype.mk` arm, which has been monomorphized away.
    let val.21 := _x.20;
    let _x.22 := 1;
    let _x.23 := Nat.add _x.15 _x.22; 
    let _x.24 := @ArrayIterator.mk ...; 
    cases val.21 : Nat
    | Option.none => recur(_x.24, b)
    | Option.some val.27 => recur(_x.24, b+val.27)

Problem analysis

1. The lambda function passed to filterMap (_f.2) is inlined by inlineApp?. Since fn has two exits (returning some _ and none respectively), inlining works such that at both exits a join point is called for the code that executes after _f.2. This way, the code after the _f.2 call is not duplicated.

2. The join point for the code after _f.2 contains three let statements, even before the case distinction on whether fn returned some _ or none. There are so many that isJpCases? .. = none, and therefore the JP-Cases optimization is not applied.

The good case

PlausibleIterStep.yield is initially not inlined. Therefore, the IR initially looks like this:

| Decidable.isTrue h =>
    let _x.38 := 1;
    let _x.39 := Nat.add _x.34 _x.38; 
    let _x.40 := @ArrayIterator.mk _ _x.35 _x.39;
    let _x.41 := @Array.getInternal _ _x.35 _x.34 ◾;
    let _x.42 := @PlausibleIterStep.yield ... _x.40 _x.41 ◾;
    cases _x.42 ...
    -- In the `yield` arm, the function passed to `filterMap` is called.

Unlike in the "bad case", the successor iterator _x.40 is needed in PlausibleIterStep.yield, which fortunately prevents floatLetIn from moving it down.

Later, the situation before applying inlineApp? looks like this:

  | Bool.true =>
    let _x.24 := 1;
    let _x.25 := Nat.add _x.20 _x.24;
    let _x.26 := @ArrayIterator.mk ...; -- here above, since `floatLetIn` cannot pass through `.yield`
    let _x.27 := Array.getInternal ...;
    let _x.28 := PlausibleIterStep.yield ... _x.26 _x.27 ...;
    cases val.29 : IterStep
    | IterStep.yield it.30 out.31 =>
        let _x.32 := _f.2 out.31; -- the function passed to `filterMap`
        -- When inlining `_f.2`, everything following becomes the continuation.
        cases val.33 : Nat
        | Option.none => recur(it.30, b)
        | Option.some val.36 => recur(it.30, b+val.36)

Why does this produce more efficient code?

In this case, the continuation directly begins with pattern matching on the return value of the function _f.2. Therefore, the JPCases optimization kicks in: Both exits of _f.2 return some _ or none, so we can apply a case-of-known-constructor optimization and eliminate the join point. The result is more efficient code.

Conclusion

In general, the optimization pipeline seems somewhat fragile, and small changes in order that depend on subtle factors make a big difference. Concretely, given this example, one could consider whether to:

• Raise the threshold for JPCases, so that more code before the match statement is allowed. (probably not a good idea)

• Tweak floatLetIn so that it does not see matches with only one match arm as a reason to float—unless that one match arm contains another match with multiple arms. In other words: floatLetIn could treat matches with only one arm as a linear continuation of control flow and only float once control flow really branches.

• Replace matches along Subtype.mk with linear code using projections, i.e., call Subtype.val, which will later be inlined to the identity function in mono.

Version

Lean 4.30.0, commit 7b3d778ab047a176130d38931794714f997a1f7b
Target: aarch64-unknown-linux-gnu

Start with the reproduction comparing Def.f and Inline.f, and enable the shown trace.Compiler output. Read the JPCases, floatLetIn, inlineApp?, and mono stages to compare how the two variants handle the join point. Done means the compiler avoids the reported efficiency regression without degrading the intended optimization behavior.