Avoiding allocation when accessing fields of a StructArray inside a Struct

General Usage Performance
1

Let me start with the MWE

using StructArrays
using BenchmarkTools

# the struct to be converted to StructArray
struct Foo1
    a::Float64
    b::Float64
end

a = [1.0, 2.0]
b = [3.0, 4.0]

foo1 = StructArray{Foo1}(a=a, b=b)

# `Bar1.foo` is type-annotated by hand
struct Bar1
    foo::StructVector{Foo1}
end

# `Bar2.foo` is type-annotated using that given by `typeof(foo1)`
struct Bar2
    foo::StructVector{Foo1, NamedTuple{(:a, :b), Tuple{Vector{Float64}, Vector{Float64}}}, Int64}
end


function access_bar(bar)
    bar.foo.a
end

@code_warntype access_bar(bar1)

#=
Output:

Body::Any
1 ā”€ %1 = Base.getproperty(bar, :foo)::StructVector
ā”‚   %2 = Base.getproperty(%1, :a)::Any
ā””ā”€ā”€      return %2
=#

@code_warntype access_bar(bar2)

#=
Body::Vector{Float64}
1 ā”€ %1 = Base.getproperty(bar, :foo)::StructVector{Foo1, NamedTuple{(:a, :b), Tuple{Vector{Float64}, Vector{Float64}}}, Int64}
ā”‚   %2 = Base.getproperty(%1, :a)::Vector{Float64}
ā””ā”€ā”€      return %2
=#

@btime access_bar(bar1)
# 52.019 ns (1 allocation: 32 bytes)

@btime access_bar(bar2)
#   13.889 ns (0 allocations: 0 bytes)

My question is, are there alternative solutions for getting rid of the allocation?

Annotating like foo::StructVector{Foo1, NamedTuple{(:a, :b), Tuple{Vector{Float64}, Vector{Float64}}}, Int64} is not so convenient as the StructArray has many fields, and Iā€™d like the code to be generic when type changes (such as for using CuArray instead).

If annotating the type is the only solution, are there short-hand aliases for foo::StructVector{Foo1, NamedTuple{(:a, :b), Tuple{Vector{Float64}, Vector{Float64}}}, Int64}?

2

After some digging, I realized that NTuple{2, Vector{Float64}} can be used.

3

This is not a concrete type, so it ends up being boxed. You can get type stability with writing about this much like so:

struct MyStruct{T <: StructVector{Foo1}}
    foo::T
end
2 Likes
4

Thanks! This is the issue.