Skip to content

Commit

Permalink
Quasi-linear indexing for reshaped reinterpretarray
Browse files Browse the repository at this point in the history
  • Loading branch information
@timholy
timholy committed Sep 29, 2020
1 parent bc2abbe commit 69a82b5
Show file tree
Hide file tree
Showing 3 changed files with 299 additions and 45 deletions.
192 changes: 183 additions & 9 deletions base/reinterpretarray.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -175,8 +175,96 @@ function check_writable(a::ReinterpretArray{T, N, S} where N) where {T,S}
end
end

IndexStyle(a::ReinterpretArray) = IndexStyle(a.parent)
IndexStyle(a::ReshapedReinterpretArray{T, N, S}) where {T, N, S} = sizeof(T) < sizeof(S) ? IndexCartesian() : IndexStyle(a.parent)
## IndexStyle specializations

# For `reinterpret(reshape, T, a)` where we're adding a channel dimension and with
# `IndexStyle(a) == IndexLinear()`, it's advantageous to retain pseudo-linear indexing.
struct IndexSCartesian2{K,N} <: IndexStyle end # K = sizeof(S) ÷ sizeof(T), a static-sized 2d cartesian iterator

IndexStyle(::Type{ReinterpretArray{T,N,S,A,false}}) where {T,N,S,A<:AbstractArray{S,N}} = IndexStyle(A)
function IndexStyle(::Type{ReinterpretArray{T,N,S,A,true}}) where {T,N,S,A<:AbstractArray{S}}
if sizeof(T) < sizeof(S)
IndexStyle(A) === IndexLinear() && return IndexSCartesian2{sizeof(S) ÷ sizeof(T),N}()
return IndexCartesian()
end
return IndexStyle(A)
end
IndexStyle(::IndexSCartesian2{K,N}, ::IndexSCartesian2{K,N}) where {K,N} = IndexSCartesian2{K,N}()

struct SCartesianIndex2{K,N} <: AbstractCartesianIndex{N}
i::Int
j::Int
end
to_index(i::SCartesianIndex2) = i

struct SCartesianIndices2{K,N,R<:AbstractUnitRange{Int}} <: AbstractMatrix{SCartesianIndex2{K}}
indices2::R
end
SCartesianIndices2{K,N}(indices2::AbstractUnitRange{Int}) where {K,N} = (@assert K::Int > 1; SCartesianIndices2{K,N,typeof(indices2)}(indices2))

eachindex(::IndexSCartesian2{K,N}, A::AbstractArray) where {K,N} = SCartesianIndices2{K,N}(eachindex(IndexLinear(), parent(A)))

size(iter::SCartesianIndices2{K}) where K = (K, length(iter.indices2))
axes(iter::SCartesianIndices2{K}) where K = (Base.OneTo(K), iter.indices2)

first(iter::SCartesianIndices2{K,N}) where {K,N} = SCartesianIndex2{K,N}(1, first(iter.indices2))
last(iter::SCartesianIndices2{K,N}) where {K,N} = SCartesianIndex2{K,N}(K, last(iter.indices2))

@inline function getindex(iter::SCartesianIndices2{K,N}, i::Int, j::Int) where {K,N}
@boundscheck checkbounds(iter, i, j)
return SCartesianIndex2{K,N}(i, iter.indices2[j])
end

function iterate(iter::SCartesianIndices2{K,N}) where {K,N}
ret = iterate(iter.indices2)
ret === nothing && return nothing
item2, state2 = ret
return SCartesianIndex2{K,N}(1, item2), (1, item2, state2)
end

function iterate(iter::SCartesianIndices2{K,N}, (state1, item2, state2)) where {K,N}
if state1 < K
item1 = state1 + 1
return SCartesianIndex2{K,N}(item1, item2), (item1, item2, state2)
end
ret = iterate(iter.indices2, state2)
ret === nothing && return nothing
item2, state2 = ret
return SCartesianIndex2{K,N}(1, item2), (1, item2, state2)
end

