Implement findall.

Author: maleadt

src/indexing.jl

@@ -43,3 +43,51 @@ function Base.getindex(xs::CuArray{T}, bools::CuArray{Bool}) where {T}
 
   return ys
 end
+
+
+## findall
+
+function Base.findall(bools::CuArray{Bool})
+    indices = cumsum(bools)
+
+    n = _getindex(indices, length(indices))
+    ys = CuArray{Int}(undef, n)
+
+    if n > 0
+        num_threads = min(n, 256)
+        num_blocks = ceil(Int, length(indices) / num_threads)
+
+        function kernel(ys::CuDeviceArray{Int}, bools, indices)
+            i = threadIdx().x + (blockIdx().x - 1) * blockDim().x
+
+            if i <= length(bools) && bools[i]
+                b = indices[i]   # new position
+                ys[b] = i
+
+            end
+
+            return
+        end
+
+        function configurator(kernel)
+            fun = kernel.fun
+            config = launch_configuration(fun)
+            blocks = cld(length(indices), config.threads)
+
+            return (threads=config.threads, blocks=blocks)
+        end
+
+        @cuda config=configurator kernel(ys, bools, indices)
+    end
+
+    unsafe_free!(indices)
+
+    return ys
+end
+
+function Base.findall(f::Function, A::CuArray)
+    bools = map(f, A)
+    ys = findall(bools)
+    unsafe_free!(bools)
+    return ys
+end

test/base.jl

@@ -359,3 +359,8 @@ end
     inds = rand(1:100, 150, 150)
     @test testf(x->permutedims(view(x, inds, :), (3, 2, 1)), rand(100, 100))
 end
+
+@testset "findall" begin
+    @test testf(x->findall(x), rand(Bool, 100))
+    @test testf(x->findall(y->y>0.5, x), rand(100))
+end