From 213ec243ac7c05eab60230323c762dc57f8a2998 Mon Sep 17 00:00:00 2001 From: William Moses Date: Sat, 13 Sep 2025 19:22:22 -0500 Subject: [PATCH 1/4] Add Enzyme to benchmark tests --- Project.toml | 1 + benchmarks/benchmarks.jl | 4 ++++ benchmarks/src/DynamicPPLBenchmarks.jl | 3 ++- 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index e636199a1..ed70f1b05 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 dfcdcc271..9362e7a4a 100644 --- a/benchmarks/benchmarks.jl +++ b/benchmarks/benchmarks.jl @@ -43,13 +43,17 @@ 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), ("Multivariate 10k", multivariate10k, :typed, :mooncake, 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. diff --git a/benchmarks/src/DynamicPPLBenchmarks.jl b/benchmarks/src/DynamicPPLBenchmarks.jl index 8c5032ace..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,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(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const), ) to_backend(x) = error("Unknown backend: $x") From a93015ba88b8b665610f5a70c1be506dafe5ee1b Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 25 Sep 2025 12:47:53 +0100 Subject: [PATCH 2/4] Add compat entries, format --- benchmarks/src/DynamicPPLBenchmarks.jl | 7 +++++-- docs/Project.toml | 1 + test/Project.toml | 1 + 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/benchmarks/src/DynamicPPLBenchmarks.jl b/benchmarks/src/DynamicPPLBenchmarks.jl index 195884ee8..dcb6ac9be 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 +using Enzyme: Enzyme export Models, make_suite, model_dimension """ @@ -37,7 +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), + :enzyme => ADTypes.AutoEnzyme(; + mode=Enzyme.set_runtime_activity(Enzyme.Reverse), + function_annotation=Enzyme.Const, + ), ) to_backend(x) = error("Unknown backend: $x") diff --git a/docs/Project.toml b/docs/Project.toml index cc0be339d..47563d9a7 100644 --- a/docs/Project.toml +++ b/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" diff --git a/test/Project.toml b/test/Project.toml index 589b150f4..0317977d4 100644 --- a/test/Project.toml +++ b/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" From a54ffd3864c863e16378780f4c54120237dec745 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 25 Sep 2025 12:49:33 +0100 Subject: [PATCH 3/4] restrict to smorgasbord for now --- benchmarks/benchmarks.jl | 3 --- 1 file changed, 3 deletions(-) diff --git a/benchmarks/benchmarks.jl b/benchmarks/benchmarks.jl index 9362e7a4a..dd1140dba 100644 --- a/benchmarks/benchmarks.jl +++ b/benchmarks/benchmarks.jl @@ -49,11 +49,8 @@ chosen_combinations = [ ("Loop univariate 10k", loop_univariate10k, :typed, :mooncake, true), ("Multivariate 10k", multivariate10k, :typed, :mooncake, 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 643dd12e6c35e5fdceb87d6a732f894925e17090 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 25 Sep 2025 13:04:27 +0100 Subject: [PATCH 4/4] add enzyme to the correct env --- Project.toml | 1 - benchmarks/Project.toml | 2 ++ 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index ed70f1b05..e636199a1 100644 --- a/Project.toml +++ b/Project.toml @@ -16,7 +16,6 @@ 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/Project.toml b/benchmarks/Project.toml index 7e7b9f7cb..97afe3532 100644 --- a/benchmarks/Project.toml +++ b/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"