Implement optimized methods for ReshapedArrays

src/ArrayIterationPlayground.jl

@@ -1,13 +1,14 @@
 module ArrayIterationPlayground
 
-import Base: getindex, setindex!, start, next, done, length, eachindex, show, parent
-using Base: ReshapedArray, linearindexing, LinearFast, LinearSlow
+import Base: getindex, setindex!, start, next, done, length, eachindex, show, parent, isless
+using Base: ReshapedArray, ReshapedIndex, linearindexing, LinearFast, LinearSlow, LinearIndexing
 using Base.PermutedDimsArrays: PermutedDimsArray
 
 export inds, index, value, stored, each, sync
 
 include("types.jl")
 include("core.jl")
+include("reshaped.jl")
 include("sparse.jl")
 
 end # module

src/core.jl

@@ -9,6 +9,10 @@ function show(io::IO, W::ArrayIndexingWrapper)
     print(io, "iteration hint over ", hint_string(W), " of a ", summary(W.data), " over the region ", W.indexes)
 end
 
+function show(io::IO, iter::ContigCartIterator)
+    print(io, "Cartesian iterator with:\n  domain ", iter.arrayrange, "\n  start  ", iter.columnrange.start, "\n  stop   ", iter.columnrange.stop)
+end
+
 parent(W::ArrayIndexingWrapper) = W.data
 
 """
@@ -83,7 +87,10 @@ each{T,N}(A::AbstractArray{T,N}, indexes::NTuple{N,IterIndex}) = each(ArrayIndex
 # Fallback definitions for each
 each{A,I,isstored}(W::ArrayIndexingWrapper{A,I,false,isstored}) = (itr = each(index(W)); ValueIterator{A,typeof(itr)}(W.data, itr))
 each{A,N,isstored}(W::ArrayIndexingWrapper{A,NTuple{N,Colon},true,isstored}) = eachindex(W.data)
-each{A,I,isstored}(W::ArrayIndexingWrapper{A,I,true,isstored}) = CartesianRange(ranges(W))
+each{A,I,isstored}(W::ArrayIndexingWrapper{A,I,true,isstored}) = _each(contiguous_index(W.indexes), W)
+
+_each(::Contiguous, W) = contiguous_iterator(W)
+_each(::Any, W) = CartesianRange(ranges(W))
 
 start(vi::ValueIterator) = start(vi.iter)
 done(vi::ValueIterator, s) = done(vi.iter, s)
@@ -182,3 +189,28 @@ checksame_storageorder(::Type{Val{false}}, A, B...) = Val{false}
 _sso(::FirstToLast, ::FirstToLast) = Val{true}
 _sso{p}(::OtherOrder{p}, ::OtherOrder{p}) = Val{true}
 _sso(::StorageOrder, ::StorageOrder) = Val{false}
+
+# indexes is contiguous if it's one of:
+#    Colon...
+#    Colon..., Union{UnitRange,Int}, Int...
+contiguous_index(I) = contiguous_index(Contiguous(), I...)
+@inline contiguous_index(c::Contiguous, ::Colon, I...) = contiguous_index(c, I...)
+@inline contiguous_index(::Contiguous, ::Any, I...) = contiguous_index(MaybeContiguous(), I...)
+@inline contiguous_index(c::MaybeContiguous, ::Int, I...) = contiguous_index(c, I...)
+@inline contiguous_index(::MaybeContiguous, ::Any, I...) = NonContiguous()
+@inline contiguous_index(::Contiguity) = Contiguous()  # won't get here for NonContiguous
+
+contiguous_iterator(W) = _contiguous_iterator(W, linearindexing(parent(W)))
+function _contiguous_iterator(W, ::LinearFast)
+    f, l = firstlast(W)
+    f:l
+end
+_contiguous_iterator(W, ::LinearSlow) = CartesianRange(ranges(W))
+
+function firstlast(W)
+    A = parent(W)
+    psz = size(A)
+    f = sub2ind(psz, map((i,j)->first(i[j]), inds(A), W.indexes)...)
+    l = sub2ind(psz, map((i,j)->last(i[j]),  inds(A), W.indexes)...)
+    f, l
+end

src/reshaped.jl

@@ -0,0 +1,37 @@
+# optimized methods for ReshapedArrays
+
+# column iteration
+start(iter::ContigCartIterator) = iter.columnrange.start
+function next(iter::ContigCartIterator, state)
+    item, newstate = next(iter.arrayrange, state)
+    ReshapedIndex(item), newstate
+end
+done(iter::ContigCartIterator, state) = isless(iter.columnrange.stop, state)
+
+function _contiguous_iterator{AA<:ReshapedArray}(W::ArrayIndexingWrapper{AA}, ::LinearSlow)
+    fi, li = firstlast(W)
+    A = parent(W)
+    f = Base.ind2sub_rs(A.mi, fi)
+    l = Base.ind2sub_rs(A.mi, li)
+    ContigCartIterator(CartesianRange(inds(parent(A))),
+                       CartesianRange(CartesianIndex(f), CartesianIndex(l)))
+end
+
+# Branching implementation
+# @inline isless{N}(I1::CartesianIndex{N}, I2::CartesianIndex{N}) = _isless(I1.I, I2.I)
+# @inline function _isless{N}(I1::NTuple{N,Int}, I2::NTuple{N,Int})
+#     i1, i2 = I1[N], I2[N]
+#     isless(i1, i2) && return true
+#     isless(i2, i1) && return false
+#     _isless(Base.front(I1), Base.front(I2))
+# end
+# _isless(::Tuple{}, ::Tuple{}) = false
+
+# Select implementation
+# @inline isless{N}(I1::CartesianIndex{N}, I2::CartesianIndex{N}) = _isless(0, I1.I, I2.I)
+# @inline function _isless{N}(ret, I1::NTuple{N,Int}, I2::NTuple{N,Int})
+#     newret = ifelse(ret==0, icmp(I1[N], I2[N]), ret)
+#     _isless(newret, Base.front(I1), Base.front(I2))
+# end
+# _isless(ret, ::Tuple{}, ::Tuple{}) = ifelse(ret==1, true, false)
+# icmp(a, b) = ifelse(isless(a,b), 1, ifelse(a==b, 0, -1))

src/types.jl

@@ -37,3 +37,15 @@ immutable FirstToLastIterator{N,AA}
     parent::AA
     itr::CartesianRange{N}
 end
+
+# Contiguous ranges
+abstract Contiguity
+immutable Contiguous <: Contiguity end
+immutable NonContiguous <: Contiguity end
+immutable MaybeContiguous <: Contiguity end  # intermediate type used in assessing contiguity
+
+# Contiguous cartesian ranges. Sometimes needed for LinearSlow arrays.
+immutable ContigCartIterator{N}
+    arrayrange::CartesianRange{N}
+    columnrange::CartesianRange{N}
+end

test/dense.jl

@@ -4,8 +4,10 @@ B = sub(A, 1:2, 1:3)
 
 @test each(index(A)) == eachindex(A) == 1:length(A)
 @test each(index(B)) == CartesianRange((1:2, 1:3))
-@test each(index(A, :, 1:2)) == CartesianRange((1:2, 1:2))
-@test each(index(A, :, 2:3)) == CartesianRange((1:2, 2:3))
+@test each(index(A, :, 1:2)) == 1:4
+@test each(index(B, :, 1:2)) == CartesianRange((1:2, 1:2))
+@test each(index(A, :, 2:3)) == 3:6
+@test each(index(B, :, 2:3)) == CartesianRange((1:2, 2:3))
 @test each(index(A, 1, :)) == CartesianRange((1, 1:3))
 
 k = 0
@@ -187,3 +189,56 @@ for (IB, IA) in sync(index(B, :, 1), index(A, :, 2))
     B[IB] = A[IA]
 end
 @test B == [A[1,2] -1; A[2,2] -1]
+
+## optimized ReshapedArray iterators
+A = reshape(1:12, 4, 3)    # LinearFast
+@test each(index(A, :, 2)) == 5:8
+@test each(index(A, 2:3, 3)) == 10:11
+A = sub(copy(reshape(1:15, 5, 3)), 1:4, :)  # LinearSlow
+a = reshape(A, 12)
+@test each(index(a, 1:4)) == CCI(CartesianRange(size(A)),
+                                 CartesianRange(CartesianIndex(1,1),CartesianIndex(4,1)))
+k = 0
+for I in each(index(a, 1:4))
+    @test a[I] == a[k+=1]
+end
+@test each(index(a, 3:9)) == CCI(CartesianRange(size(A)),
+                                 CartesianRange(CartesianIndex(3,1),CartesianIndex(1,3)))
+k = 2
+for I in each(index(a, 3:9))
+    @test a[I] == a[k+=1]
+end
+
+function sum_cols_slow!(S, A)  # slow for ReshapedArrays
+    fill!(S, 0)
+    @assert inds(S,2) == inds(A,2)
+    for j in inds(A,2)
+	tmp = S[1,j]
+	@inbounds for i in inds(A, 1)
+	    tmp += A[i,j]
+	end
+	S[1,j] = tmp
+    end
+    S
+end
+
+function sum_cols_fast!(S, A)
+    fill!(S, 0)
+    @assert inds(S,2) == inds(A,2)
+    for j in inds(A,2)
+	tmp = S[1,j]
+	@inbounds for I in each(index(A, :, j))
+	    tmp += A[I]
+	end
+	S[1,j] = tmp
+    end
+    S
+end
+
+A = rand(1000,1,999)
+B = sub(A, 1:size(A,1)-1, 1, 1:size(A,3)-1)
+R = reshape(B, (size(B,1),size(B,3)))
+@test isa(each(index(R, :, 2)), CCI)
+S1 = zeros(1,size(R,2))
+S2 = similar(S1)
+@test sum_cols_fast!(S1, R) == sum_cols_slow!(S2, R)

test/internal.jl

@@ -1,13 +1,27 @@
 A = rand(3,4)
 AIP.checksame_inds(A, A)
 @test_throws DimensionMismatch AIP.checksame_inds(A, A')
-@test isa(AIP.storageorder(A), AIP.FirstToLast)
+@test isa(@inferred(AIP.storageorder(A)), AIP.FirstToLast)
 R = reshape(A, (2,3,2))
-@test isa(AIP.storageorder(R), AIP.FirstToLast)
+@test isa(@inferred(AIP.storageorder(R)), AIP.FirstToLast)
 S = sub(A, 1:3, 1:4)
 R = reshape(S, (2,3,2))
-@test isa(AIP.storageorder(R), AIP.FirstToLast)
+@test isa(@inferred(AIP.storageorder(R)), AIP.FirstToLast)
 B = PermutedDimsArray(A, [2,1])
-@test isa(AIP.storageorder(B), AIP.OtherOrder{(2,1)})
+@test isa(@inferred(AIP.storageorder(B)), AIP.OtherOrder{(2,1)})
 R = reshape(B, (2,2,3))
-@test isa(AIP.storageorder(R), AIP.NoOrder)
+@test isa(@inferred(AIP.storageorder(R)), AIP.NoOrder)
+
+@test @inferred(AIP.contiguous_index((:, :))) == AIP.Contiguous()
+@test @inferred(AIP.contiguous_index((:, 3))) == AIP.Contiguous()
+@test @inferred(AIP.contiguous_index((3, :))) == AIP.NonContiguous()
+@test @inferred(AIP.contiguous_index((3, 3))) == AIP.Contiguous()
+@test @inferred(AIP.contiguous_index((:, :, 3))  ) == AIP.Contiguous()
+@test @inferred(AIP.contiguous_index((:, :, 1:2))) == AIP.Contiguous()
+@test @inferred(AIP.contiguous_index((:, 3, :))  ) == AIP.NonContiguous()
+@test @inferred(AIP.contiguous_index((:, 3, 1:2))) == AIP.NonContiguous()
+
+@test isless(CartesianIndex((1,1)), CartesianIndex((2,1)))
+@test isless(CartesianIndex((1,1)), CartesianIndex((1,2)))
+@test isless(CartesianIndex((2,1)), CartesianIndex((1,2)))
+@test !isless(CartesianIndex((1,2)), CartesianIndex((2,1)))

test/runtests.jl

@@ -3,6 +3,7 @@ using Base.Test
 using Base.PermutedDimsArrays: PermutedDimsArray
 
 const AIP = ArrayIterationPlayground
+const CCI = AIP.ContigCartIterator
 
 include("array_types.jl")  # just for testing