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

Make eachindex more efficient for linear arrays #10704

Merged
merged 2 commits into from
Apr 2, 2015
Merged

Make eachindex more efficient for linear arrays #10704

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
57f9e66
Make eachindex more efficient for linear arrays
mbauman Mar 31, 2015
adde889
Add a comment about the internal _inline_meta macro
mbauman Apr 2, 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
16 changes: 11 additions & 5 deletions base/abstractarray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -337,12 +337,18 @@ end

zero{T}(x::AbstractArray{T}) = fill!(similar(x), zero(T))

## iteration support for arrays as ranges ##
## iteration support for arrays ##
macro _inline_meta()
Expr(:meta, :inline)
Copy link
Member

Choose a reason for hiding this comment

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

A comment here about pushmeta! not yet being available in the bootstrap might be handy. Otherwise
there's a (small) risk we might see usages of Base.@_inline_meta popping up in package code from folks who learn by browsing through the source.

end
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])

# eachindex iterates over all indices. LinearSlow is later in bootstrap
eachindex(A::AbstractArray) = (@_inline_meta; eachindex(linearindexing(A), A))
eachindex(::LinearFast, A::AbstractArray) = 1:length(A)

start(A::AbstractArray) = _start(A,linearindexing(A))
_start(::AbstractArray,::LinearFast) = 1
next(a::AbstractArray,i) = (a[i],i+1)
done(a::AbstractArray,i) = (i > length(a))
isempty(a::AbstractArray) = (length(a) == 0)

## Conversions ##
Expand Down
5 changes: 5 additions & 0 deletions base/array.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -297,6 +297,11 @@ end

collect(itr) = collect(eltype(itr), itr)

## Iteration ##
start(A::Array) = 1
next(a::Array,i) = (a[i],i+1)
done(a::Array,i) = (i > length(a))

## Indexing: getindex ##

getindex(a::Array) = arrayref(a,1)
Expand Down
2 changes: 1 addition & 1 deletion base/dict.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -713,7 +713,7 @@ function skip_deleted(h::Dict, i)
end

start(t::Dict) = skip_deleted(t, 1)
done(t::Dict, i) = done(t.vals, i)
done(t::Dict, i) = i > length(t.vals)
Copy link
Member

Choose a reason for hiding this comment

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

I'm curious why this change is necessary...

Copy link
Member Author

Choose a reason for hiding this comment

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

This changes the iterator state to be a tuple of (IndexIterator, IteratorState) instead of just the next index. Dict was (incorrectly) passing a state to the Array iteration protocol that did not come from the Array's start or next.

Copy link
Member

Choose a reason for hiding this comment

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

Got it, thanks.

next(t::Dict, i) = ((t.keys[i],t.vals[i]), skip_deleted(t,i+1))

isempty(t::Dict) = (t.count == 0)
Expand Down
61 changes: 19 additions & 42 deletions base/multidimensional.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
### Multidimensional iterators
module IteratorsMD

import Base: eltype, length, start, _start, done, next, last, getindex, setindex!, linearindexing, min, max
import Base: eltype, length, start, done, next, last, getindex, setindex!, linearindexing, min, max, eachindex
import Base: simd_outer_range, simd_inner_length, simd_index
import Base: @nref, @ncall, @nif, @nexprs, LinearFast, LinearSlow, to_index

export CartesianIndex, CartesianRange, eachindex
export CartesianIndex, CartesianRange

# Traits for linear indexing
linearindexing(::BitArray) = LinearFast()
Expand Down Expand Up @@ -110,60 +110,37 @@ stagedfunction CartesianRange{N}(I::CartesianIndex{N})
end
CartesianRange{N}(sz::NTuple{N,Int}) = CartesianRange(CartesianIndex(sz))

stagedfunction eachindex{T,N}(A::AbstractArray{T,N})
stagedfunction eachindex{T,N}(::LinearSlow, A::AbstractArray{T,N})
startargs = fill(1, N)
stopargs = [:(size(A,$i)) for i=1:N]
:(CartesianRange(CartesianIndex{$N}($(startargs...)), CartesianIndex{$N}($(stopargs...))))
meta = Expr(:meta, :inline)
:($meta; CartesianRange(CartesianIndex{$N}($(startargs...)), CartesianIndex{$N}($(stopargs...))))
end

