Split should_rewrite_ft for call and invoke expressions, and overlay Base._unique_dims (#505)

* Split `should_rewrite_ft` for `call` and `invoke` expressions

* fix dispatch on `unique`

* fix syntax problem

* test

* fix typo

* Replace `should_rewrite_invoke` of `unique` for overlayed method on `_unique_dims`

* remove previous solution

Author: mofeing

src/Overlay.jl

@@ -142,3 +142,12 @@ end
         end
     end
 end
+
+## fixes #493
+@reactant_overlay @noinline function Base._unique_dims(A::AbstractArray, dims::Colon)
+    if use_overlayed_version(A)
+        error("Reactant doesn't have a `Base._unique_dims` with the current interpreter.")
+    else
+        Base.inferencebarrier(Base._unique_dims)(A, dims)
+    end
+end

src/utils.jl

@@ -89,7 +89,7 @@ function has_ancestor(query::Module, target::Module)
     end
 end
 
-function should_rewrite_ft(@nospecialize(ft))
+function should_rewrite_call(@nospecialize(ft))
     # Don't rewrite builtin or intrinsics
     if ft <: Core.IntrinsicFunction || ft <: Core.Builtin
         return false
@@ -178,6 +178,9 @@ function should_rewrite_ft(@nospecialize(ft))
     return true
 end
 
+# by default, same as `should_rewrite_call`
+should_rewrite_invoke(@nospecialize(ft), @nospecialize(args)) = should_rewrite_call(ft)
+
 # Avoid recursively interpreting into methods we define explicitly
 # as overloads, which we assume should handle the entirety of the
 # translation (and if not they can use call_in_reactant).
@@ -242,7 +245,7 @@ function rewrite_inst(inst, ir, interp, RT, guaranteed_error)
         end
         if ft == typeof(Core._apply_iterate)
             ft = Core.Compiler.widenconst(maybe_argextype(inst.args[3], ir))
-            if Base.invokelatest(should_rewrite_ft, ft)
+            if Base.invokelatest(should_rewrite_call, ft)
                 if RT === Union{}
                     rep = Expr(
                         :call,
@@ -256,7 +259,7 @@ function rewrite_inst(inst, ir, interp, RT, guaranteed_error)
                     return true, rep, Any
                 end
             end
-        elseif Base.invokelatest(should_rewrite_ft, ft)
+        elseif Base.invokelatest(should_rewrite_call, ft)
             if RT === Union{}
                 rep = Expr(:call, call_with_reactant, MustThrowError(), inst.args...)
                 return true, rep, Union{}
@@ -270,10 +273,13 @@ function rewrite_inst(inst, ir, interp, RT, guaranteed_error)
         omi = inst.args[1]::Core.MethodInstance
         sig = omi.specTypes
         ft = sig.parameters[1]
+        argsig = sig.parameters[2:end]
         if ft == typeof(Core.kwcall)
             ft = sig.parameters[3]
+            argsig = sig.parameters[4:end]
         end
-        if Base.invokelatest(should_rewrite_ft, ft) && !is_reactant_method(omi)
+        argsig = Core.apply_type(Core.Tuple, argsig...)
+        if Base.invokelatest(should_rewrite_invoke, ft, argsig) && !is_reactant_method(omi)
             method = omi.def::Core.Method
 
             min_world = Ref{UInt}(typemin(UInt))

test/compile.jl

@@ -128,7 +128,7 @@ end
     @test !occursin("add", repr(hlo))
 end
 
-# While a bit specific, the following is used to check for a bug in `should_rewrite_ft`
+# While a bit specific, the following is used to check for a bug in `should_rewrite_call`
 function sinusoidal_embedding(
     x::AbstractArray{T,4}, min_freq, max_freq, embedding_dims::Int
 ) where {T}
@@ -146,3 +146,9 @@ end
     x_ra = Reactant.to_rarray(rand(Float32, 1, 1, 1, 4))
     hlo = @code_hlo sinusoidal_embedding(x_ra, 0.1, 10.0, 4)
 end
+
+# test #493
+@testset "unique(::Vector{Symbol}) (#493)" begin
+    x = [:a, :b, :a]
+    @test @jit(unique(x)) == [:a, :b]
+end