Document BenchmarkToolsExt

Author: penelopeysm

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:

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"

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"]
     ],

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.

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

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,