Allow opting out of TSVI

Author: penelopeysm

docs/src/api.md

@@ -110,6 +110,12 @@ Similarly, we can revert this with [`DynamicPPL.unfix`](@ref), i.e. return the v
 DynamicPPL.unfix
 ```
 
+## Controlling threadsafe evaluation
+
+```@docs
+DynamicPPL.set_threadsafe_eval!
+```
+
 ## Predicting
 
 DynamicPPL provides functionality for generating samples from the posterior predictive distribution through the `predict` function. This allows you to use posterior parameter samples to generate predictions for unobserved data points.

src/DynamicPPL.jl

@@ -92,6 +92,7 @@ export AbstractVarInfo,
     getargnames,
     extract_priors,
     values_as_in_model,
+    set_threadsafe_eval!,
     # LogDensityFunction
     LogDensityFunction,
     # Contexts

src/model.jl

@@ -1,3 +1,24 @@
+const USE_THREADSAFE_EVAL = Ref(Threads.nthreads() > 1)
+
+"""
+    DynamicPPL.set_threadsafe_eval!(val::Bool)
+
+Enable or disable threadsafe model evaluation globally. By default, threadsafe evaluation is
+used whenever Julia is run with multiple threads.
+
+However, this is not necessary for the vast majority of DynamicPPL models. **In particular,
+use of threaded sampling with MCMCChains alone does NOT require threadsafe evaluation.**
+Threadsafe evaluation is only needed when manipulating `VarInfo` objects in parallel, e.g.
+when using `x ~ dist` statements inside `Threads.@threads` blocks.
+
+If you do not need threadsafe evaluation, disabling it can lead to significant performance
+improvements.
+"""
+function set_threadsafe_eval!(val::Bool)
+    USE_THREADSAFE_EVAL[] = val
+    return nothing
+end
+
 """
     struct Model{F,argnames,defaultnames,missings,Targs,Tdefaults,Ctx<:AbstractContext}
         f::F
@@ -863,16 +884,6 @@ function (model::Model)(rng::Random.AbstractRNG, varinfo::AbstractVarInfo=VarInf
     return first(init!!(rng, model, varinfo))
 end
 
-"""
-    use_threadsafe_eval(context::AbstractContext, varinfo::AbstractVarInfo)
-
-Return `true` if evaluation of a model using `context` and `varinfo` should
-wrap `varinfo` in `ThreadSafeVarInfo`, i.e. threadsafe evaluation, and `false` otherwise.
-"""
-function use_threadsafe_eval(context::AbstractContext, varinfo::AbstractVarInfo)
-    return Threads.nthreads() > 1
-end
-
 """
     init!!(
         [rng::Random.AbstractRNG,]
@@ -912,14 +923,14 @@ end
 
 Evaluate the `model` with the given `varinfo`.
 
-If multiple threads are available, the varinfo provided will be wrapped in a
+If threadsafe evaluation is enabled, the varinfo provided will be wrapped in a
 `ThreadSafeVarInfo` before evaluation.
 
 Returns a tuple of the model's return value, plus the updated `varinfo`
 (unwrapped if necessary).
 """
 function AbstractPPL.evaluate!!(model::Model, varinfo::AbstractVarInfo)
-    return if use_threadsafe_eval(model.context, varinfo)
+    return if DynamicPPL.USE_THREADSAFE_EVAL[]
         evaluate_threadsafe!!(model, varinfo)
     else
         evaluate_threadunsafe!!(model, varinfo)

test/threadsafe.jl

@@ -1,4 +1,26 @@
 @testset "threadsafe.jl" begin
+    @testset "set threadsafe eval" begin
+        # A dummy model that lets us see what type of VarInfo is being used for evaluation.
+        @model function find_out_varinfo_type()
+            x ~ Normal()
+            return typeof(__varinfo__)
+        end
+        model = find_out_varinfo_type()
+
+        # Check the default.
+        @test DynamicPPL.USE_THREADSAFE_EVAL[] == (Threads.nthreads() > 1)
+        # Disable it.
+        DynamicPPL.set_threadsafe_eval!(false)
+        @test DynamicPPL.USE_THREADSAFE_EVAL[] == false
+        @test !(model() <: DynamicPPL.ThreadSafeVarInfo)
+        # Enable it.
+        DynamicPPL.set_threadsafe_eval!(true)
+        @test DynamicPPL.USE_THREADSAFE_EVAL[] == true
+        @test model() <: DynamicPPL.ThreadSafeVarInfo
+        # Reset to default to avoid messing with other tests.
+        DynamicPPL.set_threadsafe_eval!(Threads.nthreads() > 1)
+    end
+
     @testset "constructor" begin
         vi = VarInfo(gdemo_default)
         threadsafe_vi = @inferred DynamicPPL.ThreadSafeVarInfo(vi)