Resolving UnionAlls with a single concrete instance in struct declarations

[brainandforce]

This is something I’ve encountered far too often, and it’s one of my biggest gripes with Julia. @Tamas_Papp explained it best:

@Tamas_Papp

To be concrete, let me give an example. The other day a student asked me for help with optimizing a calculation. Since it involved lots of small matrices, I immediately suggested using StaticArrays.SMatrix. The student then did something not unlike this MWE:

using StaticArrays

struct Foo{T,N}
    bar::SMatrix{N,N,T}
end

Now, seasoned Julia users can spot the problem quickly, which is easy to demo with

julia> baz(foo::Foo) = sum(foo.bar)
baz (generic function with 1 method)

julia> foo = Foo(SMatrix{3,3}(zeros(3,3)))
Foo{Float64, 3}([0.0 0.0 0.0; 0.0 0.0 0.0; 0.0 0.0 0.0])

julia> @code_warntype baz(foo)
MethodInstance for baz(::Foo{Float64, 3})
  from baz(foo::Foo) @ Main REPL[10]:1
Arguments
  #self#::Core.Const(Main.baz)
  foo::Foo{Float64, 3}
Body::Any
1 ─ %1 = Main.sum::Core.Const(sum)
│   %2 = Base.getproperty(foo, :bar)::SMatrix{3, 3, Float64}
│   %3 = (%1)(%2)::Any
└──      return %3

The problem is that SMatrix{3,3,Float64} is not a concrete type, it is SMatrix{3,3,Float64,9}.

I’ve encountered issues with this in the past with CliffordNumbers.jl: subtypes of AbstractCliffordNumber{Q,T} need an extra type parameter (for instance, KVector{K,Q,T,L}) where L is the length of the backing NTuple{L,T}. For this type, L = binomial(dimension(Q), K), so it’s not particularly simple arithmetic.

There’s good reason not to allow arbitrary computation in the type system, but this either becomes a pitfall for new Julia programmers, or an annoyance for those who have to leak type parameters when building a struct out of types like these.

Can there be a mechanism for the Julia compiler to resolve almost concrete types in struct declarations like SMatrix{N,N,T}, which only have a single concrete instance, to avoid the type instability and performance hit? Perhaps some sort of function, like

Base.resolve_to_concrete_type(::Type{SMatrix{M,N,T}}) = SMatrix{M,N,T,M*N}

function Base.resolve_to_concrete_type(::Type{KVector{K,Q,T}}) where {K,Q,T}
    return KVector{K,Q,T,binomial(dimension(Q), K)}
end

which a user can opt into to cover cases like this?

[Benny]

@brainandforce

SMatrix{N,N,T}, which only have a single concrete instance

Assuming you meant one instantiable concrete subtype, that’s entirely enforced by the implementation flagging or fixing an incomplete or inappropriate type like SMatrix{2,2,Int,3}, not the language itself. If we circumvent those methods with type piracy or the below hack, we can actually instantiate those inappropriate types:

julia> x = eval(Expr(:new, SMatrix{2,2,Int,6}, Tuple(1:6)))
2×2 SMatrix{2, 2, Int64, 6} with indices SOneTo(2)×SOneTo(2):
 1  3
 2  4

julia> eachindex(x), x[5], x[6]
(SOneTo(4), 5, 6)

In the general case, omitted type parameters can more reasonably become non-redundant, and just as easily as extending resolve_to_concrete_type. Generic functions in the type computation e.g. binomial, dimension can get new, even impure, methods for new input types that violate implementation invariants. Adding new constructor methods can make exceptions to invariants. Omitting a type parameter e.g. Foo{T,N,L} to Foo{T,N} by computing a field parameter SMatrix{N,N,T,N*N} essentially amounts to putting “arbitrary computation in the type system”, dangerously brittle for a dynamic runtime and still questionable for AOT compilation. Pure built-ins or intrinsics on core types instead of generic functions e.g. Base.mul_int(N, N) could be more reasonable.

[sgaure]

It’s a long story. At least from 2014: RFC: staged / meta / generated types ? · Issue #8472 · JuliaLang/julia · GitHub and has come up many times (linked to in that issue), with various suggestions for how it can be solved.

[Tamas_Papp]

@brainandforce

a mechanism for the Julia compiler to resolve almost concrete types in struct declarations like SMatrix{N,N,T}, which only have a single concrete instance, to avoid the type instability and performance hit?

Current best practice is to use a generic tool like JET to test for these instabilities, preferably in CI.

This is just one possible type instability anyway. In complex codebases you need tooling to catch them all.

[nsajko]

@brainandforce

Can there be a mechanism for the Julia compiler to resolve almost concrete types in struct declarations like SMatrix{N,N,T}, which only have a single concrete instance, to avoid the type instability and performance hit? Perhaps some sort of function

I once created and registered a package with this exact purpose, TypeCompletion.jl. Topic on Discourse:

• [ANN] TypeCompletion.jl: exploit redundancy to complete UnionAll types

Disclaimer: I have not taken a look at that package for some time. It might not be well-maintained. It might not follow some best practice. The design might not be how I would do it if I were to do it today again.

If you are interested in maintaining/developing it, I can transfer or whatever.

[ettersi]

@brainandforce

Can there be a mechanism for the Julia compiler to resolve almost concrete types in struct declarations like SMatrix{N,N,T}, which only have a single concrete instance, to avoid the type instability and performance hit?

Not quite the solution you are looking for, but I believe to a large extent you can avoid this problem by using functions as type aliases.

SMatrixFun(M, N, T) = SMatrix{M, N, T, M*N}

Unfortunately, this pattern is viral. For your example Foo struct, you now have to do this.

struct FooImpl{M}
    bar::M
end
Foo(T, N) = FooImpl{SMatrixFun(N, N, T)}

[Mason]

I prefer to use a macro for this, because then it can have the same name (but be distinguished by the @ sigil):

macro SMatrix(ex)
    Base.isexpr(ex, :braces, 3) || error("Something informative...")
    N, M, T = (ex.args)
    quote
        N = $(esc(N))
        M = $(esc(M))
        SMatrix{N, M, $(esc(T)), N*M}
    end
end

julia> @SMatrix{3, 2, Int}
SMatrix{3, 2, Int64, 6} (alias for SArray{Tuple{3, 2}, Int64, 2, 6})

Still viral though.