Safe co-iteration across an axis for 1+ arrays

[Tokazama]

This should help with #62. The main idea is that we can get the convenience of checking compatible indices like in eachindex(A...), but for specific axes. Possible concerns may be:

1. I decided to have indices(x) always return iterable indices instead of a tuple of indices. So the default behavior is like eachindex. This does two things for us:
   1. The behavior of for i in indices(args...) ... end is more predictable becaues i can always point to a single value of the array(s).
   2. It makes it very easy to incorporate/extend for other array/array-like types. If the behavior is unique to each axis then only indices(x) needs to be specified. If the behavior is unique for the array then only indices(x, i) needs to be specified.
2. I didn't implement any special behavior for combining/optimizing known range first/last/step although it was mentioned in Safe co-iteration along axes #62. I figured even a minimal range type would be meaty enough for an independent PR.

If people don't find the first point compelling enough to depart from behavior like axes I can change it. I just thought it was a really interesting idea. I'm also happy to implement a range if it's seen as immediately necessary for this PR.

[chriselrod]

I didn't implement any special behavior for combining/optimizing known range first/last/step although it was mentioned in #62. I figured even a minimal range type would be meaty enough for an independent PR.

I would like to see this.
I'd like to add that instead of the pattern matching, I'd like the iterators to also track an is-dense along the lines of #61.

Then that would provide all the information I'd need to implement the optimizations I talked about, i.e.:

1. Possibility of flattening some subset of loops, CartesianIndexing -> LinearIndexing style.
2. Realizing that stride values must necessarily be equal between arrays, to avoid spilling integer registers.

It'd also be great to get the compile time length promotion.
While LLVM would be able to figure out that if SOneTo(3) == (x::UnitRange), then x must be 1:3, LoopVectorization cannot figure that out, because it won't know that SOneTo(3) == x. Thus the consolidation of known information would really help it when mixing arrays of known and unknown dimensions.

[Tokazama]

I'd like to add that instead of the pattern matching, I'd like the iterators to also track an is-dense along the lines of #61.

Are you referring to something that does something like stridelayout(::Type{T}) -> (contig, batch, striderank) for each axis, or just flattening out the results of indices when possible?

[chriselrod]

I forgot to add axesdense to the list of returns: stridelayout(::Type{T}) -> (contig, batch, striderank, axesdense).
So that'd mean consolidating the axesdense information.
The reason I want that info from the axis iterator rather than to read it directly from the arrays is:

A = rand(10, 5);
for i in 1:3, j in 1:4
    A[i,j] *= 2
end

A is dense, but that doesn't mean these loops could be flattened.

[Tokazama]

Are you thinking of something like this?

is_dense(::Type{T}, i::Int) where {T} = nothing
is_dense(::Type{<:Array}, i::Int) = true
@inline function is_dense(::Type{<:Union{Transpose{T,A},Adjoint{T,A}}}) where {T,A<:AbstractMatrix{T}}
    return is_dense(A, (2, 1)[i])
end
@inline function is_dense(::Type{<:PermutedDimsArray{T,N,I1,I2,A}}, i::Int) where {T,N,I1,I2,A<:AbstractArray{T,N}}
    return is_dense(A, I1[i])
end

I don't really see how an iterator along each dimension in isolation can incorporate the density since it requires knowing about the previous dimensions and strides.

[chriselrod]

Something like:

is_dense(::Type{T}, i) where {T} = last(stridelayout(T))[i]

I don't really see how an iterator along each dimension in isolation can incorporate the density since it requires knowing about the previous dimensions and strides.

You're right. Given what I said:

A = rand(2,3,4);
B = PermutedDimsArray(rand(4,3,2), (3,2,1));
for k in indices((A,B),3), j in indices((A,B),2), i in indices((A,B),1)
    # do something with A[i,j,k] and B[i,j,k]
end

These each of these should be dense, even though we obviously can't flatten the loops.
However, we could flatten them if we only had A or if we only had B.

That is, is_dense by itself does not provide enough information, but I think it does in conjunction with stridelayout.
I.e., in querying whether we can flatten two loops, we need to check:

1. That they share the same order among all arrays indexed by them (provided by stridelayout)
2. That the loop with the smaller stride is_dense/covers the full axis for each of them (information must be provided by the stride-iterator).

