Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

RFC: Give AbstractArrays smart and performant indexing behaviors for free #10525

Merged
merged 5 commits into from
Jun 4, 2015
Merged

RFC: Give AbstractArrays smart and performant indexing behaviors for free #10525

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a07c6d3
Adds methods for getindex and setindex! with Colon
mbauman Mar 7, 2015
b77b026
AbstractArray: fancy getindex and setindex! fallbacks
mbauman May 5, 2015
5cb2835
Adapt AbstractArray subtypes for indexing fallbacks
mbauman May 5, 2015
9384586
Perf tests for indexing
mbauman May 6, 2015
0888a79
News for #10525
mbauman May 6, 2015
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion NEWS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,12 +192,23 @@ Library improvements

* `is_valid_char(c)` now correctly handles Unicode "non-characters", which are valid Unicode codepoints. ([#11171])

* Data-structure processing
* Array and AbstractArray improvements

* New multidimensional iterators and index types for efficient iteration over `AbstractArray`s. Array iteration should generally be written as `for i in eachindex(A) ... end` rather than `for i = 1:length(A) ... end`. ([#8432])

* New implementation of SubArrays with substantial performance and functionality improvements ([#8501]).

* AbstractArray subtypes only need to implement `size` and `getindex`
for scalar indices to support indexing; all other indexing behaviors
(including logical idexing, ranges of indices, vectors, colons, etc.) are
implemented in default fallbacks. Similarly, they only need to implement
scalar `setindex!` to support all forms of indexed assingment ([#10525]).
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Making this true in practice for general types is presumably conditional on merging #10911, but gosh darnit I am going to figure out the remaining challenges for the NTuple/Vararg unification. Status: now building julia only crashes after inference gets turned on, rather than in the 2nd file, which to me seems like progress 😄.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think I worded this in a way that it's still technically true. You just need to implement all combinations of scalar indexing (linear and otherwise) until #10911. But, yes, I'm looking forward to it!


* AbstractArrays that do not extend `similar` now return an `Array` by
default ([#10525]).

* Data-structure processing

* New `sortperm!` function for pre-allocated index arrays ([#8792]).

* Switch from `O(N)` to `O(log N)` algorithm for `dequeue!(pq, key)`
Expand Down Expand Up @@ -1417,6 +1428,7 @@ Too numerous to mention.
[#10400]: https://github.com/JuliaLang/julia/issues/10400
[#10446]: https://github.com/JuliaLang/julia/issues/10446
[#10458]: https://github.com/JuliaLang/julia/issues/10458
[#10525]: https://github.com/JuliaLang/julia/issues/10525
[#10543]: https://github.com/JuliaLang/julia/issues/10543
[#10659]: https://github.com/JuliaLang/julia/issues/10659
[#10679]: https://github.com/JuliaLang/julia/issues/10679
Expand Down
296 changes: 226 additions & 70 deletions base/abstractarray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,41 +118,52 @@ linearindexing{T<:Range}(::Type{T}) = LinearFast()
*(::LinearFast, ::LinearSlow) = LinearSlow()
*(::LinearSlow, ::LinearSlow) = LinearSlow()

# The real @inline macro is not available this early in the bootstrap, so this
# internal macro splices the meta Expr directly into the function body.
macro _inline_meta()
Expr(:meta, :inline)
end
macro _noinline_meta()
Expr(:meta, :noinline)
end

## Bounds checking ##
checkbounds(sz::Int, ::Colon) = nothing
checkbounds(sz::Int, i::Int) = 1 <= i <= sz || throw(BoundsError())
checkbounds(sz::Int, i::Real) = checkbounds(sz, to_index(i))
checkbounds(sz::Int, I::AbstractVector{Bool}) = length(I) == sz || throw(BoundsError())
checkbounds(sz::Int, r::Range{Int}) = isempty(r) || (minimum(r) >= 1 && maximum(r) <= sz) || throw(BoundsError())
checkbounds{T<:Real}(sz::Int, r::Range{T}) = checkbounds(sz, to_index(r))

function checkbounds{T <: Real}(sz::Int, I::AbstractArray{T})
_checkbounds(sz, i::Integer) = 1 <= i <= sz
_checkbounds(sz, i::Real) = 1 <= to_index(i) <= sz
_checkbounds(sz, I::AbstractVector{Bool}) = length(I) == sz
_checkbounds(sz, r::Range{Int}) = (@_inline_meta; isempty(r) || (minimum(r) >= 1 && maximum(r) <= sz))
_checkbounds{T<:Real}(sz, r::Range{T}) = (@_inline_meta; _checkbounds(sz, to_index(r)))
_checkbounds(sz, ::Colon) = true
function _checkbounds{T <: Real}(sz, I::AbstractArray{T})
@_inline_meta
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we really want to force inlining on this? EDIT: same question applies to multiple methods below (I won't annotate them).

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, that is a very good question, and it'd probably be good to review all the little performance tweaks I've been making to make sure they're necessary. This one, however, makes a huge difference. It's on the order of 10x more costly to call this function than it is to inline it. These bounds checks are effectively "inlined" into the current n-ary arrayref intrinsics, so I needed to do all sorts of these inlining tricks to match it. These sorts of optimizations are so finicky and tricky (and time consuming!) to test.

I've also played with changing the branching behavior (e.g., using & instead of &&, etc) and can get really nice-looking (and short!) LLVM code that ends up performing worse for some strange cache/pipeline/who-knows-what reason.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I really like how you wrote the function---presumably it's SIMD-able this way.

And, thanks for letting me know about the importance of inlining. I assumed anytime you cranked up a for-loop of possibly large size (but, possibly not), it wouldn't be worth inlining. Certainly I wouldn't have guessed this. Thanks for testing!

b = true
for i in I
checkbounds(sz, i)
b &= _checkbounds(sz, i)
end
b
end
# Prevent allocation of a GC frame by hiding the BoundsError in a noinline function
throw_boundserror(A, I) = (@_noinline_meta; throw(BoundsError(A, I)))

checkbounds(A::AbstractArray, I::AbstractArray{Bool}) = size(A) == size(I) || throw(BoundsError())

checkbounds(A::AbstractArray, I) = checkbounds(length(A), I)

function checkbounds(A::AbstractMatrix, I::Union(Real,Colon,AbstractArray), J::Union(Real,Colon,AbstractArray))
checkbounds(size(A,1), I)
checkbounds(size(A,2), J)
checkbounds(A::AbstractArray, I::AbstractArray{Bool}) = size(A) == size(I) || throw_boundserror(A, I)
checkbounds(A::AbstractArray, I::AbstractVector{Bool}) = length(A) == length(I) || throw_boundserror(A, I)
checkbounds(A::AbstractArray, I) = (@_inline_meta; _checkbounds(length(A), I) || throw_boundserror(A, I))
function checkbounds(A::AbstractMatrix, I::Union(Real,AbstractArray,Colon), J::Union(Real,AbstractArray,Colon))
@_inline_meta
(_checkbounds(size(A,1), I) && _checkbounds(size(A,2), J)) || throw_boundserror(A, (I, J))
end

function checkbounds(A::AbstractArray, I::Union(Real,Colon,AbstractArray), J::Union(Real,Colon,AbstractArray))
checkbounds(size(A,1), I)
checkbounds(trailingsize(A,2), J)
function checkbounds(A::AbstractArray, I::Union(Real,AbstractArray,Colon), J::Union(Real,AbstractArray,Colon))
@_inline_meta
(_checkbounds(size(A,1), I) && _checkbounds(trailingsize(A,2), J)) || throw_boundserror(A, (I, J))
end

function checkbounds(A::AbstractArray, I::Union(Real,Colon,AbstractArray)...)
function checkbounds(A::AbstractArray, I::Union(Real,AbstractArray,Colon)...)
@_inline_meta
n = length(I)
if n > 0
for dim = 1:(n-1)
checkbounds(size(A,dim), I[dim])
_checkbounds(size(A,dim), I[dim]) || throw_boundserror(A, I)
end
checkbounds(trailingsize(A,n), I[n])
_checkbounds(trailingsize(A,n), I[n]) || throw_boundserror(A, I)
end
end

Expand All @@ -178,6 +189,8 @@ similar (a::AbstractArray, T) = similar(a, T, size(a))
similar{T}(a::AbstractArray{T}, dims::Dims) = similar(a, T, dims)
similar{T}(a::AbstractArray{T}, dims::Int...) = similar(a, T, dims)
similar (a::AbstractArray, T, dims::Int...) = similar(a, T, dims)
# similar creates an Array by default
similar (a::AbstractArray, T, dims::Dims) = Array(T, dims)

function reshape(a::AbstractArray, dims::Dims)
if prod(dims) != length(a)
Expand Down Expand Up @@ -361,11 +374,7 @@ zero{T}(x::AbstractArray{T}) = fill!(similar(x), zero(T))

# While the definitions for LinearFast are all simple enough to inline on their
# own, LinearSlow's CartesianRange is more complicated and requires explicit
# inlining. The real @inline macro is not available this early in the bootstrap,
# so this internal macro splices the meta Expr directly into the function body.
macro _inline_meta()
Expr(:meta, :inline)
end
# inlining.
start(A::AbstractArray) = (@_inline_meta(); itr = eachindex(A); (itr, start(itr)))
next(A::AbstractArray,i) = (@_inline_meta(); (idx, s) = next(i[1], i[2]); (A[idx], (i[1], s)))
done(A::AbstractArray,i) = done(i[1], i[2])
Expand Down Expand Up @@ -430,19 +439,6 @@ imag{T<:Real}(x::AbstractArray{T}) = zero(x)

\(A::Number, B::AbstractArray) = B ./ A

## Indexing: getindex ##

getindex(t::AbstractArray, i::Real) = error("indexing not defined for ", typeof(t))

# linear indexing with a single multi-dimensional index
function getindex(A::AbstractArray, I::AbstractArray)
x = similar(A, size(I))
for i in eachindex(I)
x[i] = A[I[i]]
end
return x
end

# index A[:,:,...,i,:,:,...] where "i" is in dimension "d"
# TODO: more optimized special cases
slicedim(A::AbstractArray, d::Integer, i) =
Expand Down Expand Up @@ -490,42 +486,202 @@ function circshift{T,N}(a::AbstractArray{T,N}, shiftamts)
a[(I::NTuple{N,Vector{Int}})...]
end

## Indexing: setindex! ##

# 1-d indexing is assumed defined on subtypes
setindex!(t::AbstractArray, x, i::Real) =
error("setindex! not defined for ",typeof(t))
setindex!(t::AbstractArray, x) = throw(MethodError(setindex!, (t, x)))

## Indexing: handle more indices than dimensions if "extra" indices are 1

# Don't require vector/matrix subclasses to implement more than 1/2 indices,
# respectively, by handling the extra dimensions in AbstractMatrix.

function getindex(A::AbstractVector, i1,i2,i3...)
if i2*prod(i3) != 1
throw(BoundsError())
## Approach:
# We only define one fallback method on getindex for all argument types.
# That dispatches to an (inlined) internal _getindex function, where the goal is
# to transform the indices such that we can call the only getindex method that
# we require AbstractArray subtypes must define, either:
# getindex(::T, ::Int) # if linearindexing(T) == LinearFast()
# getindex(::T, ::Int, ::Int, #=...ndims(A) indices...=#) if LinearSlow()
# Unfortunately, it is currently impossible to express the latter method for
# arbitrary dimensionalities. We could get around that with ::CartesianIndex{N},
# but that isn't as obvious and would require that the function be inlined to
# avoid allocations. If the subtype hasn't defined those methods, it goes back
# to the _getindex function where an error is thrown to prevent stack overflows.
#
# We use the same scheme for unsafe_getindex, with the exception that we can
# fallback to the safe version if the subtype hasn't defined the required
# unsafe method.

function getindex(A::AbstractArray, I...)
@_inline_meta
_getindex(linearindexing(A), A, I...)
end
function unsafe_getindex(A::AbstractArray, I...)
@_inline_meta
_unsafe_getindex(linearindexing(A), A, I...)
end
## Internal defitions
_getindex(::LinearFast, A::AbstractArray) = (@_inline_meta; getindex(A, 1))
_getindex(::LinearSlow, A::AbstractArray) = (@_inline_meta; _getindex(A, 1))
_unsafe_getindex(::LinearFast, A::AbstractArray) = (@_inline_meta; unsafe_getindex(A, 1))
_unsafe_getindex(::LinearSlow, A::AbstractArray) = (@_inline_meta; _unsafe_getindex(A, 1))
_getindex(::LinearIndexing, A::AbstractArray, I...) = error("indexing $(typeof(A)) with types $(typeof(I)) is not supported")
_unsafe_getindex(::LinearIndexing, A::AbstractArray, I...) = error("indexing $(typeof(A)) with types $(typeof(I)) is not supported")

## LinearFast Scalar indexing
_getindex(::LinearFast, A::AbstractArray, I::Int) = error("indexing not defined for ", typeof(A))
function _getindex(::LinearFast, A::AbstractArray, I::Real...)
@_inline_meta
# We must check bounds for sub2ind; so we can then call unsafe_getindex
checkbounds(A, I...)
unsafe_getindex(A, sub2ind(size(A), to_index(I)...))
end
_unsafe_getindex(::LinearFast, A::AbstractArray, I::Int) = (@_inline_meta; getindex(A, I))
function _unsafe_getindex(::LinearFast, A::AbstractArray, I::Real...)
@_inline_meta
unsafe_getindex(A, sub2ind(size(A), to_index(I)...))
end

# LinearSlow Scalar indexing
@generated function _getindex{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, I::Real...)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suspect we want to get rid of the typing on I and the calls to to_index. In my view, for an AbstractArray with LinearSlow, this policy should be handled by the array type itself. Interpolations.jl will thank you for it 😄.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I definitely agree, but I want to hear more (from you or @tlycken) about how Interpolations would like to use Dual or other non-Real indices. I think these definitions are correct, but that we'll want to change the nonscalar Union(Real, AbstractVector, Colon) to be more accepting.

Does it make sense to use linear indexing with non-real numbers? They would still need to implement arithmetic and div/rem for ind2sub and sub2ind, so perhaps these definitions could be widened to Number, but I hesitate to go further.

For nonscalar indexing, right now all indices must support Base.index_shape, index_lengths, and unsafe_getindex. The tricky one here is indexing - we could have the shape and length default to 1. We could also make to_index default to being a no-op for unknown types.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any time one's using LinearSlow and the number of indexes is fewer than the number of dimensions, I think it makes sense to call to_index on the final index before calling ind2sub. But when the number of indexes matches or exceeds the dimensionality, I think one should leave to_index out of it.

Yes, we could broaden that to Union(Number, AbstractVector, Colon). Will that/can that satisfy the constraints you mention in your last paragraph?

N = length(I)
if N == AN
:(error("indexing not defined for ", typeof(A)))
elseif N > AN
# Drop trailing ones
Isplat = Expr[:(to_index(I[$d])) for d = 1:AN]
Osplat = Expr[:(to_index(I[$d]) == 1) for d = AN+1:N]
quote
$(Expr(:meta, :inline))
(&)($(Osplat...)) || throw_boundserror(A, I)
getindex(A, $(Isplat...))
end
else
# Expand the last index into the appropriate number of indices
Isplat = Expr[:(to_index(I[$d])) for d = 1:N-1]
i = 0
for d=N:AN
push!(Isplat, :(s[$(i+=1)]))
end
sz = Expr(:tuple)
sz.args = Expr[:(size(A, $d)) for d=N:AN]
quote
$(Expr(:meta, :inline))
# ind2sub requires all dimensions to be nonzero, so checkbounds first
checkbounds(A, I...)
s = ind2sub($sz, to_index(I[$N]))
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This call to to_index is fine, however.

unsafe_getindex(A, $(Isplat...))
end
end
A[i1]
end
function getindex(A::AbstractMatrix, i1,i2,i3,i4...)
if i3*prod(i4) != 1
throw(BoundsError())
@generated function _unsafe_getindex{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, I::Real...)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it be worth putting these into a for (funcname, exprs) in ((:_getindex, :_unsafe_getindex), ...) @eval begin ... block? Only if you think it's better, of course.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I actually had that at one point, but decided it'd be better to start verbosely and then condense later if it ends up making sense. I'll play with it.

N = length(I)
if N == AN
Isplat = Expr[:(to_index(I[$d])) for d = 1:N]
:($(Expr(:meta, :inline)); getindex(A, $(Isplat...)))
elseif N > AN
# Drop trailing dimensions (unchecked)
Isplat = Expr[:(to_index(I[$d])) for d = 1:AN]
quote
$(Expr(:meta, :inline))
unsafe_getindex(A, $(Isplat...))
end
else
# Expand the last index into the appropriate number of indices
Isplat = Expr[:(to_index(I[$d])) for d = 1:N-1]
for d=N:AN
push!(Isplat, :(s[$(d-N+1)]))
end
sz = Expr(:tuple)
sz.args = Expr[:(size(A, $d)) for d=N:AN]
quote
$(Expr(:meta, :inline))
s = ind2sub($sz, to_index(I[$N]))
unsafe_getindex(A, $(Isplat...))
end
end
A[i1,i2]
end

function setindex!(A::AbstractVector, x, i1,i2,i3...)
if i2*prod(i3) != 1
throw(BoundsError())
## Setindex! is defined similarly. We first dispatch to an internal _setindex!
# function that allows dispatch on array storage
function setindex!(A::AbstractArray, v, I...)
@_inline_meta
_setindex!(linearindexing(A), A, v, I...)
end
function unsafe_setindex!(A::AbstractArray, v, I...)
@_inline_meta
_unsafe_setindex!(linearindexing(A), A, v, I...)
end
## Internal defitions
_setindex!(::LinearFast, A::AbstractArray, v) = (@_inline_meta; setindex!(A, v, 1))
_setindex!(::LinearSlow, A::AbstractArray, v) = (@_inline_meta; _setindex!(A, v, 1))
_unsafe_setindex!(::LinearFast, A::AbstractArray, v) = (@_inline_meta; unsafe_setindex!(A, v, 1))
_unsafe_setindex!(::LinearSlow, A::AbstractArray, v) = (@_inline_meta; _unsafe_setindex!(A, v, 1))
_setindex!(::LinearIndexing, A::AbstractArray, v, I...) = error("indexing $(typeof(A)) with types $(typeof(I)) is not supported")
_unsafe_setindex!(::LinearIndexing, A::AbstractArray, v, I...) = error("indexing $(typeof(A)) with types $(typeof(I)) is not supported")

## LinearFast Scalar indexing
_setindex!(::LinearFast, A::AbstractArray, v, I::Int) = error("indexed assignment not defined for ", typeof(A))
function _setindex!(::LinearFast, A::AbstractArray, v, I::Real...)
@_inline_meta
# We must check bounds for sub2ind; so we can then call unsafe_setindex!
checkbounds(A, I...)
unsafe_setindex!(A, v, sub2ind(size(A), to_index(I)...))
end
_unsafe_setindex!(::LinearFast, A::AbstractArray, v, I::Int) = (@_inline_meta; setindex!(A, v, I))
function _unsafe_setindex!(::LinearFast, A::AbstractArray, v, I::Real...)
@_inline_meta
unsafe_setindex!(A, v, sub2ind(size(A), to_index(I)...))
end

# LinearSlow Scalar indexing
@generated function _setindex!{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, v, I::Real...)
N = length(I)
if N == AN
:(error("indexed assignment not defined for ", typeof(A)))
elseif N > AN
# Drop trailing ones
Isplat = Expr[:(to_index(I[$d])) for d = 1:AN]
Osplat = Expr[:(to_index(I[$d]) == 1) for d = AN+1:N]
quote
$(Expr(:meta, :inline))
(&)($(Osplat...)) || throw_boundserror(A, I)
setindex!(A, v, $(Isplat...))
end
else
# Expand the last index into the appropriate number of indices
Isplat = Expr[:(to_index(I[$d])) for d = 1:N-1]
i = 0
for d=N:AN
push!(Isplat, :(s[$(i+=1)]))
end
sz = Expr(:tuple)
sz.args = Expr[:(size(A, $d)) for d=N:AN]
quote
$(Expr(:meta, :inline))
checkbounds(A, I...)
s = ind2sub($sz, to_index(I[$N]))
unsafe_setindex!(A, v, $(Isplat...))
end
end
A[i1] = x
end
function setindex!(A::AbstractMatrix, x, i1,i2,i3,i4...)
if i3*prod(i4) != 1
throw(BoundsError())
@generated function _unsafe_setindex!{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, v, I::Real...)
N = length(I)
if N == AN
Isplat = Expr[:(to_index(I[$d])) for d = 1:N]
:(setindex!(A, v, $(Isplat...)))
elseif N > AN
# Drop trailing dimensions (unchecked)
Isplat = Expr[:(to_index(I[$d])) for d = 1:AN]
quote
$(Expr(:meta, :inline))
unsafe_setindex!(A, v, $(Isplat...))
end
else
# Expand the last index into the appropriate number of indices
Isplat = Expr[:(to_index(I[$d])) for d = 1:N-1]
for d=N:AN
push!(Isplat, :(s[$(d-N+1)]))
end
sz = Expr(:tuple)
sz.args = Expr[:(size(A, $d)) for d=N:AN]
quote
$(Expr(:meta, :inline))
s = ind2sub($sz, to_index(I[$N]))
unsafe_setindex!(A, v, $(Isplat...))
end
end
A[i1,i2] = x
end

## get (getindex with a default value) ##
Expand Down
Loading