Properly set within_autodiff (#442)

src/Enzyme.jl

@@ -314,6 +314,7 @@ 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

@@ -63,6 +63,18 @@ function set_reactant_abi(
     if f === Reactant.call_with_reactant
         arginfo2 = ArgInfo(fargs isa Nothing ? nothing : fargs[2:end], argtypes[2:end])
         return abstract_call(interp, arginfo2::ArgInfo, si, sv, max_methods)
+    elseif !(interp.within_autodiff_rewrite) && f === overload_autodiff
+        interp′ = Enzyme.Compiler.Interpreter.EnzymeInterpreter(
+            interp; within_autodiff_rewrite=true
+        )
+        return Base.@invoke abstract_call_known(
+            interp′::Enzyme.Compiler.Interpreter.EnzymeInterpreter,
+            f,
+            arginfo,
+            si,
+            sv,
+            max_methods,
+        )
     end
 
     return Base.@invoke abstract_call_known(
@@ -78,7 +90,9 @@ end
 @static if Enzyme.GPUCompiler.HAS_INTEGRATED_CACHE
     struct ReactantCacheToken end
 
-    function ReactantInterpreter(; world::UInt=Base.get_world_counter())
+    function ReactantInterpreter(;
+        world::UInt=Base.get_world_counter(), within_autodiff=false
+    )
         return Enzyme.Compiler.Interpreter.EnzymeInterpreter(
             ReactantCacheToken(),
             REACTANT_METHOD_TABLE,
@@ -87,15 +101,17 @@ end
             false,            #=reverse_rules=#
             false,            #=inactive_rules=#
             false,            #=broadcast_rewrite=#
-            false,            #=within_autodiff_rewrite=#
+            within_autodiff,  #=within_autodiff_rewrite=#
             set_reactant_abi,
         )
     end
 else
     const REACTANT_CACHE = Enzyme.GPUCompiler.CodeCache()
 
     function ReactantInterpreter(;
-        world::UInt=Base.get_world_counter(), code_cache=REACTANT_CACHE
+        world::UInt=Base.get_world_counter(),
+        code_cache=REACTANT_CACHE,
+        within_autodiff=false,
     )
         return Enzyme.Compiler.Interpreter.EnzymeInterpreter(
             REACTANT_CACHE,
@@ -105,7 +121,7 @@ else
             false,            #=reverse_rules=#
             false,            #=inactive_rules=#
             false,            #=broadcast_rewrite=#
-            false,            #=within_autodiff_rewrite=#
+            within_autodiff,  #=within_autodiff_rewrite=#
             set_reactant_abi,
         )
     end

src/TracedUtils.jl

@@ -270,6 +270,7 @@ function make_mlir_fn(
     args_in_result::Symbol=:all,
     construct_function_without_args::Bool=false,
     do_transpose=true,
+    within_autodiff=false,
     input_shardings=nothing,  # This is not meant to be used by the user.
     output_shardings=nothing, # This is not meant to be used by the user.
     runtime=nothing,
@@ -329,9 +330,19 @@ function make_mlir_fn(
         process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_pad, inv_map)
 
         if isempty(kwargs)
-            Reactant.call_with_reactant(f, traced_args...)
+            if within_autodiff
+                Reactant.call_with_reactant_within_autodiff(f, traced_args...)
+            else
+                Reactant.call_with_reactant(f, traced_args...)
+            end
         else
-            Reactant.call_with_reactant(Core.kwcall, kwargs, f, traced_args...)
+            if within_autodiff
+                Reactant.call_with_reactant_within_autodiff(
+                    Core.kwcall, kwargs, f, traced_args...
+                )
+            else
+                Reactant.call_with_reactant(Core.kwcall, kwargs, f, traced_args...)
+            end
         end
     finally
         MLIR.IR.deactivate!(fnbody)

src/utils.jl

@@ -4,6 +4,7 @@ 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
@@ -632,7 +633,9 @@ function call_with_reactant_generator(
         ))
     end
 
-    interp = ReactantInterpreter(; world)
+    interp = ReactantInterpreter(;
+        world, within_autodiff=self == typeof(Reactant.call_with_reactant_within_autodiff)
+    )
 
     min_world = Ref{UInt}(typemin(UInt))
     max_world = Ref{UInt}(typemax(UInt))
@@ -876,6 +879,10 @@ end
     $(Expr(:meta, :generated_only))
     return $(Expr(:meta, :generated, call_with_reactant_generator))
 end
+@eval function call_with_reactant_within_autodiff($REDUB_ARGUMENTS_NAME...)
+    $(Expr(:meta, :generated_only))
+    return $(Expr(:meta, :generated, call_with_reactant_generator))
+end
 
 @static if isdefined(Core, :BFloat16)
     nmantissa(::Type{Core.BFloat16}) = 7

test/autodiff.jl

@@ -126,6 +126,21 @@ end
     @test res1[1] ≈ ores1[1]
 end
 
+function error_not_within_autodiff()
+    !Enzyme.within_autodiff() && error("Not within autodiff")
+    return nothing
+end
+
+fwd_within_autodiff(Mode, RT) = Enzyme.autodiff(Mode, error_not_within_autodiff, RT)
+
+@testset "within_autodiff" begin
+    @test_throws ErrorException error_not_within_autodiff()
+    @test fwd_within_autodiff(Forward, Const) == ()
+
+    @test_throws ErrorException @jit error_not_within_autodiff()
+    @test (@jit fwd_within_autodiff(Forward, Const)) == ()
+end
+
 function gw(z)
     return Enzyme.gradient(Forward, sum, z; chunk=Val(1))
 end