Make threadsafe evaluation opt-in

Author: penelopeysm

HISTORY.md

@@ -9,12 +9,41 @@
 This version provides a reimplementation of `LogDensityFunction` that provides performance improvements on the order of 2–10× for both model evaluation as well as automatic differentiation.
 Exact speedups depend on the model size: larger models have less significant speedups because the bulk of the work is done in calls to `logpdf`.
 
-For more information about how this is accomplished, please see https://github.com/TuringLang/DynamicPPL.jl/pull/1113 as well as the `src/fasteval.jl` file, which contains extensive comments.
+For more information about how this is accomplished, please see https://github.com/TuringLang/DynamicPPL.jl/pull/1113 as well as the `src/logdensityfunction.jl` file, which contains extensive comments.
 
 As a result of this change, `LogDensityFunction` no longer stores a VarInfo inside it.
 In general, if `ldf` is a `LogDensityFunction`, it is now only valid to access `ldf.model` and `ldf.adtype`.
 If you were previously relying on this behaviour, you will need to store a VarInfo separately.
 
+#### Threadsafe evaluation
+
+DynamicPPL models are by default no longer thread-safe.
+If you have threading in a model, you **must** now manually mark it as so, using:
+
+```julia
+@model f() = ...
+model = f()
+model = setthreadsafe(model, true)
+```
+
+It used to be that DynamicPPL would 'automatically' enable thread-safe evaluation if Julia was launched with more than one thread (i.e., by checking `Threads.nthreads() > 1`).
+
+The problem with this approach is that it sacrifices a huge amount of performance.
+Furthermore, it is not actually the correct approach: just because Julia has multiple threads does not mean that a particular model actually requires threadsafe evaluation.
+
+**A model requires threadsafe evaluation if, and only if, the VarInfo object used inside the model is manipulated in parallel.**
+This can occur if any of the following are inside `Threads.@threads` or other concurrency functions / macros:
+
+  - tilde-statements
+  - calls to `@addlogprob!`
+  - any direct manipulation of the special `__varinfo__` variable
+
+If you have none of these inside threaded blocks, then you do not need to mark your model as threadsafe.
+**Notably, the following do not require threadsafe evaluation:**
+
+  - Using threading for anything that does not involve VarInfo. For example, you can calculate a log-probability in parallel, and then add it using `@addlogprob!` outside of the threaded block. This does not require threadsafe evaluation.
+  - Sampling with `AbstractMCMC.MCMCThreads()`.
+
 #### Parent and leaf contexts
 
 The `DynamicPPL.NodeTrait` function has been removed.

docs/src/api.md

