automated tracing of ProbProgTrace and Constraint structs

Author: sbrantq

src/Reactant.jl

@@ -246,11 +246,13 @@ include("Tracing.jl")
 include("Compiler.jl")
 
 include("Overlay.jl")
-include("probprog/ProbProg.jl")
 
 # Serialization
 include("serialization/Serialization.jl")
 
+# ProbProg
+include("probprog/ProbProg.jl")
+
 using .Compiler: @compile, @code_hlo, @code_mhlo, @jit, @code_xla, traced_getfield, compile
 export ConcreteRArray,
     ConcreteRNumber,

src/probprog/Modeling.jl

@@ -92,7 +92,7 @@ function simulate_(rng::AbstractRNG, f::Function, args::Vararg{Any,Nargs}) where
     seed_buffer = only(rng.seed.data).buffer
     GC.@preserve seed_buffer begin
         t, _, _ = compiled_fn(rng, f, args...)
-        trace = from_trace_tensor(t)
+        trace = ProbProgTrace(t)
     end
 
     return trace, trace.weight
@@ -146,37 +146,30 @@ end
 
 # Gen-like helper function.
 function generate_(
-    rng::AbstractRNG,
-    f::Function,
-    args::Vararg{Any,Nargs};
-    constraint::Constraint=Constraint(),
+    rng::AbstractRNG, constraint::Constraint, f::Function, args::Vararg{Any,Nargs}
 ) where {Nargs}
     trace = nothing
 
-    constraint_ptr = ConcreteRNumber(reinterpret(UInt64, pointer_from_objref(constraint)))
-
     constrained_addresses = extract_addresses(constraint)
 
-    function wrapper_fn(rng, constraint_ptr, args...)
-        return generate(rng, f, args...; constraint_ptr, constrained_addresses)
-    end
-
-    compiled_fn = @compile optimize = :probprog wrapper_fn(rng, constraint_ptr, args...)
+    compiled_fn = @compile optimize = :probprog generate(
+        rng, constraint, f, args...; constrained_addresses
+    )
 
     seed_buffer = only(rng.seed.data).buffer
     GC.@preserve seed_buffer constraint begin
-        t, _, _ = compiled_fn(rng, constraint_ptr, args...)
-        trace = from_trace_tensor(t)
+        t, _, _ = compiled_fn(rng, constraint, f, args...)
+        trace = ProbProgTrace(t)
     end
 
     return trace, trace.weight
 end
 
 function generate(
     rng::AbstractRNG,
+    constraint,
     f::Function,
     args::Vararg{Any,Nargs};
-    constraint_ptr::TracedRNumber,
     constrained_addresses::Set{Address},
 ) where {Nargs}
     args = (rng, args...)
