add Enzyme benchmarks and incorporate some CompatHelper bits (#1056)

* Add Enzyme to benchmark tests

* Add compat entries, format

* restrict to smorgasbord for now

* add enzyme to the correct env

---------

Co-authored-by: William Moses <[email protected]>

Author: penelopeysm

benchmarks/Project.toml

@@ -7,6 +7,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
@@ -23,6 +24,7 @@ ADTypes = "1.14.0"
 BenchmarkTools = "1.6.0"
 Distributions = "0.25.117"
 DynamicPPL = "0.37"
+Enzyme = "0.13"
 ForwardDiff = "0.10.38, 1"
 LogDensityProblems = "2.1.2"
 Mooncake = "0.4"

benchmarks/benchmarks.jl

@@ -43,6 +43,7 @@ chosen_combinations = [
     ("Smorgasbord", smorgasbord_instance, :simple_dict, :forwarddiff, true),
     ("Smorgasbord", smorgasbord_instance, :typed, :reversediff, true),
     ("Smorgasbord", smorgasbord_instance, :typed, :mooncake, true),
+    ("Smorgasbord", smorgasbord_instance, :typed, :enzyme, true),
     ("Loop univariate 1k", loop_univariate1k, :typed, :mooncake, true),
     ("Multivariate 1k", multivariate1k, :typed, :mooncake, true),
     ("Loop univariate 10k", loop_univariate10k, :typed, :mooncake, true),

benchmarks/src/DynamicPPLBenchmarks.jl

@@ -13,7 +13,7 @@ using StableRNGs: StableRNG
 
 include("./Models.jl")
 using .Models: Models
-
+using Enzyme: Enzyme
 export Models, make_suite, model_dimension
 
 """
@@ -37,6 +37,10 @@ const SYMBOL_TO_BACKEND = Dict(
     :reversediff => ADTypes.AutoReverseDiff(; compile=false),
     :reversediff_compiled => ADTypes.AutoReverseDiff(; compile=true),
     :mooncake => ADTypes.AutoMooncake(; config=nothing),
+    :enzyme => ADTypes.AutoEnzyme(;
+        mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
+        function_annotation=Enzyme.Const,
+    ),
 )
 
 to_backend(x) = error("Unknown backend: $x")

docs/Project.toml

@@ -24,5 +24,6 @@ FillArrays = "0.13, 1"
 ForwardDiff = "0.10, 1"
 JET = "0.9, 0.10"
 LogDensityProblems = "2"
+MarginalLogDensities = "0.4"
 MCMCChains = "5, 6, 7"
 StableRNGs = "1"

test/Project.toml

@@ -44,6 +44,7 @@ EnzymeCore = "0.6 - 0.8"
 ForwardDiff = "0.10.12, 1"
 JET = "0.9 - 0.10.6"
 LogDensityProblems = "2"
+MarginalLogDensities = "0.4"
 MCMCChains = "7.2.1"
 MacroTools = "0.5.6"
 OrderedCollections = "1"