Merge pull request #597 from JuliaGPU/tb/findfirst_oob

Fix OOB in find_first kernel.

Author: maleadt

src/indexing.jl

@@ -156,9 +156,9 @@ function Base.findfirst(vals::CuArray, xs::CuArray)
             Ipre = Rpre[I[1]]
             Ipost = Rpost[I[2]]
 
-            @inbounds if xs[Ipre, i, Ipost] == vals[Ipre, Ipost]
+            @inbounds if xs[Ipre, i, Ipost] == vals[Ipre, 1, Ipost]
                 full_index = LinearIndices(xs)[Ipre, i, Ipost]   # atomic_min only works with integers
-                reduced_index = LinearIndices(indices)[Ipre, Ipost] # FIXME: @atomic doesn't handle array ref with CartesianIndices
+                reduced_index = LinearIndices(indices)[Ipre, 1, Ipost] # FIXME: @atomic doesn't handle array ref with CartesianIndices
                 CUDAnative.@atomic indices[reduced_index] = min(indices[reduced_index], full_index)
             end
         end

test/base.jl

@@ -366,6 +366,13 @@ end
   let x = rand(Float32, 10, 10)
       @test findmax(x) == findmax(CuArray(x))
       @test findmax(x; dims=1) == Array.(findmax(CuArray(x); dims=1))
+      @test findmax(x; dims=2) == Array.(findmax(CuArray(x); dims=2))
+  end
+  let x = rand(Float32, 10, 10, 10)
+      @test findmax(x) == findmax(CuArray(x))
+      @test findmax(x; dims=1) == Array.(findmax(CuArray(x); dims=1))
+      @test findmax(x; dims=2) == Array.(findmax(CuArray(x); dims=2))
+      @test findmax(x; dims=3) == Array.(findmax(CuArray(x); dims=3))
   end
 
   let x = rand(Float32, 100)
@@ -375,6 +382,13 @@ end
   let x = rand(Float32, 10, 10)
       @test findmin(x) == findmin(CuArray(x))
       @test findmin(x; dims=1) == Array.(findmin(CuArray(x); dims=1))
+      @test findmin(x; dims=2) == Array.(findmin(CuArray(x); dims=2))
+  end
+  let x = rand(Float32, 10, 10, 10)
+      @test findmin(x) == findmin(CuArray(x))
+      @test findmin(x; dims=1) == Array.(findmin(CuArray(x); dims=1))
+      @test findmin(x; dims=2) == Array.(findmin(CuArray(x); dims=2))
+      @test findmin(x; dims=3) == Array.(findmin(CuArray(x); dims=3))
   end
 end