I’m using a multiple dispatch on an abstract type (AbstractBody) to switch between two behaviours. I need to store a performance critical function in one case, but don’t need to store anything in the other case.
I’ve implemented a version which avoids ambiguous function type allocations, but I lose this specification when I use a struct to hold any AbstractBody.
abstract type AbstractBody end
struct NoBody<:AbstractBody end
struct FuncBody{F<:Function}<:AbstractBody
sdf::F
end
struct SimAmbiguous
body::AbstractBody # this works even when there is `NoBody`
end
struct SimUnambiguous{F<:Function}
body::FuncBody{F} # This locks me into needing a function
end
function circle(ambiguous; radius = 8.0)
sdf(x) = norm2(x) - radius
ambiguous && return SimAmbiguous(FuncBody(sdf))
return SimUnambiguous(FuncBody(sdf))
end
ambiguous = circle(true) # slow
unambiguous = circle(false) # fast
julia> @btime $ambiguous.body.sdf(x) setup=(x=rand(3));
73.024 ns (2 allocations: 32 bytes)
julia> @btime $unambiguous.body.sdf(x) setup=(x=rand(3));
10.431 ns (0 allocations: 0 bytes)
Do I need to give up on the abstract type approach? If so, how do I handle when there is NoBody to use in SimUnambiguous? Some default function?
I’m not sure why this wasn’t working before, maybe something to do with the benchmark. Here’s a more complete example which works fine.
using BenchmarkTools
using LinearAlgebra: norm2
abstract type AbstractBody end
abstract type NoBody <: AbstractBody end
struct FuncBody{F1<:Function,F2<:Function} <: AbstractBody
sdf::F1
map::F2
function FuncBody(sdf,map=(x,t)->x)
comp(x,t) = sdf(map(x,t),t)
new{typeof(comp),typeof(map)}(comp, map)
end
end
struct SimAmbiguous{A<:AbstractArray,B<:AbstractBody}
a::A
body::B
end
function nsphere(radius = 8.0, n = 2)
sdf(x,t) = norm2(x) - radius
map(x,t) = (x[1]-=radius*(1-cos(t/radius)); x)
return SimAmbiguous(zeros(ntuple(i->Int(radius), n)), FuncBody(sdf,map))
end
measure!(sim) = measure!(sim.a, sim.body)
function measure!(a::Array{Float64,N}, body; t=0) where N
x = zeros(N) # here's the single allocation.
for I in CartesianIndices(a)
@. x = I.I-1.4
a[I] = body.sdf(x,t)
end
end
function measure!(a::Array{Float64,N}, body::NoBody; t=0) where N end
julia> @btime measure!(sim) setup=(sim=nsphere(2^6,2))
76.061 μs (1 allocation: 96 bytes)
The one allocation is a small price to pay for the speed up relative to using tuples for the sdf.