eltype{I}(::Type{CartesianRange{I}}) = I
eltype{I}(::CartesianRange{I}) = I

stagedfunction start{I}(iter::CartesianRange{I})
N=length(I)
finishedex = Expr(:(||), [:(iter.stop[$i] < iter.start[$i]) for i = 1:N]...)
:(return $finishedex, iter.start)
end

stagedfunction _start{T,N}(A::AbstractArray{T,N}, ::LinearSlow)
args = fill(1, N)
:(return isempty(A), CartesianIndex{$N}($(args...)))
end

# Prevent an ambiguity warning
next(R::StepRange, state::(Bool, CartesianIndex{1})) = (index=state[2]; return R[index], (index[1]==length(R), CartesianIndex{1}(index[1]+1)))
next{T}(R::UnitRange{T}, state::(Bool, CartesianIndex{1})) = (index=state[2]; return R[index], (index[1]==length(R), CartesianIndex{1}(index[1]+1)))
done(R::StepRange, state::(Bool, CartesianIndex{1})) = state[1]
done(R::UnitRange, state::(Bool, CartesianIndex{1})) = state[1]

stagedfunction next{T,N}(A::AbstractArray{T,N}, state::(Bool, CartesianIndex{N}))
I = state[2]
finishedex = (N==0 ? true : :(newindex[$N] > size(A, $N)))
start(iter::CartesianRange) = iter.start
stagedfunction next{I<:CartesianIndex}(iter::CartesianRange{I}, state)
N = length(I)
meta = Expr(:meta, :inline)
quote
$meta
index=state[2]
@inbounds v = A[index]
newindex=@nif $N d->(index[d] < size(A, d)) d->@ncall($N, $I, k->(k>d ? index[k] : k==d ? index[k]+1 : 1))
finished=$finishedex
v, (finished,newindex)
index=state
@nif $N d->(index[d] < iter.stop[d]) d->(@nexprs($N, k->(ind_k = ifelse(k>=d, index[k] + (k==d), iter.start[k]))))
newindex = @ncall $N $I ind
index, newindex
end
end
stagedfunction next{I<:CartesianIndex}(iter::CartesianRange{I}, state::(Bool, I))
stagedfunction done{I<:CartesianIndex}(iter::CartesianRange{I}, state)
N = length(I)
finishedex = (N==0 ? true : :(newindex[$N] > iter.stop[$N]))
meta = Expr(:meta, :inline)
quote
$meta
index=state[2]
newindex=@nif $N d->(index[d] < iter.stop[d]) d->@ncall($N, $I, k->(k>d ? index[k] : k==d ? index[k]+1 : iter.start[k]))
finished=$finishedex
index, (finished,newindex)
end
:(state[$N] > iter.stop[$N])
end

done{T,N}(A::AbstractArray{T,N}, state::(Bool, CartesianIndex{N})) = state[1]
done{I<:CartesianIndex}(iter::CartesianRange{I}, state::(Bool, I)) = state[1]
# 0-d cartesian ranges are special-cased to iterate once and only once
start{I<:CartesianIndex{0}}(iter::CartesianRange{I}) = false
next{I<:CartesianIndex{0}}(iter::CartesianRange{I}, state) = iter.start, true
done{I<:CartesianIndex{0}}(iter::CartesianRange{I}, state) = state

Copy link
Contributor

Choose a reason for hiding this comment

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

That's indeed much nicer than the old trick with finished.

