From bbc82eff97e580dd10f2221aa37d4af2ec25069c Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Tue, 25 Mar 2025 18:53:57 +0000 Subject: [PATCH 1/3] Move code to BenchmarkToolsExt --- Project.toml | 3 + benchmarks/{src => }/Models.jl | 0 benchmarks/benchmarks.jl | 37 +++++++-- .../DynamicPPLBenchmarkToolsExt.jl | 82 ++++++------------- src/DynamicPPL.jl | 7 +- 5 files changed, 63 insertions(+), 66 deletions(-) rename benchmarks/{src => }/Models.jl (100%) rename benchmarks/src/DynamicPPLBenchmarks.jl => ext/DynamicPPLBenchmarkToolsExt.jl (50%) diff --git a/Project.toml b/Project.toml index d5185d727..e026c3121 100644 --- a/Project.toml +++ b/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" diff --git a/benchmarks/src/Models.jl b/benchmarks/Models.jl similarity index 100% rename from benchmarks/src/Models.jl rename to benchmarks/Models.jl diff --git a/benchmarks/benchmarks.jl b/benchmarks/benchmarks.jl index 89b65d2de..d971e8472 100644 --- a/benchmarks/benchmarks.jl +++ b/benchmarks/benchmarks.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,14 +69,14 @@ 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 @@ -65,7 +84,7 @@ 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 diff --git a/benchmarks/src/DynamicPPLBenchmarks.jl b/ext/DynamicPPLBenchmarkToolsExt.jl similarity index 50% rename from benchmarks/src/DynamicPPLBenchmarks.jl rename to ext/DynamicPPLBenchmarkToolsExt.jl index 4c73bf355..bfd79065c 100644 --- a/benchmarks/src/DynamicPPLBenchmarks.jl +++ b/ext/DynamicPPLBenchmarkToolsExt.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 diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl index 50fe0edc7..d27720a6a 100644 --- a/src/DynamicPPL.jl +++ b/src/DynamicPPL.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 From c88e6f1796b1a8ed9ada3e7dd02e3e940e4311f6 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Tue, 25 Mar 2025 19:21:07 +0000 Subject: [PATCH 2/3] Document BenchmarkToolsExt --- HISTORY.md | 21 +++++++++++++++++++++ docs/Project.toml | 3 +++ docs/make.jl | 10 ++++++++-- docs/src/api.md | 18 ++++++++++++++++++ ext/DynamicPPLBenchmarkToolsExt.jl | 17 ++++++++++------- src/DynamicPPL.jl | 2 ++ 6 files changed, 62 insertions(+), 9 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 3ea8071f3..07977cc2e 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,5 +1,26 @@ # DynamicPPL Changelog +## 0.36.0 + +### BenchmarkTools extension + +DynamicPPL now contains a BenchmarkTools extension. +If you have both packages loaded, then it exports a single function `make_benchmark_suite`, which returns a `BenchmarkTools.BenchmarkGroup` object. +Running this will give you information about the time taken to evaluate the log density of the model ("evaluation"), as well as the time taken to evaluate its gradient. + +Note that benchmarking both of these is in general much easier since the changes in 0.35.0. +If you just want to run a single model, the easiest way is to do this: + +```julia +@model f() = ... +ldf = LogDensityFunction(f(); adtype=AutoMyBackend()) +params = ldf.varinfo[:] +@btime LogDensityProblems.logdensity($ldf, params) +@btime LogDensityProblems.logdensity_and_gradient($ldf, params) +``` + +The `make_benchmark_suite` function is essentially a nice wrapper around this. + ## 0.35.5 Several internal methods have been removed: diff --git a/docs/Project.toml b/docs/Project.toml index fa57f2c1c..074be5eac 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,9 +1,11 @@ [deps] Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" DocumenterMermaid = "a078cd44-4d9c-4618-b545-3ab9d77f9177" +DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" @@ -13,6 +15,7 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" [compat] Accessors = "0.1" +BenchmarkTools = "1" DataStructures = "0.18" Distributions = "0.25" Documenter = "1" diff --git a/docs/make.jl b/docs/make.jl index c69b72fb8..b887e8a3e 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -9,8 +9,10 @@ using DynamicPPL: AbstractPPL # consistent with that. using Distributions using DocumenterMermaid -# load MCMCChains package extension to make `predict` available + +# To get docstrings from package extensions using MCMCChains +using BenchmarkTools # Doctest setup DocMeta.setdocmeta!( @@ -22,7 +24,11 @@ makedocs(; # The API index.html page is fairly large, and violates the default HTML page size # threshold of 200KiB, so we double that. format=Documenter.HTML(; size_threshold=2^10 * 400), - modules=[DynamicPPL, Base.get_extension(DynamicPPL, :DynamicPPLMCMCChainsExt)], + modules=[ + DynamicPPL, + Base.get_extension(DynamicPPL, :DynamicPPLMCMCChainsExt), + Base.get_extension(DynamicPPL, :DynamicPPLBenchmarkToolsExt), + ], pages=[ "Home" => "index.md", "API" => "api.md", "Internals" => ["internals/varinfo.md"] ], diff --git a/docs/src/api.md b/docs/src/api.md index 9c8249c97..670d3adae 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -245,6 +245,24 @@ DynamicPPL.TestUtils.update_values!! DynamicPPL.TestUtils.test_values ``` +## Benchmarking Utilities + +If you have `BenchmarkTools` loaded, this function will be available: + +```@docs +DynamicPPL.make_benchmark_suite +``` + +For more fine-grained control over this, you can construct a [`LogDensityFunction`](@ref) yourself and run something along the lines of: + +```julia +# set up your model and varinfo here +ldf = LogDensityFunction(model, varinfo; adtype=adtype) +params = varinfo[:] +@benchmark LogDensityProblems.logdensity($ldf, params) +@benchmark LogDensityProblems.logdensity_with_gradient($ldf, params) +``` + ## Debugging Utilities DynamicPPL provides a few methods for checking validity of a model-definition. diff --git a/ext/DynamicPPLBenchmarkToolsExt.jl b/ext/DynamicPPLBenchmarkToolsExt.jl index bfd79065c..4f5dac767 100644 --- a/ext/DynamicPPLBenchmarkToolsExt.jl +++ b/ext/DynamicPPLBenchmarkToolsExt.jl @@ -16,14 +16,17 @@ using Random: Random Create a benchmark suite for `model` using the selected varinfo type and AD backend. Available varinfo choices: - • `:untyped` → uses `VarInfo()` - • `:typed` → uses `VarInfo(model)` - • `:simple_namedtuple` → uses `SimpleVarInfo{Float64}(model())` - • `:simple_dict` → builds a `SimpleVarInfo{Float64}` from a Dict (pre-populated with the model’s outputs) +- `:untyped` → uses `VarInfo()` +- `:typed` → uses `VarInfo(model)` +- `:simple_namedtuple` → uses `SimpleVarInfo{Float64}(model())` +- `:simple_dict` → builds a `SimpleVarInfo{Float64}` from a Dict (pre-populated with the model’s outputs) + +Note that to run AD with a given backend, you will also need to load the backend yourself. For example, +if `adtype` is `AutoForwardDiff()`, you will need to `import ForwardDiff`. `islinked` determines whether to link the VarInfo for evaluation. """ -function make_benchmark_suite( +function DynamicPPL.make_benchmark_suite( rng::Random.AbstractRNG, model::Model, varinfo_choice::Symbol, @@ -67,10 +70,10 @@ function make_benchmark_suite( return suite end -function make_benchmark_suite( +function DynamicPPL.make_benchmark_suite( model::Model, varinfo_choice::Symbol, adtype::Symbol, islinked::Bool ) - return make_benchmark_suite( + return DynamicPPL.make_benchmark_suite( Random.default_rng(), model, varinfo_choice, adtype, islinked ) end diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl index d27720a6a..901e80666 100644 --- a/src/DynamicPPL.jl +++ b/src/DynamicPPL.jl @@ -125,6 +125,8 @@ export AbstractVarInfo, value_iterator_from_chain, check_model, check_model_and_trace, + # Benchmarking + make_benchmark_suite, # Deprecated. @logprob_str, @prob_str, From e9966ffadcac00a83eaa09bd699aec6aed810cb6 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Tue, 25 Mar 2025 19:26:13 +0000 Subject: [PATCH 3/3] Update CI benchmark script --- benchmarks/benchmarks.jl | 48 ++++++++++++++++++++++++++-------------- 1 file changed, 32 insertions(+), 16 deletions(-) diff --git a/benchmarks/benchmarks.jl b/benchmarks/benchmarks.jl index d971e8472..fc064e0a6 100644 --- a/benchmarks/benchmarks.jl +++ b/benchmarks/benchmarks.jl @@ -3,6 +3,7 @@ using Pkg Pkg.develop(; path=joinpath(@__DIR__, "..")) using DynamicPPL: DynamicPPL, make_benchmark_suite, VarInfo +using ADTypes using BenchmarkTools: @benchmark, median, run using PrettyTables: PrettyTables, ft_printf using ForwardDiff: ForwardDiff @@ -48,6 +49,21 @@ lda_instance = begin Models.lda(2, d, w) end +# AD types setup +fd = AutoForwardDiff() +rd = AutoReverseDiff() +mc = AutoMooncake(; config=nothing) +""" + get_adtype_shortname(adtype::ADTypes.AbstractADType) + +Get the package name that corresponds to the the AD backend `adtype`. Only used +for pretty-printing. +""" +get_adtype_shortname(::AutoMooncake) = "Mooncake" +get_adtype_shortname(::AutoForwardDiff) = "ForwardDiff" +get_adtype_shortname(::AutoReverseDiff{false}) = "ReverseDiff" +get_adtype_shortname(::AutoReverseDiff{true}) = "ReverseDiff:Compiled" + # Specify the combinations to test: # (Model Name, model instance, VarInfo choice, AD backend, linked) chosen_combinations = [ @@ -55,22 +71,22 @@ chosen_combinations = [ "Simple assume observe", Models.simple_assume_observe(randn(StableRNG(23))), :typed, - :forwarddiff, + fd, false, ), - ("Smorgasbord", smorgasbord_instance, :typed, :forwarddiff, false), - ("Smorgasbord", smorgasbord_instance, :simple_namedtuple, :forwarddiff, true), - ("Smorgasbord", smorgasbord_instance, :untyped, :forwarddiff, true), - ("Smorgasbord", smorgasbord_instance, :simple_dict, :forwarddiff, true), - ("Smorgasbord", smorgasbord_instance, :typed, :reversediff, true), - ("Smorgasbord", smorgasbord_instance, :typed, :mooncake, true), - ("Loop univariate 1k", loop_univariate1k, :typed, :mooncake, true), - ("Multivariate 1k", multivariate1k, :typed, :mooncake, true), - ("Loop univariate 10k", loop_univariate10k, :typed, :mooncake, true), - ("Multivariate 10k", multivariate10k, :typed, :mooncake, true), - ("Dynamic", Models.dynamic(), :typed, :mooncake, true), - ("Submodel", Models.parent(randn(StableRNG(23))), :typed, :mooncake, true), - ("LDA", lda_instance, :typed, :reversediff, true), + ("Smorgasbord", smorgasbord_instance, :typed, fd, false), + ("Smorgasbord", smorgasbord_instance, :simple_namedtuple, fd, true), + ("Smorgasbord", smorgasbord_instance, :untyped, fd, true), + ("Smorgasbord", smorgasbord_instance, :simple_dict, fd, true), + ("Smorgasbord", smorgasbord_instance, :typed, rd, true), + ("Smorgasbord", smorgasbord_instance, :typed, mc, true), + ("Loop univariate 1k", loop_univariate1k, :typed, mc, true), + ("Multivariate 1k", multivariate1k, :typed, mc, true), + ("Loop univariate 10k", loop_univariate10k, :typed, mc, true), + ("Multivariate 10k", multivariate10k, :typed, mc, true), + ("Dynamic", Models.dynamic(), :typed, mc, true), + ("Submodel", Models.parent(randn(StableRNG(23))), :typed, mc, true), + ("LDA", lda_instance, :typed, rd, true), ] # Time running a model-like function that does not use DynamicPPL, as a reference point. @@ -83,7 +99,7 @@ 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" + @info "Running benchmark for $model_name / $varinfo_choice / $(get_adtype_shortname(adbackend))" suite = make_benchmark_suite(StableRNG(23), model, varinfo_choice, adbackend, islinked) results = run(suite) eval_time = median(results["evaluation"]).time @@ -95,7 +111,7 @@ for (model_name, model, varinfo_choice, adbackend, islinked) in chosen_combinati ( model_name, model_dimension(model, islinked), - string(adbackend), + get_adtype_shortname(adbackend), string(varinfo_choice), islinked, relative_eval_time,