@@ -42,6 +42,13 @@ The context of a model can be set using [`contextualize`](@ref):
 contextualize
 ```
 
+Some models require threadsafe evaluation (see https://turinglang.org/docs/THIS_DOESNT_EXIST_YET for more information on when this is necessary).
+If this is the case, one must enable threadsafe evaluation for a model:
+
+```@docs
+setthreadsafe
+```
+
 ## Evaluation
 
 With [`rand`](@ref) one can draw samples from the prior distribution of a [`Model`](@ref).

src/DynamicPPL.jl

@@ -90,6 +90,7 @@ export AbstractVarInfo,
     Model,
     getmissings,
     getargnames,
+    setthreadsafe,
     extract_priors,
     values_as_in_model,
     # evaluation

src/compiler.jl

@@ -301,7 +301,7 @@ function model(mod, linenumbernode, expr, warn)
     modeldef = build_model_definition(expr)
 
     # Generate main body
-    modeldef[:body] = generate_mainbody(mod, modeldef[:body], warn)
+    modeldef[:body] = generate_mainbody(mod, modeldef[:body], warn, false)
 
     return build_output(modeldef, linenumbernode)
 end
@@ -346,36 +346,64 @@ Generate the body of the main evaluation function from expression `expr` and arg
 If `warn` is true, a warning is displayed if internal variables are used in the model
 definition.
 """
-generate_mainbody(mod, expr, warn) = generate_mainbody!(mod, Symbol[], expr, warn)
+generate_mainbody(mod, expr, warn, warned_about_threads_threads) =
+    generate_mainbody!(mod, Symbol[], expr, warn, warned_about_threads_threads)
 
-generate_mainbody!(mod, found, x, warn) = x
-function generate_mainbody!(mod, found, sym::Symbol, warn)
+generate_mainbody!(mod, found, x, warn, warned_about_threads_threads) = x
+function generate_mainbody!(mod, found, sym::Symbol, warn, warned_about_threads_threads)
     if warn && sym in INTERNALNAMES && sym ∉ found
         @warn "you are using the internal variable `$sym`"
         push!(found, sym)
     end
 
     return sym
 end
-function generate_mainbody!(mod, found, expr::Expr, warn)
+function generate_mainbody!(mod, found, expr::Expr, warn, warned_about_threads_threads)
     # Do not touch interpolated expressions
     expr.head === :$ && return expr.args[1]
 
+    # Flag to determine whether we've issued a warning for threadsafe macros Note that this
+    # detection is not fully correct. We can only detect the presence of a macro that has
+    # the symbol `Threads.@threads`, however, we can't detect if that *is actually*
+    # Threads.@threads from Base.Threads.
+
     # Do we don't want escaped expressions because we unfortunately
     # escape the entire body afterwards.
-    Meta.isexpr(expr, :escape) && return generate_mainbody(mod, found, expr.args[1], warn)
+    Meta.isexpr(expr, :escape) && return generate_mainbody(
+        mod, found, expr.args[1], warn, warned_about_threads_threads
+    )
 
     # If it's a macro, we expand it
     if Meta.isexpr(expr, :macrocall)
-        return generate_mainbody!(mod, found, macroexpand(mod, expr; recursive=true), warn)
+        if expr.args[1] == Expr(:., :Threads, QuoteNode(Symbol("@threads"))) &&
+            !warned_about_threads_threads
+            warned_about_threads_threads = true
+            @warn (
+                "It looks like you are using `Threads.@threads` in your model definition." *
+                "\n\nNote that since version 0.39 of DynamicPPL, threadsafe evaluation of models is disabled by default." *
+                " If you need it, you will need to explicitly enable it by creating the model, and then running `model = setthreadsafe(model, true)`." *
+                "\n\nAvoiding threadsafe evaluation can often lead to significant performance improvements. Please see https://turinglang.org/docs/THIS_PAGE_DOESNT_EXIST_YET for more details of when threadsafe evaluation is actually required."
+            )
+        end
+        return generate_mainbody!(
+            mod,
+            found,
+            macroexpand(mod, expr; recursive=true),
+            warn,
+            warned_about_threads_threads,
+        )
     end
 
     # Modify dotted tilde operators.
     args_dottilde = getargs_dottilde(expr)
     if args_dottilde !== nothing
         L, R = args_dottilde
         return generate_mainbody!(
-            mod, found, Base.remove_linenums!(generate_dot_tilde(L, R)), warn
+            mod,
+            found,
+            Base.remove_linenums!(generate_dot_tilde(L, R)),
+            warn,
+            warned_about_threads_threads,
         )
     end
 
@@ -385,8 +413,8 @@ function generate_mainbody!(mod, found, expr::Expr, warn)
         L, R = args_tilde
         return Base.remove_linenums!(
             generate_tilde(
-                generate_mainbody!(mod, found, L, warn),
-                generate_mainbody!(mod, found, R, warn),
+                generate_mainbody!(mod, found, L, warn, warned_about_threads_threads),
+                generate_mainbody!(mod, found, R, warn, warned_about_threads_threads),
             ),
         )
     end
@@ -397,13 +425,19 @@ function generate_mainbody!(mod, found, expr::Expr, warn)
         L, R = args_assign
         return Base.remove_linenums!(
             generate_assign(
-                generate_mainbody!(mod, found, L, warn),
-                generate_mainbody!(mod, found, R, warn),
+                generate_mainbody!(mod, found, L, warn, warned_about_threads_threads),
+                generate_mainbody!(mod, found, R, warn, warned_about_threads_threads),
             ),
         )
     end
 
-    return Expr(expr.head, map(x -> generate_mainbody!(mod, found, x, warn), expr.args)...)
+    return Expr(
+        expr.head,
+        map(
+            x -> generate_mainbody!(mod, found, x, warn, warned_about_threads_threads),
+            expr.args,
+        )...,
+    )
 end
 
 function generate_assign(left, right)

src/debug_utils.jl

@@ -424,8 +424,11 @@ function check_model_and_trace(
     # Perform checks before evaluating the model.
     issuccess = check_model_pre_evaluation(model)
 
-    # Force single-threaded execution.
-    _, varinfo = DynamicPPL.evaluate_threadunsafe!!(model, varinfo)
+    # TODO(penelopeysm): Implement merge, etc. for DebugAccumulator, and then perform a
+    # check on the merged accumulator, rather than checking it in the accumulate_assume
+    # calls. That way we can also support multi-threaded evaluation and use `evaluate!!`
+    # here instead of `_evaluate!!`.
+    _, varinfo = DynamicPPL._evaluate!!(model, varinfo)
 
     # Perform checks after evaluating the model.
     debug_acc = DynamicPPL.getacc(varinfo, Val(_DEBUG_ACC_NAME))

src/model.jl

@@ -1,5 +1,5 @@
 """
-    struct Model{F,argnames,defaultnames,missings,Targs,Tdefaults,Ctx<:AbstractContext}
+    struct Model{F,argnames,defaultnames,missings,Targs,Tdefaults,Ctx<:AbstractContext,Threaded}
         f::F
         args::NamedTuple{argnames,Targs}
         defaults::NamedTuple{defaultnames,Tdefaults}
@@ -17,6 +17,10 @@ An argument with a type of `Missing` will be in `missings` by default. However,
 non-traditional use-cases `missings` can be defined differently. All variables in `missings`
 are treated as random variables rather than observations.
 
+The `Threaded` type parameter indicates whether the model requires threadsafe evaluation
+(i.e., whether the model contains statements which modify the internal VarInfo that are
+executed in parallel). By default, this is set to `false`.
+
 The default arguments are used internally when constructing instances of the same model with
 different arguments.
 
@@ -33,8 +37,9 @@ julia> Model{(:y,)}(f, (x = 1.0, y = 2.0), (x = 42,)) # with special definition
 Model{typeof(f),(:x, :y),(:x,),(:y,),Tuple{Float64,Float64},Tuple{Int64}}(f, (x = 1.0, y = 2.0), (x = 42,))
 ```
 """
-struct Model{F,argnames,defaultnames,missings,Targs,Tdefaults,Ctx<:AbstractContext} <:
-       AbstractProbabilisticProgram
+struct Model{
+    F,argnames,defaultnames,missings,Targs,Tdefaults,Ctx<:AbstractContext,Threaded
+} <: AbstractProbabilisticProgram
     f::F
     args::NamedTuple{argnames,Targs}
     defaults::NamedTuple{defaultnames,Tdefaults}
@@ -46,13 +51,13 @@ struct Model{F,argnames,defaultnames,missings,Targs,Tdefaults,Ctx<:AbstractConte
     Create a model with evaluation function `f` and missing arguments overwritten by
     `missings`.
     """
