Skip to content

Orthogonalize re-indexing for FastSubArrays #53369

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 2 commits into from
Feb 18, 2024
Merged

Orthogonalize re-indexing for FastSubArrays #53369

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8454477
Orthogonalize re-indexing for FastSubArrays
jishnub Feb 17, 2024
3143648
Change Integer index to Int in _reindexlinear
jishnub Feb 17, 2024
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
86 changes: 28 additions & 58 deletions base/subarray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -320,47 +320,46 @@ 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::Int) = 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::Int) = 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
const OneBasedRanges = Union{OneTo{Int}, UnitRange{Int}, Slice{OneTo{Int}}, IdentityUnitRange{OneTo{Int}}}
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
Expand All @@ -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}
Expand All @@ -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

Expand Down