Split should_rewrite_ft for call and invoke expressions, and overlay Base._unique_dims

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
@@ -171,6 +171,20 @@ function should_rewrite_ft(@nospecialize(ft))
     return true
 end
 
+# by default, same as `should_rewrite_call`
+function should_rewrite_invoke(@nospecialize(ft), @nospecialize(args))
+    Core.println(Core.stderr, "[should_rewrite_invoke] ft = $ft, parameters = $(args)")
+    return should_rewrite_call(ft)
+end
+
+function should_rewrite_invoke(::typeof(Base.unique), @nospecialize(args))
+    Core.println(Core.stderr, "[should_rewrite_invoke] Base.unique catched!")
+    if args === Tuple{Vector{Symbol}}
+        return false
+    end
+    return true
+end
+
 # 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).
@@ -231,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,
@@ -245,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{}
@@ -259,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}