stagedfunction length{I<:CartesianIndex}(iter::CartesianRange{I})
N = length(I)
Expand Down Expand Up @@ -502,7 +479,7 @@ function copy!{T,N}(dest::AbstractArray{T,N}, src::AbstractArray{T,N})
break
end
end
if samesize
if samesize && linearindexing(dest) == linearindexing(src)
for I in eachindex(dest)
@inbounds dest[I] = src[I]
end
Expand Down
4 changes: 2 additions & 2 deletions base/reducedim.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,7 +209,7 @@ function _mapreducedim!{T,N}(f, op, R::AbstractArray, A::AbstractArray{T,N})
end
else
sizeR = CartesianIndex{N}(size(R))
@inbounds @simd for IA in eachindex(A)
@inbounds @simd for IA in CartesianRange(size(A))
IR = min(IA, sizeR)
R[IR] = op(R[IR], f(A[IA]))
end
Expand Down Expand Up @@ -306,7 +306,7 @@ function findminmax!{T,N}(f, Rval, Rind, A::AbstractArray{T,N})
end
else
sizeR = CartesianIndex(size(Rval))
@inbounds for IA in eachindex(A)
@inbounds for IA in CartesianRange(size(A))
IR = min(sizeR, IA)
k += 1
tmpAv = A[IA]
Expand Down
30 changes: 16 additions & 14 deletions doc/stdlib/arrays.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,29 +33,31 @@ Basic functions

.. function:: eachindex(A)

Creates an iterable object for visiting each multi-dimensional index of the AbstractArray ``A``. Example for a 2-d array::
Creates an iterable object for visiting each index of an AbstractArray ``A`` in an efficient manner. For array types that have opted into fast linear indexing (like ``Array``), this is simply the range ``1:length(A)``. For other array types, this returns a specialized Cartesian range to efficiently index into the array with indices specified for every dimension. Example for a sparse 2-d array::

julia> A = rand(2,3)
2x3 Array{Float64,2}:
0.960084 0.629326 0.625155
0.432588 0.955903 0.991614
julia> A = sprand(2, 3, 0.5)
2x3 sparse matrix with 4 Float64 entries:
[1, 1] = 0.598888
[1, 2] = 0.0230247
[1, 3] = 0.486499
[2, 3] = 0.809041

julia> for iter in eachindex(A)
@show iter.I_1, iter.I_2
@show A[iter]
end
@show iter.I_1, iter.I_2
@show A[iter]
end
(iter.I_1,iter.I_2) = (1,1)
A[iter] = 0.9600836263003063
A[iter] = 0.5988881393454597
(iter.I_1,iter.I_2) = (2,1)
A[iter] = 0.4325878255452178
A[iter] = 0.0
(iter.I_1,iter.I_2) = (1,2)
A[iter] = 0.6293256402775211
A[iter] = 0.02302469881746183
(iter.I_1,iter.I_2) = (2,2)
A[iter] = 0.9559027084099654
A[iter] = 0.0
(iter.I_1,iter.I_2) = (1,3)
A[iter] = 0.6251548453735303
A[iter] = 0.4864987874354343
(iter.I_1,iter.I_2) = (2,3)
A[iter] = 0.9916142534546522
A[iter] = 0.8090413606455655

.. function:: countnz(A)

Expand Down
36 changes: 19 additions & 17 deletions test/arrayops.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -993,7 +993,7 @@ I2 = CartesianIndex((-1,5,2))

@test length(I1) == 3

a = zeros(2,3)
a = spzeros(2,3)
@test CartesianRange(size(a)) == eachindex(a)
a[CartesianIndex{2}(2,3)] = 5
@test a[2,3] == 5
Expand All @@ -1015,22 +1015,24 @@ indexes = collect(R)
@test length(R) == 12

r = 2:3
state = start(eachindex(r))
@test !done(r, state)
_, state = next(r, state)
@test !done(r, state)
val, state = next(r, state)
@test done(r, state)
@test val == 3
r = 2:3:8
state = start(eachindex(r))
@test !done(r, state)
_, state = next(r, state)
_, state = next(r, state)
@test !done(r, state)
val, state = next(r, state)
@test val == 8
@test done(r, state)
itr = eachindex(r)
state = start(itr)
@test !done(itr, state)
_, state = next(itr, state)
@test !done(itr, state)
val, state = next(itr, state)
@test done(itr, state)
@test r[val] == 3
r = sparse(collect(2:3:8))
itr = eachindex(r)
state = start(itr)
@test !done(itr, state)
_, state = next(itr, state)
_, state = next(itr, state)
@test !done(itr, state)
val, state = next(itr, state)
@test r[val] == 8
@test done(itr, state)


#rotates
Expand Down