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)
endFor 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
endThis 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