`Broadcast.broadcast_shape` for BlockedUnitRanges fails inference

[charleskawczynski]

MWE:

using BlockArrays
shape1 = (BlockArrays._BlockedUnitRange((2,)),);
shape2 = (BlockArrays._BlockedUnitRange((2,)),);
Base.Broadcast._bcs(shape1, shape2)
# @code_warntype Base.Broadcast._bcs1(shape1[1], shape2[1]) # lower level, but internals
@code_warntype Base.Broadcast.broadcast_shape(shape1, shape2)

[jishnub]

This comes from

BlockArrays.jl/src/blockbroadcast.jl

Line 37 in ec8c7f5

Base.Broadcast.axistype(a::T, b::T) where T<:BlockedUnitRange = length(b) == 1 ? a : combine_blockaxes(a, b)

Perhaps we may remove the special-casing, as the performance impact in the 1-term union will be minimal

[charleskawczynski]

Yeah, I tried fixing this. I'll open the PR for convenience (#312). First, I thought that the core issue was that union is used in sortedunion, and union does not preserve tuple types:

julia> union((1,), (2,))
2-element Vector{Int64}:
 1
 2

However, even after fixing that in the PR (which borrows some functions in TupleTools.jl and defines union on Tuples to return a Tuple), the result is still type unstable because the tuple length depends on the values in the tuple:

So, while this could fix the type preservation, it won't actually fix the type instability. If you have ideas about how to make Broadcast.broadcast_shape type stable / inferrable, that'd be great

[jishnub]

I think it's better for this to return a vector, as the non-trivial branch does. Since unique works on values, the length can't be known at compile time. I wonder if we may just remove the if-else here? I haven't checked if this changes the results

[charleskawczynski]

Yeah, that does fix it. I guess it's worth it. The general case of calling combine_blockaxes is either going to be type unstable in the case of Tuples (leading to inference triggers), or allocating heap allocating from union(a,b). Right now, it's a mixture, so it allocates and it's type-unstable.

[charleskawczynski]

I'll update the PR

[charleskawczynski]

Ok, the PR is updated.

[charleskawczynski]

Thanks for the tip @jishnub!

[charleskawczynski]

Actually, I prefer #313 if that's okay 🙂

[charleskawczynski]

So, it seems that length(b) == 1 ? a is needed for upstream tests to pass.

I updated #313, and basically specialized it so that the tests pass, but I have a feeling that I'm introducing inconsistent behavior between blocklasts()::Tuple and blocklasts()::Vector cases.

I think one potential solution would be to put the length of BlockedUnitRange in the type space:

struct BlockedUnitRange{CS,L} <: AbstractUnitRange{Int}
    first::Int
    lasts::CS
    global _BlockedUnitRange(f, cs::CS) where CS = new{CS, len(f, cs)}(f, cs)
end

len(f, l) = isempty(l) ? 0 : Integer(last(l)-f+1)
...
length(a::BlockedUnitRange{CS, L}) where {CS, L} = L

which will make the call to length (in combine_blockaxes) compile-time known, and remove the type instability. However, that may cause other issues. And I'm now seeing that things like DefaultBlockAxis may make this change more intrusive.

Thoughts? jishnub