SROA for array allocation

[oxinabox]

Consider the situation where you have a function that does two things, and returns both outputs, but then in the outer function you only use one of them.
We would like the optimizer to remove the work that is being done for the part you don't use.
And it can do that on nightly, but not it is a subpart of a function that is being called (not even an inlined one). Only if the function being eleminated is directly being called.

Consider the following code:

using BenchmarkTools
# Effect inference can't seem to workout that this doesn't throw
# It does work out that it is effect_free, and terminates_globally
Base.@assume_effects :nothrow foo(x::Int) = [x]
bar(x) = foo(x), 2x

function bar_last_manual(x)
    _ = foo(x)
    return 2x
end
@ballocated bar_last_manual(1)
@code_warntype optimize=true bar_last_manual(1)


function bar_last_inlined(x)
    _, ret = @inline bar(x)
    return ret
end
@ballocated bar_last_inlined(1)
@code_warntype optimize=true bar_last_inlined(1)

we can see that bar_last_manual acts as expected, it eliminated the allocation leaving just the multiplication

julia> @ballocated bar_last_manual(1)
0

julia> @code_warntype optimize=true bar_last_manual(1)
MethodInstance for bar_last_manual(::Int64)
  from bar_last_manual(x) @ Main REPL[18]:1
Arguments
  #self#::Core.Const(bar_last_manual)
  x::Int64
Body::Int64
1 ─ %1 = Base.mul_int(2, x)::Int64
└──      return %1

On the other hand if we do it via an inlined called to foo, it does not:

julia> @ballocated bar_last_inlined(1)
64

julia> @code_warntype optimize=true bar_last_inlined(1)
MethodInstance for bar_last_inlined(::Int64)
  from bar_last_inlined(x) @ Main REPL[21]:1
Arguments
  #self#::Core.Const(bar_last_inlined)
  x::Int64
Locals
  val::Tuple{Vector{Int64}, Int64}
  @_4::Int64
  ret::Int64
Body::Int64
1 ─ %1 = $(Expr(:foreigncall, :(:jl_alloc_array_1d), Vector{Int64}, svec(Any, Int64), 0, :(:ccall), Vector{Int64}, 1, 1))::Vector{Int64}
└──      goto #3 if not true
2 ─      Base.arrayset(false, %1, x, 1)
3 ┄      goto #4
4 ─      goto #5
5 ─ %6 = Base.mul_int(2, x)::Int64
└──      goto #6
6 ─      return %6

This is sad.

Solving this would make batched forward-mode in Diffractor much easier.
Because it would mean we can effectively just run the pushforward part of the frules if we don't use the primal output.

[Seelengrab]

we can see that bar_last_manual acts as expected, it eliminated the allocation leaving just the multiplication

I think it eliminated the entirety of the call to foo since the entire result of that call is ignored. Your example just happens to then return the exact same thing that would be returned in the second element of the tuple returned by foo. So I don't think it's true that we eliminate parts of an inner call now (or at least this example doesn't show this happening).

[oxinabox]

So I don't think it's true that we eliminate parts of an inner call now

Yes. That's the issue. We don't do that. We should do that.
The example illustrates that foo can be eliminated. But not if it's a constitute part of a function being called (not even if that function is inlined).

[Seelengrab]

Ah, gotcha - then I misunderstood what you wrote in your initial post:

We would like the optimizer to remove the work that is being done for the part you don't use.
And it can do that on nightly, but not it is a subpart of a function that is being called (not even an inlined one). Only if the function being eleminated is directly being called.

which reads to me like we already do partial function elimination.

[aviatesk]

This is a tricky problem. I think we have two options:

1. We can mark foo as @noinline function: Then it isn't inlined into bar and we can eliminate the call to it when the return value is unused since we know foo is is_removable_if_unused (with the help of the @assume_effects annotation). When it's inlined, however, we need to prove that [x] is is_removable_if_unused instead, but the compiler can't prove it since the inlined body of [x] includes e.g. arrayset, that itself isn't effect_free.

2. We could implement SROA pass for array: If we define foo(x) = Ref(x) we can eliminate it no matter if it's inlined or not, since SROA works nicely with inlined code. Currently our SROA works only for structs, but I guess we can have a simplified version of optimizer: simple array SROA #43909 for optimizing array allocations like this.

Option 2. would be a pure compiler improvement, so it may be nice to have on its own but it may take some time for us to implement it right. If you require an immediate solution, option 1 could be the way to go.

[oxinabox]

oh, it's not the function call that is blocking it.
It is the fact that it was assigned to a tuple, then split up.
vs being split up immediately.

See

function bar_last_manual2(x)
    both = (foo(x), 2x)
    _, ret = both
    return ret
end
@ballocated bar_last_manual2(1)
@code_warntype optimize=true bar_last_manual2(1)

Which is just like the one with the inlined function call:

julia> @ballocated bar_last_manual2(1)
64

julia> @code_warntype optimize=true bar_last_manual2(1)
MethodInstance for bar_last_manual2(::Int64)
  from bar_last_manual2(x) @ Main REPL[10]:1
Arguments
  #self#::Core.Const(bar_last_manual2)
  x::Int64
Locals
  @_3::Int64
  ret::Int64
  both::Tuple{Vector{Int64}, Int64}
Body::Int64
1 ─ %1 = $(Expr(:foreigncall, :(:jl_alloc_array_1d), Vector{Int64}, svec(Any, Int64), 0, :(:ccall), Vector{Int64}, 1, 1))::Vector{Int64}
└──      goto #3 if not true
2 ─      Base.arrayset(false, %1, x, 1)
3 ┄      goto #4
4 ─      goto #5
5 ─ %6 = Base.mul_int(2, x)::Int64
└──      return %6