WithAlloc

This package implements a very small extension to Bumper.jl. Bumper strongly enourages (almost enforces) that it is used purely from within @no_escape blocks. Bumper-allocating an array in a function and passing it back to the caller should generally be avoided. This results in a common pattern:

@no_escape begin 
   # determine the type and size a required array
   T, N = determine_array(x1, x2, x3)
   # preallocate some arrays
   A = @alloc(T, N)
   # do a computation on A 
   calculate_something!(A, x1, x2, x3)
end 

The goal of WithAlloc.jl is to replace the above 3 lines with

@no_escape begin 
   A = @withalloc calculate_something!(x1, x2, x3)
end

Preliminary Documentation

For now, there are just a few simple use-case examples. Proper documentation will follow once the packages has been tested a bit and there is some agreement it will be long-term useful.

using WithAlloc, LinearAlgebra, Bumper 

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

# we wrap mul! into a new function so we don't become pirates...
mymul!(A, B, C) = mul!(A, B, C)

# tell `WithAlloc` how to allocate memory for `mymul!`
WithAlloc.whatalloc(::typeof(mymul!), 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(mymul!, B, C)
   A2 = @alloc(A2_alloc_info...)
   mymul!(A2, B, C)

   @show A2 ≈ A1
end

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

# ------------------------------------------------------------------------

# Multiple arrays is handled via tuples: 

B = randn(5,10)
C = randn(10, 3)
D = randn(10, 5)
A1 = B * C 
A2 = B * D

mymul2!(A1, A2, B, C, D) = mul!(A1, B, C), mul!(A2, B, D)

function WithAlloc.whatalloc(::typeof(mymul2!), B, C, D) 
   T1 = promote_type(eltype(B), eltype(C)) 
   T2 = promote_type(eltype(B), eltype(D))
   return ( (T1, size(B, 1), size(C, 2)), 
            (T2, size(B, 1), size(D, 2)) )
end

@no_escape begin 
   A1b, A2b = WithAlloc.@withalloc mymul2!(B, C, D)
   @show A1 ≈ A1b, A2 ≈ A2b   # true, true 
end

This approach should become non-allocating, which we can quickly check.

using WithAlloc, LinearAlgebra, Bumper 

mymul!(A, B, C) = mul!(A, B, C)

WithAlloc.whatalloc(::typeof(mymul!), B, C) = 
          (promote_type(eltype(B), eltype(C)), size(B, 1), size(C, 2))

nalloc = let B = randn(5,10), C = randn(10, 3)
   @allocated sum( @withalloc mymul!(B, C) )
end          

@show nalloc  # 0