MasonProtter / Bumper.jl

Bring Your Own Stack
MIT License
152 stars 6 forks source link

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

Open artemsolod opened 9 months ago

artemsolod commented 9 months ago

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 commented 9 months ago

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 commented 9 months ago

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 commented 9 months ago

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