Alternate ways to deal with arrays of subtypes of abstract type

New to Julia here. I am currently working on a raymarching program which supports different types of primitives. I thus have a Primitive abstract type which has subtypes, for example Sphere and Box. For each primitive subtype, I have a function sdf which takes it and a point and returns a float value. The obvious way to work with a scene would be a vector of primitive, something lie Vector{T} where T<:Primitive, but I then need to define sdf on this scene in such that it returns the minimum value of sdf for all elements of the array. This is the simple implementation I have

sdf(p::Vec3, sphere::Sphere) = norm(sphere.c - p) - sphere.r #Example implementation of sdf

function sdf(p::Vec3, objects::Vector{T}) where T<:Primitive
   n = length(objects)
   n == 1 && return sdf(p, objects[1])
   d = sdf(p, objects[1])
   for i ∈ 2:n
       d = min(d, sdf(p, objects[i]))
   end
   return d
end

The problem is that this function is highly type unstable. Obviously, objects[i] can only be determined to be Primitive, but the calls to sdf are all determined to be of type Any, which propagates all the way through the program as this method is the core of the loop. All the individual sdf methods are type stable. My question is if there is a way to indicate that all implementations of sdf will return a Float64, or in general if there is a way to stabilize this type of pattern? The best would be a way to keep working with vectors of primitives, but if there is a completely different way to implement this type of behaviour, I am open.

First of all, welcome to our community! :confetti_ball:

Second, I see two ways out there, the first uses the structure you already have (just need a little discipline) and the other changes the basics.

  1. Your problem is one of the few cases in which is advantageous to annotate the return type at call site. Julia is probably inferring sdf to be of type any because sdf has too many methods and the type passed is unknown at compile-time. It could be smart enough to see that every method returns a Float64 but to avoid having very long compilation times the compiler has an internal threshold that makes it gave up in the inference if there is too many methods. Your code with the types annotated would look like this:
function sdf(p::Vec3, objects::Vector{T}) where T<:Primitive
   n = length(objects)
   n == 1 && return (sdf(p, objects[1]) :: Float64)
   d = sdf(p, objects[1]) :: Float64
   for i ∈ 2:n
       d = min(d, sdf(p, objects[i]) :: Float64)
   end
   return d
end
  1. The problem I see is that sdf is a very lightweight function, so even if its return is inferred correctly you yet have the problem that its is a lightweight function being dynamic dispatched (considerable overhead). Unfortunately, the efficient solution I have to suggest you is not pretty (or modular). It consists in not having multiple Primitive types, instead you should have a single concrete type with a tag field indicating which primitive it represents, and enough fields that it can be used to model any of your primitives (note that if all your fields are Float64 you can just have as many fields as the object with most fields and reuse them with different meaning for each distinct primitive). The sdf is then a single method, with an if ... elseif ... end that manually selects what to be done.

If you could avoid the abstractly typed container at all, this would be the best. But probably this is not possible in your case.

1 Like

Hi, I’m sorry: your MWE is a bit too abstract for me to validate this claim. It would be great if you could share

  • the definition of Primitive and Sphere
  • maybe a second primitive (Rectangle?)
  • an example call to sdf which exhibits the problem

Thanks!

Here is a minimal implementation of what I have

abstract type Primitive end

const Float = Float64
const Scene = Vector{T} where T <: Primitive #Type alias for scene

struct Sphere <: Primitive
    c::Vec3
    r::Float
end
Sphere() = Sphere{Lambert}(Vec3(0), 1, Lambert()) #Default Debugging Sphere

struct Plane <: Primitive
    n::Vec3
    h::Float
end
Plane() = Plane{Lambert}(Vec3(0, 1, 0), 0) #Default Debugging Plane

sdf(p::Vec3, sphere::Sphere)    = norm(sphere.c - p) - sphere.r
sdf(p::Vec3, plane::Plane)      = p⋅plane.n + plane.h

The operations on Vec3 seem to be fully stable from what I see because I have used it previously with no issues. The file is quite long so I haven’t included it, but if it is necessary, I’ll add it.

function sdfUnion(p::Vec3, objects::Scene)
    n = length(objects)
    n == 1 && return sdf(p, objects[1])
    d = sdf(p, objects[1])
    for i ∈ 2:n
        d = min(d, sdf(p, objects[i]))
    end
    return d
end

Then, a simple call to
sdfUnion(Vec3(0., 0., 0.), [Sphere(), Plane()])
shows the problems I am talking about. I hope this can help and if any more information is necessary, I’ll try to add it. Thanks

1 Like

