Merge remote-tracking branch 'origin/mt/fixes_threaded' into fixes_threaded

Author: devmotion

src/compat/ad.jl

@@ -9,16 +9,3 @@ ZygoteRules.@adjoint function push!(
 )
     return push!(vi, vn, r, dist, gidset), _ -> nothing
 end
-
-# Multithreaded evaluation is not compatible with Zygote.
-ZygoteRules.@adjoint function (model::Model)(
-    vi::AbstractVarInfo,
-    spl::AbstractSampler,
-    ctx::AbstractContext
-)
-    function evaluate(vi, spl, ctx)
-        return evaluate_singlethreaded(model, vi, spl, ctx)
-    end
-    return ZygoteRules.pullback(evaluate, vi, spl, ctx)
-end
-

src/model.jl

@@ -154,7 +154,7 @@ function evaluate_multithreaded(model, varinfo, sampler, context)
     end
     wrapper = ThreadSafeVarInfo(varinfo)
     result = model.f(model, wrapper, sampler, context)
-    acclogp!(varinfo, sum(wrapper.logps))
+    setlogp!(varinfo, getlogp(wrapper))
     return result
 end
 

src/threadsafe.jl

@@ -9,29 +9,33 @@ struct ThreadSafeVarInfo{V<:AbstractVarInfo,L} <: AbstractVarInfo
     logps::L
 end
 function ThreadSafeVarInfo(vi::AbstractVarInfo)
-    return ThreadSafeVarInfo(vi, [zero(getlogp(vi)) for _ in 1:Threads.nthreads()])
+    return ThreadSafeVarInfo(vi, [Ref(zero(getlogp(vi))) for _ in 1:Threads.nthreads()])
 end
 ThreadSafeVarInfo(vi::ThreadSafeVarInfo) = vi
 
 # Instead of updating the log probability of the underlying variables we
 # just update the array of log probabilities.
 function acclogp!(vi::ThreadSafeVarInfo, logp)
-    vi.logps[Threads.threadid()] += logp
+    vi.logps[Threads.threadid()][] += logp
     return vi
 end
 
 # The current log probability of the variables has to be computed from
 # both the wrapped variables and the thread-specific log probabilities.
-getlogp(vi::ThreadSafeVarInfo) = getlogp(vi.varinfo) + sum(vi.logps)
+getlogp(vi::ThreadSafeVarInfo) = getlogp(vi.varinfo) + sum(getindex, vi.logps)
 
 # TODO: Make remaining methods thread-safe.
 
 function resetlogp!(vi::ThreadSafeVarInfo)
-    fill!(vi.logps, zero(getlogp(vi)))
+    for x in vi.logps
+        x[] = zero(x[])
+    end
     return resetlogp!(vi.varinfo)
 end
 function setlogp!(vi::ThreadSafeVarInfo, logp)
-    fill!(vi.logps, zero(logp))
+    for x in vi.logps
+        x[] = zero(x[])
+    end
     return setlogp!(vi.varinfo, logp)
 end
 

test/compat/ad.jl

@@ -62,6 +62,6 @@ using Tracker
 
     y, back = Zygote.pullback(logp_model, x)
     @test y ≈ lp
-    @test back(1) ≈ grad
+    @test back(1)[1] ≈ grad
 end