refactored mh inference steps with new calling convention enforced

Author: sbrantq

src/ProbProg.jl

@@ -13,6 +13,7 @@ using ..Reactant:
 using ..Compiler: @jit, @compile
 using Enzyme
 using Base: ReentrantLock
+using Random
 
 mutable struct ProbProgTrace
     fn::Union{Nothing,Function}
@@ -21,13 +22,18 @@ mutable struct ProbProgTrace
     retval::Any
     weight::Any
     subtraces::Dict{Symbol,Any}
+    rng::Union{Nothing,AbstractRNG}
 
     function ProbProgTrace(fn::Function, args::Tuple)
-        return new(fn, args, Dict{Symbol,Any}(), nothing, nothing, Dict{Symbol,Any}())
+        return new(
+            fn, args, Dict{Symbol,Any}(), nothing, nothing, Dict{Symbol,Any}(), nothing
+        )
     end
 
     function ProbProgTrace()
-        return new(nothing, (), Dict{Symbol,Any}(), nothing, nothing, Dict{Symbol,Any}())
+        return new(
+            nothing, (), Dict{Symbol,Any}(), nothing, nothing, Dict{Symbol,Any}(), nothing
+        )
     end
 end
 
@@ -587,6 +593,10 @@ function simulate(rng::AbstractRNG, f::Function, args::Vararg{Any,Nargs}) where
 
     trace = unsafe_pointer_to_objref(Ptr{Any}(Array(trace)[1]))
 
+    trace.fn = f
+    trace.args = args
+    trace.rng = rng
+
     return trace, trace.weight
 end
 
