From 9bb1b8e04ec2243b445ebba3304c6058fbceb0eb Mon Sep 17 00:00:00 2001 From: William Moses Date: Sat, 13 Sep 2025 19:22:22 -0500 Subject: [PATCH 1/3] Add enzyme to benchmark tests --- benchmarks/benchmarks.jl | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/benchmarks/benchmarks.jl b/benchmarks/benchmarks.jl index b733d810c..d0a6c7ee8 100644 --- a/benchmarks/benchmarks.jl +++ b/benchmarks/benchmarks.jl @@ -43,13 +43,21 @@ 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), + ("Loop univariate 1k", loop_univariate1k, :typed, :enzyme, true), ("Multivariate 1k", multivariate1k, :typed, :mooncake, true), + ("Multivariate 1k", multivariate1k, :typed, :enzyme, true), ("Loop univariate 10k", loop_univariate10k, :typed, :mooncake, true), + ("Loop univariate 10k", loop_univariate10k, :typed, :enzyme, true), ("Multivariate 10k", multivariate10k, :typed, :mooncake, true), + ("Multivariate 10k", multivariate10k, :typed, :enzyme, true), ("Dynamic", Models.dynamic(), :typed, :mooncake, true), + ("Dynamic", Models.dynamic(), :typed, :enzyme, true), ("Submodel", Models.parent(randn(rng)), :typed, :mooncake, true), + ("Submodel", Models.parent(randn(rng)), :typed, :enzyme, true), ("LDA", lda_instance, :typed, :reversediff, true), + ("LDA", lda_instance, :typed, :enzyme, true), ] # Time running a model-like function that does not use DynamicPPL, as a reference point. From 54f94b839331e396e24a65a2b8d9a94bb7f4d6f3 Mon Sep 17 00:00:00 2001 From: William Moses Date: Sat, 13 Sep 2025 22:37:48 -0500 Subject: [PATCH 2/3] Add Enzyme as an automatic differentiation backend --- benchmarks/src/DynamicPPLBenchmarks.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/benchmarks/src/DynamicPPLBenchmarks.jl b/benchmarks/src/DynamicPPLBenchmarks.jl index 8c5032ace..213c6a47c 100644 --- a/benchmarks/src/DynamicPPLBenchmarks.jl +++ b/benchmarks/src/DynamicPPLBenchmarks.jl @@ -37,6 +37,7 @@ const SYMBOL_TO_BACKEND = Dict( :reversediff => ADTypes.AutoReverseDiff(; compile=false), :reversediff_compiled => ADTypes.AutoReverseDiff(; compile=true), :mooncake => ADTypes.AutoMooncake(; config=nothing), + :enzyme => ADTypes.AutoEnzyme(; config=nothing), ) to_backend(x) = error("Unknown backend: $x") From 043c3a3f16acb6df0e5faae7bb20fe5a727d784a Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Mon, 22 Sep 2025 18:35:15 -0500 Subject: [PATCH 3/3] address comments --- Project.toml | 1 + benchmarks/benchmarks.jl | 4 ---- benchmarks/src/DynamicPPLBenchmarks.jl | 4 ++-- 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/Project.toml b/Project.toml index 5f11cba3f..e0062c9ea 100644 --- a/Project.toml +++ b/Project.toml @@ -16,6 +16,7 @@ ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" diff --git a/benchmarks/benchmarks.jl b/benchmarks/benchmarks.jl index d0a6c7ee8..3857bfccd 100644 --- a/benchmarks/benchmarks.jl +++ b/benchmarks/benchmarks.jl @@ -45,13 +45,9 @@ chosen_combinations = [ ("Smorgasbord", smorgasbord_instance, :typed, :mooncake, true), ("Smorgasbord", smorgasbord_instance, :typed, :enzyme, true), ("Loop univariate 1k", loop_univariate1k, :typed, :mooncake, true), - ("Loop univariate 1k", loop_univariate1k, :typed, :enzyme, true), ("Multivariate 1k", multivariate1k, :typed, :mooncake, true), - ("Multivariate 1k", multivariate1k, :typed, :enzyme, true), ("Loop univariate 10k", loop_univariate10k, :typed, :mooncake, true), - ("Loop univariate 10k", loop_univariate10k, :typed, :enzyme, true), ("Multivariate 10k", multivariate10k, :typed, :mooncake, true), - ("Multivariate 10k", multivariate10k, :typed, :enzyme, true), ("Dynamic", Models.dynamic(), :typed, :mooncake, true), ("Dynamic", Models.dynamic(), :typed, :enzyme, true), ("Submodel", Models.parent(randn(rng)), :typed, :mooncake, true), diff --git a/benchmarks/src/DynamicPPLBenchmarks.jl b/benchmarks/src/DynamicPPLBenchmarks.jl index 213c6a47c..195884ee8 100644 --- a/benchmarks/src/DynamicPPLBenchmarks.jl +++ b/benchmarks/src/DynamicPPLBenchmarks.jl @@ -13,7 +13,7 @@ using StableRNGs: StableRNG include("./Models.jl") using .Models: Models - +import Enzyme export Models, make_suite, model_dimension """ @@ -37,7 +37,7 @@ const SYMBOL_TO_BACKEND = Dict( :reversediff => ADTypes.AutoReverseDiff(; compile=false), :reversediff_compiled => ADTypes.AutoReverseDiff(; compile=true), :mooncake => ADTypes.AutoMooncake(; config=nothing), - :enzyme => ADTypes.AutoEnzyme(; config=nothing), + :enzyme => ADTypes.AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const), ) to_backend(x) = error("Unknown backend: $x")