Avoid ambiguous function type allocation without losing abstract type

General Usage Performance
,
1

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?

2

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.

3

Cool, but what happened to NoBody? Did you just drop it or does it work fine despite the ambiguity?

4

Just for curiosity, couldn’t this be written as:

struct SimAmbiguous{A<:AbstractArray,B<:AbstractBody}
    a::A
    body::B
end

?
seems safer from a type stability point of view, if it works.

5

Yes, it works fine. I’ve put it back into the example.

6

It does! Good idea.