Adding the type annotations seem to help, with the compiler infering that d is of type Float64. However, benchmarking it is seems to not be the biggest bottleneck, as the performace only improves by about 10%. The problem with multiple dispatch seems to be the cause of the biggest bottleneck.
In the past, for example implementing this in GLSL, I have done what you suggest, with a single big object type, so I could implement this if necessary. I just figured Julia being much more high-level might have a way to deal with this issue.

Thanks a lot. First thing I notice is the difference in the signatures of the original MWE

function sdf(p::Vec3, objects::Vector{T}) where T<:Primitive

where we have a homogeneous vector of primitives and

function sdfUnion(p::Vec3, objects::Scene)

where we have a heterogeneous one. But we probably can work around this. Just to be on the safe side: how is Vec3 defined? Is it Vec3 = SVector{3, Float64}?

Thanks for your help!

No worries.

I am not sure what you mean by the different signature. Since Scene is an alias for Vector{T} where T <: Primitive, shouldn’t the two be the same? Or is there a difference?

I learnt too late about static vectors, so Vec3 is a homemade type. It’s not too complex and like previously said I already used it without issues, so it shouldn’t pose problems, but if it’s better I could migrate to SVector{3, Float64} instead.
Here is the Vec3 file, which is loaded in before the rest. Again Float is an alias for Float64 just to simplify a little bit

struct Vec3 <: Number
    x::Float
    y::Float
    z::Float
end
Vec3(v::Real) = (vf = convert(Float, v); Vec3(vf, vf, vf))
Vec3(v::Vector{Real}) = Vec3(convert(Float, v[1]), convert(Float, v[2]), convert(Float, v[3]))

const ZERO = Vec3(0, 0, 0)
const ORIGIN = ZERO
Base.zero(::Vec3) = ZERO
const I = Vec3(1, 0, 0)
const J = Vec3(0, 1, 0)
const K = Vec3(0, 0, 1)

function Base.getindex(v::Vec3, i::Integer)
    i == 1 && return v.x
    i == 2 && return v.y
    i == 3 && return v.z
    throw(BoundsError("cannot access a 3-element vector at position $i"))
end
Base.firstindex(::Vec3) = 1
Base.lastindex(::Vec3) = 3

Base.iterate(v::Vec3, state=1) = state > 3 ? nothing : (v[state], state+1)
Base.IteratorSize(::Type{Vec3}) = HasLength(3)
Base.IteratorEltype(::Type{Vec3}) = HasElType()
Base.eltype(::Type{Vec3}) = Float
Base.length(::Vec3) = 3

Base.:+(v1::Vec3, v2::Vec3) = Vec3(v1.x+v2.x, v1.y+v2.y, v1.z+v2.z)
Base.:+(v::Vec3, c::Real) = Vec3(v.x+c, v.y+c, v.z+c)
Base.:+(c::Real, v::Vec3) = v+c

Base.:-(v1::Vec3) = Vec3(-v1.x, -v2.x, -v3.x)
Base.:-(v1::Vec3, v2::Vec3) = Vec3(v1.x-v2.x, v1.y-v2.y, v1.z-v2.z)
Base.:-(v::Vec3, c::Real) = Vec3(v.x-c, v.y-c, v.z-c)
Base.:-(c::Real, v::Vec3) = v-c

Base.:*(v1::Vec3, v2::Vec3) = Vec3(v1.x*v2.x, v1.y*v2.y, v1.z*v2.z)
Base.:*(v::Vec3, c::Real) = Vec3(c*v.x, c*v.y, c*v.z)
Base.:*(c::Real, v::Vec3) = v*c

Base.:/(v1::Vec3, v2::Vec3) = Vec3(v1.x/v2.x, v1.y/v2.y, v1.z/v2.z)
Base.:/(v::Vec3, c::Real) = Vec3(v.x/c, v.y/c, v.z/c)

Base.:^(v1::Vec3, v2::Vec3) = Vec3(v1.x^v2.x, v1.y^v2.y, v1.z^v2.z)
Base.:^(v::Vec3, c::Real) = Vec3(v.x^c, v.y^c, v.z^c)

LinearAlgebra.:⋅(v1::Vec3, v2::Vec3) = v1.x*v2.x + v1.y*v2.y + v1.z*v2.z
×(v1::Vec3, v2::Vec3) = Vec3(
    v1.y*v2.z - v1.z*v2.y,
    v1.z*v2.x - v1.x*v2.z,
    v1.x*v2.y - v1.y*v2.x
)

LinearAlgebra.norm(v::Vec3) = sqrt(v⋅v)
normSq(v::Vec3) = v⋅v
LinearAlgebra.normalize(v::Vec3) = v/norm(v)

