use NoCache to improve set_to_zero!! performance with Mooncake

sunxd3 · sunxd3 · commit e93458c98f6b · 2025-07-08T08:25:35.000+01:00
diff --git a/HISTORY.md b/HISTORY.md
@@ -1,5 +1,9 @@
 # DynamicPPL Changelog
 
+## Unreleased
+
+Improved performance for some models with Mooncake.jl by using `NoCache` with `Mooncake.set_to_zero!!` for DynamicPPL types.
+
 ## 0.36.14
 
 Added compatibility with AbstractPPL@0.12.
diff --git a/ext/DynamicPPLMooncakeExt.jl b/ext/DynamicPPLMooncakeExt.jl
@@ -1,9 +1,110 @@
 module DynamicPPLMooncakeExt
 
+__precompile__(false)
+
 using DynamicPPL: DynamicPPL, istrans
 using Mooncake: Mooncake
+import Mooncake: set_to_zero!!
+using Mooncake: NoTangent, Tangent, MutableTangent, NoCache, set_to_zero_internal!!
 
 # This is purely an optimisation.
 Mooncake.@zero_adjoint Mooncake.DefaultCtx Tuple{typeof(istrans),Vararg}
 
+"""
+Check if a tangent has the expected structure for a given type.
+"""
+function has_expected_structure(
+    x, expected_type::Type{<:Union{Tangent,MutableTangent}}, expected_fields
+)
+    x isa expected_type || return false
+    hasfield(typeof(x), :fields) || return false
+
+    fields = x.fields
+    if expected_fields isa Tuple
+        # Exact match required
+        propertynames(fields) == expected_fields || return false
+    else
+        # All expected fields must be present
+        all(f in propertynames(fields) for f in expected_fields) || return false
+    end
+
+    return true
+end
+
+"""
+Check if a tangent corresponds to a DynamicPPL.LogDensityFunction
+"""
+function is_dppl_ldf_tangent(x)
+    has_expected_structure(x, Tangent, (:model, :varinfo, :context, :adtype, :prep)) ||
+        return false
+
+    fields = x.fields
+    is_dppl_varinfo_tangent(fields.varinfo) || return false
+    is_dppl_model_tangent(fields.model) || return false
+
+    return true
+end
+
+"""
+Check if a tangent corresponds to a DynamicPPL.VarInfo
+"""
+function is_dppl_varinfo_tangent(x)
+    return has_expected_structure(x, Tangent, (:metadata, :logp, :num_produce))
+end
+
+"""
+Check if a tangent corresponds to a DynamicPPL.Model
+"""
+function is_dppl_model_tangent(x)
+    return has_expected_structure(x, Tangent, (:f, :args, :defaults, :context))
+end
+
+"""
+Check if a MutableTangent corresponds to DynamicPPL.Metadata
+"""
+function is_dppl_metadata_tangent(x)
+    return has_expected_structure(
+        x, MutableTangent, (:idcs, :vns, :ranges, :vals, :dists, :orders, :flags)
+    )
+end
+
+"""
+Check if a model function tangent represents a closure.
+"""
+function is_closure_model(model_f_tangent)
+    model_f_tangent isa MutableTangent && return true
+
+    if model_f_tangent isa Tangent && hasfield(typeof(model_f_tangent), :fields)
+        # Check if any field is a MutableTangent with PossiblyUninitTangent{Any}
+        for (_, fval) in pairs(model_f_tangent.fields)
+            if fval isa MutableTangent &&
+                hasfield(typeof(fval), :fields) &&
+                hasfield(typeof(fval.fields), :contents) &&
+                fval.fields.contents isa Mooncake.PossiblyUninitTangent{Any}
+                return true
+            end
+        end
+    end
+
+    return false
+end
+
+function Mooncake.set_to_zero!!(x)
+    # Check for DynamicPPL types and use NoCache for better performance
+    if is_dppl_ldf_tangent(x)
+        # Special handling for LogDensityFunction to detect closures
+        model_f_tangent = x.fields.model.fields.f
+        cache = is_closure_model(model_f_tangent) ? IdDict{Any,Bool}() : NoCache()
+        return set_to_zero_internal!!(cache, x)
+    elseif is_dppl_varinfo_tangent(x) ||
+        is_dppl_model_tangent(x) ||
+        is_dppl_metadata_tangent(x)
+        # These types can always use NoCache
+        return set_to_zero_internal!!(NoCache(), x)
+    else
+        # Use the original implementation with IdDict for all other types
+        return set_to_zero_internal!!(IdDict{Any,Bool}(), x)
+    end
+end
+
 end # module
