Proposed Interface for "Automatic" Allocations

[cortner]

I've written a small experimental extension to Bumper.jl I currently call WithAlloc.jl, which substitutes

A = @alloc # figure out how to allocate A 
B = @alloc # figure out how to allocate B 
dosomething!(A, B, x1, x2, x3) 

with

A, B = @withalloc dosomething!(x1, x2, x3) 

by specifying how dosomething! wants its outputs allocated. I am interested in

• feedback about the general idea
• comments of the concrete approach, potential bugs etc?
• is this a sufficiently useful general pattern that I should make a PR in Bumper instead of maintaining it as an extension?
• I am open-minded about alternatives, naming suggestions, pretty much anything ...

Thank you.

Example

using WithAlloc, LinearAlgebra, Bumper 

# simple allocating operation
B = randn(5,10)
C = randn(10, 3)
A1 = B * C

# tell `WithAlloc` how to allocate memory for `mymul!`
WithAlloc.whatalloc(::typeof(mul!), B, C) = 
          (promote_type(eltype(B), eltype(C)), size(B, 1), size(C, 2))

# the "naive use" of automated pre-allocation could look like this: 
# This is essentially the code that the macro @withalloc generates
@no_escape begin 
   A2_alloc_info = WithAlloc.whatalloc(mul!, B, C)
   A2 = @alloc(A2_alloc_info...)
   mul!(A2, B, C)
   @show A2 ≈ A1
end

# but the same pattern will be repreated over and over so ... 
@no_escape begin 
   A3 = @withalloc mul!(B, C)
   @show A3 ≈ A1 
end