Base.min(v::Vec3, a::Real) = Vec3(min(v.x, a), min(v.y, a), min(v.z, a))
Base.min(a::Real, v::Vec3) = Vec3(min(v.x, a), min(v.y, a), min(v.z, a))
Base.min(v1::Vec3, v2::Vec3) = Vec3(min(v1.x, v2.x), min(v1.y, v2.y), min(v1.z, v2.z))
Base.max(v::Vec3, a::Real) = Vec3(max(v.x, a), max(v.y, a), max(v.z, a))
Base.max(a::Real, v::Vec3) = Vec3(max(v.x, a), max(v.y, a), max(v.z, a))
Base.max(v1::Vec3, v2::Vec3) = Vec3(max(v1.x, v2.x), max(v1.y, v2.y), max(v1.z, v2.z))
Base.clamp(v::Vec3, a::Real, b::Real) = Vec3(
    min(max(v.x, a), b),
    min(max(v.y, a), b),
    min(max(v.z, a), b)
)
Base.clamp(v::Vec3, a::Vec3, b::Vec3) = Vec3(
    min(max(v.x, a.x), b.x),
    min(max(v.y, a.y), b.y),
    min(max(v.z, a.z), b.z)
)
Base.mod(v::Vec3, a::Real) = Vec3(mod(v.x, a), mod(v.y, a), mod(v.z, a))
Base.mod(a::Real, v::Vec3) = Vec3(mod(v.x, a), mod(v.y, a), mod(v.z, a))
Base.mod(v1::Vec3, v2::Vec3) = Vec3(mod(v1.x, v2.x), mod(v1.y, v2.y), mod(v1.z, v2.z))

Base.abs(v::Vec3) = Vec3(abs(v.x), abs(v.y), abs(v.z))

Also, thank you very much for the help

Disregard my comment: I was mistaken. Here is what I have

using StaticArrays
using LinearAlgebra
using BenchmarkTools

const Vec3 = SVector{3, Float64}

abstract type Primitive end

const Scene = Vector{T} where T <: Primitive #Type alias for scene

struct Sphere <: Primitive
    c::Vec3
    r::Float64
end
Sphere() = Sphere(Vec3(0., 0., 0.), 1.) #Default Debugging Sphere

struct Plane <: Primitive
    n::Vec3
    h::Float64
end
Plane() = Plane(Vec3(0., 1., 0.), 0.) #Default Debugging Plane

sdf(p::Vec3, sphere::Sphere)    = norm(sphere.c - p) - sphere.r
sdf(p::Vec3, plane::Plane)      = p⋅plane.n + plane.h

function sdfUnion(p::Vec3, objects::Scene)
    n = length(objects)
    n == 1 && return sdf(p, objects[1])
    d = sdf(p, objects[1])
    for i ∈ 2:n
        d = min(d, sdf(p, objects[i]))
    end
    return d
end

@btime sdfUnion(vec3, scene) setup=(vec3 = Vec3(0., 0., 0.); scene = Primitive[Sphere(), Plane()])

yielding

  13.527 ns (0 allocations: 0 bytes)

which looks pretty good to me. Only thing I seem to remember is that certain compiler optimizations stop working in such cases if the number of sub types gets larger than a small limit (could be 3 or 4?), but then manual dispatch could help. Could this be the problem? Or do we need to investigate Vec3?

1 Like

Testing different scenarios, I think the main problem is the limitation of optimizations. Increasing to three possible subtypes of Primitive doesn’t change the performance, 4 and up does increase the time to about 80ns instead of ~15.
Changing the optimization level in Juno doesn’t seem to change this, so it might be a hard limit on the compiler.
In this case, what I might have to do is mix a little bit with @Henrique_Becker’s answer by bundling together primitives which have the same kinds of fields, like Plane and Sphere to try and reduce the number of subtypes while keeping the structure a bit cleaner.

Also, switching between SVector{3, Float64} and my implementation of Vec3 doesn’t seem to change anything performance-wise, so that should be ok.

Knowing this and combining it with type annotations on sdfUnion to make it type stable, I think I should be able to get the performance to good levels, at least for now.

Anyways, thank you very much for your help

1 Like

Without extending the MWE I can only guess that manual dispatch enumerating all primitives like

function sdf_dispatch(p, o)
    if p isa Sphere
        sphere = o
        sdf(sphere, sphere)
    elseif p isa Plane
        plane = o
        sdf(sphere, plane)
    else
        return 0.0
    end
end

could help. There is also GitHub - jlapeyre/ManualDispatch.jl: Avoid method dispatch at runtime which I haven’t had a chance to evaluate yet.

i’d recommend keeping types homogenous. In my implementation of Raytracing in a weekend, I’ve used that trick to keep dynamic dispatch to a minimum. I’ve noticed a considerable speedup, even for small scenes, simply because those functions are called so often that minimizing that dispatch is really worth it.

3 Likes

Just for reference, this kind of thing was discussed in other threads:

2 Likes