Fix type instability when using custom indices

[barche]

This adds a new method indicestype to overcome a type instability when using non-conventional indices. Example:

using Compat
using Base.Test

using CustomUnitRanges: filename_for_zerorange
include(filename_for_zerorange)

immutable ZeroView{ScalarT,N} <: AbstractArray{ScalarT,N}
  a::Array{ScalarT,N}
end

@compat Base.IndexStyle{ScalarT,N}(::Type{ZeroView{ScalarT,N}}) = IndexLinear()
Base.indices{ScalarT,N}(v::ZeroView{ScalarT,N}) = map(ZeroRange, size(v.a))
Base.getindex{ScalarT}(v::ZeroView{ScalarT,1}, i::Integer) = v.a[i+1]
Base.setindex!{ScalarT}(v::ZeroView{ScalarT,1}, value, i::Integer) = (v.a[i+1] = value)
Base.getindex{ScalarT,N}(v::ZeroView{ScalarT,N}, i::Integer) = v.a[i] # linearindices are always 1-based for ND arrays
Base.setindex!{ScalarT,N}(v::ZeroView{ScalarT,N}, value, i::Integer) = (v.a[i] = value)

function benchmark_fill_vec(A)
  for i in linearindices(A)
    A[i] = i
  end
end

function benchmark_fill_mat(A)
  for i in indices(A,1)
    for j in indices(A,2)
      A[i,j] = i*(j+1)
    end
  end
end

const bench_size = 1000000
const ncols = 3
v_bench = ZeroView(zeros(bench_size))
A_bench = ZeroView(zeros(bench_size,ncols))

println("vector timings:")
@time benchmark_fill_vec(v_bench)
@time benchmark_fill_vec(v_bench)
@time benchmark_fill_vec(v_bench)

@test v_bench[0] == 0
@test v_bench[bench_size-1] == bench_size-1

println("matrix timings, linear indexing:")
@time benchmark_fill_vec(A_bench)
@time benchmark_fill_vec(A_bench)
@time benchmark_fill_vec(A_bench)

@test A_bench[0,0] == 1
@test A_bench[bench_size-1,2] == bench_size*ncols

println("matrix timings, [i,j] indexing:")
@time benchmark_fill_mat(A_bench)
@time benchmark_fill_mat(A_bench)
@time benchmark_fill_mat(A_bench)

@test A_bench[0,0] == 0
@test A_bench[bench_size-1,2] == (bench_size-1)*ncols

This gives a time of 0.18 s with 16 M allocations on the benchmark_fill_mat function. After the fix and specializing indicestype as follows:

Base.indicestype{ScalarT,N}(::Type{ZeroView{ScalarT,N}},d) = ZeroRange{Int}

Then the new timing is 0.02 s with 4 allocations.

Context:
https://discourse.julialang.org/t/multidimensional-arrays-with-0-based-indices/3325/7

cc @timholy

[timholy]

This might be worth having, but first let me ask: can't you just specialize indices(A::ZeroView, d)?

Since you're checking performance, it's worth noting that your benchmark_fill_mat traverses A in out-of-cache order. You might consider reversing the order of the two loops. Note that for I in CartesianRange(indices(A)) should always give you the proper order. However, #9080 will lead to a penalty, so if you're really trying to get maximal performance consider emulating this function. (You don't need the complicated ΔI stuff, just use it as an example of creating N loops from indices without paying the #9080 price.) Perhaps also add @inbounds.

The reason I'm mentioning these performance tips is that there's some chance (especially for high dimensions) that a ZeroView will be slightly faster than conventional julia arrays due to the simpler form of sub2ind; it would be interesting to document whether this is true and whether it's enough of an effect for anyone to care about. (I know @JaredCrean2 is interested in this.) If you want to try to see how big an effect it is, you might consider filling with a constant rather than something more complicated, just to isolate the access pattern itself.

[barche]

My current workaround is to specialize indices(A,d), but I feel this is a bit more complicated because you have to remember to properly implement the conditional on N. Maybe that's ok though, especially if it would be possible to write indices(A) in terms of indices(A,d). That way no extra API is needed.

Thanks for the tip on CartesianRange, I had forgotten about that.

[timholy]

That's a legitimate concern, and there are other places where I've had to be careful about constructing the correct type. So having some kind of trait for this might be the right thing to do.

