For the simple case of a single pre-allocation @withalloc is non-allocating. But for more than one pre-allocation there is a very small allocation that I cannot track down and don't understand.
using WithAlloc, LinearAlgebra, Bumper, BenchmarkTools
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
function alloctest1(B, C, D)
@no_escape begin
a1, a2 = WithAlloc.whatalloc(mymul2!, B, C, D)
A1 = @alloc(a1...)
A2 = @alloc(a2...)
mymul2!(A1, A2, B, C, D)
# this change makes it again allocating:
# A12 = ( (@alloc(a1...)), (@alloc(a2...)) )
# mymul2!(A12..., B, C, D)
sum(A1) + sum(A2)
end
end
function alloctest2(B, C, D)
@no_escape begin
A1, A2 = @withalloc mymul2!(B, C, D)
sum(A1) + sum(A2)
end
end
B = randn(5,10); C = randn(10, 3); D = randn(10, 5)
@btime alloctest1($B, $C, $D)
@btime alloctest2($B, $C, $D)
For the simple case of a single pre-allocation
@withalloc
is non-allocating. But for more than one pre-allocation there is a very small allocation that I cannot track down and don't understand.