@@ -193,7 +186,7 @@ function generate(
 
     constraint_val = MLIR.IR.result(
         MLIR.Dialects.builtin.unrealized_conversion_cast(
-            [TracedUtils.get_mlir_data(constraint_ptr)]; outputs=[constraint_ty]
+            [TracedUtils.get_mlir_data(constraint)]; outputs=[constraint_ty]
         ),
         1,
     )

src/probprog/Types.jl

@@ -1,4 +1,5 @@
 using Base: ReentrantLock
+using ..Reactant: AbstractConcreteNumber, AbstractConcreteArray
 
 mutable struct ProbProgTrace
     choices::Dict{Symbol,Any}
@@ -9,6 +10,9 @@ mutable struct ProbProgTrace
     function ProbProgTrace()
         return new(Dict{Symbol,Any}(), nothing, nothing, Dict{Symbol,Any}())
     end
+    function ProbProgTrace(x::Union{AbstractConcreteNumber,AbstractConcreteArray})
+        return convert(ProbProgTrace, x)
+    end
 end
 
 struct Address
@@ -42,6 +46,9 @@ mutable struct Constraint <: AbstractDict{Address,Any}
 
     Constraint() = new(Dict{Address,Any}())
     Constraint(d::Dict{Address,Any}) = new(d)
+    function Constraint(x::Union{AbstractConcreteNumber,AbstractConcreteArray})
+        return convert(Constraint, x)
+    end
 end
 
 Base.getindex(c::Constraint, k::Address) = c.dict[k]

src/probprog/Utils.jl

@@ -3,14 +3,20 @@ using ..Reactant:
     TracedUtils,
     Ops,
     TracedRArray,
+    TracedRNumber,
     Compiler,
     OrderedIdDict,
-    make_tracer,
     TracedToTypes,
-    TracedTrack,
     TracedType,
-    TracedSetPath
-import ..Reactant: promote_to
+    TracedTrack,
+    TracedSetPath,
+    ConcreteToTraced,
+    AbstractConcreteArray,
+    XLA,
+    Sharding,
+    to_number
+import ..Reactant: promote_to, make_tracer
+import ..Compiler: donate_argument!
 
 """
     process_probprog_function(f, args, op_name)
@@ -171,30 +177,59 @@ function process_probprog_outputs(
     return traced_result
 end
 
-to_trace_tensor(t::ProbProgTrace) = promote_to(TracedRArray{UInt64,0}, t)
+function promote_to(::Type{TracedRArray{UInt64,0}}, t::Union{ProbProgTrace,Constraint})
+    return Ops.fill(reinterpret(UInt64, pointer_from_objref(t)), Int64[])
+end
 
-function from_trace_tensor(trace_tensor)
-    while !isready(trace_tensor)
+function Base.convert(
+    ::Type{T}, x::AbstractConcreteArray
+) where {T<:Union{ProbProgTrace,Constraint}}
+    while !isready(x)
         yield()
     end
-    return unsafe_pointer_to_objref(Ptr{Any}(Array(trace_tensor)[1]))::ProbProgTrace
+    return unsafe_pointer_to_objref(Ptr{Any}(collect(x)[1]))::T
 end
 
-function promote_to(::Type{TracedRArray{UInt64,0}}, t::ProbProgTrace)
-    ptr = reinterpret(UInt64, pointer_from_objref(t))
-    return Ops.fill(ptr, Int64[])
+function Base.convert(
+    ::Type{T}, x::AbstractConcreteNumber
+) where {T<:Union{ProbProgTrace,Constraint}}
+    while !isready(x)
+        yield()
+    end
+    return unsafe_pointer_to_objref(Ptr{Any}(to_number(x)))::T
 end
 
-to_constraint_tensor(c::Constraint) = promote_to(TracedRArray{UInt64,0}, c)
-
-function from_constraint_tensor(constraint_tensor)
-    while !isready(constraint_tensor)
-        yield()
+function Base.getproperty(t::Union{ProbProgTrace,Constraint}, s::Symbol)
+    if s === :data
+        return ConcreteRNumber(reinterpret(UInt64, pointer_from_objref(t))).data
+    else
+        return getfield(t, s)
     end
-    return unsafe_pointer_to_objref(Ptr{Any}(Array(constraint_tensor)[1]))::Constraint
 end
 
-function promote_to(::Type{TracedRArray{UInt64,0}}, c::Constraint)
-    ptr = reinterpret(UInt64, pointer_from_objref(c))
-    return Ops.fill(ptr, Int64[])
+function donate_argument!(
+    ::Any, ::Union{ProbProgTrace,Constraint}, ::Int, ::Any, ::Any
+)
+    return nothing
+end
+
+Base.@nospecializeinfer function make_tracer(
+    seen,
+    @nospecialize(prev::Union{ProbProgTrace,Constraint}),
+    @nospecialize(path),
+    mode;
+    @nospecialize(sharding = Sharding.NoSharding()),
+    kwargs...,
+)
+    if mode == ConcreteToTraced
+        haskey(seen, prev) && return seen[prev]::TracedRNumber{UInt64}
+        result = TracedRNumber{UInt64}((path,), nothing)
+        seen[prev] = result
+        return result
+    elseif mode == TracedToTypes
+        push!(path, typeof(prev))
+        return nothing
+    else
+        error("Unsupported mode for $(typeof(prev)): $mode")
+    end
 end

test/probprog/generate.jl

@@ -34,7 +34,7 @@ end
         rng = ReactantRNG(seed)
         μ = Reactant.ConcreteRNumber(0.0)
         σ = Reactant.ConcreteRNumber(1.0)
-        trace, weight = ProbProg.generate_(rng, model, μ, σ, shape)
+        trace, weight = ProbProg.generate_(rng, ProbProg.Constraint(), model, μ, σ, shape)
         @test mean(trace.retval[1]) ≈ 0.0 atol = 0.05 rtol = 0.05
     end
 
@@ -47,7 +47,7 @@ end
 
         constraint = ProbProg.Constraint(:s => (fill(0.1, shape),))
 
-        trace, weight = ProbProg.generate_(rng, model, μ, σ, shape; constraint)
+        trace, weight = ProbProg.generate_(rng, constraint, model, μ, σ, shape)
 
         @test trace.choices[:s][1] == constraint[ProbProg.Address(:s)][1]
 
@@ -72,7 +72,7 @@ end
             :u => :y => (fill(0.3, shape),),
         )
 
-        trace, weight = ProbProg.generate_(rng, nested_model, μ, σ, shape; constraint)
+        trace, weight = ProbProg.generate_(rng, constraint, nested_model, μ, σ, shape)
 
         @test trace.choices[:s][1] == fill(0.1, shape)
         @test trace.subtraces[:t].choices[:x][1] == fill(0.2, shape)
@@ -108,33 +108,24 @@ end
 
         constrained_addresses = ProbProg.extract_addresses(constraint1)
 
-        constraint_ptr1 = Reactant.ConcreteRNumber(
-            reinterpret(UInt64, pointer_from_objref(constraint1))
+        compiled_fn = @compile optimize = :probprog ProbProg.generate(
+            rng, constraint1, model, μ, σ, shape; constrained_addresses
         )
 
-        wrapper_fn(rng, constraint_ptr, μ, σ) = ProbProg.generate(
-            rng, model, μ, σ, shape; constraint_ptr, constrained_addresses
-        )
-
-        compiled_fn = @compile optimize = :probprog wrapper_fn(rng, constraint_ptr1, μ, σ)
-
         trace1 = nothing
         seed_buffer = only(rng.seed.data).buffer
         GC.@preserve seed_buffer constraint1 begin
-            trace1, _ = compiled_fn(rng, constraint_ptr1, μ, σ)
-            trace1 = ProbProg.from_trace_tensor(trace1)
+            trace1, _ = compiled_fn(rng, constraint1, model, μ, σ, shape)
+            trace1 = ProbProg.ProbProgTrace(trace1)
         end
 
         constraint2 = ProbProg.Constraint(:s => (fill(0.2, shape),))
-        constraint_ptr2 = Reactant.ConcreteRNumber(
-            reinterpret(UInt64, pointer_from_objref(constraint2))
-        )
 
         trace2 = nothing
         seed_buffer = only(rng.seed.data).buffer
         GC.@preserve seed_buffer constraint2 begin
-            trace2, _ = compiled_fn(rng, constraint_ptr2, μ, σ)
-            trace2 = ProbProg.from_trace_tensor(trace2)
+            trace2, _ = compiled_fn(rng, constraint2, model, μ, σ, shape)
+            trace2 = ProbProg.ProbProgTrace(trace2)
         end
 
         @test trace1.choices[:s][1] != trace2.choices[:s][1]

test/probprog/hmc.jl

@@ -36,16 +36,10 @@ function hmc_program(
     num_steps,
     mass,
     initial_momentum,
-    constraint_ptr,
+    constraint,
     constrained_addresses,
 )
-    t, _, _ = ProbProg.generate(
-        rng,
-        model,
-        xs;
-        constraint_ptr=constraint_ptr,
-        constrained_addresses=constrained_addresses,
-    )
+    t, _, _ = ProbProg.generate(rng, constraint, model, xs; constrained_addresses)
 
     t, accepted, _ = ProbProg.hmc(
         rng,
@@ -73,7 +67,6 @@ end
         :param_a => ([0.0],), :param_b => ([0.0],), :ys_a => (ys_a,), :ys_b => (ys_b,)
     )
     constrained_addresses = ProbProg.extract_addresses(obs)
-    constraint_ptr = ConcreteRNumber(reinterpret(UInt64, pointer_from_objref(obs)))
 
     step_size = ConcreteRNumber(0.001)
     num_steps_compile = ConcreteRNumber(1000)
@@ -89,7 +82,7 @@ end
         num_steps_compile,
         mass,
         initial_momentum,
-        constraint_ptr,
+        obs,
         constrained_addresses,
     )
     @test contains(repr(code), "enzyme_probprog_get_flattened_samples_from_trace")
@@ -106,7 +99,7 @@ end
             num_steps_compile,
             mass,
             initial_momentum,
-            constraint_ptr,
+            obs,
             constrained_addresses,
         )
     end
@@ -116,18 +109,18 @@ end
     trace = nothing
     GC.@preserve seed_buffer obs begin
         run_time_s = @elapsed begin
-            trace_ptr, _ = compiled_fn(
+            trace, _ = compiled_fn(
                 rng,
                 model,
                 xs,
                 step_size,
                 num_steps_run,
                 mass,
                 initial_momentum,
-                constraint_ptr,
+                obs,
                 constrained_addresses,
             )
-            trace = ProbProg.from_trace_tensor(trace_ptr)
+            trace = ProbProg.ProbProgTrace(trace)
         end
         println("HMC run time: $(round(run_time_s * 1000, digits=2)) ms")
     end