SimdLoop.simd_outer_range(iter::SCartesianIndices2) = iter.indices2
SimdLoop.simd_inner_length(::SCartesianIndices2{K}, ::Any) where K = K
@inline function SimdLoop.simd_index(::SCartesianIndices2{K,N}, Ilast::Int, I1::Int) where {K,N}
SCartesianIndex2{K,N}(I1+1, Ilast)
end

_maybe_reshape(::IndexSCartesian2, A::ReshapedReinterpretArray, I...) = A

# fallbacks
function _getindex(::IndexSCartesian2, A::AbstractArray{T,N}, I::Vararg{Int, N}) where {T,N}
@_propagate_inbounds_meta
getindex(A, I...)
end
function _setindex!(::IndexSCartesian2, A::AbstractArray{T,N}, v, I::Vararg{Int, N}) where {T,N}
@_propagate_inbounds_meta
setindex!(A, v, I...)
end
# fallbacks for array types that use "pass-through" indexing (e.g., `IndexStyle(A) = IndexStyle(parent(A))`)
# but which don't handle SCartesianIndex2
function _getindex(::IndexSCartesian2, A::AbstractArray{T,N}, ind::SCartesianIndex2) where {T,N}
@_propagate_inbounds_meta
I = _to_subscript_indices(A, ind.i, ind.j)
getindex(A, I...)
end
function _setindex!(::IndexSCartesian2, A::AbstractArray{T,N}, v, ind::SCartesianIndex2) where {T,N}
@_propagate_inbounds_meta
I = _to_subscript_indices(A, ind.i, ind.j)
setindex!(A, v, I...)
end


## AbstractArray interface

parent(a::ReinterpretArray) = a.parent
dataids(a::ReinterpretArray) = dataids(a.parent)
Expand Down Expand Up @@ -231,6 +319,19 @@ end
_getindex_ra(a, inds[1], tail(inds))
end

@inline @propagate_inbounds function getindex(a::ReshapedReinterpretArray{T,N,S}, ind::SCartesianIndex2) where {T,N,S}
check_readable(a)
n = sizeof(S) ÷ sizeof(T)
t = Ref{NTuple{n,T}}()
s = Ref{S}(a.parent[ind.j])
GC.@preserve t s begin
tptr = Ptr{UInt8}(unsafe_convert(Ref{T}, t))
sptr = Ptr{UInt8}(unsafe_convert(Ref{S}, s))
_memcpy!(tptr, sptr, sizeof(S))
end
return t[][ind.i]
end

@inline _memcpy!(dst, src, n) = ccall(:memcpy, Cvoid, (Ptr{UInt8}, Ptr{UInt8}, Csize_t), dst, src, n)

@inline @propagate_inbounds function _getindex_ra(a::NonReshapedReinterpretArray{T,N,S}, i1::Int, tailinds::TT) where {T,N,S,TT}
Expand Down Expand Up @@ -292,9 +393,17 @@ end
if sizeof(T) > sizeof(S)
# Extra dimension in the parent array
n = sizeof(T) ÷ sizeof(S)
for i = 1:n
s[] = a.parent[i, i1, tailinds...]
_memcpy!(tptr + (i-1)*sizeof(S), sptr, sizeof(S))
if isempty(tailinds) && IndexStyle(a.parent) === IndexLinear()
offset = n * (i1 - firstindex(a))
for i = 1:n
s[] = a.parent[i + offset]
_memcpy!(tptr + (i-1)*sizeof(S), sptr, sizeof(S))
end
else
for i = 1:n
s[] = a.parent[i, i1, tailinds...]
_memcpy!(tptr + (i-1)*sizeof(S), sptr, sizeof(S))
end
end
else
# No extra dimension
Expand Down Expand Up @@ -334,6 +443,20 @@ end
_setindex_ra!(a, v, inds[1], tail(inds))
end

