Add docstrings and tests

Author: devmotion

src/model.jl

@@ -133,21 +133,42 @@ function (model::Model)(
     context::AbstractContext = DefaultContext()
 )
     if Threads.nthreads() == 1
-        return evaluate_singlethreaded(rng, model, varinfo, sampler, context)
+        return evaluate_threadunsafe(rng, model, varinfo, sampler, context)
     else
-        return evaluate_multithreaded(rng, model, varinfo, sampler, context)
+        return evaluate_threadsafe(rng, model, varinfo, sampler, context)
     end
 end
 
-function evaluate_singlethreaded(rng, model, varinfo, sampler, context)
+"""
+    evaluate_threadunsafe(rng, model, varinfo, sampler, context)
+
+Evaluate the `model` without wrapping `varinfo` inside a `ThreadSafeVarInfo`.
+
+If the `model` makes use of Julia's multithreading this will lead to undefined behaviour.
+This method is not exposed and supposed to be used only internally in DynamicPPL.
+
+See also: [`evaluate_threadsafe`](@ref)
+"""
+function evaluate_threadunsafe(rng, model, varinfo, sampler, context)
     resetlogp!(varinfo)
     if has_eval_num(sampler)
         sampler.state.eval_num += 1
     end
     return model.f(rng, model, varinfo, sampler, context)
 end
 
-function evaluate_multithreaded(rng, model, varinfo, sampler, context)
+"""
+    evaluate_threadsafe(rng, model, varinfo, sampler, context)
+
+Evaluate the `model` with `varinfo` wrapped inside a `ThreadSafeVarInfo`.
+
+With the wrapper, Julia's multithreading can be used for assume statements in the `model`
+but parallel sampling will lead to undefined behaviour.
+This method is not exposed and supposed to be used only internally in DynamicPPL.
+
+See also: [`evaluate_threadunsafe`](@ref)
+"""
+function evaluate_threadsafe(rng, model, varinfo, sampler, context)
     resetlogp!(varinfo)
     if has_eval_num(sampler)
         sampler.state.eval_num += 1

test/threadsafe.jl

@@ -39,6 +39,7 @@
         x = rand(10_000)
 
         @model function wthreads(x)
+            global vi_ = _varinfo
             x[1] ~ Normal(0, 1)
             Threads.@threads for i in 2:length(x)
                 x[i] ~ Normal(x[i-1], 1)
@@ -48,21 +49,28 @@
         vi = VarInfo()
         wthreads(x)(vi)
         lp_w_threads = getlogp(vi)
+        if Threads.nthreads() == 1
+            @test vi_ isa VarInfo
+        else
+            @test vi_ isa DynamicPPL.ThreadSafeVarInfo
+        end
 
         println("With `@threads`:")
         println("  default:")
         @time wthreads(x)(vi)
 
-        # Ensure that we use `ThreadSafeVarInfo`.
+        # Ensure that we use `ThreadSafeVarInfo` to handle multithreaded assume statements.
+        DynamicPPL.evaluate_threadsafe(Random.GLOBAL_RNG, wthreads(x), vi,
+                                       SampleFromPrior(), DefaultContext())
         @test getlogp(vi) ≈ lp_w_threads
-        DynamicPPL.evaluate_multithreaded(Random.GLOBAL_RNG, wthreads(x), vi,
-                                          SampleFromPrior(), DefaultContext())
+        @test vi_ isa DynamicPPL.ThreadSafeVarInfo
 
-        println("  evaluate_multithreaded:")
-        @time DynamicPPL.evaluate_multithreaded(Random.GLOBAL_RNG, wthreads(x), vi,
-                                                SampleFromPrior(), DefaultContext())
+        println("  evaluate_threadsafe:")
+        @time DynamicPPL.evaluate_threadsafe(Random.GLOBAL_RNG, wthreads(x), vi,
+                                             SampleFromPrior(), DefaultContext())
 
         @model function wothreads(x)
+            global vi_ = _varinfo
             x[1] ~ Normal(0, 1)
             for i in 2:length(x)
                 x[i] ~ Normal(x[i-1], 1)
@@ -72,6 +80,11 @@
         vi = VarInfo()
         wothreads(x)(vi)
         lp_wo_threads = getlogp(vi)
+        if Threads.nthreads() == 1
+            @test vi_ isa VarInfo
+        else
+            @test vi_ isa DynamicPPL.ThreadSafeVarInfo
+        end
 
         println("Without `@threads`:")
         println("  default:")
@@ -80,12 +93,13 @@
         @test lp_w_threads ≈ lp_wo_threads
 
         # Ensure that we use `VarInfo`.
-        DynamicPPL.evaluate_singlethreaded(Random.GLOBAL_RNG, wothreads(x), vi,
-                                           SampleFromPrior(), DefaultContext())
+        DynamicPPL.evaluate_threadunsafe(Random.GLOBAL_RNG, wothreads(x), vi,
+                                         SampleFromPrior(), DefaultContext())
         @test getlogp(vi) ≈ lp_w_threads
+        @test vi_ isa VarInfo
 
-        println("  evaluate_singlethreaded:")
-        @time DynamicPPL.evaluate_singlethreaded(Random.GLOBAL_RNG, wothreads(x), vi,
-                                                 SampleFromPrior(), DefaultContext())
+        println("  evaluate_threadunsafe:")
+        @time DynamicPPL.evaluate_threadunsafe(Random.GLOBAL_RNG, wothreads(x), vi,
+                                               SampleFromPrior(), DefaultContext())
     end
 end