feat: sorted search overloads (#1325)

* feat: sorted search overloads

* feat: add reverse overloads

* fix: searchsorted last

* test: searchsorted routines

Author: avik-pal

src/TracedRArray.jl

@@ -21,6 +21,11 @@ using ..TracedUtils: TracedUtils, get_mlir_data, set_mlir_data!, materialize_tra
 using ReactantCore: ReactantCore
 using GPUArraysCore: GPUArraysCore, @allowscalar
 
+__lt(::Base.Order.ForwardOrdering, a, b) = isless.(a, b)
+__lt(o::Base.Order.ReverseOrdering, a, b) = __lt(o.fwd, b, a)
+__lt(o::Base.Order.By, a, b) = __lt(o.order, o.by.(a), o.by.(b))
+__lt(o::Base.Order.Lt, a, b) = o.lt.(a, b)
+
 ReactantCore.is_traced(::TracedRArray, seen) = true
 ReactantCore.is_traced(::TracedRArray) = true
 
@@ -943,20 +948,26 @@ function overloaded_stack(dims::Union{Integer,Colon}, xs)
 end
 
 # sort
-function Base.sort(x::AnyTracedRArray; alg=missing, order=missing, kwargs...)
-    return sort!(copy(x); alg, order, kwargs...)
+function Base.sort(x::AnyTracedRArray; alg=missing, kwargs...)
+    return sort!(copy(x); alg, kwargs...)
 end
-function Base.sort(x::AnyTracedRVector; alg=missing, order=missing, kwargs...)
-    return sort!(copy(x); alg, order, kwargs...)
+function Base.sort(x::AnyTracedRVector; alg=missing, kwargs...)
+    return sort!(copy(x); alg, kwargs...)
 end
 
 function Base.sort!(
-    x::AnyTracedRVector; lt=isless, by=identity, rev::Bool=false, alg=missing, order=missing
+    x::AnyTracedRVector;
+    lt=isless,
+    by=identity,
+    rev::Bool=false,
+    alg=missing,
+    order=Base.Order.Forward,
 )
     @assert alg === missing "Reactant doesn't support `alg` kwarg for `sort!`"
-    @assert order === missing "Reactant doesn't support `order` kwarg for `sort!`"
 
-    comparator = rev ? (a, b) -> !lt(by(a), by(b)) : (a, b) -> lt(by(a), by(b))
+    ordering = Base.ord(lt, by, rev, order)
+    comparator = (a, b) -> __lt(ordering, a, b)
+
     res = only(Ops.sort(materialize_traced_array(x); comparator, dimension=1))
     set_mlir_data!(x, get_mlir_data(res))
     return x
@@ -969,22 +980,23 @@ function Base.sort!(
     by=identity,
     rev::Bool=false,
     alg=missing,
-    order=missing,
+    order=Base.Order.Forward,
 )
     @assert alg === missing "Reactant doesn't support `alg` kwarg for `sort!`"
-    @assert order === missing "Reactant doesn't support `order` kwarg for `sort!`"
 
-    comparator = rev ? (a, b) -> !lt(by(a), by(b)) : (a, b) -> lt(by(a), by(b))
+    ordering = Base.ord(lt, by, rev, order)
+    comparator = (a, b) -> __lt(ordering, a, b)
+
     res = only(Ops.sort(materialize_traced_array(x); dimension=dims, comparator))
     set_mlir_data!(x, get_mlir_data(res))
     return x
 end
 
-function Base.sortperm(x::AnyTracedRArray; alg=missing, order=missing, kwargs...)
-    return sortperm!(similar(x, Int), x; alg, order, kwargs...)
+function Base.sortperm(x::AnyTracedRArray; alg=missing, kwargs...)
+    return sortperm!(similar(x, Int), x; alg, kwargs...)
 end
-function Base.sortperm(x::AnyTracedRVector; alg=missing, order=missing, kwargs...)
-    return sortperm!(similar(x, Int), x; alg, order, dims=1, kwargs...)
+function Base.sortperm(x::AnyTracedRVector; alg=missing, kwargs...)
+    return sortperm!(similar(x, Int), x; alg, dims=1, kwargs...)
 end
 
 function Base.sortperm!(
@@ -995,18 +1007,18 @@ function Base.sortperm!(
     by=identity,
     rev::Bool=false,
     alg=missing,
-    order=missing,
+    order=Base.Order.Forward,
 ) where {N}
     if dims === nothing
         @assert ndims(x) == 1
         dims = 1
     end
 
     @assert alg === missing "Reactant doesn't support `alg` kwarg for `sortperm!`"
-    @assert order === missing "Reactant doesn't support `order` kwarg for `sortperm!`"
 
-    comparator =
-        rev ? (a, b, i1, i2) -> !lt(by(a), by(b)) : (a, b, i1, i2) -> lt(by(a), by(b))
+    ordering = Base.ord(lt, by, rev, order)
+    comparator = (a, b, i1, i2) -> __lt(ordering, a, b)
+
     idxs = Ops.constant(collect(LinearIndices(x)))
     _, res = Ops.sort(materialize_traced_array(x), idxs; dimension=dims, comparator)
     set_mlir_data!(ix, get_mlir_data(res))
@@ -1346,4 +1358,50 @@ function scan_impl!(
     return output
 end
 
+function Base.searchsortedfirst(
+    v::AnyTracedRVector, x, lo::T, hi::T, o::Base.Ordering
+) where {T<:Integer}
+    return sum(T.(__lt(o, v[lo:hi], x)); init=lo)
+end
+
+function Base.searchsortedlast(
+    v::AnyTracedRVector, x, lo::T, hi::T, o::Base.Ordering
+) where {T<:Integer}
+    return sum(T.(.!(__lt(o, x, v[lo:hi]))); init=lo - 1)
+end
+
+function Base.searchsorted(
+    v::AnyTracedRVector, x, lo::T, hi::T, o::Base.Ordering
+) where {T<:Integer}
+    firstidx = searchsortedfirst(v, x, lo, hi, o)
+    lastidx = searchsortedlast(v, x, lo, hi, o)
+    return Reactant.TracedRNumberOverrides.TracedUnitRange(firstidx, lastidx)
+end
+
+function Base.reverse(
+    v::AnyTracedRVector{T}, start::Integer, stop::Integer=lastindex(v)
+) where {T}
+    v[start:stop] = reverse!(v[start:stop])
+    return v
+end
+
+function Base.reverse!(
+    v::AnyTracedRVector{T}, start::Integer, stop::Integer=lastindex(v)
+) where {T}
+    reverse!(view(v, start:stop))
+    return v
+end
+
+function Base.reverse!(v::AnyTracedRVector{T}) where {T}
+    v_mat = materialize_traced_array(v)
+    copyto!(v, Ops.reverse(v_mat; dimensions=1))
+    return v
+end
+
+function Base._reverse!(a::AnyTracedRArray{T,N}, dims::NTuple{M,Int}) where {T,N,M}
+    a_mat = materialize_traced_array(a)
+    copyto!(a, Ops.reverse(a_mat; dimensions=dims))
+    return a
+end
+
 end

test/basic.jl

@@ -1295,3 +1295,35 @@ accum_fn(x, y) = abs2(x) + abs2(y)
         end ≈ accumulate(accum_fn, b; dims=3, init=0.0f0)
     end
 end
+
+sameunitrange(x, y) = first(x) == first(y) && last(x) == last(y)
+
+@testset "searchsorted" begin
+    x = [1, 2, 4, 5, 5, 7]
+    x_ra = Reactant.to_rarray(x)
+
+    @testset "searchsortedfirst" begin
+        @testset for val in (4, 5, 3, 9, 0)
+            @test @jit(searchsortedfirst(x_ra, val)) == searchsortedfirst(x, val)
+            @test @jit(searchsortedfirst(x_ra, ConcreteRNumber(val))) ==
+                searchsortedfirst(x, val)
+        end
+    end
+
+    @testset "searchsortedlast" begin
+        @testset for val in (4, 5, 3, 9, 0)
+            @test @jit(searchsortedlast(x_ra, val)) == searchsortedlast(x, val)
+            @test @jit(searchsortedlast(x_ra, ConcreteRNumber(val))) ==
+                searchsortedlast(x, val)
+        end
+    end
+
+    @testset "searchsorted" begin
+        @testset for val in (4, 5, 3, 9, 0)
+            @test sameunitrange(@jit(searchsorted(x_ra, val)), searchsorted(x, val))
+            @test sameunitrange(
+                @jit(searchsorted(x_ra, ConcreteRNumber(val))), searchsorted(x, val)
+            )
+        end
+    end
+end