-    function Model{missings}(
+    function Model{missings,Threaded}(
         f::F,
         args::NamedTuple{argnames,Targs},
         defaults::NamedTuple{defaultnames,Tdefaults},
         context::Ctx=DefaultContext(),
-    ) where {missings,F,argnames,Targs,defaultnames,Tdefaults,Ctx}
-        return new{F,argnames,defaultnames,missings,Targs,Tdefaults,Ctx}(
+    ) where {missings,F,argnames,Targs,defaultnames,Tdefaults,Ctx,Threaded}
+        return new{F,argnames,defaultnames,missings,Targs,Tdefaults,Ctx,Threaded}(
             f, args, defaults, context
         )
     end
@@ -71,18 +76,27 @@ model with different arguments.
     args::NamedTuple{argnames,Targs},
     defaults::NamedTuple{kwargnames,Tkwargs},
     context::AbstractContext=DefaultContext(),
+    threadsafe::Bool=false,
 ) where {F,argnames,Targs,kwargnames,Tkwargs}
     missing_args = Tuple(
         name for (name, typ) in zip(argnames, Targs.types) if typ <: Missing
     )
     missing_kwargs = Tuple(
         name for (name, typ) in zip(kwargnames, Tkwargs.types) if typ <: Missing
     )
-    return :(Model{$(missing_args..., missing_kwargs...)}(f, args, defaults, context))
+    return :(Model{$(missing_args..., missing_kwargs...),threadsafe}(
+        f, args, defaults, context
+    ))
 end
 
