I have a function that allocates a large matrix H each time.
function f(x::Vector{T}) where T <: Real
H = Matrix{Complex{T}}(undef, 2, 3)
# do stuff...
H
end
I want pre-allocate H to speed up my code, but sometimes T is Float64, sometimes it’s a dual number or maybe something else.
I’m guessing this is a fairly common problem, but I didn’t know how to search for it. Is there a “best practice” way of pre-allocating arrays for each type and dispatching for the appropriate one?
The only way I know of is to make a type with the same static parameters, eg.
struct fData{T}
H::Matrix{Complex{T}}
function fData{T}() where {T}
H = Matrix{Complex{T}}(undef, 2, 3)
return new(H)
end
end
function f(data::fData{T}, x::Vector{T}) where T <: Real
H = data.h
# do stuff
end
# in the caller, T must be known from somewhere
data = fData{T}()
x = rand(n) # however x gets created
for i=1:1000
f(x, data)
end
This approach effectively pushes the problem out one level, because data is allocated outside the for loop. Its gets more difficult the more levels you have though.
function f(x::Vector{T}) where T <: Real
H = similar(x)
f!(H, x)
H
end
function f!(out::Vector{T}, in::Vector{T}) where T <: Real
... # your code here
end
Edit: sorry, didn’t notice H has a different type.
You can use eltype(x) in caller script or function - and replace similar(x) with the type you need.
I see this pattern a lot, but it implicitly relies on T being closed under f. This is a reasonable assumption, but can quickly break down when relying eg on multiple AD packages at the same time.
The fundamental problem is that in general, it is very hard to preallocate for an operation
d = f(a, b, c)
as
f!(d, a, b, c)
when a, b and c can be sufficiently generic, because predicting the type of d can be tricky.
My current approach to this problem is to
initially avoid it, not think about preallocation,
the first line of optimizations is to use immutables, eg StaticArrays, and hope for the best,
if I decide I really need it, use some kind of a heuristic to determine the element type.
For the last one, something like
T = typeof(g(one(eltype(a)), one(eltybe(b)), one(eltype(c)))
d = Matrix{T}(undef, I_know, the_dimensions)
where g somehow reflects the elementary operations for what is going on in f. But this can be brittle as the <: Real interface is not formally specified and packages can deviate from expectations.