Pre-allocating differently typed arrays for multiple dispatch


#1

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?


#2

Typically you know T outside the big loop (or something) and then you can just allocate it there and pass it into f.


#3

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.


#4

We had some discussion about this issue over at https://github.com/JuliaRobotics/RigidBodyDynamics.jl/issues/363 . Perhaps that kind of approach would work for you?


#5

Yeah, that approach has been working quite well. I’ve since abstracted out the key parts into an AbstractTypeDict:

You could just copy the AbstractTypeDict type, make a concrete subtype and implement:

  • a constructor
  • a valuetype method
  • a makevalue method

(Just pattern match StateCache or one of the other subtypes in that file).


#6
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.


#7

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

  1. initially avoid it, not think about preallocation,
  2. the first line of optimizations is to use immutables, eg StaticArrays, and hope for the best,
  3. 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.