Type instability of surfpt_nearby for abstract Shape

[wsshin]

Currently surfpt_nearby is type-stable for every concrete subtype of Shape{N} and returns NTuple{2,SVector{N,Float64}}. However, when surfpt_nearby is called with an abstract Shape{N} object, it is type-unstable. For example, consider the following code:

julia> VERSION
v"0.6.1-pre.0"

julia> using GeometryPrimitives, StaticArrays

julia> svec = Vector{Shape{3}}(2);

julia> function test(p, svec::Vector{Shape{3}})
           rsum = @SVector zeros(3)
           nsum = @SVector zeros(3)
           for s = svec
               r, n = surfpt_nearby(p, s)
               rsum += r
               nsum += n
           end
           return rsum, nsum
       end
test (generic function with 1 method)

The above test function is type-unstable, because the second argument s of surfpt_nearby is an element of Vector{Shape{3}} and thus typed Shape{3}:

julia> @code_warntype test(@SVector(zeros(3)), svec)
Variables:
  ...
Body:
  begin
      ...
      SSAValue(2) = (Main.surfpt_nearby)(p::SVector{3,Float64}, s::GeometryPrimitives.Shape{3,D} where D)::Any
      ...
  end::Tuple{Any,Any}

This makes sense. When surfpt_nearby is called with an abstract Shape object, the following method defined in GeometryPrimitives.jl is called:

surfpt_nearby(x::AbstractVector, o::Shape{N}) where {N} = surfpt_nearby(SVector{N}(x), o)

Because the specific type of o is not known, the compiler cannot further infer the return type in this case.

To resolve this type instability, I tried to change the above method definition using return type annotation:

(surfpt_nearby(x::AbstractVector, o::Shape{N})::NTuple{2,SVector{N,Float64}}) where {N} = surfpt_nearby(SVector{N}(x), o)

I thought this would remove the type instability, because now the compiler would know the return type of the above method. Strangely, the type instability remains:

julia> @code_warntype test(@SVector(zeros(3)), svec)
Variables:
  ...
Body:
  begin
      ...
      SSAValue(2) = (Main.surfpt_nearby)(p::SVector{3,Float64}, s::GeometryPrimitives.Shape{3,D} where D)::Any
      ...
  end::Tuple{Any,Any}

I'm not sure which method of surfpt_nearby is being called here. To me, the only method with a matching argument types is the one defined in GeometryPrimitives.jl, which is now annotated with a return type. Still, the compiler fails to pick up this annotated return type.

In fact, if we call surfpt_nearby with a regular Vector instead of an SVector in the first argument p, the return type annotation works as intended and test becomes type-stable:

julia> @code_warntype test(zeros(3), svec)
Variables:
  ...
Body:
  begin
      ...
      SSAValue(2) = (Main.surfpt_nearby)(p::Array{Float64,1}, s::GeometryPrimitives.Shape{3,D} where D)::Tuple{SVector{3,Float64},SVector{3,Float64}}
      ...
  end::Tuple{SVector{3,Float64},SVector{3,Float64}}

I don't understand why the return type annotation does not kick in for SVector p. How can we resolve this problem? Any suggestions?

[wsshin]

I think now I found the source of this instability. I will spend a few more hours to examine this finding to make sure it is really the source of the instability, and will report back here.

[wsshin]

So it turned out that this instability was related with #6, where we introduced the extra L parameter in Box and Ellipsoid to stably type their projection matrices as SMatrix{N,N,Float64,L}.

In the above unstable test function, we called surfpt_nearby(p, s) inside a for loop with s being an entry of svec::Vector{Shape{3}}. Now, surfpt_nearby for each concrete Shape type is defined to be stable, returning NTuple{2,SVector{N,Float64}}. Even so, when an abstract s::Shape{3} is passed to surfpt_nearby, the the compiler gets only N = 3 and not L = 9. This means that the projection matrices p::SMatrix{N,N,Float64,L} of Box and Ellipsoid cannot be fully typed, leading to unstable surfpt_nearby that internally uses p.

So, to remove the above type instability, we need to embed the information of L when creating an array of abstract Shapes. In other words, we should define Shape to take L as

abstract type Shape{N,L,D} end

