add support for ignore_derivatives (#2547)

Author: vchuravy

Project.toml

@@ -1,7 +1,7 @@
 name = "Enzyme"
 uuid = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>"]
-version = "0.13.80"
+version = "0.13.81"
 
 [deps]
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
@@ -44,7 +44,7 @@ BFloat16s = "0.2, 0.3, 0.4, 0.5"
 CEnum = "0.4, 0.5"
 ChainRulesCore = "1"
 DynamicPPL = "0.35, 0.36, 0.37"
-EnzymeCore = "0.8.13"
+EnzymeCore = "0.8.14"
 Enzyme_jll = "0.0.201"
 GPUArraysCore = "0.1.6, 0.2"
 GPUCompiler = "1.6.2"

docs/src/notebooks/ignore_derivatives.jl

@@ -31,6 +31,9 @@ begin
     )
 end
 
+# ╔═╡ 23a8503f-3c68-4523-aebe-a4ce4575a02b
+import Enzyme: ignore_derivatives
+
 # ╔═╡ df72e42f-7eec-476f-8ce5-72b09f620005
 md"""
 # Reproducing "Stabilizing backpropagation through time to learn complex physics" 
@@ -116,51 +119,6 @@ end
 # ╔═╡ be852753-126d-42fa-a55c-c907f5dce99d
 plot_gradientfield(N, S, x₀, y, θ₁, θ₂, n)
 
-# ╔═╡ 0b6d0456-1f94-479f-b690-89ad7fc61e44
-begin
-    @noinline function ignore_derivatives(x::T) where {T}
-        return Core.inferencebarrier(x)::T
-    end
-
-    function EnzymeRules.forward(
-            config,
-            ::Const{typeof(ignore_derivatives)},
-            A, x::Duplicated
-        )
-        return Enzyme.make_zero(x.val)
-    end
-
-    function EnzymeRules.augmented_primal(
-            config,
-            ::Const{typeof(ignore_derivatives)},
-            FA, x
-        )
-        primal = EnzymeRules.needs_primal(config) ? x.val : nothing
-        if x isa Active
-            shadow = nothing
-        else
-            shadow = Enzyme.make_zero(x.val)
-        end
-
-        return EnzymeRules.AugmentedReturn(primal, shadow, nothing)
-    end
-    function EnzymeRules.reverse(
-            config,
-            ::Const{typeof(ignore_derivatives)},
-            dret::Active, tape, x::Active
-        )
-        return (Enzyme.make_zero(x.val),)
-    end
-
-    function EnzymeRules.reverse(
-            config,
-            ::Const{typeof(ignore_derivatives)},
-            ::Type{<:Duplicated}, tape, x::Duplicated
-        )
-        return (nothing,)
-    end
-end
-
 # ╔═╡ 873e7792-99a1-4472-92c2-6fc32e2889fa
 N_stop(xᵢ, θ) = θ[1] * ignore_derivatives(xᵢ^2) + θ[2] * ignore_derivatives(xᵢ)
 
@@ -170,6 +128,7 @@ plot_gradientfield(N_stop, S, x₀, y, θ₁, θ₂, n)
 # ╔═╡ Cell order:
 # ╠═b72e9218-81ba-11f0-1eba-5bd949c7ade4
 # ╠═9f5c0822-a19a-4c63-95e7-d2f066a7440f
+# ╠═23a8503f-3c68-4523-aebe-a4ce4575a02b
 # ╠═a4453d23-6e31-451f-b2cd-97346accac82
 # ╠═bd0352c3-1b3c-42f5-ab93-7ca4cb67b9ad
 # ╟─df72e42f-7eec-476f-8ce5-72b09f620005
@@ -182,6 +141,5 @@ plot_gradientfield(N_stop, S, x₀, y, θ₁, θ₂, n)
 # ╠═45ee18f4-d6d3-40f4-bbc0-04cbd3b7b840
 # ╠═ae6a671d-1559-4bff-af6e-78d2b54db020
 # ╠═be852753-126d-42fa-a55c-c907f5dce99d
-# ╠═0b6d0456-1f94-479f-b690-89ad7fc61e44
 # ╠═873e7792-99a1-4472-92c2-6fc32e2889fa
 # ╠═d71a22cc-c1f3-4425-8a6f-442a0bc4f215

lib/EnzymeCore/Project.toml

@@ -1,7 +1,7 @@
 name = "EnzymeCore"
 uuid = "f151be2c-9106-41f4-ab19-57ee4f262869"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>"]
-version = "0.8.13"
+version = "0.8.14"
 
 [compat]
 Adapt = "3, 4"
@@ -15,6 +15,7 @@ AdaptExt = "Adapt"
 
 [extras]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 
 [weakdeps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

lib/EnzymeCore/src/EnzymeCore.jl

@@ -6,7 +6,7 @@ export Const, Active, Duplicated, DuplicatedNoNeed, BatchDuplicated, BatchDuplic
 export MixedDuplicated, BatchMixedDuplicated
 export DefaultABI, FFIABI, InlineABI, NonGenABI
 export BatchDuplicatedFunc
-export within_autodiff
+export within_autodiff, ignore_derivatives
 export needs_primal
 
 function batch_size end
@@ -620,6 +620,27 @@ Returns true if within autodiff, otherwise false.
     return false
 end
 
+"""
+    ignore_derivatives(x::T)::T
+
+Behaves like the `identity` function, but disconnects the "shadow"
+associated with `x`. This has the effect of preventing any derivatives
+from being propagated through `x`.
+
+!!! compat "Enzyme 0.13.74"
+    Support for `ignore_derivatives` was added in Enzyme 0.13.74.
+"""
+@generated function ignore_derivatives(x::T) where {T}
+    name = "extern __enzyme_ignore_derivatives." * string(T)
+    return quote
+        if EnzymeCore.within_autodiff()
+            return ccall($name, llvmcall, $T, ($T,), x)
+        else
+            return x
+        end
+    end
+end
+
 """
     set_err_if_func_written(::Mode)
 

