Make benchmarks not depend on TuringBenchmarking.jl

Author: mhauru

benchmarks/Project.toml

@@ -3,9 +3,12 @@ uuid = "d94a1522-c11e-44a7-981a-42bf5dc1a001"
 version = "0.1.0"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
-TuringBenchmarking = "0db1332d-5c25-4deb-809f-459bc696f94f"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"

benchmarks/benchmarks.jl

@@ -62,22 +62,12 @@ end
 results_table = Tuple{String,String,String,Bool,Float64,Float64}[]
 
 for (model_name, model, varinfo_choice, adbackend, islinked) in chosen_combinations
-    suite = make_suite(model, varinfo_choice, adbackend)
+    suite = make_suite(model, varinfo_choice, adbackend, islinked)
     results = run(suite)
-    result_key = islinked ? "linked" : "standard"
-
-    eval_time = median(results["evaluation"][result_key]).time
+    eval_time = median(results["evaluation"]).time
     relative_eval_time = eval_time / reference_time
-
-    grad_group = results["gradient"]
-    if isempty(grad_group)
-        relative_ad_eval_time = NaN
-    else
-        grad_backend_key = first(keys(grad_group))
-        ad_eval_time = median(grad_group[grad_backend_key][result_key]).time
-        relative_ad_eval_time = ad_eval_time / eval_time
-    end
-
+    ad_eval_time = median(results["gradient"]).time
+    relative_ad_eval_time = ad_eval_time / eval_time
     push!(
         results_table,
         (

benchmarks/src/DynamicPPLBenchmarks.jl

@@ -1,16 +1,39 @@
 module DynamicPPLBenchmarks
 
 using DynamicPPL: VarInfo, SimpleVarInfo, VarName
-using BenchmarkTools: BenchmarkGroup
-using TuringBenchmarking: make_turing_suite
+using BenchmarkTools: BenchmarkGroup, @benchmarkable
+using DynamicPPL: DynamicPPL
+using ADTypes: ADTypes
+using LogDensityProblems: LogDensityProblems
+
+# Load some the default backends to trigger conditional loading.
+using ForwardDiff: ForwardDiff
+using ReverseDiff: ReverseDiff
 
 include("./Models.jl")
 using .Models: Models
 
 export Models, make_suite
 
+# Utility functions for representing AD backends using symbols.
+# Copied from TuringBenchmarking.jl.
+const SYMBOL_TO_BACKEND = Dict(
+    :forwarddiff => ADTypes.AutoForwardDiff(; chunksize=0),
+    :reversediff => ADTypes.AutoReverseDiff(; compile=false),
+    :reversediff_compiled => ADTypes.AutoReverseDiff(; compile=true),
+    :mooncake => ADTypes.AutoMooncake(; config=nothing),
+)
+
+to_backend(x) = error("Unknown backend: $x")
+to_backend(x::ADTypes.AbstractADType) = x
+function to_backend(x::Union{AbstractString,Symbol})
+    k = Symbol(lowercase(string(x)))
+    haskey(SYMBOL_TO_BACKEND, k) || error("Unknown backend: $x")
+    return SYMBOL_TO_BACKEND[k]
+end
+
 """
-    make_suite(model, varinfo_choice::Symbol, adbackend::Symbol)
+    make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::Bool)
 
 Create a benchmark suite for `model` using the selected varinfo type and AD backend.
 Available varinfo choices:
@@ -20,8 +43,10 @@ Available varinfo choices:
   • `:simple_dict`       → builds a `SimpleVarInfo{Float64}` from a Dict (pre-populated with the model’s outputs)
 
 The AD backend should be specified as a Symbol (e.g. `:forwarddiff`, `:reversediff`, `:zygote`).
+
+`islinked` determines whether to link the VarInfo for evaluation.
 """
-function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol)
+function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::Bool)
     suite = BenchmarkGroup()
 
     vi = if varinfo_choice == :untyped
@@ -40,14 +65,32 @@ function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol)
         error("Unknown varinfo choice: $varinfo_choice")
     end
 
-    # Add the AD benchmarking suite.
-    suite = make_turing_suite(
-        model;
-        adbackends=[adbackend],
-        varinfo=vi,
-        check_grads=true,
-        error_on_failed_backend=true,
-    )
+    adbackend = to_backend(adbackend)
+    context = DynamicPPL.DefaultContext()
+
+    if islinked
+        vi = DynamicPPL.link(vi, model)
+    end
+
+    # We construct `LogDensityFunction` using different values
+    # than the ones we're going to use for the test. Some of the AD backends
+    # compile the tape upon `LogDensityFunction` construction, and we want to
+    # evaluate using inputs different from those that the tape was compiled for.
+    f = DynamicPPL.LogDensityFunction(model, vi, context; adtype=adbackend)
+
+    # The parameters at which we evaluate f.
+    θ = if islinked
+        randn(length(vi[:]))
+    else
+        rand(Vector, model)
+    end
+
+    # Run once to trigger compilation.
+    LogDensityProblems.logdensity_and_gradient(f, θ)
+    suite["gradient"] = @benchmarkable $(LogDensityProblems.logdensity_and_gradient)($f, $θ)
+
+    # Also benchmark just standard model evaluation because why not.
+    suite["evaluation"] = @benchmarkable $(DynamicPPL.evaluate!!)($model, $vi, $context)
 
     return suite
 end

benchmarks/src/Models.jl

@@ -17,7 +17,7 @@ using Distributions:
     logpdf,
     product_distribution,
     truncated
-using DynamicPPL: @model, to_submodel
+using DynamicPPL: DynamicPPL, @model, to_submodel
 using LinearAlgebra: cholesky
 
 export simple_assume_observe_non_model,