-function Model(f, args::NamedTuple, context::AbstractContext=DefaultContext(); kwargs...)
-    return Model(f, args, NamedTuple(kwargs), context)
+function Model(
+    f,
+    args::NamedTuple,
+    context::AbstractContext=DefaultContext(),
+    threadsafe=false;
+    kwargs...,
+)
+    return Model(f, args, NamedTuple(kwargs), context, threadsafe)
 end
 
 """
@@ -91,8 +105,10 @@ end
 Return a new `Model` with the same evaluation function and other arguments, but
 with its underlying context set to `context`.
 """
-function contextualize(model::Model, context::AbstractContext)
-    return Model(model.f, model.args, model.defaults, context)
+function contextualize(
+    model::Model{F,A,D,M,Ta,Td,Ctx,Threaded}, context::AbstractContext
+) where {F,A,D,M,Ta,Td,Ctx,Threaded}
+    return Model(model.f, model.args, model.defaults, context, Threaded)
 end
 
 """
@@ -105,6 +121,31 @@ function setleafcontext(model::Model, context::AbstractContext)
     return contextualize(model, setleafcontext(model.context, context))
 end
 
+"""
+    setthreadsafe(model::Model, threadsafe::Bool)
+
+Returns a new `Model` with its threadsafe flag set to `threadsafe`.
+
+Threadsafe evaluation allows for parallel execution of model statements that mutate the
+internal `VarInfo` object. For example, this is needed if tilde-statements are nested inside
+`Threads.@threads` or similar constructs.
+
+It is not needed for generic multithreaded operations that don't involve VarInfo. For
+example, calculating a log-likelihood term in parallel and then calling `@addlogprob!`
+outside of the parallel region is safe without needing to set `threadsafe=true`.
+
+It is also not needed for multithreaded sampling with AbstractMCMC's `MCMCThreads()`.
+"""
+function setthreadsafe(
+    model::Model{F,A,D,M,Ta,Td,Ctx,Threaded}, threadsafe::Bool
+) where {F,A,D,M,Ta,Td,Ctx,Threaded}
+    return if Threaded == threadsafe
+        model
+    else
+        Model{M,threadsafe}(model.f, model.args, model.defaults, model.context)
+    end
+end
+
 """
     model | (x = 1.0, ...)
 
@@ -863,16 +904,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,]
@@ -944,40 +975,14 @@ If multiple threads are available, the varinfo provided will be wrapped in a
 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)
-        evaluate_threadsafe!!(model, varinfo)
-    else
-        evaluate_threadunsafe!!(model, varinfo)
-    end
-end
-
-"""
-    evaluate_threadunsafe!!(model, varinfo)
-
-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!!(model, varinfo)
+function AbstractPPL.evaluate!!(
+    model::Model{F,A,D,M,Ta,Td,Ctx,false}, varinfo::AbstractVarInfo
+) where {F,A,D,M,Ta,Td,Ctx}
     return _evaluate!!(model, resetaccs!!(varinfo))
 end
-
-"""
-    evaluate_threadsafe!!(model, varinfo, context)
-
-Evaluate the `model` with `varinfo` wrapped inside a `ThreadSafeVarInfo`.
-
-With the wrapper, Julia's multithreading can be used for observe 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!!(model, varinfo)
+function AbstractPPL.evaluate!!(
+    model::Model{F,A,D,M,Ta,Td,Ctx,true}, varinfo::AbstractVarInfo
+) where {F,A,D,M,Ta,Td,Ctx}
     wrapper = ThreadSafeVarInfo(resetaccs!!(varinfo))
     result, wrapper_new = _evaluate!!(model, wrapper)
     # TODO(penelopeysm): If seems that if you pass a TSVI to this method, it