diff --git a/test/ext/DynamicPPLMooncakeExt.jl b/test/ext/DynamicPPLMooncakeExt.jl
@@ -1,5 +1,195 @@
+using DynamicPPL
+using Distributions
+using Random
+using Test
+using StableRNGs
+using Mooncake: NoCache, set_to_zero!!, set_to_zero_internal!!, zero_tangent
+using DynamicPPL.TestUtils.AD: @be
+using Statistics: median
+
+# Define models globally to avoid closure issues
+@model function test_model1(x)
+    s ~ InverseGamma(2, 3)
+    m ~ Normal(0, sqrt(s))
+    return x .~ Normal(m, sqrt(s))
+end
+
+@model function test_model2(x, y)
+    τ ~ Gamma(1, 1)
+    σ ~ InverseGamma(2, 3)
+    μ ~ Normal(0, τ)
+    x .~ Normal(μ, σ)
+    return y .~ Normal(μ, σ)
+end
+
 @testset "DynamicPPLMooncakeExt" begin
-    Mooncake.TestUtils.test_rule(
-        StableRNG(123456), istrans, VarInfo(); unsafe_perturb=true, interface_only=true
-    )
+    @testset "istrans rule" begin
+        Mooncake.TestUtils.test_rule(
+            StableRNG(123456), istrans, VarInfo(); unsafe_perturb=true, interface_only=true
+        )
+    end
+
+    @testset "set_to_zero!! optimization" begin
+        # Test with a real DynamicPPL model
+        model = test_model1([1.0, 2.0, 3.0])
+        vi = VarInfo(Random.default_rng(), model)
+        ldf = LogDensityFunction(model, vi, DefaultContext())
+        tangent = zero_tangent(ldf)
+
+        # Test that set_to_zero!! works correctly
+        result = set_to_zero!!(deepcopy(tangent))
+        @test result isa typeof(tangent)
+
+        # Test with metadata - verify structure exists
+        if hasfield(typeof(tangent.fields.varinfo.fields), :metadata)
+            metadata = tangent.fields.varinfo.fields.metadata
+            @test !isnothing(metadata)
+        end
+    end
+
+    @testset "NoCache optimization correctness" begin
+        # Test that set_to_zero!! uses NoCache for DynamicPPL types
+        model = test_model1([1.0, 2.0, 3.0])
+        vi = VarInfo(Random.default_rng(), model)
+        ldf = LogDensityFunction(model, vi, DefaultContext())
+        tangent = zero_tangent(ldf)
+
+        # Modify some values
+        if hasfield(typeof(tangent.fields.model.fields), :args) &&
+            hasfield(typeof(tangent.fields.model.fields.args), :x)
+            x_tangent = tangent.fields.model.fields.args.x
+            if !isempty(x_tangent)
+                x_tangent[1] = 5.0
+            end
+        end
+
+        # Call set_to_zero!! and verify it works
+        set_to_zero!!(tangent)
+
+        # Check that values are zeroed
+        if hasfield(typeof(tangent.fields.model.fields), :args) &&
+            hasfield(typeof(tangent.fields.model.fields.args), :x)
+            x_tangent = tangent.fields.model.fields.args.x
+            if !isempty(x_tangent)
+                @test x_tangent[1] == 0.0
+            end
+        end
+    end
+
+    @testset "Performance improvement" begin
+        # Test with DEMO_MODELS if available
+        if isdefined(DynamicPPL.TestUtils, :DEMO_MODELS) &&
+            !isempty(DynamicPPL.TestUtils.DEMO_MODELS)
+            model = DynamicPPL.TestUtils.DEMO_MODELS[1]
+        else
+            # Fallback to our test model
+            model = test_model1([1.0, 2.0, 3.0, 4.0])
+        end
+
+        vi = VarInfo(Random.default_rng(), model)
+        ldf = LogDensityFunction(model, vi, DefaultContext())
+        tangent = zero_tangent(ldf)
+
+        # Run benchmarks
+        result_iddict = @be begin
+            cache = IdDict{Any,Bool}()
+            set_to_zero_internal!!(cache, tangent)
+        end
+
+        result_nocache = @be set_to_zero!!(tangent)
+
+        # Extract median times
+        time_iddict = median(result_iddict).time
+        time_nocache = median(result_nocache).time
+
+        # We expect NoCache to be faster
+        speedup = time_iddict / time_nocache
+        @test speedup > 1.5  # Conservative expectation - should be ~4x
+
+        @info "Performance improvement" speedup time_iddict_μs = time_iddict / 1000 time_nocache_μs =
+            time_nocache / 1000
+    end
+
+    @testset "Aliasing safety" begin
+        # Test with aliased data
+        shared_data = [1.0, 2.0, 3.0]
+        model = test_model2(shared_data, shared_data)  # x and y are the same array
+        vi = VarInfo(Random.default_rng(), model)
+        ldf = LogDensityFunction(model, vi, DefaultContext())
+        tangent = zero_tangent(ldf)
+
+        # Check that aliasing is preserved in tangent
+        if hasfield(typeof(tangent.fields.model.fields), :args)
+            args = tangent.fields.model.fields.args
+            if hasfield(typeof(args), :x) && hasfield(typeof(args), :y)
+                @test args.x === args.y  # Aliasing should be preserved
+
+                # Modify via x
+                if !isempty(args.x)
+                    args.x[1] = 10.0
+                    @test args.y[1] == 10.0  # Should also change y
+                end
+
+                # Zero and check both are zeroed
+                # Since x and y are aliased, zeroing one zeros both
+                set_to_zero!!(tangent)
+                if !isempty(args.x)
+                    @test args.x[1] == 0.0
+                    @test args.y[1] == 0.0
+                end
+            end
+        end
+    end
+
+    @testset "Closure handling" begin
+        # Test that closure models are correctly handled
+
+        # Create closure model (captures environment, has circular references)
+        function create_closure_model()
+            local_var = 42
+            @model function closure_model(x)
+                s ~ InverseGamma(2, 3)
+                m ~ Normal(0, sqrt(s))
+                return x .~ Normal(m, sqrt(s))
+            end
+            return closure_model
+        end
+
+        closure_fn = create_closure_model()
+        model_closure = closure_fn([1.0, 2.0, 3.0])
+        vi_closure = VarInfo(Random.default_rng(), model_closure)
+        ldf_closure = LogDensityFunction(model_closure, vi_closure, DefaultContext())
+        tangent_closure = zero_tangent(ldf_closure)
+
+        # Test that it works without stack overflow
+        @test_nowarn set_to_zero!!(deepcopy(tangent_closure))
+
+        # Compare with global model (no closure)
+        model_global = test_model1([1.0, 2.0, 3.0])
+        vi_global = VarInfo(Random.default_rng(), model_global)
+        ldf_global = LogDensityFunction(model_global, vi_global, DefaultContext())
+        tangent_global = zero_tangent(ldf_global)
+
+        # Verify model.f tangent types differ
+        f_tangent_closure = tangent_closure.fields.model.fields.f
+        f_tangent_global = tangent_global.fields.model.fields.f
+
+        @test f_tangent_global isa Mooncake.NoTangent  # Global function
+        @test f_tangent_closure isa Mooncake.Tangent   # Closure function
+
+        # Performance comparison
+        time_global = @elapsed for _ in 1:100
+            set_to_zero!!(tangent_global)
+        end
+
+        time_closure = @elapsed for _ in 1:100
+            set_to_zero!!(tangent_closure)
+        end
+
+        # Global should be faster (uses NoCache)
+        @test time_global < time_closure
+
+        @info "Closure handling" time_global_ms = time_global * 1000 time_closure_ms =
+            time_closure * 1000
+    end
 end
