Skip to content

RFC: Refactor indexing to use Varargs and ReshapedArrays #16251

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

Merged
merged 8 commits into from
May 12, 2016
Merged

RFC: Refactor indexing to use Varargs and ReshapedArrays #16251

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
e082468
Fix testing setindex! -- make sure we are using the proper dimensiona…
mbauman May 7, 2016
3440bf5
Refactor abstract scalar indexing to use Varargs
mbauman May 7, 2016
51b4bc1
Refactor abstract nonscalar indexing to use Varargs
mbauman May 7, 2016
7f2d845
Refactor SubArray indexing to use Varargs
mbauman May 7, 2016
a6954cd
Rewrite array interface doc in terms of Vararg
mbauman May 7, 2016
a071314
Use a lispy recursive function for reshape(A, Val{N})
mbauman May 8, 2016
bd38ee9
Enforce index/parent consistency in constructor
mbauman May 10, 2016
e5fec99
Use dispatch to avoid a branch in first(::Tuple)
mbauman May 10, 2016
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
125 changes: 62 additions & 63 deletions base/abstractarray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -414,14 +414,11 @@ pointer{T}(x::AbstractArray{T}, i::Integer) = (@_inline_meta; unsafe_convert(Ptr
# 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 require the type A{T,N} <: AbstractArray{T,N} to define; either:
# getindex(::A, ::Int) # if linearindexing(A) == LinearFast() OR
# getindex{T,N}(::A{T,N}, ::Vararg{Int, N}) # if LinearSlow()
# If the subtype hasn't defined the required method, it falls back to the
# _getindex function again where an error is thrown to prevent stack overflows.

function getindex(A::AbstractArray, I...)
@_propagate_inbounds_meta
Expand All @@ -432,62 +429,62 @@ function unsafe_getindex(A::AbstractArray, I...)
@inbounds r = getindex(A, I...)
r
end
## Internal defitions
# 0-dimensional indexing is defined to prevent ambiguities. LinearFast is easy:
_getindex(::LinearFast, A::AbstractArray) = (@_propagate_inbounds_meta; getindex(A, 1))
# But LinearSlow must take into account the dimensionality of the array:
_getindex{T}(::LinearSlow, A::AbstractArray{T,0}) = error("indexing not defined for ", typeof(A))
_getindex(::LinearSlow, A::AbstractVector) = (@_propagate_inbounds_meta; getindex(A, 1))
_getindex(l::LinearSlow, A::AbstractArray) = (@_propagate_inbounds_meta; _getindex(l, A, 1))

## Internal definitions
_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
## LinearFast Scalar indexing: canonical method is one Int
_getindex(::LinearFast, A::AbstractArray, ::Int) = error("indexing not defined for ", typeof(A))
_getindex(::LinearFast, A::AbstractArray, i::Real) = (@_propagate_inbounds_meta; getindex(A, to_index(i)))
function _getindex{T,N}(::LinearFast, A::AbstractArray{T,N}, I::Vararg{Real,N})
Copy link
Member

Choose a reason for hiding this comment

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

My impression is that this is still slow without @inline, because of the lack of call specialization (which will hopefully happen someday, see #16159).

Copy link
Member

Choose a reason for hiding this comment

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

Oh, nvm, next line 😦

# We must check bounds for sub2ind; so we can then use @inbounds
@_inline_meta
J = to_indexes(I...)
@boundscheck checkbounds(A, J...)
@inbounds r = getindex(A, sub2ind(size(A), J...))
r
end
function _getindex(::LinearFast, A::AbstractArray, I::Real...) # TODO: DEPRECATE FOR #14770
@_inline_meta
J = to_indexes(I...)
@boundscheck checkbounds(A, J...)
@inbounds r = getindex(A, sub2ind(size(A), J...))
r
end

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

## LinearSlow Scalar indexing: Canonical method is full dimensionality of Ints
_getindex{T,N}(::LinearSlow, A::AbstractArray{T,N}, ::Vararg{Int, N}) = error("indexing not defined for ", typeof(A))
_getindex{T,N}(::LinearSlow, A::AbstractArray{T,N}, I::Vararg{Real, N}) = (@_propagate_inbounds_meta; getindex(A, to_indexes(I...)...))
function _getindex(::LinearSlow, A::AbstractArray, i::Real)
# ind2sub requires all dimensions to be > 0; may as well just check bounds
@_inline_meta
@boundscheck checkbounds(A, i)
@inbounds r = getindex(A, ind2sub(size(A), to_index(i))...)
r
end
@generated function _getindex{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, I::Real...) # TODO: DEPRECATE FOR #14770
N = length(I)
if N == AN
if all(x->x===Int, I)
:(error("indexing not defined for ", typeof(A)))
else
:(@_propagate_inbounds_meta; getindex(A, to_indexes(I...)...))
end
elseif N > AN
if N > AN
# Drop trailing ones
Isplat = Expr[:(I[$d]) for d = 1:AN]
Osplat = Expr[:(to_index(I[$d]) == 1) for d = AN+1:N]
quote
# We only check the trailing ones, so just propagate @inbounds state
@_propagate_inbounds_meta
@boundscheck (&)($(Osplat...)) || throw_boundserror(A, I)
getindex(A, $(Isplat...))
end
else
# Expand the last index into the appropriate number of indices
Isplat = Expr[:(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]
szcheck = Expr[:(size(A, $d) > 0) for d=N:AN]
sz.args = Expr[:(size(A, $d)) for d=max(N,1):AN]
szcheck = Expr[:(size(A, $d) > 0) for d=max(N,1):AN]
last_idx = N > 0 ? :(to_index(I[$N])) : 1
quote
# ind2sub requires all dimensions to be > 0:
@_propagate_inbounds_meta
@boundscheck (&)($(szcheck...)) || throw_boundserror(A, I)
s = ind2sub($sz, to_index(I[$N]))
getindex(A, $(Isplat...))
getindex(A, $(Isplat...), ind2sub($sz, $last_idx)...)
end
end
end
Expand All @@ -504,34 +501,40 @@ function unsafe_setindex!(A::AbstractArray, v, I...)
r
end
## Internal defitions
_setindex!(::LinearFast, A::AbstractArray, v) = (@_propagate_inbounds_meta; setindex!(A, v, 1))
_setindex!{T}(::LinearSlow, A::AbstractArray{T,0}, v) = error("indexing not defined for ", typeof(A))
_setindex!(::LinearSlow, A::AbstractVector, v) = (@_propagate_inbounds_meta; setindex!(A, v, 1))
_setindex!(l::LinearSlow, A::AbstractArray, v) = (@_propagate_inbounds_meta; _setindex!(l, A, v, 1))

_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
_setindex!(::LinearFast, A::AbstractArray, v, ::Int) = error("indexed assignment not defined for ", typeof(A))
_setindex!(::LinearFast, A::AbstractArray, v, i::Real) = (@_propagate_inbounds_meta; setindex!(A, v, to_index(i)))
function _setindex!{T,N}(::LinearFast, A::AbstractArray{T,N}, v, I::Vararg{Real,N})
# We must check bounds for sub2ind; so we can then use @inbounds
@_inline_meta
J = to_indexes(I...)
@boundscheck checkbounds(A, J...)
@inbounds r = setindex!(A, v, sub2ind(size(A), J...))
r
end
function _setindex!(::LinearFast, A::AbstractArray, v, I::Real...) # TODO: DEPRECATE FOR #14770
@_inline_meta
J = to_indexes(I...)
@boundscheck checkbounds(A, J...)
@inbounds r = setindex!(A, v, sub2ind(size(A), J...))
r
end

# LinearSlow Scalar indexing
@generated function _setindex!{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, v, I::Real...)
_setindex!{T,N}(::LinearSlow, A::AbstractArray{T,N}, v, ::Vararg{Int, N}) = error("indexed assignment not defined for ", typeof(A))
_setindex!{T,N}(::LinearSlow, A::AbstractArray{T,N}, v, I::Vararg{Real, N}) = (@_propagate_inbounds_meta; setindex!(A, v, to_indexes(I...)...))
function _setindex!(::LinearSlow, A::AbstractArray, v, i::Real)
# ind2sub requires all dimensions to be > 0; may as well just check bounds
@_inline_meta
@boundscheck checkbounds(A, i)
@inbounds r = setindex!(A, v, ind2sub(size(A), to_index(i))...)
r
end
@generated function _setindex!{T,AN}(::LinearSlow, A::AbstractArray{T,AN}, v, I::Real...) # TODO: DEPRECATE FOR #14770
N = length(I)
if N == AN
if all(x->x===Int, I)
:(error("indexing not defined for ", typeof(A)))
else
:(@_propagate_inbounds_meta; setindex!(A, v, to_indexes(I...)...))
end
elseif N > AN
if N > AN
# Drop trailing ones
Isplat = Expr[:(I[$d]) for d = 1:AN]
Osplat = Expr[:(to_index(I[$d]) == 1) for d = AN+1:N]
Expand All @@ -544,19 +547,15 @@ end
else
# Expand the last index into the appropriate number of indices
Isplat = Expr[:(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]
szcheck = Expr[:(size(A, $d) > 0) for d=N:AN]
sz.args = Expr[:(size(A, $d)) for d=max(N,1):AN]
szcheck = Expr[:(size(A, $d) > 0) for d=max(N,1):AN]
last_idx = N > 0 ? :(to_index(I[$N])) : 1
quote
@_propagate_inbounds_meta
# ind2sub requires all dimensions to be > 0:
@_propagate_inbounds_meta
@boundscheck (&)($(szcheck...)) || throw_boundserror(A, I)
s = ind2sub($sz, to_index(I[$N]))
setindex!(A, v, $(Isplat...))
setindex!(A, v, $(Isplat...), ind2sub($sz, $last_idx)...)
end
end
end
Expand Down
4 changes: 2 additions & 2 deletions base/array.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -308,7 +308,7 @@ done(a::Array,i) = i == length(a)+1

# This is more complicated than it needs to be in order to get Win64 through bootstrap
getindex(A::Array, i1::Real) = arrayref(A, to_index(i1))
getindex(A::Array, i1::Real, i2::Real, I::Real...) = arrayref(A, to_index(i1), to_index(i2), to_indexes(I...)...)
getindex(A::Array, i1::Real, i2::Real, I::Real...) = arrayref(A, to_index(i1), to_index(i2), to_indexes(I...)...) # TODO: REMOVE FOR #14770

# Faster contiguous indexing using copy! for UnitRange and Colon
function getindex(A::Array, I::UnitRange{Int})
Expand Down Expand Up @@ -337,7 +337,7 @@ end

## Indexing: setindex! ##
setindex!{T}(A::Array{T}, x, i1::Real) = arrayset(A, convert(T,x)::T, to_index(i1))
setindex!{T}(A::Array{T}, x, i1::Real, i2::Real, I::Real...) = arrayset(A, convert(T,x)::T, to_index(i1), to_index(i2), to_indexes(I...)...)
setindex!{T}(A::Array{T}, x, i1::Real, i2::Real, I::Real...) = arrayset(A, convert(T,x)::T, to_index(i1), to_index(i2), to_indexes(I...)...) # TODO: REMOVE FOR #14770

# These are redundant with the abstract fallbacks but needed for bootstrap
function setindex!(A::Array, x, I::AbstractVector{Int})
Expand Down
123 changes: 29 additions & 94 deletions base/multidimensional.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,10 +251,25 @@ index_shape_dim(A, dim, ::Colon) = (trailingsize(A, dim),)
# ambiguities for AbstractArray subtypes. See the note in abstractarray.jl

# Note that it's most efficient to call checkbounds first, and then to_index
@inline function _getindex(l::LinearIndexing, A::AbstractArray, I::Union{Real, AbstractArray, Colon}...)
@inline function _getindex{T,N}(l::LinearIndexing, A::AbstractArray{T,N}, I::Vararg{Union{Real, AbstractArray, Colon},N})
@boundscheck checkbounds(A, I...)
_unsafe_getindex(l, A, I...)
end
# Explicitly allow linear indexing with one non-scalar index
@inline function _getindex(l::LinearIndexing, A::AbstractArray, i::Union{Real, AbstractArray, Colon})
@boundscheck checkbounds(A, i)
_unsafe_getindex(l, _maybe_linearize(l, A), i)
end
# But we can speed up LinearSlow arrays by reshaping them to vectors:
_maybe_linearize(::LinearFast, A::AbstractArray) = A
_maybe_linearize(::LinearSlow, A::AbstractVector) = A
_maybe_linearize(::LinearSlow, A::AbstractArray) = reshape(A, length(A))
Copy link
Member

Choose a reason for hiding this comment

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

👍


@inline function _getindex{N}(l::LinearIndexing, A::AbstractArray, I::Vararg{Union{Real, AbstractArray, Colon},N}) # TODO: DEPRECATE FOR #14770
@boundscheck checkbounds(A, I...)
_unsafe_getindex(l, reshape(A, Val{N}), I...)
end

@generated function _unsafe_getindex(::LinearIndexing, A::AbstractArray, I::Union{Real, AbstractArray, Colon}...)
N = length(I)
quote
Expand All @@ -268,8 +283,6 @@ end
end

# logical indexing optimization - don't use find (within to_index)
# This is inherently a linear operation in the source, but we could potentially
# use fast dividing integers to speed it up.
function _unsafe_getindex(::LinearIndexing, src::AbstractArray, I::AbstractArray{Bool})
shape = index_shape(src, I)
dest = similar(src, shape)
Expand All @@ -294,7 +307,7 @@ end
$(Expr(:meta, :inline))
D = eachindex(dest)
Ds = start(D)
idxlens = index_lengths(src, I...) # TODO: unsplat?
idxlens = index_lengths(src, I...)
@nloops $N i d->(1:idxlens[d]) d->(@inbounds j_d = getindex(I[d], i_d)) begin
d, Ds = next(D, Ds)
@inbounds dest[d] = @ncall $N getindex src j
Expand All @@ -311,10 +324,21 @@ end
# before redispatching to the _unsafe_batchsetindex!
_iterable(v::AbstractArray) = v
_iterable(v) = repeated(v)
@inline function _setindex!(l::LinearIndexing, A::AbstractArray, x, J::Union{Real,AbstractArray,Colon}...)
@inline function _setindex!{T,N}(l::LinearIndexing, A::AbstractArray{T,N}, x, J::Vararg{Union{Real,AbstractArray,Colon},N})
@boundscheck checkbounds(A, J...)
_unsafe_setindex!(l, A, x, J...)
end
@inline function _setindex!(l::LinearIndexing, A::AbstractArray, x, j::Union{Real,AbstractArray,Colon})
@boundscheck checkbounds(A, j)
_unsafe_setindex!(l, _maybe_linearize(l, A), x, j)
A
end
@inline function _setindex!{N}(l::LinearIndexing, A::AbstractArray, x, J::Vararg{Union{Real, AbstractArray, Colon},N}) # TODO: DEPRECATE FOR #14770
@boundscheck checkbounds(A, J...)
_unsafe_setindex!(l, reshape(A, Val{N}), x, J...)
A
end

@inline function _unsafe_setindex!(::LinearIndexing, A::AbstractArray, x, J::Union{Real,AbstractArray,Colon}...)
_unsafe_batchsetindex!(A, _iterable(x), to_indexes(J...)...)
end
Expand Down Expand Up @@ -434,95 +458,6 @@ for (f, fmod, op) = ((:cummin, :_cummin!, :min), (:cummax, :_cummax!, :max))
@eval ($f)(A::AbstractArray) = ($f)(A, 1)
end

## SubArray index merging
# A view created like V = A[2:3:8, 5:2:17] can later be indexed as V[2:7],
# creating a new 1d view.
# In such cases we have to collapse the 2d space spanned by the ranges.
#
# API:
# merge_indexes(V, indexes::NTuple, index)
# indexes encodes the view's trailing indexes into the parent array,
# and index encodes the subset of these elements that we'll select.
#
# It returns a CartesianIndex or array of CartesianIndexes.

# Checking 'in' a range is fast -- so check all possibilities and keep the good ones
@generated function merge_indexes{N}(V, indexes::NTuple{N}, index::Union{Colon, Range})
# There may be a vector of cartesian indices in the passed indexes... which
# makes the number of indices more than N. Since we pre-allocate the array
# of CartesianIndexes, we need to figure out how big to make it
M = 0
for T in indexes.parameters
T <: CartesianIndex ? (M += length(T)) : (M += 1)
end
index_length_expr = index <: Colon ? Symbol("Istride_", N+1) : :(length(index))
quote
Cartesian.@nexprs $N d->(I_d = indexes[d])
dimlengths = Cartesian.@ncall $N index_lengths_dim V.parent length(V.indexes)-N+1 I
Istride_1 = 1 # strides of the indexes to merge
Cartesian.@nexprs $N d->(Istride_{d+1} = Istride_d*dimlengths[d])
idx_len = $(index_length_expr)
if idx_len < 0.1*$(Symbol("Istride_", N+1)) # this has not been carefully tuned
return merge_indexes_div(V, indexes, index, dimlengths)
end
Cartesian.@nexprs $N d->(counter_d = 1) # counter_0 is the linear index
k = 0
merged = Array(CartesianIndex{$M}, idx_len)
Cartesian.@nloops $N i d->(1:dimlengths[d]) d->(counter_{d-1} = counter_d + (i_d-1)*Istride_d; @inbounds idx_d = I_d[i_d]) begin
if counter_0 in index # this branch is elided for ::Colon
@inbounds merged[k+=1] = Cartesian.@ncall $N CartesianIndex{$M} idx
end
end
merged
end
end

# mapping getindex across the parent and subindices rapidly gets too big to
# automatically inline, but it is crucial that it does so to avoid allocations
# Unlike SubArray's reindex, merge_indexes doesn't drop any indices.
@inline inlinemap(f, t::Tuple, s::Tuple) = (f(t[1], s[1]), inlinemap(f, tail(t), tail(s))...)
inlinemap(f, t::Tuple{}, s::Tuple{}) = ()
inlinemap(f, t::Tuple{}, s::Tuple) = ()
inlinemap(f, t::Tuple, s::Tuple{}) = ()

# Otherwise, we fall back to the slow div/rem method, using ind2sub.
@inline merge_indexes{N}(V, indexes::NTuple{N}, index) =
merge_indexes_div(V, indexes, index, index_lengths_dim(V.parent, length(V.indexes)-N+1, indexes...))

@inline merge_indexes_div{N}(V, indexes::NTuple{N}, index::Real, dimlengths) =
CartesianIndex(inlinemap(getindex, indexes, ind2sub(dimlengths, index)))
merge_indexes_div{N}(V, indexes::NTuple{N}, index::AbstractArray, dimlengths) =
reshape([CartesianIndex(inlinemap(getindex, indexes, ind2sub(dimlengths, i))) for i in index], size(index))
merge_indexes_div{N}(V, indexes::NTuple{N}, index::Colon, dimlengths) =
[CartesianIndex(inlinemap(getindex, indexes, ind2sub(dimlengths, i))) for i in 1:prod(dimlengths)]

# Merging indices is particularly difficult in the case where we partially linearly
# index through a multidimensional array. It's easiest if we can simply reduce the
# partial indices to a single linear index into the parent index array.
function merge_indexes{N}(V, indexes::NTuple{N}, index::Tuple{Colon, Vararg{Colon}})
shape = index_shape(indexes[1], index...)
reshape(merge_indexes(V, indexes, :), (shape[1:end-1]..., shape[end]*prod(index_lengths_dim(V.parent, length(V.indexes)-length(indexes)+2, tail(indexes)...))))
end
@inline merge_indexes{N}(V, indexes::NTuple{N}, index::Tuple{Real, Vararg{Real}}) = merge_indexes(V, indexes, sub2ind(size(indexes[1]), index...))
# In general, it's a little trickier, but we can use the product iterator
# if we replace colons with ranges. This can be optimized further.
function merge_indexes{N}(V, indexes::NTuple{N}, index::Tuple)
I = replace_colons(V, indexes, index)
shp = index_shape(indexes[1], I...) # index_shape does no bounds checking
dimlengths = index_lengths_dim(V.parent, length(V.indexes)-N+1, indexes...)
sz = size(indexes[1])
reshape([CartesianIndex(inlinemap(getindex, indexes, ind2sub(dimlengths, sub2ind(sz, i...)))) for i in product(I...)], shp)
end
@inline replace_colons(V, indexes, I) = replace_colons_dim(V, indexes, 1, I)
@inline replace_colons_dim(V, indexes, dim, I::Tuple{}) = ()
@inline replace_colons_dim(V, indexes, dim, I::Tuple{Colon}) =
(1:trailingsize(indexes[1], dim)*prod(index_lengths_dim(V.parent, length(V.indexes)-length(indexes)+2, tail(indexes)...)),)
@inline replace_colons_dim(V, indexes, dim, I::Tuple{Colon, Vararg{Any}}) =
(1:size(indexes[1], dim), replace_colons_dim(V, indexes, dim+1, tail(I))...)
@inline replace_colons_dim(V, indexes, dim, I::Tuple{Any, Vararg{Any}}) =
(I[1], replace_colons_dim(V, indexes, dim+1, tail(I))...)


cumsum(A::AbstractArray, axis::Integer=1) = cumsum!(similar(A, Base._cumsum_type(A)), A, axis)
cumsum!(B, A::AbstractArray) = cumsum!(B, A, 1)
cumprod(A::AbstractArray, axis::Integer=1) = cumprod!(similar(A), A, axis)
Expand Down
Loading