Properly set within_autodiff (#442)

src/Interpreter.jl

@@ -44,6 +44,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(
@@ -59,7 +71,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,
@@ -68,14 +82,17 @@ end
             false,            #=reverse_rules=#
             false,            #=inactive_rules=#
             false,            #=broadcast_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,
@@ -85,6 +102,7 @@ else
             false,            #=reverse_rules=#
             false,            #=inactive_rules=#
             false,            #=broadcast_rewrite=#
+            within_autodiff,  #=within_autodiff_rewrite=#
             set_reactant_abi,
         )
     end
@@ -238,7 +256,7 @@ function overload_autodiff(
     primargs = ((v.val for v in args)...,)
 
     fnwrap, func2, traced_result, result, seen_args, ret, linear_args, in_tys, linear_results = TracedUtils.make_mlir_fn(
-        primf, primargs, (), string(f) * "_autodiff", false
+        primf, primargs, (), string(f) * "_autodiff", false; within_autodiff=true
     )
 
     activity = Int32[]

src/TracedUtils.jl

@@ -94,6 +94,7 @@ function make_mlir_fn(
     no_args_in_result::Bool=false,
     construct_function_without_args::Bool=false,
     do_transpose=true,
+    within_autodiff=false,
 )
     if sizeof(typeof(f)) != 0 || f isa Base.BroadcastFunction
         return (
@@ -180,8 +181,11 @@ function make_mlir_fn(
                 arg.mlir_data = row_maj_arg
             end
         end
-
-        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
     finally
         MLIR.IR.deactivate!(fnbody)
     end

src/utils.jl

@@ -4,6 +4,7 @@ function apply(f, args...; kwargs...)
 end
 
 function call_with_reactant end
+function call_with_reactant_within_autodiff end
 
 function maybe_argextype(@nospecialize(x), src)
     return try
@@ -483,7 +484,11 @@ function call_with_reactant_generator(
         MethodError($REDUB_ARGUMENTS_NAME[1], $REDUB_ARGUMENTS_NAME[2:end], $world)
     ))
 
-    interp = ReactantInterpreter(; world)
+    if self == typeof(Reactant.call_with_reactant_within_autodiff)
+        interp = ReactantInterpreter(; world, within_autodiff=true)
+    else
+        interp = ReactantInterpreter(; world, within_autodiff=false)
+    end
 
     min_world = Ref{UInt}(typemin(UInt))
     max_world = Ref{UInt}(typemax(UInt))
@@ -728,3 +733,7 @@ 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

test/autodiff.jl

@@ -63,6 +63,21 @@ fwd(Mode, RT, x, y) = Enzyme.autodiff(Mode, square, RT, Duplicated(x, y))
     @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