Possibility of tranforming existing functions to use Bumper.jl

[artemsolod]

Is it in principle possible to have some sort of a macro that would take a function and replace its inner calls to Vector{T}(undef, n) (and similar) with @alloc(T, n)? To me this appears possible after the calls are inlined and possibly some escape analysis applied but I have a very limited understanding of the problem.
So the macro could do something like:

function mapsum(f, x)
    arr = Vector{Float64}(undef, length(x))
    arr .= f.(x)
    return sum(arr)
end

transforms into

function mapsum_bumpered(f, x)
   @no_escape begin
        arr = @alloc(Float64, length(x))
        arr .= f.(x)
        ans = sum(arr)
    end
   return ans
end

Thanks!

[MasonProtter]

This isn't quite the same, but there is the package https://github.com/ericphanson/AllocArrays.jl which does something similar. I must say though, the idea makes me a bit uncomfortable. Without escape analysis available at code generation time, we can't really know if it's safe to do this transformation or not.

[artemsolod]

Thank you for the answer! AllocArrays.jl looks interesting, will try it out.

I agree the approach looks unsafe but I find it would help my workflow and speed up running a bunch of parallel simulations.

It's interesting also if some approaches can be done without full escape analysis. Like considering a function safe for Bumper transformation if it only returns isbits types. Or maybe inlining and then wrapping all remaining return statements in copy_if_bumpallocated that could tell if something is bump-allocated and copy to memory when it is.

Meanwhile, feel free to close the issue if that is something not relevant for Bumper.jl atm.

[MasonProtter]

I'll leave this open as a somewhat far-off speculative potential feature.