Skip to content
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.

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

Add missing optimization for iterate(::LogicalIndex{<:CartesianIndex,<:BitArray}) #43448

Merged
merged 3 commits into from
Feb 16, 2022

Conversation

N5N3
Copy link
Member

@N5N3 N5N3 commented Dec 17, 2021

This PR aims to make view(A, A .> 0) faster when IndexStyle(A) isa IndexCartesian.
Some benchmark:

a = randn(128, 32);
A = view(a, axes(a)...);
A2 = view(reshape(a,32,32,4),1:32,1:32,1:4);
p = a .> 0;

@btime view($a, $p); # 1.780 μs (4 allocations: 16.56 KiB)
@btime view($A, $p); # 16.700 μs (2 allocations: 32.92 KiB)
@btime view($A2, reshape($p, size($A2))); # 13.500 μs (4 allocations: 49.47 KiB)

# After This PR
Base.collect(L::Base.LogicalIndex{Int,<:BitArray}) = findall(vec(L.mask))
Base.collect(L::Base.LogicalIndex{CartesianIndex{N},BitArray{N}}) where {N} = N >= 2 ? findall(L.mask) : [i for i in L]

@btime view($a, $p); # 1.710 μs (6 allocations: 16.66 KiB)
@btime view($A, $p); # 2.522 μs (2 allocations: 32.92 KiB)
@btime view($A2, reshape($p, size($A2))); # 3.638 μs (4 allocations: 49.47 KiB)

borrow code from findall
@N5N3
Copy link
Member Author

N5N3 commented Dec 18, 2021

Profile shows that the slow speed for IndexCartesian comes from the unoptimized fallback.
So I add the missing optimization with findall as a reference.
This will also make A[A.>0] = B faster.

@N5N3 N5N3 changed the title make collect(::LogicalIndex) faster by call findall Add missing optimization for iterate(::LogicalIndex{<:CartesianIndex,<:BitArray}) Dec 18, 2021
@N5N3
Copy link
Member Author

N5N3 commented Dec 20, 2021

Update the Benchmark

using Random, BenchmarkTools
Random.seed!(0)
a = randn(128, 32);
A = view(a, axes(a)...);
A2 = view(reshape(a,32,32,4),1:32,1:32,1:4);
p = a .> 0;
@btime view($a, $p); # 1.630 μs  -->  1.360 μs
@btime view($A, $p); # 16.000 μs --> 2.500 μs
@btime view($A2, reshape($p, size($A2))); # 13.100 μs --> 3.638 μs

Co-authored-by: Kristoffer Carlsson <[email protected]>
@ViralBShah
Copy link
Member

Is this good to merge?

@N5N3
Copy link
Member Author

N5N3 commented Feb 16, 2022

We already have

julia/test/abstractarray.jl

Lines 458 to 462 in f090992

mask = bitrand(shape)
@testset "test logical indexing" begin
@test B[mask] == A[mask] == B[findall(mask)] == A[findall(mask)] == LinearIndices(mask)[findall(mask)]
@test B[vec(mask)] == A[vec(mask)] == LinearIndices(mask)[findall(mask)]
mask1 = bitrand(size(A, 1))

So the added optimization should have been tested.

@KristofferC KristofferC reopened this Feb 16, 2022
@KristofferC KristofferC added the status:merge me PR is reviewed. Merge when all tests are passing label Feb 16, 2022
@ViralBShah ViralBShah merged commit d85108a into JuliaLang:master Feb 16, 2022
@N5N3 N5N3 deleted the fasterLogicalIndex branch February 16, 2022 14:40
antoine-levitt pushed a commit to antoine-levitt/julia that referenced this pull request Feb 17, 2022
…,<:BitArray})` (JuliaLang#43448)

* make `collect(::LogicalIndex)` faster by call `findall`

* borrow code from findall

borrow code from findall

* Update base/multidimensional.jl

Co-authored-by: Kristoffer Carlsson <[email protected]>

Co-authored-by: Kristoffer Carlsson <[email protected]>
@DilumAluthge DilumAluthge removed the status:merge me PR is reviewed. Merge when all tests are passing label Feb 19, 2022
LilithHafner pushed a commit to LilithHafner/julia that referenced this pull request Feb 22, 2022
…,<:BitArray})` (JuliaLang#43448)

* make `collect(::LogicalIndex)` faster by call `findall`

* borrow code from findall

borrow code from findall

* Update base/multidimensional.jl

Co-authored-by: Kristoffer Carlsson <[email protected]>

Co-authored-by: Kristoffer Carlsson <[email protected]>
LilithHafner pushed a commit to LilithHafner/julia that referenced this pull request Mar 8, 2022
…,<:BitArray})` (JuliaLang#43448)

* make `collect(::LogicalIndex)` faster by call `findall`

* borrow code from findall

borrow code from findall

* Update base/multidimensional.jl

Co-authored-by: Kristoffer Carlsson <[email protected]>

Co-authored-by: Kristoffer Carlsson <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
performance Must go faster
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

5 participants