Reducing allocations in quadgk within loop

General Usage Performance
8

I played around more with the solution suggested by @PeterSimon and figured out that many of the allocations in the original code actually stem from passing around the function vrad_fun as argument, and not from quadgk itself. In the example provided here, it is not obvious why this would be necessary, but in my actual use-case I need vrad_fun to be passable as argument and not defined in global scope.

Reading more about it in the manual and in the thread Memory allocation with function passed as arguments, I realized that I need to ensure specialization for the function type by some where{V} (see code below). With this, I managed to get the same amount of allocations as in PeterSimon’s reply, but keeping the vrad_fun as function argument. For completeness, and the small chance it might help someone in the future, I post my solution:

using SpecialFunctions
using QuadGK
using BenchmarkTools

Lsum = 2
v_arr = zeros(100,Lsum+1)
alpha_arr = rand(100)

function vrad_fun(r)
    return exp(-r^2)
end

function precompute_varr!(v_arr,alpha_arr,Lsum,vrad_fun::V) where {V}
    buf = alloc_segbuf()
    for n = 0:Lsum
        for k=1:lastindex(alpha_arr)
            v_arr[k,n+1] = vrad_integration2(vrad_fun,alpha_arr[k],n,buf)
        end
    end
    return v_arr
end

function integrand(r,alpha,n,vrad_fun)
    return vrad_fun(r)*r^(2*n+2)*exp(-alpha*r^2)
end

function vrad_integration2(vrad_fun,alpha,n,buf) #where {V}
    val = quadgk(r -> integrand(r,alpha,n,vrad_fun),0,Inf;segbuf=buf)[1]
end

v_arr2 = zeros(100,Lsum+1);
@btime precompute_varr!($v_arr2,$alpha_arr,$Lsum,vrad_fun)
# 1.001 ms (2 allocations: 368 bytes)
2 Likes