@inline @propagate_inbounds function setindex!(a::ReshapedReinterpretArray{T,N,S}, v, ind::SCartesianIndex2) where {T,N,S}
check_writable(a)
v = convert(T, v)::T
t = Ref{T}(v)
s = Ref{S}(a.parent[ind.j])
GC.@preserve t s begin
tptr = Ptr{UInt8}(unsafe_convert(Ref{T}, t))
sptr = Ptr{UInt8}(unsafe_convert(Ref{S}, s))
_memcpy!(sptr + (ind.i-1)*sizeof(T), tptr, sizeof(T))
end
a.parent[ind.j] = s[]
return a
end

@inline @propagate_inbounds function _setindex_ra!(a::NonReshapedReinterpretArray{T,N,S}, v, i1::Int, tailinds::TT) where {T,N,S,TT}
v = convert(T, v)::T
# Make sure to match the scalar reinterpret if that is applicable
Expand Down Expand Up @@ -407,13 +530,21 @@ end
GC.@preserve t s begin
tptr = Ptr{UInt8}(unsafe_convert(Ref{T}, t))
sptr = Ptr{UInt8}(unsafe_convert(Ref{S}, s))
if sizeof(T) >= sizeof(S) == 0
if sizeof(T) >= sizeof(S)
if sizeof(T) > sizeof(S)
# Extra dimension in the parent array
n = sizeof(T) ÷ sizeof(S)
for i = 1:n
_memcpy!(sptr, tptr + (i-1)*sizeof(S), sizeof(S))
a.parent[i, i1, tailinds...] = s[]
if isempty(tailinds) && IndexStyle(a.parent) === IndexLinear()
offset = n * (i1 - firstindex(a))
for i = 1:n
_memcpy!(sptr, tptr + (i-1)*sizeof(S), sizeof(S))
a.parent[i + offset] = s[]
end
else
for i = 1:n
_memcpy!(sptr, tptr + (i-1)*sizeof(S), sizeof(S))
a.parent[i, i1, tailinds...] = s[]
end
end
else
# No extra dimension
Expand Down Expand Up @@ -523,3 +654,46 @@ using .Iterators: Stateful
end
return true
end

# Reductions with IndexSCartesian2

function _mapreduce(f::F, op::OP, style::IndexSCartesian2{K}, A::AbstractArrayOrBroadcasted) where {F,OP,K}
inds = eachindex(style, A)
n = size(inds)[2]
if n == 0
return mapreduce_empty_iter(f, op, A, IteratorEltype(A))
else
return mapreduce_impl(f, op, A, first(inds), last(inds))
end
end

@noinline function mapreduce_impl(f::F, op::OP, A::AbstractArrayOrBroadcasted,
ifirst::SCI, ilast::SCI, blksize::Int) where {F,OP,SCI<:SCartesianIndex2{K}} where K
if ifirst.j + blksize > ilast.j
# sequential portion
@inbounds a1 = A[ifirst]
@inbounds a2 = A[SCI(2,ifirst.j)]
v = op(f(a1), f(a2))
@simd for i = ifirst.i + 2 : K
@inbounds ai = A[SCI(i,ifirst.j)]
v = op(v, f(ai))
end
# Remaining columns
for j = ifirst.j+1 : ilast.j
@simd for i = 1:K
@inbounds ai = A[SCI(i,j)]
v = op(v, f(ai))
end
end
return v
else
# pairwise portion
jmid = (ifirst.j + ilast.j) >> 1
v1 = mapreduce_impl(f, op, A, ifirst, SCI(K,jmid), blksize)
v2 = mapreduce_impl(f, op, A, SCI(1,jmid+1), ilast, blksize)
return op(v1, v2)
end
end

mapreduce_impl(f::F, op::OP, A::AbstractArrayOrBroadcasted, ifirst::SCartesianIndex2, ilast::SCartesianIndex2) where {F,OP} =
mapreduce_impl(f, op, A, ifirst, ilast, pairwise_blocksize(f, op))
Loading

0 comments on commit 69a82b5

Please sign in to comment.