However, the d part is a bit of a concern. If it can have different types for different dimensions (and I've seen cases where that seems like the right thing to do), then this interface won't be inferrable because d is a value.

[barche]

Would using Val be a viable option here, i.e

indicestype{ScalarT,N,D}(::Type{Array{ScalarT,N}},::Type{Val{D}})

[andyferris]

Personally, I would rather prefer that functions in base like broadcast never needed to call this function in this (second/negative) case... I.e. If custom types neglect to implement indexing operations other than linear or exactly N dimensional Cartesian, then it would be great if everything respected that and just worked (this surprised me in some point in the past).

[barche]

That would indeed simplify things. It's what I'm doing right now as a workaround. So far I noticed this breaks display(A), but I'm sure that's fixable.

[timholy]

Yes, it seems that Val or similar would be necessary for this.

I guess the main thing I'd consider is whether indicestype has more widespread use. If it reduces complexity in other code, I'd be less equivocal in my enthusiasm.

[barche]

I can't really think of any other use, to be honest. Following on @andyferris 's comment, I wonder what would happen if the branch on d were simply removed, so that calling indices(A,d) with d> N would be an error? It seems this would at least allow removing indices1 and make the current PR unnecessary.

[timholy]

That's an interesting suggestion, and given that indices have different types this is probably the right way to go. It does introduce a difference with size (which doesn't have this problem since it always returns an Int), and the size(A, d) construct is occasionally useful:

function foo(A::AbstractVecOrMat)
    m, n = size(A, 1), size(A, 2)
    ...

However, this also causes some performance gotchas (e.g., JuliaMath/Interpolations.jl#154) because of the branch.

I wonder if it would be better to have a more composable solution. For example, two that work now are

function foo(A::AbstractVecOrMat)
    sz = size(A)
    m, n = sz[1], length(sz) == 2 ? sz[2] : 1
    ...

or

function foo(A::AbstractVecOrMat)
    m, n = Base.fill_to_length(size(A), 1, Val{2})
    ...

The indices variant would be a little more complicated,

function foo(A::AbstractVecOrMat)
    inds = indices(A)
    inds1, inds2 = Base.fill_to_length(inds, oftype(inds[1], Base.OneTo(1)), Val{2})
    ...

Making all this easier to use seems like a worthwhile exercise.

[andyferris]

As a bit of a indices-noob - how do implied singleton dimensions work in the indices world? E.g. does/can/should broadcast treat anything other than OneTo(1) as expandable?

If so, then I'm wondering if there is no well-defined answer for what that implied index is for all AbstractArray(even if we allow for its size to be 1, which seems reasonable independent of the index type, if it is 1-based or zero-based, or whatever).

[andyferris]

OTOH a type like NoIndex could be the implied singleton index for n > N as well as the first dimension of a RowVector (It seems to me that the first index of a row vector should broadcast correctly with any e.g. vector).

This is a tricky one, but I think "missing" early dimensions are a thing that we should try to support correctly.

[andyferris]

(Cc @mbauman @jiahao w.r.t. earlier discussions of tensors/indices)

[barche]

I must admit I didn't consider broadcasting at all. The implied indexing does seem counter-intuitive, with or without this PR:

# Column vector:
v = ZeroView([1,2,3])
v[0] # gives 1
v[0,0] # bounds error at abstractarray.jl:372
v[0,1] # gives 1
# Row vector:
v = ZeroView([1 2 3])
v[0] # bounds error, because linear indices are 1-based
v[0,0] # gives 1
v[1,0] # bounds error

The second bounds error is because linear indices are 1-based and has nothing to do with this PR.

The first one comes from line 404 or 408 in abstractarray.jl I think, and it might be fixable (if deemed necessary) in the context of the current PR by replacing the 1:1 with one(indicestype(... although this would require passing along information about the indices type to checkbounds_indices

[mbauman]

I don't have any good answers here, but I can similarly point at my own failed experiments at trying to use custom indices types to support the Axis-propagation behavior in AxisArrays (JuliaArrays/AxisArrays.jl#58 and JuliaArrays/AxisArrays.jl#81).

That second attempt is failing due to type assertions that work around these sorts of instabilities within broadcasting code.

In the context of AxisArrays, I think the "right" behavior is to replace the singleton axes with the broadcasted one. I'm not sure how generalizable that choice is, but it implies to me that any length-1 dimension can support broadcasting, regardless of its indices/axes.

[andyferris]

I think the "right" behavior is to replace the singleton axes with the broadcasted one

My feeling is that is correct, but...

but it implies to me that any length-1 dimension can support broadcasting, regardless of its indices/axes.

... note that this sometimes might be ambiguous - if I broadcast two vectors with indices (0:0,) and (2:2,), what are the output indices?

(It's definitely a bit of a strange case, but if the length of the vectors are set at run-time, then I'm worried we won't be able to avoid odd behavior in a small number of corner cases when lengths of indices happen to equal 1).

[barche]

This got fixed along the road to 1.0, it seems!