Move code to BenchmarkToolsExt

Author: penelopeysm

Project.toml

@@ -25,6 +25,7 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [weakdeps]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -34,6 +35,7 @@ MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 
 [extensions]
+DynamicPPLBenchmarkToolsExt = ["BenchmarkTools"]
 DynamicPPLChainRulesCoreExt = ["ChainRulesCore"]
 DynamicPPLEnzymeCoreExt = ["EnzymeCore"]
 DynamicPPLForwardDiffExt = ["ForwardDiff"]
@@ -47,6 +49,7 @@ AbstractMCMC = "5"
 AbstractPPL = "0.10.1"
 Accessors = "0.1"
 BangBang = "0.4.1"
+BenchmarkTools = "1.6.0"
 Bijectors = "0.13.18, 0.14, 0.15"
 ChainRulesCore = "1"
 Compat = "4"

benchmarks/src/Models.jl renamed to benchmarks/Models.jl

@@ -2,23 +2,42 @@ using Pkg
 # To ensure we benchmark the local version of DynamicPPL, dev the folder above.
 Pkg.develop(; path=joinpath(@__DIR__, ".."))
 
-using DynamicPPLBenchmarks: Models, make_suite, model_dimension
+using DynamicPPL: DynamicPPL, make_benchmark_suite, VarInfo
 using BenchmarkTools: @benchmark, median, run
 using PrettyTables: PrettyTables, ft_printf
+using ForwardDiff: ForwardDiff
+using Mooncake: Mooncake
+using ReverseDiff: ReverseDiff
 using StableRNGs: StableRNG
 
-rng = StableRNG(23)
+include("Models.jl")
+
+"""
+    model_dimension(model, islinked)
+
+Return the dimension of `model`, accounting for linking, if any.
+"""
+function model_dimension(model, islinked)
+    vi = VarInfo()
+    model(vi)
+    if islinked
+        vi = DynamicPPL.link(vi, model)
+    end
+    return length(vi[:])
+end
 
 # Create DynamicPPL.Model instances to run benchmarks on.
-smorgasbord_instance = Models.smorgasbord(randn(rng, 100), randn(rng, 100))
+smorgasbord_instance = Models.smorgasbord(
+    randn(StableRNG(23), 100), randn(StableRNG(23), 100)
+)
 loop_univariate1k, multivariate1k = begin
-    data_1k = randn(rng, 1_000)
+    data_1k = randn(StableRNG(23), 1_000)
     loop = Models.loop_univariate(length(data_1k)) | (; o=data_1k)
     multi = Models.multivariate(length(data_1k)) | (; o=data_1k)
     loop, multi
 end
 loop_univariate10k, multivariate10k = begin
-    data_10k = randn(rng, 10_000)
+    data_10k = randn(StableRNG(23), 10_000)
     loop = Models.loop_univariate(length(data_10k)) | (; o=data_10k)
     multi = Models.multivariate(length(data_10k)) | (; o=data_10k)
     loop, multi
