Skip to content

Commit

Permalink
speed up logical indexing by bitarray
Browse files Browse the repository at this point in the history
  • Loading branch information
chethega committed Oct 20, 2018
1 parent 2e2b14c commit 1f81b7b
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 11 deletions.
1 change: 1 addition & 0 deletions base/bitarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@ IndexStyle(::Type{<:BitArray}) = IndexLinear()
const _msk64 = ~UInt64(0)
@inline _div64(l) = l >> 6
@inline _mod64(l) = l & 63
@inline _BLSR(x)= x & (x-1) #zeros the last set bit. Has native instruction on many archs. needed in multidimensional.jl
@inline _msk_end(l::Integer) = _msk64 >>> _mod64(-l)
@inline _msk_end(B::BitArray) = _msk_end(length(B))
num_bit_chunks(n::Int) = _div64(n+63)
Expand Down
27 changes: 16 additions & 11 deletions base/multidimensional.jl
Original file line number Diff line number Diff line change
Expand Up @@ -523,19 +523,24 @@ end
L.mask[idx] && return (idx, s)
end
end
# When wrapping a BitArray, lean heavily upon its internals -- this is a common
# case. Just use the Int index and count as its state.
@inline function iterate(L::LogicalIndex{Int,<:BitArray}, s=(0,1))
s[2] > length(L) && return nothing
i, n = s
# When wrapping a BitArray, lean heavily upon its internals.
@inline function iterate(L::Base.LogicalIndex{Int,<:BitArray})
L.sum == 0 && return nothing
Bc = L.mask.chunks
while true
if Bc[_div64(i)+1] & (UInt64(1)<<_mod64(i)) != 0
i += 1
return (i, (i, n+1))
end
i += 1
return iterate(L::Base.LogicalIndex{Int,<:BitArray}, (1, @inbounds Bc[1]))
end
@propagate_inbounds function iterate(L::Base.LogicalIndex{Int,<:BitArray}, s)
Bc = L.mask.chunks
i1, c = s

while c==0
i1 == length(Bc) && return nothing
i1 += 1
c = Bc[i1]
end
tz = trailing_zeros(c) + 1
c = _BLSR(c)
return ((i1-1)<<6 + tz, (i1, c))
end

@inline checkbounds(::Type{Bool}, A::AbstractArray, I::LogicalIndex{<:Any,<:AbstractArray{Bool,1}}) =
Expand Down

0 comments on commit 1f81b7b

Please sign in to comment.