ArrayInterface.size vs StaticArrays.Size

[cscherrer]

Hi,

I'd like to get a better understanding of the intending idioms for working with arrays, especially the case where we want different methods for the static vs dynamic case.

StaticArrays.Size solves the problem like this:

julia> Size(zeros(SMatrix{3, 4}))
Size(3, 4)

julia> Size(zeros(3, 4))
Size(StaticArrays.Dynamic(), StaticArrays.Dynamic())

Though we might have some missing results, it's nice that the size computation itself is always static. Working with types is also convenient:

julia> Size(typeof(zeros(SMatrix{3, 4})))
Size(3, 4)

julia> Size(typeof(zeros(3, 4)))
Size(StaticArrays.Dynamic(), StaticArrays.Dynamic())

By contrast, here's ArrayInterface.size:

julia> ArrayInterface.size(zeros(SMatrix{3, 4}))
(static(3), static(4))

julia> ArrayInterface.size(zeros(3, 4))
(3, 4)

That last result might seem convenient, but computing this requires dynamic dispatch, which seems to defeat the purpose. And it fails on types:

julia> ArrayInterface.size(typeof(zeros(3, 4)))
ERROR: MethodError: no method matching size(::Type{Matrix{Float64}})
Closest candidates are:
  size(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:524
  size(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}, ::Integer) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:523
  size(::Union{LinearAlgebra.Cholesky, LinearAlgebra.CholeskyPivoted}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/cholesky.jl:442
  ...
Stacktrace:
 [1] axes
   @ ./abstractarray.jl:89 [inlined]
 [2] axes
   @ ~/.julia/packages/ArrayInterface/CmPZc/src/axes.jl:166 [inlined]
 [3] size(a::Type{Matrix{Float64}})
   @ ArrayInterface ~/.julia/packages/ArrayInterface/CmPZc/src/size.jl:18
 [4] top-level scope
   @ REPL[95]:1

julia> ArrayInterface.size(typeof(zeros(3, 4)))
ERROR: MethodError: no method matching size(::Type{Matrix{Float64}})
Closest candidates are:
  size(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:524
  size(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}, ::Integer) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:523
  size(::Union{LinearAlgebra.Cholesky, LinearAlgebra.CholeskyPivoted}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/cholesky.jl:442
  ...
Stacktrace:
 [1] axes
   @ ./abstractarray.jl:89 [inlined]
 [2] axes
   @ ~/.julia/packages/ArrayInterface/CmPZc/src/axes.jl:166 [inlined]
 [3] size(a::Type{Matrix{Float64}})
   @ ArrayInterface ~/.julia/packages/ArrayInterface/CmPZc/src/size.jl:18
 [4] top-level scope
   @ REPL[95]:1

Is ArrayInterface.size intended to be used for building generated functions or for something else? In what cases does it have benefits over StaticArrays.Size?

[Tokazama]

How does this require dynamic dispatch? All the statically known info is computed in the type domain and inlined so that it is known at compile time. If you want to operate on just the type you can use ArrayInterface.known_size.

For reference, here are some related methods:

• known_length(x)
• known_first(x)
• known_last(x)
• known_step(x)

[cscherrer]

Say I'm building a mul(A::AbstractMatrix, x::AbstractVector), and I dispatch using ArrayInterface.size. It seems to me this would only be entirely static when all sizes are known statically.

known_size seems to partially replicate Size, but the result is not type-level, so it still can't be used directly for dispatch.

I think I'm just missing some context on intended usage. Do you know of examples where these methods make it easier to build fast code than just using StaticArrays.Size?

[cscherrer]

Oh, maybe something like

struct Sized{S,T}
    x::T
end

Sized(x::T) where {T} = Sized{ArrayInterface.known_size(T), T}(x)

?

Then dispatch is easier:

julia> f(::Sized{(3,)}) = "three"
f (generic function with 1 method)

julia> f(x) = "not three"
f (generic function with 2 methods)

julia> x = Sized(SizedVector{3}([1,2,3]))
Sized{(3,), SizedVector{3, Int64, Vector{Int64}}}([1, 2, 3])

julia> f(x)
"three"

[Tokazama]

If you want a single type to dispatch on you could do VecSize{S} = Tuple{StaticInt{S}} and MatSize{S1,S2} = Tuple{StaticInt{S1},StaticInt{S2}}.

If you don't know how many dimensions it is and all you know is that you have a type T = Tuple{Vararg{StaticInt}}, then you can get the regular tuple with Static.known(T).

[cscherrer]

Ok, how would you decide whether to take this approach or to use StaticArrays.Size?

[Tokazama]

I wouldn't ever use StaticArrays.Size. I think some maintainers for StaticArrays.jl agree. I honestly don't see any benefit to Size.

[cscherrer]

Ah, there's the missing context. That's really helpful, thanks!

[Tokazama]

Of course! Happy to help

[chriselrod]

and I dispatch using ArrayInterface.size. It seems to me this would only be entirely static when all sizes are known statically.

Not sure what you meant here, but ArrayInterface.size shouldn't be resulting in dynamic dispatches.

julia> using StrideArrays, Static, ArrayInterface

julia> A = StrideArray{Float64}(undef, (static(2), 3, static(4), 5));

julia> ArrayInterface.size(A)
(static(2), 3, static(4), 5)

julia> typeof(ans)
Tuple{StaticInt{2}, Int64, StaticInt{4}, Int64}

julia> Static.known(typeof(ArrayInterface.size(A)))
(2, nothing, 4, nothing)

You can dispatch on a Tuple{StaticInt{2}, Int, StaticInt{4}, Int} just fine, and the type is known at compile time. You may have to be careful indexing into it, i.e. the indexes should be known at compile time.

I'll echo Tokazama's points.

Also, map tends to be type stable when working with heterogenous tuples, while broadcast isn't. My code working with tuples thus tends to be a mix of

1. map
2. first and Base.tail in recursive functions, where each call pops off the first remaining element, processes it, and then calls itself using Base.tail as the arg.
3. @generated

[cscherrer]

Not sure what you meant here, but ArrayInterface.size shouldn't be resulting in dynamic dispatches.

I think it depends how you use it - I'm just trying to figure out the right idioms.

Say you're building a mul(m::AbstractMatrix, y::Abstractvector). You might try something like

function mul(m::AbstractMatrix, v::Abstractvector)
    m_size = ArrayInterface.size(m)
    v_size = ArrayInterface.size(v)
    _mul(m, v, m_size, v_size)

where _mul has some generated methods for static sizes. It seemed to me if m or v have any dynamic sizes this would require dynamic dispatch. Maybe the compiler is smart enough to work around this?

Then, if you're building a generated function, ArrayInterface.size doesn't help, because it fails on types. It's helpful to know about known_size.

Basically what I'm looking for is a design pattern for writing functions that split into generic methods vs static methods with generated functions.

In the multivariate normal example we're looking at, I'm currently doing something like

logdensity(d::MvNormal{(:ω,), <:Tuple{<:Cholesky}}, y) = logdensity_mvnormal_ω(getfield(d.ω, :factors), y)

logdensity_mvnormal_ω(UL::Union{UpperTriangular, LowerTriangular}, y::AbstractVector) = logdensity_mvnormal_ω(KnownSize(UL), y)

where KnownSize is

struct KnownSize{S, T}
    value::T
end

KnownSize(x::T) where {T} = KnownSize{Tuple{ArrayInterface.known_size(T)...}, T}(x)

This gives a size wrapper around the value. Dispatch is easy, you can keep the same argument slots, and it's easy to unwrap. But maybe there's a cleaner way to do this?

[Tokazama]

It seemed to me if m or v have any dynamic sizes this would require dynamic dispatch. Maybe the compiler is smart enough to work around this?

It still shouldn't be dynamic dispatch (although this is helping me realize how obtuse some of the code currently is without more detailed documentation).

If you look at related tests, you can see that it's all fully inferred.

I would just call logdensity_mvnormal_ω(ArrayInterface.size(UL), y). Then if you want a unique function for each combination of "staticness" define:

logdensity_mvnormal_ω(s::Tuple{Int,Int}, y) = ...
logdensity_mvnormal_ω(s::Tuple{StaticInt{S1},Int}, y) where {S1} = ...
logdensity_mvnormal_ω(s::Tuple{Int,StaticInt{S2}}, y) where {S2} = ...
logdensity_mvnormal_ω(s::Tuple{StaticInt{S1},StaticInt{S2}}, y) where {S1,S2} = ...

Depending on what you need to do you might want to make any of the last three @generated

You could also do something like

@generated function logdensity_mvnormal_ω(s::Tuple{S1,S2}, y) where {S1,S2}
    s1 = known(S1) === nothing ? :(getfield(s, 1)) : known(S1)
    s2 = known(S2) === nothing ? :(getfield(s, 2)) : known(S2)

    ...do stuff
end

[cscherrer]

Ok, this is making some more sense. Thank you for your help!

[chriselrod]

Then, if you're building a generated function, ArrayInterface.size doesn't help, because it fails on types. It's helpful to know about known_size.

Looking at your earlier example, something like this should also work:

using Static, ArrayInterface
function mul(m::AbstractMatrix, v::Abstractvector)
    m_size = ArrayInterface.size(m)
    v_size = ArrayInterface.size(v)
    _mul(m, v, m_size, v_size)
end
@generated function _mul(m, v, m_size::MS, v_size::VS) where {MS,VS}
    # `ms` and `vs` are tuples, with values for known information and `nothing` when the sizes were dynamic
    ms = known(MS)
    vs = known(VS)
    m_sizes = Any[]
    q = quote end
    # build the expression. If `ms[i] === nothing`, you'd use `:(m_size[$i])` for the size,
    # otherwise you can use `ms[i]` or manually unroll the expression by `ms[i]`.
end

I agree that this could use more documentation.

Ideas on how to make this more convenient are welcome.

And to repeat, we do want to keep all the ArrayInterface.size calls outside of generated functions to avoid possible world age issues.

[cscherrer]

I've thought about this some more while working on https://github.com/cscherrer/MultivariateMeasures.jl. There are two main cases that come up. First, we may have full control over the generic implementation. That's like this:

Looking at your earlier example, something like this should also work:

using Static, ArrayInterface
function mul(m::AbstractMatrix, v::Abstractvector)
    m_size = ArrayInterface.size(m)
    v_size = ArrayInterface.size(v)
    _mul(m, v, m_size, v_size)
end
@generated function _mul(m, v, m_size::MS, v_size::VS) where {MS,VS}
    # `ms` and `vs` are tuples, with values for known information and `nothing` when the sizes were dynamic
    ms = known(MS)
    vs = known(VS)
    m_sizes = Any[]
    q = quote end
    # build the expression. If `ms[i] === nothing`, you'd use `:(m_size[$i])` for the size,
    # otherwise you can use `ms[i]` or manually unroll the expression by `ms[i]`.
end

Let's say m is i by j, and v has length j. Some or all of these might be known statically. So maybe we have a strange case where

ArrayInterface.size(m) == (StaticInt(5), 4)
ArrayInterface.size(v) == (StaticInt(4),)

Making this fast should be no problem, but my earlier concern is that even asking for ArrayInterface.size(m) would (I thought) trigger a dynamic dispatch to get its column count.

The second situation is where a generic mul(m::AbstractMatrix, v::Abstractvector) already exists, and we need to add optimized methods for static cases. Is there a nice way to do this?

[Tokazama]

Dispatch will not be dynamic, because the operations that need to be performed will be known at compile time. Finding the column count would be computed dynamically though. For example, ArrayInterface.size(ones(3)') would lower to code like this at compile time: Static.StaticInt(1), Base.arraysize(x, 1).

The second example is problematic because it really does require community buy in to traits or directly dispatching on the array types. This is where we'll probably need libraries that redefine these functions outside of the standard library, which is already happening with packages like Octavio.jl. I think this sort of development is probably a necessary step for evolving LinearAlgebra further.

[cscherrer]

ArrayInterface.size(ones(3)') would lower to code like this at compile time: Static.StaticInt(1), Base.arraysize(x, 1).

This is a great example! It really helps to see what's computed at compile time. Could this be added to the docs?

The second example is problematic because it really does require community buy in to traits or directly dispatching on the array types. This is where we'll probably need libraries that redefine these functions outside of the standard library, which is already happening with packages like Octavio.jl. I think this sort of development is probably a necessary step for evolving LinearAlgebra further.

Good point. How about for designing packages around this idea? I've heard from a few devs that they'd like fewer dependencies in MeasureTheory.jl. So I'd like to have generic methods here, and high-performance implementations (with dependencies like ArrayInterface and LoopVectorization) in MultivariateMeasures.jl. The challenge is that MeasureTheory doesn't "know" about about MultivariateMeasures.

That's one example, but I'd guess a situation where specialized implementations add to something more generic could be (or become) very common.

[Tokazama]

This is a great example! It really helps to see what's computed at compile time. Could this be added to the docs?

Can do.

How about for designing packages around this idea?

I think that having a common set of traits is a necessary starting point. Fortunately, ArrayInterface.jl is relatively light weight and widely used (looks like MeasureTheory.jl already indirectly depends on it). But there are a lot of options when creating a generic design that permits high performance methods to step in later. For example, ArrayInterface.can_avx allows generic buy in for LoopVectorization.jl without a dependency.

You could have a trait defined in MeasureTheory that dispatches to a high performance trait in MultivariateMeasures

MeasureTheory.fxn(x) = fxn(MTTrait(x), x)
MeasureTheory.fxn(::MTTrait, x) = ...

MultivariateMeasures.fxn(::FastMTTrait, x) = ...

[Tokazama]

@chriselrod , what do you think about replacing ArrayInterface.size with a type like ArrayInterface.Size that wraps a tuple of optionally static ints? Something like this:

struct Size{N,S<:Tuple}
    size::S
end

Base.size(x::Size) = map(Int, x.size)

Base.length(x::Size) = Int(prod(x.size))

const Length{L} = Size{1,Tuple{L}}
Length(x::CanonicalInt) = Size((x,))
@inline function Length(x)
    len = known_length(x)
    if len === missing
        return Length(length(x))
    else
        return Length(static(len))
    end
end
Size(A::Array) = Size(size(A))
Size(A::Array, dim::Int) = Length(size(A, dim))

...which would replace ArrayInterface.size and ArrayInterface.static_length.

[chriselrod]

What is the benefit of that?

[Tokazama]

Its a bit more user friendly than having a module specific definition of size. We can also dispatch on it similar(::Abstract device, ::Type{T}, ::Size).

Personally, I think it'd be nicer than what we have now but if you don't like it I'm fine with leaving things the way they are.

[ChrisRackauckas]

No longer has the silly naming that confuses everyone, and no longer in this repo.