ACEsuit / WithAlloc.jl

A simple Bumper convenience extension
MIT License
1 stars 0 forks source link

@withalloc allocates but @withalloc1 does not #3

Closed cortner closed 4 months ago

cortner commented 4 months ago

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)
cortner commented 4 months ago

fixed by commit e58c6171d11d9feb1ec1aae9fd938fd4d77e1cf2