cleanup and fix

Author: jumerckx

src/Enzyme.jl

@@ -321,7 +321,6 @@ function overload_autodiff(
         argprefix,
         resprefix,
         resargprefix,
-        within_autodiff=true,
     )
     (; result, linear_args, in_tys, linear_results) = mlir_fn_res
     fnwrap = mlir_fn_res.fnwrapped

src/Interpreter.jl

@@ -90,17 +90,17 @@ end
     struct ReactantCacheToken end
 
     function ReactantInterpreter(;
-        world::UInt=Base.get_world_counter(), within_autodiff=false
+        world::UInt=Base.get_world_counter()
     )
         return Enzyme.Compiler.Interpreter.EnzymeInterpreter(
             ReactantCacheToken(),
             REACTANT_METHOD_TABLE,
             world,
-            false,            #=forward_rules=#
-            false,            #=reverse_rules=#
-            false,            #=inactive_rules=#
-            false,            #=broadcast_rewrite=#
-            within_autodiff,  #=within_autodiff_rewrite=#
+            false,  #=forward_rules=#
+            false,  #=reverse_rules=#
+            false,  #=inactive_rules=#
+            false,  #=broadcast_rewrite=#
+            false,  #=within_autodiff_rewrite=#
             set_reactant_abi,
         )
     end
@@ -110,17 +110,16 @@ else
     function ReactantInterpreter(;
         world::UInt=Base.get_world_counter(),
         code_cache=REACTANT_CACHE,
-        within_autodiff=false,
     )
         return Enzyme.Compiler.Interpreter.EnzymeInterpreter(
             REACTANT_CACHE,
             REACTANT_METHOD_TABLE,
             world,
-            false,            #=forward_rules=#
-            false,            #=reverse_rules=#
-            false,            #=inactive_rules=#
-            false,            #=broadcast_rewrite=#
-            within_autodiff,  #=within_autodiff_rewrite=#
+            false,  #=forward_rules=#
+            false,  #=reverse_rules=#
+            false,  #=inactive_rules=#
+            false,  #=broadcast_rewrite=#
+            false,  #=within_autodiff_rewrite=#
             set_reactant_abi,
         )
     end

src/Overlay.jl

@@ -30,7 +30,13 @@ end
 @reactant_overlay @noinline function Enzyme.autodiff(
     rmode::Enzyme.Mode, f::FA, rt::Type{A}, args::Vararg{Annotation,Nargs}
 ) where {FA<:Annotation,A<:Annotation,Nargs}
-    return overload_autodiff(rmode, f, rt, args...)
+    original_within_autodiff = WITHIN_AUTODIFF[]
+    try
+        WITHIN_AUTODIFF[] = true
+        return overload_autodiff(rmode, f, rt, args...)
+    finally
+        WITHIN_AUTODIFF[] = original_within_autodiff
+    end
 end
 
 @reactant_overlay function EnzymeCore.ignore_derivatives(args...)

src/utils.jl

@@ -19,7 +19,6 @@ function apply(f::F, args...; kwargs...) where {F}
 end
 
 function call_with_reactant end
-function call_with_reactant_within_autodiff end
 
 function maybe_argextype(@nospecialize(x), src)
     return try