lib/EnzymeCore/test/runtests.jl

@@ -41,3 +41,7 @@ end
 @testset "within_autodiff" begin
     @test !EnzymeCore.within_autodiff()
 end
+
+@testset "ignore_derivatives" begin
+    @test EnzymeCore.ignore_derivatives(3) == 3
+end

src/Enzyme.jl

@@ -53,6 +53,7 @@ import EnzymeCore:
     set_strong_zero,
     clear_strong_zero,
     within_autodiff,
+    ignore_derivatives,
     WithPrimal,
     NoPrimal,
     needs_primal,

src/compiler.jl

@@ -2604,6 +2604,29 @@ function enzyme!(
     if DumpPostWrap[]
         API.EnzymeDumpModuleRef(mod.ref)
     end
+
+    # Rewrite enzyme_ignore_derivatives functions to the identity of their first argument.
+    to_delete = LLVM.Function[]
+    for fn in functions(mod)
+        if startswith(name(fn), "__enzyme_ignore_derivatives")
+            push!(to_delete, fn)
+            to_delete_inst = LLVM.CallInst[]
+            for u in LLVM.uses(fn)
+                ci = LLVM.user(u)
+                @assert isa(ci, LLVM.CallInst)
+                LLVM.replace_uses!(ci, operands(ci)[1])
+                push!(to_delete_inst, ci)
+            end
+            for ci in to_delete_inst
+                LLVM.erase!(ci)
+            end
+        end
+    end
+    for fn in to_delete
+        LLVM.erase!(fn)
+    end
+    LLVM.verify(mod)
+
     API.EnzymeLogicErasePreprocessedFunctions(logic)
     adjointfname = adjointf == nothing ? nothing : LLVM.name(adjointf)
     augmented_primalfname =

test/ignore_derivatives.jl

@@ -0,0 +1,25 @@
+using Test
+using Enzyme
+import Enzyme: ignore_derivatives
+
+@testset "ignore_derivatives" begin
+    @test autodiff(Enzyme.Forward, ignore_derivatives, Duplicated(1.0, 2.0)) == (0.0,)
+    @test autodiff(Enzyme.Reverse, ignore_derivatives, Active(1.0)) == ((0.0,),)
+end
+
+N(xᵢ, θ) = θ[1] * xᵢ^2 + θ[2] * xᵢ
+N_stop(xᵢ, θ) = θ[1] * ignore_derivatives(xᵢ^2) + θ[2] * ignore_derivatives(xᵢ)
+
+@testset "simulate with ignore_derivatives" begin
+    x₀ = -0.3
+    θ = (-4.0, 4.0)
+
+    dθ = MixedDuplicated(θ, Ref(Enzyme.make_zero(θ)))
+    @test Enzyme.autodiff(Enzyme.Reverse, N, Active(x₀), dθ) == ((6.4, nothing),)
+    @test dθ.dval[] == (0.09, -0.3)
+
+
+    dθ = MixedDuplicated(θ, Ref(Enzyme.make_zero(θ)))
+    @test Enzyme.autodiff(Enzyme.Reverse, N_stop, Active(x₀), dθ) == ((0.0, nothing),)
+    @test dθ.dval[] == (0.09, -0.3)
+end