Bumper.jl icon indicating copy to clipboard operation
Bumper.jl copied to clipboard

Possibility of tranforming existing functions to use `Bumper.jl`

Open artemsolod opened this issue 1 year ago • 3 comments

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!

artemsolod avatar Oct 26 '23 14:10 artemsolod

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.

MasonProtter avatar Nov 07 '23 15:11 MasonProtter

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.

artemsolod avatar Nov 07 '23 16:11 artemsolod

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

MasonProtter avatar Nov 07 '23 16:11 MasonProtter