# Faster array indexing with a logical array

**URL:** <https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364>\
**Category:** New to Julia\
**Tags:** array\
**Created:** [September 27, 2020, 3:24pm UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364 "2020-09-27T15:24:04Z")\
**Posts on this page:** 13\
**Page:** 1

<div class="post-metadata">

**Author:** ![bhaveshshrimali](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bhaveshshrimali/32/13233_2.png) [@bhaveshshrimali](https://discourse.julialang.org/u/bhaveshshrimali)\
**Post date:** [September 27, 2020, 3:24pm UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/1 "2020-09-27T15:24:04Z")

</div>

Hi everyone,

Pretty basic question…

I was experimenting with indexing arrays with logical arrays. For instance, say I have an array shaped `(10000,2)` and I need to extract all the indices where the sum across the columns is less than a threshold. What would be the best _Julian_ way to do it? I tried list comprehension but that is slower than the method below

MWE:

```julia
using BenchmarkTools 

# what I could think of ?
function getindicesVectorized(x::Array)
    return x[getindex.(findall(sum(x, dims=2) .<= 1), 1), :];
end

function getindicesListComp(x::Array)
    ind = findall(sum(x, dims=2) .<= 1)
    temp = [x[ind[i][1], :] for i in 1:length(ind)]
    return hcat(temp...)';
end

points = rand(10000,2)
julia> @btime getindicesVectorized($points)
  32.201 μs (15 allocations: 279.70 KiB)

```

It seems to be fast and readable (see equivalent numpy version), but I was wondering if there would be a faster way to do this ?

What I would do in python naively

```nohighlight
import numpy as np
from timeit import timeit
points = np.random.random((10000,2))
In [8]: %timeit points[np.sum(points, axis=1) <= 1, :]
212 µs ± 2.73 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

```

Thanks in advance 🙂 !!

---

<div class="post-metadata">

**Author:** ![Elrod](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/elrod/32/22461_2.png) [@Elrod](https://discourse.julialang.org/u/Elrod)\
**Post date:** [September 27, 2020, 4:24pm UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/2 "2020-09-27T16:24:32Z")

</div>

I could write a version that’s fast if you have AVX512 and probably okay otherwise.  
I should try to come up with a map-like API so I can put it in LoopVectorization next to vfilter.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [September 27, 2020, 4:27pm UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/3 "2020-09-27T16:27:57Z")

</div>

You don’t need `findall`, you can index with the boolean array from broadcasting `.<` (as in the python version). But even better is to avoid creating the sum anywhere:

```julia
julia> gb1(x) = x[vec(sum(x,dims=2) .< 1), :];

julia> gb2(x) = x[(@views x[:,1] .+ x[:,2] .< 1), :];

julia> @btime gb1($points);
  30.618 μs (13 allocations: 160.41 KiB)

julia> @btime gb2($points);
  20.707 μs (5 allocations: 82.00 KiB)

julia> @btime getindicesVectorized($points); # to compare
  37.253 μs (15 allocations: 275.02 KiB)

```

---

<div class="post-metadata">

**Author:** ![Elrod](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/elrod/32/22461_2.png) [@Elrod](https://discourse.julialang.org/u/Elrod)\
**Post date:** [September 27, 2020, 4:47pm UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/4 "2020-09-27T16:47:06Z")

</div>

AVX512-optimized version:

```julia
using VectorizationBase, SIMDPirates

function getindscompress(X::AbstractMatrix{T}) where {T<:Base.HWReal}
    M, N = size(X)
    @assert N == 2
    Y = similar(X)
    W, Wshift = VectorizationBase.pick_vector_width_shift(T)
    m = 0; my = 0
    ptrX = stridedpointer(X); ptrY = stridedpointer(Y);
    for _ ∈ 1:(M >>> Wshift)
        x₁ = vload(ptrX, (_MM{W}(m),0))
        x₂ = vload(ptrX, (_MM{W}(m),1))
        mask = (x₁ + x₂) <= 1
        SIMDPirates.compressstore!(gep(ptrY, (my,)), extract_data(x₁), mask)
        SIMDPirates.compressstore!(gep(ptrY, (my+M,)), extract_data(x₂), mask)
        my += count_ones(mask)
        m += W
    end
    if (M & (W -1)) > 0
        x₁ = vload(ptrX, (_MM{W}(m),0))
        x₂ = vload(ptrX, (_MM{W}(m),1))
        mask = ((x₁ + x₂) <= 1) & VectorizationBase.mask(T, M)
        SIMDPirates.compressstore!(gep(ptrY, (my,)), extract_data(x₁), mask)
        SIMDPirates.compressstore!(gep(ptrY, (my+M,)), extract_data(x₂), mask)        
    end
    view(Y, 1:my, :)
end

```

Benchmarks:

```julia
julia> X = rand(1000,2);

julia> getindicesVectorized(X) == getindscompress(X)
true

julia> @benchmark getindicesVectorized($X)
BenchmarkTools.Trial:
  memory estimate: 32.83 KiB
  allocs estimate: 11
  --------------
  minimum time: 4.207 μs (0.00% GC)
  median time: 4.600 μs (0.00% GC)
  mean time: 5.130 μs (9.27% GC)
  maximum time: 1.259 ms (97.40% GC)
  --------------
  samples: 10000
  evals/sample: 7

julia> @benchmark getindscompress($X)
BenchmarkTools.Trial:
  memory estimate: 15.75 KiB
  allocs estimate: 1
  --------------
  minimum time: 681.347 ns (0.00% GC)
  median time: 881.409 ns (0.00% GC)
  mean time: 1.011 μs (10.37% GC)
  maximum time: 37.658 μs (88.79% GC)
  --------------
  samples: 10000
  evals/sample: 121

julia> @benchmark gb1($X)
BenchmarkTools.Trial:
  memory estimate: 20.69 KiB
  allocs estimate: 11
  --------------
  minimum time: 2.971 μs (0.00% GC)
  median time: 3.364 μs (0.00% GC)
  mean time: 3.723 μs (8.39% GC)
  maximum time: 1.178 ms (92.85% GC)
  --------------
  samples: 10000
  evals/sample: 8

julia> @benchmark gb2($X)
BenchmarkTools.Trial:
  memory estimate: 12.55 KiB
  allocs estimate: 4
  --------------
  minimum time: 1.847 μs (0.00% GC)
  median time: 2.149 μs (0.00% GC)
  mean time: 2.397 μs (7.13% GC)
  maximum time: 975.513 μs (92.04% GC)
  --------------
  samples: 10000
  evals/sample: 10

```

Using 10\_000 points:

```julia
julia> points = rand(10000,2);

julia> @btime getindicesVectorized($points);
  30.121 μs (15 allocations: 279.58 KiB)

julia> @btime getindscompress($points);
  4.958 μs (2 allocations: 156.33 KiB)

julia> @btime gb1($points);
  23.194 μs (13 allocations: 162.22 KiB)

julia> @btime gb2($points);
  14.622 μs (5 allocations: 83.81 KiB)

```

> [@bhaveshshrimali](#):
>
> What would be the best _Julian_ way to do it?

To me, _Julian_ means winning all the benchmarks. 😉

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [September 27, 2020, 4:57pm UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/5 "2020-09-27T16:57:25Z")

</div>

> [@bhaveshshrimali](#):
>
> I need to extract all the indices where the sum across the columns

It looks to me like you’re extracting the values, not the indices.

---

<div class="post-metadata">

**Author:** ![bhaveshshrimali](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bhaveshshrimali/32/13233_2.png) [@bhaveshshrimali](https://discourse.julialang.org/u/bhaveshshrimali)\
**Post date:** [September 27, 2020, 7:01pm UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/6 "2020-09-27T19:01:49Z")

</div>

> [@Elrod](#):
>
> I could write a version that’s fast if you have AVX512 and probably okay otherwise.

I don’t think so, but frankly I don’t know if I have some kind of vector extensions instructions. The processor I have is [Intel Xeon(R)](https://ark.intel.com/content/www/us/en/ark/products/97468/intel-xeon-processor-e3-1535m-v6-8m-cache-3-10-ghz.html). And simply looking up under Instruction Set Extension, I only see AVX2.

I can read more about it as I don’t have background in architecture 🙂

A bit off the topic, but is there a way to check if my CPU supports AVX512 (I am on Windows) ?

---

<div class="post-metadata">

**Author:** ![bhaveshshrimali](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bhaveshshrimali/32/13233_2.png) [@bhaveshshrimali](https://discourse.julialang.org/u/bhaveshshrimali)\
**Post date:** [September 27, 2020, 7:07pm UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/7 "2020-09-27T19:07:25Z")

</div>

The timings are great Chris! I think this and the previous problem that I posted above both have significant performance gain when doing SIMD operations. I will read more on it.

The idea is to be able to use this in Finite Element Assembly, where surely one can take advantage of [LoopVectorization.jl](https://chriselrod.github.io/LoopVectorization.jl/latest/). Thanks!

Sadly, for the current problem, I don’t gain much by using it. ☹

```julia
julia> @btime getindscompress($points)
  34.701 μs (2 allocations: 156.33 KiB)

```

> [@mcabbott](#):
>
> ```julia
> julia> gb2(x) = x[(@views x[:,1] .+ x[:,2] .< 1), :];
> 
> ```

Thanks @mcabbott!! 🙂

---

<div class="post-metadata">

**Author:** ![bhaveshshrimali](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bhaveshshrimali/32/13233_2.png) [@bhaveshshrimali](https://discourse.julialang.org/u/bhaveshshrimali)\
**Post date:** [September 27, 2020, 7:09pm UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/8 "2020-09-27T19:09:59Z")

</div>

You are right, I mean’t extracting the elements of the array at _those_ indices. 🙂

---

<div class="post-metadata">

**Author:** ![Elrod](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/elrod/32/22461_2.png) [@Elrod](https://discourse.julialang.org/u/Elrod)\
**Post date:** [September 28, 2020, 1:07am UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/9 "2020-09-28T01:07:13Z")

</div>

> [@bhaveshshrimali](#):
>
> A bit off the topic, but is there a way to check if my CPU supports AVX512 (I am on Windows) ?

If the Xeon CPU you linked had AVX512, it would’ve been listed. You can also use:

```julia
julia> using VectorizationBase; VectorizationBase.AVX512F
true

```

But between the link and the fact that it wasn’t faster on your computer, we can already be quite sure that it’ll be `false`.

> [@bhaveshshrimali](#):
>
> The timings are great Chris! I think this and the previous problem that I posted above both have significant performance gain when doing SIMD operations. I will read more on it.
> 
> The idea is to be able to use this in Finite Element Assembly, where surely one can take advantage of [LoopVectorization.jl](https://chriselrod.github.io/LoopVectorization.jl/latest/). Thanks!
> 
> Sadly, for the current problem, I don’t gain much by using it.

The code needs AVX512 to be fast, because AVX512F (F for “Foundation”, common to all AVX512 CPUs) added a “compressed store” instruction that’s perfect for accelerating things like this.  
This will be emulated if you don’t have AVX512 so that the code still works, but that’s of course much slower.

---

<div class="post-metadata">

**Author:** ![Huite](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/huite/32/14034_2.png) [@Huite](https://discourse.julialang.org/u/Huite)\
**Post date:** [September 28, 2020, 7:08am UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/10 "2020-09-28T07:08:49Z")

</div>

If you happen to have Python and numba installed, running `numba -s` is works pretty well:

```julia
System info:
--------------------------------------------------------------------------------
__Time Stamp__
2020-07-10 09:46:04.439609

__Hardware Information__
Machine : AMD64
CPU Name : skylake
CPU count : 8
CPU Features : 
64bit adx aes avx avx2 bmi bmi2 clflushopt cmov cx16 f16c fma fsgsbase invpcid
lzcnt mmx movbe pclmul popcnt prfchw rdrnd rdseed sahf sgx sse sse2 sse3 sse4.1
sse4.2 ssse3 xsave xsavec xsaveopt xsaves

__OS Information__
Platform : Windows-10-10.0.17134-SP0
Release : 10
System Name : Windows
Version : 10.0.17134
OS specific info : 1010.0.17134SP0

[...omitted]

```

Code here:  
[https://github.com/numba/numba/blob/5b9d29b441a014793e3ad38f8cce2a16144a56ef/numba/misc/numba\_sysinfo.py#L281](https://github.com/numba/numba/blob/5b9d29b441a014793e3ad38f8cce2a16144a56ef/numba/misc/numba_sysinfo.py#L281)

I couldn’t find a similar Julia utility by briefly googling (but might be worthwhile?).

---

<div class="post-metadata">

**Author:** ![Elrod](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/elrod/32/22461_2.png) [@Elrod](https://discourse.julialang.org/u/Elrod)\
**Post date:** [September 28, 2020, 7:30am UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/11 "2020-09-28T07:30:47Z")

</div>

> [@Huite](#):
>
> I couldn’t find a similar Julia utility by briefly googling (but might be worthwhile?).

There is [CpuId.jl](https://github.com/m-j-w/CpuId.jl).

```julia
julia> using CpuId

julia> cpufeaturetable()
  Cpu Feature Description
  ––––––––––– ––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
  ACPI Thermal monitor and software controlled clock facilities (MSR)
  ADX Intel ADX (Multi-Precision Add-Carry Instruction Extensions)
  AES AES encryption instruction set
  AHF64 LAHF and SAHF in PM64
  APIC APIC on-chip (Advanced Programmable Interrupt Controller)
  AVX 256bit Advanced Vector Extensions, AVX
  AVX2 SIMD 256bit Advanced Vector Extensions 2
  AVX512BW AVX-512 Byte and Word Instructions
  AVX512CD AVX-512 Conflict Detection Instructions
  AVX512DQ AVX-512 Doubleword and Quadword Instructions
  AVX512F AVX-512 Foundation
  AVX512VL AVX-512 Vector Length Extensions
  BMI1 Bit Manipulation Instruction Set 1

```

Or

```julia
julia> using LLVM

julia> s = unsafe_string(LLVM.API.LLVMGetHostCPUFeatures())
"+sse2,+cx16,+sahf,-tbm,-avx512ifma,-sha,-gfni,-fma4,-vpclmulqdq,+prfchw,+bmi2,-cldemote,+fsgsbase,-ptwrite,+xsavec,+popcnt,+mpx,+aes,-avx512bitalg,-movdiri,+xsaves,-avx512er,+avx512vnni,-avx512vpopcntdq,-pconfig,+clwb,+avx512f,-clzero,-pku,+mmx,-lwp,-rdpid,-xop,+rdseed,-waitpkg,-movdir64b,-sse4a,+avx512bw,+clflushopt,+xsave,-avx512vbmi2,+64bit,+avx512vl,+invpcid,+avx512cd,+avx,-vaes,+cx8,+fma,-rtm,+bmi,-enqcmd,+rdrnd,-mwaitx,+sse4.1,+sse4.2,+avx2,+fxsr,-wbnoinvd,+sse,+lzcnt,+pclmul,-prefetchwt1,+f16c,+ssse3,-sgx,-shstk,+cmov,-avx512vbmi,-avx512bf16,+movbe,+xsaveopt,+avx512dq,+adx,-avx512pf,+sse3"

julia> filter(f -> occursin("avx512", f), split(s, ','))
14-element Array{SubString{String},1}:
 "-avx512ifma"
 "-avx512bitalg"
 "-avx512er"
 "+avx512vnni"
 "-avx512vpopcntdq"
 "+avx512f"
 "+avx512bw"
 "-avx512vbmi2"
 "+avx512vl"
 "+avx512cd"
 "-avx512vbmi"
 "-avx512bf16"
 "+avx512dq"
 "-avx512pf"

julia> Libc.free(s); # don't leak memory

```

A `+` means that it has the feature, while a `-` indicates it doesn’t.  
So, for example, this CPU has AVX512F, but not AVX512ER (accurate reciprocals or exponentiation) or AVX512BF16 (for BF16 support).  
This is basically what VectorizationBase does to define variables for each of these as `true`/`false`.

```julia
julia> using VectorizationBase

julia> VectorizationBase.AVX512IFMA # integer fused multiply add
false

julia> VectorizationBase.AVX512VNNI # vector neural net instructions
true

```

Note that LLVM is more complete. CpuId.jl seems to be missing VNNI, for example.

---

<div class="post-metadata">

**Author:** ![RoyiAvital](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/royiavital/32/571_2.png) [@RoyiAvital](https://discourse.julialang.org/u/RoyiAvital)\
**Post date:** [September 28, 2020, 7:40am UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/12 "2020-09-28T07:40:30Z")

</div>

> A bit off the topic, but is there a way to check if my CPU supports AVX512 (I am on Windows) ?

Use [CPU-Z](https://www.cpuid.com/softwares/cpu-z.html).  
On its `CPU` tab you’ll have:

![image](https://global.discourse-cdn.com/julialang/original/3X/c/9/c9e3933d139271228ac155c12a9b8b28e336bb4f.png)

The `Instructions` section shows what’s your CPU support.  
If you see `AVX512F` there, then it is supported (You won’t see as yours doesn’t support it).

---

<div class="post-metadata">

**Author:** ![bhaveshshrimali](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bhaveshshrimali/32/13233_2.png) [@bhaveshshrimali](https://discourse.julialang.org/u/bhaveshshrimali)\
**Post date:** [September 28, 2020, 9:36pm UTC](https://discourse.julialang.org/t/faster-array-indexing-with-a-logical-array/47364/13 "2020-09-28T21:36:54Z")

</div>

Thanks Chris!

Indeed it doesn’t have AVX512

```julia
julia> cpufeaturetable()
  Cpu Feature Description
  ––––––––––– ––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
  ACPI Thermal monitor and software controlled clock facilities (MSR)
  ADX Intel ADX (Multi-Precision Add-Carry Instruction Extensions)  
  AES AES encryption instruction set
  AHF64 LAHF and SAHF in PM64
  APIC APIC on-chip (Advanced Programmable Interrupt Controller)     
  AVX 256bit Advanced Vector Extensions, AVX
  AVX2 SIMD 256bit Advanced Vector Extensions 2
  BMI1 Bit Manipulation Instruction Set 1
  BMI2 Bit Manipulation Instruction Set 2
  .....

```