speed up logical indexing by bitarray (#29746)

* speed up logical indexing by bitarray * changed inbounds / OOB behavior to match array iterators switched spelling of _blsr * Update base/multidimensional.jl Co-Authored-By: chethega <[email protected]> (cherry picked from commit 44f2563)
JuliaLang · Feb 11, 2019 · 2de628f · 2de628f
1 parent a97f1e6
commit 2de628f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 11 deletions.
diff --git a/base/bitarray.jl b/base/bitarray.jl
@@ -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)

diff --git a/base/multidimensional.jl b/base/multidimensional.jl
@@ -520,19 +520,23 @@ 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, (1, @inbounds Bc[1]))
+end
+@inline function iterate(L::Base.LogicalIndex{Int,<:BitArray}, s)
+    Bc = L.mask.chunks
+    i1, c = s
+    while c==0
+        i1 % UInt >= length(Bc) % UInt && return nothing
+        i1 += 1
+        @inbounds 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}}) =