@@ -702,7 +712,7 @@ function generate(
     trace = nothing
 
     constraint_ptr = ConcreteRNumber(reinterpret(UInt64, pointer_from_objref(constraint)))
-    constrained_symbols = collect(keys(constraint))
+    constrained_symbols = Set(keys(constraint))
 
     function wrapper_fn(rng, constraint_ptr, args...)
         return generate_internal(rng, f, args...; constraint_ptr, constrained_symbols)
@@ -719,6 +729,10 @@ function generate(
 
     trace = unsafe_pointer_to_objref(Ptr{Any}(Array(trace)[1]))
 
+    trace.fn = f
+    trace.args = args
+    trace.rng = rng
+
     return trace, trace.weight
 end
 
@@ -727,7 +741,7 @@ function generate_internal(
     f::Function,
     args::Vararg{Any,Nargs};
     constraint_ptr::TracedRNumber,
-    constrained_symbols::Vector{Symbol},
+    constrained_symbols::Set{Symbol},
 ) where {Nargs}
     argprefix::Symbol = gensym("generatearg")
     resprefix::Symbol = gensym("generateresult")
@@ -947,4 +961,80 @@ end
 
 get_choices(trace::ProbProgTrace) = trace.choices
 
+const Selection = Set{Symbol}
+select(syms::Symbol...) = Set(syms)
+choicemap() = Constraint()
+const CompiledFnCache = Dict{Tuple{Type,Set{Symbol}},Any}
+
+function metropolis_hastings(
+    trace::ProbProgTrace,
+    sel::Selection;
+    compiled_cache::Union{Nothing,CompiledFnCache}=nothing,
+)
+    if trace.fn === nothing || trace.rng === nothing
+        error("MH requires a trace with fn and rng recorded (use generate to create trace)")
+    end
+
+    constraints = Dict{Symbol,Any}()
+    constrained_symbols = Set{Symbol}()
+
+    for (sym, val) in trace.choices
+        if !(sym in sel)
+            constraints[sym] = val
+            push!(constrained_symbols, sym)
+        end
+    end
+
+    cache_key = (typeof(trace.fn), constrained_symbols)
+
+    compiled_fn = nothing
+    if compiled_cache !== nothing
+        compiled_fn = get(compiled_cache, cache_key, nothing)
+    end
+
+    if compiled_fn === nothing
+        function wrapper_fn(rng, constraint_ptr, args...)
+            return generate_internal(
+                rng, trace.fn, args...; constraint_ptr, constrained_symbols
+            )
+        end
+
+        constraint_ptr = ConcreteRNumber(
+            reinterpret(UInt64, pointer_from_objref(constraints))
+        )
+
+        compiled_fn = @compile optimize = :probprog wrapper_fn(
+            trace.rng, constraint_ptr, trace.args...
+        )
+
+        if compiled_cache !== nothing
+            compiled_cache[cache_key] = compiled_fn
+        end
+    end
+
+    constraint_ptr = ConcreteRNumber(reinterpret(UInt64, pointer_from_objref(constraints)))
+
+    old_gc_state = GC.enable(false)
+    new_trace_ptr = nothing
+    try
+        new_trace_ptr, _, _ = compiled_fn(trace.rng, constraint_ptr, trace.args...)
+    finally
+        GC.enable(old_gc_state)
+    end
+
+    new_trace = unsafe_pointer_to_objref(Ptr{Any}(Array(new_trace_ptr)[1]))
+
+    new_trace.fn = trace.fn
+    new_trace.args = trace.args
+    new_trace.rng = trace.rng
+
+    log_alpha = new_trace.weight - trace.weight
+
+    if log(rand()) < log_alpha
+        return (new_trace, true)
+    else
+        return (trace, false)
+    end
+end
+
 end

test/probprog/linear_regression.jl

@@ -1,78 +1,83 @@
 using Reactant, Test, Random
-using Reactant: ProbProg
+using Reactant: ProbProg, ReactantRNG
 
 # Reference: https://www.gen.dev/docs/stable/getting_started/linear_regression/
 
 normal(rng, μ, σ, shape) = μ .+ σ .* randn(rng, shape)
 normal_logpdf(x, μ, σ, _) = -sum(log.(σ)) - sum((μ .- x) .^ 2) / (2 * σ^2)
 
-function my_model(seed, xs)
-    rng = Random.default_rng()
-    Random.seed!(rng, seed)
-
+function my_model(rng, xs)
     slope = ProbProg.sample(
-        normal, rng, 0.0, 2.0, (1,); symbol=:slope, logpdf=normal_logpdf
+        rng, normal, 0.0, 2.0, (1,); symbol=:slope, logpdf=normal_logpdf
     )
     intercept = ProbProg.sample(
-        normal, rng, 0.0, 10.0, (1,); symbol=:intercept, logpdf=normal_logpdf
+        rng, normal, 0.0, 10.0, (1,); symbol=:intercept, logpdf=normal_logpdf
     )
 
     ys = ProbProg.sample(
-        normal,
         rng,
+        normal,
         slope .* xs .+ intercept,
         1.0,
         (length(xs),);
         symbol=:ys,
         logpdf=normal_logpdf,
     )
 
-    return rng.seed, ys
+    return ys
 end
 
 function my_inference_program(xs, ys, num_iters)
     xs_r = Reactant.to_rarray(xs)
 
-    constraints = ProbProg.choicemap()
-    constraints[:ys] = [ys]
+    constraint = ProbProg.choicemap()
+    constraint[:ys] = [ys]
 
     seed = Reactant.to_rarray(UInt64[1, 4])
+    rng = ReactantRNG(seed)
 
-    trace, _ = ProbProg.generate(my_model, seed, xs_r; constraints)
-    trace.args = (trace.retval[1], trace.args[2:end]...) # TODO: this is a temporary hack
+    trace, _ = ProbProg.generate(rng, my_model, xs_r; constraint)
+
+    compiled_cache = ProbProg.CompiledFnCache()
 
     for i in 1:num_iters
-        trace, _ = ProbProg.metropolis_hastings(trace, ProbProg.select(:slope))
-        trace, _ = ProbProg.metropolis_hastings(trace, ProbProg.select(:intercept))
-        choices = ProbProg.get_choices(trace)
-        # @show i, choices[:slope], choices[:intercept]
+        trace, _ = ProbProg.metropolis_hastings(
+            trace, ProbProg.select(:slope); compiled_cache
+        )
+        trace, _ = ProbProg.metropolis_hastings(
+            trace, ProbProg.select(:intercept); compiled_cache
+        )
     end
 
     choices = ProbProg.get_choices(trace)
-    return (choices[:slope], choices[:intercept])
+    return (Array(choices[:slope][1])[1], Array(choices[:intercept][1])[1])
 end
 
 @testset "linear_regression" begin
     @testset "simulate" begin
         seed = Reactant.to_rarray(UInt64[1, 4])
-        Random.seed!(42) # For Julia side RNG
+        rng = ReactantRNG(seed)
 
         xs = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
         xs_r = Reactant.to_rarray(xs)
 
-        trace = ProbProg.simulate(my_model, seed, xs_r)
+        trace, _ = ProbProg.simulate(rng, my_model, xs_r)
 
         @test haskey(trace.choices, :slope)
         @test haskey(trace.choices, :intercept)
         @test haskey(trace.choices, :ys)
     end
 
     @testset "inference" begin
+        Random.seed!(1) # For Julia side RNG
         xs = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
         ys = [8.23, 5.87, 3.99, 2.59, 0.23, -0.66, -3.53, -6.91, -7.24, -9.90]
 
-        slope, intercept = my_inference_program(xs, ys, 5)
+        slope, intercept = my_inference_program(xs, ys, 10000)
+
+        @show slope, intercept
 
-        # @show slope, intercept
+        @test slope ≈ -2.0 rtol = 0.05
+        @test intercept ≈ 10.0 rtol = 0.05
     end
 end