Array of union causing 10x slowdown

I have a type-stability question. In my 2D axisymmetric ray tracing code, I define an array of surfaces that has the type Union{Cone, Plane}, like

surfaces = Matrix{Union{Cone, Plane}}(undef, 100, 2)

All of them should in principle be cones, but due to small errors from a discretization earlier on in the code, there might be a few adjacent points that end up defining a plane, thus the Union. Which elements are Plane exactly is not known until the geometry is created.

For context, the type hierarchy goes like:

abstract type Reflection end
abstract type Surface end

struct ReflectionDiffuse  <: Reflection end
struct ReflectionSpecular <: Reflection end

struct Cone{T<:Reflection} <: Surface
    reflection::T
end

struct Plane{T<:Reflection} <: Surface
    reflection::T
end

In the part of the code that checks the intersection, I’m looping over all surfaces (looping over i and j) and checking for intersections. The function call is something like

check_intersection!(surfaces[i, j], otherargs...)

It’s important to dispatch on the correct check_intersection! function based on the type of the surface as it then determines the calculation of other things like the normal and tangent vectors later on.

The problem is that a @code_warntype analysis reveals red-colored text in this function call in Union{Cone, Plane}, and it’s causing a 10x slowdown in performance and a huge number of allocations (millions to billions). I think I know why it’s type-unstable - the code is boxing all elements of surfaces - but I’m not sure how to make this type-stable.

This related question on StackOverflow suggests to create a Union and claims that since the Union is small, performance should still be good…

maybe Allocation and slow down when # of types involved increase (slower than C++ virtual methods)

Yeah that’s roughly what I’m looking for. Looks like this way of doing things isn’t ideal then.

But I’m also recalling discussion about updates to Julia where small Unions now have fairly good performance, whereas before, Unions were a performance-killer. What gives? Or is the 10x slowdown still “decent” performance compared to if that Union were, say, a union of a dozen types?

Cone and Plane are abstract types. Only unions of isbits types are fast.
If you are using 64 bits floats then you probably want the type to be Union{Cone{Float64},Plane{Float64}}.

6 Likes

Thank you, this helped. The field is not a float, but a subtype of Reflection. What would be the idiomatic way to write this function? Is it using typeof(), as in the following:

function foo(refl)
    T = typeof(refl)
    surfaces = Matrix{Union{Plane{T}, Cone{T}}}(undef, 200, 2)
    return surfaces
end

The question is if you need the parameterization on Reflection. You have two reflection types, so you could make an enum with two entries out of those. That would remove the type parameter of Cone and Plane, making the union small and therefore fast.

1 Like

This looks fine and idiomatic to me. Some people might replace the first two lines with function foo(refl::T) where T (or ... where T<:Reflection if you want to restrict inputs to Reflection, but the subsequent code will throw for other types anyway). The effect is identical so the choice is stylistic. Personally, which of the forms I choose depends on the situation.

My only remark is that you don’t use refl (only it’s type) so you could write this function to take the type rather than a value of the type, if that seemed more sensible to you. If you did that, I would write the signature as function foo(::Type{T}) where T (or ... where T<:Reflection}).

1 Like