fix: iteration ordering for findfirst (#1537)

* fix: iteration ordering for findfirst

* Update src/TracedRArray.jl

* fix: return -1

* test: make it deterministic

* chore: remove unused deps

* Update src/Ops.jl

Author: avik-pal

src/Ops.jl

@@ -1275,6 +1275,30 @@ end
     return (; values, indices)
 end
 
+# Taken from https://github.com/JuliaGPU/GPUArrays.jl/blob/49a339c63a50f1a00ac84844675bcb3a11070cb0/src/host/indexing.jl#L193
+@noinline function findfirst(
+    x::TracedRArray{Bool,N};
+    dimension::Integer=N,
+    location=mlir_stacktrace("findfirst", @__FILE__, @__LINE__),
+) where {N}
+    return reduce(
+        TracedRArray[
+            x, iota(Int64, collect(Int64, size(x)); iota_dimension=dimension, location)
+        ],
+        TracedRNumber[
+            Reactant.TracedUtils.promote_to(TracedRNumber{Bool}, false),
+            Reactant.TracedUtils.promote_to(TracedRNumber{Int64}, typemax(Int64)),
+        ],
+        [dimension],
+        function (x, i, y, j)
+            cond_val = x | y
+            idx = ifelse(x, ifelse(i < j, i, j), ifelse(y, j, typemax(Int64)))
+            return cond_val, idx
+        end;
+        location,
+    )[2] .+ 1
+end
+
 @noinline function argmax(
     x::TracedRArray{T,N};
     dimension::Integer=N,

src/TracedRArray.jl

@@ -1114,16 +1114,17 @@ Base.argmax(x::AnyTracedRArray; kwargs...) = findmax(identity, x; kwargs...)[2]
 Base.findfirst(x::AnyTracedRArray) = findfirst(identity, x)
 Base.findlast(x::AnyTracedRArray) = findlast(identity, x)
 
+# FIXME: we need to conditionally return `nothing` here if idx < 0
 function Base.findfirst(f::Function, x::AnyTracedRArray)
-    fA = materialize_traced_array(vec(f.(x)))
-    (; indices) = Ops.top_k(fA, 1)
-    return @allowscalar indices[1]
+    idx = Ops.findfirst(materialize_traced_array(vec(f.(x))))
+    return TracedRNumber{Int}((), idx.mlir_data)
 end
 
+# FIXME: we need to conditionally return `nothing` here if idx < 0
 function Base.findlast(f::Function, x::AnyTracedRArray)
     fA = Ops.reverse(materialize_traced_array(vec(f.(x))); dimensions=[1])
-    (; indices) = Ops.top_k(fA, 1)
-    return length(x) - @allowscalar(indices[1]) + 1
+    idx = Ops.findfirst(fA)
+    return length(x) + 1 - TracedRNumber{Int}((), idx.mlir_data)
 end
 
 Base.findmin(x::AnyTracedRVector) = findmin(identity, x; dims=1)

test/sorting.jl

@@ -171,7 +171,11 @@ end
 end
 
 @testset "findfirst / findlast" begin
-    x = rand(Bool, 3, 4)
+    x = Bool[
+        0 0 0 0
+        1 0 1 0
+        0 1 0 1
+    ]
     x_ra = Reactant.to_rarray(x)
 
     ffirstlinindices(x) = LinearIndices(x)[findfirst(x)]
@@ -182,7 +186,11 @@ end
     @test ffirstlinindices(x) == @jit(findfirst(x_ra))
     @test flastlinindices(x) == @jit(findlast(x_ra))
 
-    x = rand(1:256, 3, 4)
+    x = Int64[
+        3 5 7 9
+        4 6 7 8
+        5 7 8 9
+    ]
     x_ra = Reactant.to_rarray(x)
 
     @test ffirstlinindices(iseven, x) == @jit(findfirst(iseven, x_ra))