@@ -34,7 +53,7 @@ end
 chosen_combinations = [
     (
         "Simple assume observe",
-        Models.simple_assume_observe(randn(rng)),
+        Models.simple_assume_observe(randn(StableRNG(23))),
         :typed,
         :forwarddiff,
         false,
@@ -50,22 +69,22 @@ chosen_combinations = [
     ("Loop univariate 10k", loop_univariate10k, :typed, :mooncake, true),
     ("Multivariate 10k", multivariate10k, :typed, :mooncake, true),
     ("Dynamic", Models.dynamic(), :typed, :mooncake, true),
-    ("Submodel", Models.parent(randn(rng)), :typed, :mooncake, true),
+    ("Submodel", Models.parent(randn(StableRNG(23))), :typed, :mooncake, true),
     ("LDA", lda_instance, :typed, :reversediff, true),
 ]
 
 # Time running a model-like function that does not use DynamicPPL, as a reference point.
 # Eval timings will be relative to this.
 reference_time = begin
-    obs = randn(rng)
+    obs = randn(StableRNG(23))
     median(@benchmark Models.simple_assume_observe_non_model(obs)).time
 end
 
 results_table = Tuple{String,Int,String,String,Bool,Float64,Float64}[]
 
 for (model_name, model, varinfo_choice, adbackend, islinked) in chosen_combinations
     @info "Running benchmark for $model_name"
-    suite = make_suite(model, varinfo_choice, adbackend, islinked)
+    suite = make_benchmark_suite(StableRNG(23), model, varinfo_choice, adbackend, islinked)
     results = run(suite)
     eval_time = median(results["evaluation"]).time
     relative_eval_time = eval_time / reference_time

benchmarks/benchmarks.jl

@@ -1,54 +1,18 @@
-module DynamicPPLBenchmarks
+module DynamicPPLBenchmarkToolsExt
 
-using DynamicPPL: VarInfo, SimpleVarInfo, VarName
+using DynamicPPL:
+    DynamicPPL, ADTypes, LogDensityProblems, Model, VarInfo, SimpleVarInfo, VarName
 using BenchmarkTools: BenchmarkGroup, @benchmarkable
-using DynamicPPL: DynamicPPL
-using ADTypes: ADTypes
-using LogDensityProblems: LogDensityProblems
-
-using ForwardDiff: ForwardDiff
-using Mooncake: Mooncake
-using ReverseDiff: ReverseDiff
-using StableRNGs: StableRNG
-
-include("./Models.jl")
-using .Models: Models
-
-export Models, make_suite, model_dimension
+using Random: Random
 
 """
-    model_dimension(model, islinked)
-
-Return the dimension of `model`, accounting for linking, if any.
-"""
-function model_dimension(model, islinked)
-    vi = VarInfo()
-    model(vi)
-    if islinked
-        vi = DynamicPPL.link(vi, model)
-    end
-    return length(vi[:])
-end
-
-# Utility functions for representing AD backends using symbols.
-# Copied from TuringBenchmarking.jl.
-const SYMBOL_TO_BACKEND = Dict(
-    :forwarddiff => ADTypes.AutoForwardDiff(),
-    :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, islinked::Bool)
+    make_benchmark_suite(
+        [rng::Random.AbstractRNG,]
+        model::Model,
+        varinfo_choice::Symbol,
+        adtype::ADTypes.AbstractADType,
+        islinked::Bool
+    )
 
 Create a benchmark suite for `model` using the selected varinfo type and AD backend.
 Available varinfo choices:
@@ -57,13 +21,15 @@ Available varinfo choices:
   • `:simple_namedtuple` → uses `SimpleVarInfo{Float64}(model())`
   • `: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, islinked::Bool)
-    rng = StableRNG(23)
-
+function make_benchmark_suite(
+    rng::Random.AbstractRNG,
+    model::Model,
+    varinfo_choice::Symbol,
+    adtype::ADTypes.AbstractADType,
+    islinked::Bool,
+)
     suite = BenchmarkGroup()
 
     vi = if varinfo_choice == :untyped
@@ -82,14 +48,13 @@ function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::
         error("Unknown varinfo choice: $varinfo_choice")
     end
 
-    adbackend = to_backend(adbackend)
     context = DynamicPPL.DefaultContext()
 
     if islinked
         vi = DynamicPPL.link(vi, model)
     end
 
-    f = DynamicPPL.LogDensityFunction(model, vi, context; adtype=adbackend)
+    f = DynamicPPL.LogDensityFunction(model, vi, context; adtype=adtype)
     # The parameters at which we evaluate f.
     θ = vi[:]
 
@@ -102,5 +67,12 @@ function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::
 
     return suite
 end
+function make_benchmark_suite(
+    model::Model, varinfo_choice::Symbol, adtype::Symbol, islinked::Bool
+)
+    return make_benchmark_suite(
+        Random.default_rng(), model, varinfo_choice, adtype, islinked
+    )
+end
 
-end # module
+end

benchmarks/src/DynamicPPLBenchmarks.jl renamed to ext/DynamicPPLBenchmarkToolsExt.jl

@@ -217,8 +217,11 @@ if isdefined(Base.Experimental, :register_error_hint)
     end
 end
 
-# Standard tag: Improves stacktraces
-# Ref: https://www.stochasticlifestyle.com/improved-forwarddiff-jl-stacktraces-with-package-tags/
+# DynamicPPLForwardDiffExt
+# Improves stacktraces, see https://www.stochasticlifestyle.com/improved-forwarddiff-jl-stacktraces-with-package-tags/
 struct DynamicPPLTag end
 
+# DynamicPPLBenchmarkToolsExt
+function make_benchmark_suite end
+
 end # module