Skip to content

Commit abb4942

Browse files
jishnubN5N3
andauthored
Specialize iterate for IndexLinear arrays (#58635)
This adds a specialized `iterate` method for linearly indexed arrays. Firstly, we explicitly evaluate the iteration over the range that corresponds to the linear indices. Secondly, we return `nothing` for an out-of-bounds state, instead of throwing a `BoundsError`. This also lets us annotate the actual indexing with `@inbounds`. On master ```julia julia> A = rand(1000,1000); v = view(A, :); v2 = view(A, 1:2:lastindex(A)); julia> using LinearAlgebra julia> @Btime norm(Iterators.map(splat(-), zip($A, $A))); 434.982 μs (0 allocations: 0 bytes) julia> @Btime norm(Iterators.map(splat(-), zip($v, $v))); 2.080 ms (0 allocations: 0 bytes) # v"1.13.0-DEV.709" 510.835 μs (0 allocations: 0 bytes) # this PR ``` The performance of iterating over the contiguous `view` becomes comparable to that of the `Array`. Fixes #43295 without the need to disable bounds-checking. --------- Co-authored-by: N5N3 <2642243996@qq.com>
1 parent 66d3cf5 commit abb4942

File tree

3 files changed

+20
-2
lines changed

3 files changed

+20
-2
lines changed

base/abstractarray.jl

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1234,11 +1234,19 @@ oneunit(x::AbstractMatrix{T}) where {T} = _one(oneunit(T), x)
12341234
# While the definitions for IndexLinear are all simple enough to inline on their
12351235
# own, IndexCartesian's CartesianIndices is more complicated and requires explicit
12361236
# inlining.
1237-
function iterate(A::AbstractArray, state=(eachindex(A),))
1237+
iterate_starting_state(A) = iterate_starting_state(A, IndexStyle(A))
1238+
iterate_starting_state(A, ::IndexLinear) = firstindex(A)
1239+
iterate_starting_state(A, ::IndexStyle) = (eachindex(A),)
1240+
iterate(A::AbstractArray, state = iterate_starting_state(A)) = _iterate(A, state)
1241+
function _iterate(A::AbstractArray, state::Tuple)
12381242
y = iterate(state...)
12391243
y === nothing && return nothing
12401244
A[y[1]], (state[1], tail(y)...)
12411245
end
1246+
function _iterate(A::AbstractArray, state::Integer)
1247+
checkbounds(Bool, A, state) || return nothing
1248+
@inbounds(A[state]), state + one(state)
1249+
end
12421250

12431251
isempty(a::AbstractArray) = (length(a) == 0)
12441252

test/abstractarray.jl

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2286,3 +2286,13 @@ end
22862286
@test_throws "no method matching $Int(::$Infinity)" similar(ones(2), OneToInf())
22872287
end
22882288
end
2289+
2290+
@testset "iterate for linear indexing" begin
2291+
A = [1 2; 3 4]
2292+
v = view(A, :)
2293+
@test sum(x for x in v) == sum(A)
2294+
v = view(A, 1:2:lastindex(A))
2295+
@test sum(x for x in v) == sum(A[1:2:end])
2296+
v2 = view(A, Base.IdentityUnitRange(1:length(A)))
2297+
@test sum(x for x in v2) == sum(A)
2298+
end

test/stacktraces.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,7 @@ struct F49231{a,b,c,d,e,f,g} end
248248
stacktrace(catch_backtrace())
249249
end
250250
str = sprint(Base.show_backtrace, st, context = (:limit=>true, :stacktrace_types_limited => Ref(false), :color=>true, :displaysize=>(50,105)))
251-
@test contains(str, "[5] \e[0m\e[1mcollect_to!\e[22m\e[0m\e[1m(\e[22m\e[90mdest\e[39m::\e[0mVector\e[90m{…}\e[39m, \e[90mitr\e[39m::\e[0mBase.Generator\e[90m{…}\e[39m, \e[90moffs\e[39m::\e[0m$Int, \e[90mst\e[39m::\e[0mTuple\e[90m{…}\e[39m\e[0m\e[1m)\e[22m\n\e[90m")
251+
@test contains(str, "[5] \e[0m\e[1mcollect_to!\e[22m\e[0m\e[1m(\e[22m\e[90mdest\e[39m::\e[0mVector\e[90m{…}\e[39m, \e[90mitr\e[39m::\e[0mBase.Generator\e[90m{…}\e[39m, \e[90moffs\e[39m::\e[0m$Int, \e[90mst\e[39m::\e[0m$Int\e[0m\e[1m)\e[22m\n\e[90m")
252252

253253
st = try
254254
F49231{Vector,Val{'}'},Vector{Vector{Vector{Vector}}},Tuple{Int,Int,Int,Int,Int,Int,Int},Int,Int,Int}()(1,2,3)

0 commit comments

Comments
 (0)