From 84544778f041aee8cfaddeb5c77a38f1881718e4 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Sat, 17 Feb 2024 22:31:50 +0530 Subject: [PATCH 1/2] Orthogonalize re-indexing for FastSubArrays --- base/subarray.jl | 86 ++++++++++++++++-------------------------------- 1 file changed, 28 insertions(+), 58 deletions(-) diff --git a/base/subarray.jl b/base/subarray.jl index eca06fa3eacff..e131bdd7d5b8e 100644 --- a/base/subarray.jl +++ b/base/subarray.jl @@ -320,22 +320,35 @@ end # But SubArrays with fast linear indexing pre-compute a stride and offset FastSubArray{T,N,P,I} = SubArray{T,N,P,I,true} +# We define a convenience functions to compute the shifted parent index +# This differs from reindex as this accepts the view directly, instead of its indices +@inline _reindexlinear(V::FastSubArray, i::Integer) = V.offset1 + V.stride1*i +@inline _reindexlinear(V::FastSubArray, i::AbstractUnitRange{Int}) = V.offset1 .+ V.stride1 .* i + function getindex(V::FastSubArray, i::Int) @inline @boundscheck checkbounds(V, i) - @inbounds r = V.parent[V.offset1 + V.stride1*i] + @inbounds r = V.parent[_reindexlinear(V, i)] r end -# We can avoid a multiplication if the first parent index is a Colon or AbstractUnitRange, -# or if all the indices are scalars, i.e. the view is for a single value only -FastContiguousSubArray{T,N,P,I<:Union{Tuple{Union{Slice, AbstractUnitRange}, Vararg{Any}}, - Tuple{Vararg{ScalarIndex}}}} = SubArray{T,N,P,I,true} -function getindex(V::FastContiguousSubArray, i::Int) + +# For vector views with linear indexing, we disambiguate to favor the stride/offset +# computation as that'll generally be faster than (or just as fast as) re-indexing into a range. +function getindex(V::FastSubArray{<:Any, 1}, i::Int) @inline @boundscheck checkbounds(V, i) - @inbounds r = V.parent[V.offset1 + i] + @inbounds r = V.parent[_reindexlinear(V, i)] r end + +# We can avoid a multiplication if the first parent index is a Colon or AbstractUnitRange, +# or if all the indices are scalars, i.e. the view is for a single value only +FastContiguousSubArray{T,N,P,I<:Union{Tuple{Union{Slice, AbstractUnitRange}, Vararg{Any}}, + Tuple{Vararg{ScalarIndex}}}} = SubArray{T,N,P,I,true} + +@inline _reindexlinear(V::FastContiguousSubArray, i::Integer) = V.offset1 + i +@inline _reindexlinear(V::FastContiguousSubArray, i::AbstractUnitRange{Int}) = V.offset1 .+ i + # parents of FastContiguousSubArrays may support fast indexing with AbstractUnitRanges, # so we may just forward the indexing to the parent # This may only be done for non-offset ranges, as the result would otherwise have offset axes @@ -343,24 +356,10 @@ const OneBasedRanges = Union{OneTo{Int}, UnitRange{Int}, Slice{OneTo{Int}}, Iden function getindex(V::FastContiguousSubArray, i::OneBasedRanges) @inline @boundscheck checkbounds(V, i) - @inbounds r = V.parent[V.offset1 .+ i] + @inbounds r = V.parent[_reindexlinear(V, i)] r end -# For vector views with linear indexing, we disambiguate to favor the stride/offset -# computation as that'll generally be faster than (or just as fast as) re-indexing into a range. -function getindex(V::FastSubArray{<:Any, 1}, i::Int) - @inline - @boundscheck checkbounds(V, i) - @inbounds r = V.parent[V.offset1 + V.stride1*i] - r -end -function getindex(V::FastContiguousSubArray{<:Any, 1}, i::Int) - @inline - @boundscheck checkbounds(V, i) - @inbounds r = V.parent[V.offset1 + i] - r -end @inline getindex(V::FastContiguousSubArray, i::Colon) = getindex(V, to_indices(V, (:,))...) # Indexed assignment follows the same pattern as `getindex` above @@ -373,40 +372,23 @@ end function setindex!(V::FastSubArray, x, i::Int) @inline @boundscheck checkbounds(V, i) - @inbounds V.parent[V.offset1 + V.stride1*i] = x + @inbounds V.parent[_reindexlinear(V, i)] = x V end -function setindex!(V::FastContiguousSubArray, x, i::Int) +function setindex!(V::FastSubArray{<:Any, 1}, x, i::Int) @inline @boundscheck checkbounds(V, i) - @inbounds V.parent[V.offset1 + i] = x + @inbounds V.parent[_reindexlinear(V, i)] = x V end + function setindex!(V::FastSubArray, x, i::AbstractUnitRange{Int}) @inline @boundscheck checkbounds(V, i) - @inbounds V.parent[V.offset1 .+ V.stride1 .* i] = x - V -end -function setindex!(V::FastContiguousSubArray, x, i::AbstractUnitRange{Int}) - @inline - @boundscheck checkbounds(V, i) - @inbounds V.parent[V.offset1 .+ i] = x + @inbounds V.parent[_reindexlinear(V, i)] = x V end -function setindex!(V::FastSubArray{<:Any, 1}, x, i::Int) - @inline - @boundscheck checkbounds(V, i) - @inbounds V.parent[V.offset1 + V.stride1*i] = x - V -end -function setindex!(V::FastContiguousSubArray{<:Any, 1}, x, i::Int) - @inline - @boundscheck checkbounds(V, i) - @inbounds V.parent[V.offset1 + i] = x - V -end @inline setindex!(V::FastSubArray, x, i::Colon) = setindex!(V, x, to_indices(V, (i,))...) function isassigned(V::SubArray{T,N}, I::Vararg{Int,N}) where {T,N} @@ -418,25 +400,13 @@ end function isassigned(V::FastSubArray, i::Int) @inline @boundscheck checkbounds(Bool, V, i) || return false - @inbounds r = isassigned(V.parent, V.offset1 + V.stride1*i) - r -end -function isassigned(V::FastContiguousSubArray, i::Int) - @inline - @boundscheck checkbounds(Bool, V, i) || return false - @inbounds r = isassigned(V.parent, V.offset1 + i) + @inbounds r = isassigned(V.parent, _reindexlinear(V, i)) r end function isassigned(V::FastSubArray{<:Any, 1}, i::Int) @inline @boundscheck checkbounds(Bool, V, i) || return false - @inbounds r = isassigned(V.parent, V.offset1 + V.stride1*i) - r -end -function isassigned(V::FastContiguousSubArray{<:Any, 1}, i::Int) - @inline - @boundscheck checkbounds(Bool, V, i) || return false - @inbounds r = isassigned(V.parent, V.offset1 + i) + @inbounds r = isassigned(V.parent, _reindexlinear(V, i)) r end From 3143648d84c3b1a9ec00dad779ffb3f40953a843 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Sat, 17 Feb 2024 22:52:33 +0530 Subject: [PATCH 2/2] Change Integer index to Int in _reindexlinear --- base/subarray.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/base/subarray.jl b/base/subarray.jl index e131bdd7d5b8e..396f7c52bd77a 100644 --- a/base/subarray.jl +++ b/base/subarray.jl @@ -322,7 +322,7 @@ end FastSubArray{T,N,P,I} = SubArray{T,N,P,I,true} # We define a convenience functions to compute the shifted parent index # This differs from reindex as this accepts the view directly, instead of its indices -@inline _reindexlinear(V::FastSubArray, i::Integer) = V.offset1 + V.stride1*i +@inline _reindexlinear(V::FastSubArray, i::Int) = V.offset1 + V.stride1*i @inline _reindexlinear(V::FastSubArray, i::AbstractUnitRange{Int}) = V.offset1 .+ V.stride1 .* i function getindex(V::FastSubArray, i::Int) @@ -346,7 +346,7 @@ end FastContiguousSubArray{T,N,P,I<:Union{Tuple{Union{Slice, AbstractUnitRange}, Vararg{Any}}, Tuple{Vararg{ScalarIndex}}}} = SubArray{T,N,P,I,true} -@inline _reindexlinear(V::FastContiguousSubArray, i::Integer) = V.offset1 + i +@inline _reindexlinear(V::FastContiguousSubArray, i::Int) = V.offset1 + i @inline _reindexlinear(V::FastContiguousSubArray, i::AbstractUnitRange{Int}) = V.offset1 .+ i # parents of FastContiguousSubArrays may support fast indexing with AbstractUnitRanges,