[Tokazama]

This sounds a lot like what Slice is for. Would it make sense to wrap whatever we get from indices into something like Base.Slice or should we have a unique type that's "co-slicing"?

[chriselrod]

Okay, if that's the understood meaning of Slice.

Simply doing that, and then checking stridelayout to see if all the sliced dimensions are dense sounds like a good approach to me.
Meaning indices only has to accumulate statically known first/last/step.

[chriselrod]

I'm curious, what is the reason to prefer getfield(r, :start) over r.start?
I've always used the latter in my own code.

[Tokazama]

Sorry, it's just a habit from working with types that have overloaded the getproperty method for metadata like tables. I can change this, as it's probably easier to read r.start.

[chriselrod]

Maybe Base.udiv_int if we can guarantee the checks will pass?

[Tokazama]

The code for calculating the length on StepRange is what I'm basing this one and I should probably expand this out to also include this sort of code.

[Tokazama]

One last question and I think I can get something more polished together. Base.Slice only accepts AbstractUnitRange. Do you need OptinallyStaticRange to be capable of having different step sizes?

[chriselrod]

Maybe add

struct OptionallyStaticUnitRange{T,F,L} <: AbstractUnitRange{T}
    start::F
    stop::L
end

maybe_slice(r) = r
function maybe_slice(r::OptionallyStaticRange{T,F,::Val{1},L}) where {T,F,L}
    Base.Slice(OptionallyStaticUnitRange{T,F,L}(first(r), last(r)))
end

Then you could use maybe_slice(reduce(_pick_range, inds)).

Although...I think you're right.
It would probably be better to just make it a UnitRange. I'm having a hard time coming up with any examples where the step size could be something else?

[Tokazama]

I would really love if we had a common interface where someone could do this with a new type

function Base.length(r::NewRangeType{T}) where{T}
    if isempty(r)
        return zero(T)
    else
        return ArrayInterface.unsafe_length(first(r), step(r), last(r))
    end
end

But that turns into a lot of code very quickly and I should probably save that for a different PR. Would it be okay if I just do a unit range thing for now and then we can see how everyone feels about a more aggressive range focused PR later that could include OptionallyStaticStepRange?

[chriselrod]

I would really love if we had a common interface where someone could do this with a new type

Do you mean, define an interface such that length will work in that manner?
Defining your own abstract type that they could subtype (along with defining first, etc) would be easiest, otherwise wouldn't you need a PR on base Julia to do this with traits without committing type piracy? Or am I missing something?

The OptionallyStaticStepRange is fine.
Are we going to make it so that the type promises to be a full slice? Or do you want to still wrap the returns in Base.Slice?

[Tokazama]

Do you mean, define an interface such that length will work in that manner?
Defining your own abstract type that they could subtype (along with defining first, etc) would be easiest, otherwise wouldn't you need a PR on base Julia to do this with traits without committing type piracy? Or am I missing something?

I think it would ultimately need to be in base but there are a lot of little problems with how ranges are implemented in base and we'd need to start by getting the known_* methods there first.

Are we going to make it so that the type promises to be a full slice? Or do you want to still wrap the returns in Base.Slice?

I think it might be good to still wrap the return in Base.Slice because then things like getindex(A, indices(...)) can know that it represents and inbounds collection.

[chriselrod]

I think this looks good. Are you happy with it?

[Tokazama]

With this final change I think it's good. Now you can use indices((A1, A2)) like eachindex(A1, A2) but it will give you a slice so you know it spans the full span of each array.

[timholy]

This computes the indices of the indices, rather than the indices of x. You're implicitly assuming that axes(x) is idempotent. The older implementations of OffsetArrays would break this, for example, although the community seems to have widely settled on idempotency as a useful characteristic for the axes.

[Tokazama]

In terms of iterating values it should theoretically be the same, but it's not entirely idempotent because indices(x::OneTo) = Slice(OptionallyStaticUnitRange(Val(1), last(x))). If indices(x::OffsetArray, d) = indices(axes(x, d)) doesn't produce something that is appropriately offset then it could do something like indices(x::OffsetArray, d) = indices(range(x.offset[d], stop = size(x, d)).