and define svec as svec = Vector{Shape{3,9}}(...). Having to tag the extra type parameter L = 9 like this whenever creating an array of Shapes is a bit ugly, but I cannot think of a better way of eliminating the instability. I will create a PR implementing this solution.

(Another possibility is to define surfpt_nearby of Box and Ellipsoid separately for each N = 1,2,3 and to type-assert the projection matrix p as SMatrix{1,1,Float64,1}, SMatrix{2,2,Float64,4}, and SMatrix{3,3,Float64,9}, respectively. But this means I need to define any functions using p (e.g., Base.in and bounds) for individual values of N. This is doable with @generated, but I'm worried if this approach would be a bit convoluted and make the code less readable for other users.)

[wsshin]

It turns out that this issue hasn't been fully resolved yet. Strangely, I still get type instability in the return type of surfpt_nearby when the GeometryPrimitives module has four or more concrete subtypes of Shape. If I keep only three subtypes in the module (by commenting out Sphere, for example), then the return type of surfpt_nearby becomes stable again.

Maybe a Julia bug? Filed an issue at JuliaLang/julia#23210.

[wsshin]

According to JuliaLang/julia#23210, when a function is called with abstract arguments, Julia seems to give up type inference if the function has more than four method definitions. The comments in base/inference.jl reads

function abstract_call_gf_by_type(f::ANY, atype::ANY, sv::InferenceState)
    tm = _topmod(sv)
    # don't consider more than N methods. this trades off between
    # compiler performance and generated code performance.
    # typically, considering many methods means spending lots of time
    # obtaining poor type information.
    # It is important for N to be >= the number of methods in the error()
    # function, so we can still know that error() is always Bottom.
    # here I picked 4.

From this comment, I figured that I shouldn't abuse calling a function with abstract arguments. So, as for surfpt_nearby and normal, I was able to avoid type instability by replacing the following code here

surfpt_nearby(x::AbstractVector, o::Shape{N}) where {N} = surfpt_nearby(SVector{N}(x), o)
normal(x::AbstractVector, o::Shape) = surfpt_nearby(x, o)[2]

with more concrete versions:

surfpt_nearby(x::AbstractVector, b::Box{N}) where {N} = surfpt_nearby(SVector{N}(x), b)
surfpt_nearby(x::AbstractVector, s::Cylinder{N}) where {N} = surfpt_nearby(SVector{N}(x), s)
surfpt_nearby(x::AbstractVector, b::Ellipsoid{N}) where {N} = surfpt_nearby(SVector{N}(x), b)
surfpt_nearby(x::AbstractVector, s::Sphere{N}) where {N} = surfpt_nearby(SVector{N}(x), s)

normal(x::AbstractVector, b::Box) = surfpt_nearby(x, b)[2]
normal(x::AbstractVector, s::Cylinder) = surfpt_nearby(x, s)[2]
normal(x::AbstractVector, b::Ellipsoid) = surfpt_nearby(x, b)[2]
normal(x::AbstractVector, s::Sphere) = surfpt_nearby(x, s)[2]

I verified that this trick works even if I add more concrete subtypes of Shape to the module, so it doesn't seem to be a solution limited to the four concrete Shapes currently implemented.

Unfortunately, the same trick doesn't work for Base.in. In other words, replacing

Base.in(x::AbstractVector, o::Shape{N}) where {N} = SVector{N}(x) in o

with

Base.in(x::AbstractVector, b::Box{N}) where {N} = in(SVector{N}(x), b)
Base.in(x::AbstractVector, s::Cylinder{N}) where {N} = in(SVector{N}(x), s)
Base.in(x::AbstractVector, b::Ellipsoid{N}) where {N} = in(SVector{N}(x), b)
Base.in(x::AbstractVector, s::Sphere{N}) where {N} = in(SVector{N}(x), s)

does not make calling in for abstract Shape stable. This is a serious problem, because it means whenever we call Base.in on an abstract Shape (e.g., findin in kdtree.jl) we need to type-assert the result.

Why does Base.in behave differently from non-Base functions surfpt_nearby and normal? I asked this question with MWE at https://discourse.julialang.org/t/extending-base-in-type-stably/5341.

[wsshin]

Closed via #26.