generate op fixup: replacing constrained_symbols with constrained_addresses. Address struct are now the key of Constraint dicts

Author: sbrantq

src/probprog/FFI.jl

@@ -157,7 +157,14 @@ function getSampleFromConstraint(
     shape_ptr_array = unsafe_wrap(Array, shape_ptr_array, num_samples)
     sample_ptr_array = unsafe_wrap(Array, sample_ptr_array, num_samples)
 
-    tostore = get(constraint, symbol, nothing)
+    tostore = get(constraint, Address(symbol), nothing)
+
+    if tostore === nothing
+        @ccall printf(
+            "No constraint found for symbol: %s\n"::Cstring, string(symbol)::Cstring
+        )::Cvoid
+        return nothing
+    end
 
     for i in 1:num_samples
         ndims = ndims_array[i]
@@ -228,6 +235,37 @@ function getSampleFromConstraint(
     return nothing
 end
 
+function getSubconstraint(
+    constraint_ptr_ptr::Ptr{Ptr{Any}},
+    symbol_ptr_ptr::Ptr{Ptr{Any}},
+    subconstraint_ptr_ptr::Ptr{Ptr{Any}},
+)
+    constraint = unsafe_pointer_to_objref(unsafe_load(constraint_ptr_ptr))::Constraint
+    symbol = unsafe_pointer_to_objref(unsafe_load(symbol_ptr_ptr))::Symbol
+
+    subconstraint = Constraint()
+
+    for (key, value) in constraint
+        if key.path[1] == symbol
+            @assert isa(key, Address) "Expected Address type for constraint key"
+            @assert length(key.path) > 1 "Expected composite address with length > 1"
+            tail_address = Address(key.path[2:end])
+            subconstraint[tail_address] = value
+        end
+    end
+
+    if isempty(subconstraint)
+        @ccall printf(
+            "No subconstraint found for symbol: %s\n"::Cstring, string(symbol)::Cstring
+        )::Cvoid
+        return nothing
+    end
+
+    _keepalive!(subconstraint)
+    unsafe_store!(subconstraint_ptr_ptr, pointer_from_objref(subconstraint))
+    return nothing
+end
+
 function __init__()
     init_trace_ptr = @cfunction(initTrace, Cvoid, (Ptr{Ptr{Any}},))
     @ccall MLIR.API.mlir_c.EnzymeJaXMapSymbol(
@@ -297,5 +335,12 @@ function __init__()
         get_sample_from_constraint_ptr::Ptr{Cvoid},
     )::Cvoid
 
+    get_subconstraint_ptr = @cfunction(
+        getSubconstraint, Cvoid, (Ptr{Ptr{Any}}, Ptr{Ptr{Any}}, Ptr{Ptr{Any}})
+    )
+    @ccall MLIR.API.mlir_c.EnzymeJaXMapSymbol(
+        :enzyme_probprog_get_subconstraint::Cstring, get_subconstraint_ptr::Ptr{Cvoid}
+    )::Cvoid
+
     return nothing
 end

src/probprog/Modeling.jl

@@ -349,15 +349,16 @@ function generate(
     rng::AbstractRNG,
     f::Function,
     args::Vararg{Any,Nargs};
-    constraint::Constraint=Dict{Symbol,Any}(),
+    constraint::Constraint=Constraint(),
 ) where {Nargs}
     trace = nothing
 
     constraint_ptr = ConcreteRNumber(reinterpret(UInt64, pointer_from_objref(constraint)))
-    constrained_symbols = Set(keys(constraint))
+
+    constrained_addresses = _extract_addresses(constraint)
 
     function wrapper_fn(rng, constraint_ptr, args...)
-        return generate_internal(rng, f, args...; constraint_ptr, constrained_symbols)
+        return generate_internal(rng, f, args...; constraint_ptr, constrained_addresses)
     end
 
     compiled_fn = @compile optimize = :probprog wrapper_fn(rng, constraint_ptr, args...)
@@ -378,12 +379,18 @@ function generate(
     return trace, trace.weight
 end
 
+function _extract_addresses(constraint::Constraint)
+    addresses = Set(keys(constraint))
+
+    return addresses
+end
+
 function generate_internal(
     rng::AbstractRNG,
     f::Function,
     args::Vararg{Any,Nargs};
     constraint_ptr::TracedRNumber,
-    constrained_symbols::Set{Symbol},
+    constrained_addresses::Set{Address},
 ) where {Nargs}
     argprefix::Symbol = gensym("generatearg")
     resprefix::Symbol = gensym("generateresult")
@@ -441,15 +448,19 @@ function generate_internal(
         1,
     )
 
-    constrained_symbols_attr = MLIR.IR.Attribute[]
-    for sym in constrained_symbols
-        addr = reinterpret(UInt64, pointer_from_objref(sym))
-        push!(
-            constrained_symbols_attr,
-            @ccall MLIR.API.mlir_c.enzymeSymbolAttrGet(
-                MLIR.IR.context()::MLIR.API.MlirContext, addr::UInt64
-            )::MLIR.IR.Attribute
-        )
+    constrained_addresses_attr = MLIR.IR.Attribute[]
+    for address in constrained_addresses
+        address_attr = MLIR.IR.Attribute[]
+        for sym in address.path
+            sym_addr = reinterpret(UInt64, pointer_from_objref(sym))
+            push!(
+                address_attr,
+                @ccall MLIR.API.mlir_c.enzymeSymbolAttrGet(
+                    MLIR.IR.context()::MLIR.API.MlirContext, sym_addr::UInt64
+                )::MLIR.IR.Attribute
+            )
+        end
+        push!(constrained_addresses_attr, MLIR.IR.Attribute(address_attr))
     end
 
     trace_ty = @ccall MLIR.API.mlir_c.enzymeTraceTypeGet(
@@ -464,7 +475,7 @@ function generate_internal(
         weight=weight_ty,
         outputs=out_tys,
         fn=fn_attr,
-        constrained_symbols=MLIR.IR.Attribute(constrained_symbols_attr),
+        constrained_addresses=MLIR.IR.Attribute(constrained_addresses_attr),
     )
 
     for (i, res) in enumerate(linear_results)

src/probprog/ProbProg.jl

@@ -19,7 +19,7 @@ include("Modeling.jl")
 include("Inference.jl")
 include("Display.jl")
 
-export ProbProgTrace, Constraint, Selection, CompiledFnCache
+export ProbProgTrace, Constraint, Selection, CompiledFnCache, Address
 export get_choices, select, choicemap, with_compiled_cache
 
 export sample, call, simulate, generate

src/probprog/Types.jl

@@ -22,26 +22,68 @@ mutable struct ProbProgTrace
     end
 end
 
-const Constraint = Dict{Symbol,Any}
+struct Address
+    path::Vector{Symbol}
+
+    Address(path::Vector{Symbol}) = new(path)
+end
+Address(sym::Symbol) = Address([sym])
+Address(syms::Symbol...) = Address([syms...])
+
+Base.:(==)(a::Address, b::Address) = a.path == b.path
+Base.hash(a::Address, h::UInt) = hash(a.path, h)
+
+mutable struct Constraint <: AbstractDict{Address,Any}
+    dict::Dict{Address,Any}
+
+    function Constraint(pairs::Pair...)
+        dict = Dict{Address,Any}()
+        for pair in pairs
+            symbols = Symbol[]
+            current = pair
+            while isa(current, Pair) && isa(current.first, Symbol)
+                push!(symbols, current.first)
+                current = current.second
+            end
+            dict[Address(symbols...)] = current
+        end
+        return new(dict)
+    end
+
+    Constraint() = new(Dict{Address,Any}())
+    Constraint(d::Dict{Address,Any}) = new(d)
+end
+
+Base.getindex(c::Constraint, k::Address) = c.dict[k]
+Base.setindex!(c::Constraint, v, k::Address) = (c.dict[k] = v)
+Base.delete!(c::Constraint, k::Address) = delete!(c.dict, k)
+Base.keys(c::Constraint) = keys(c.dict)
+Base.values(c::Constraint) = values(c.dict)
+Base.iterate(c::Constraint) = iterate(c.dict)
+Base.iterate(c::Constraint, state) = iterate(c.dict, state)
+Base.length(c::Constraint) = length(c.dict)
+Base.isempty(c::Constraint) = isempty(c.dict)
+Base.haskey(c::Constraint, k::Address) = haskey(c.dict, k)
+Base.get(c::Constraint, k::Address, default) = get(c.dict, k, default)
+
 const Selection = Set{Symbol}
 const CompiledFnCache = Dict{Tuple{Type,Set{Symbol}},Any}
 
-const _trace_ref_lock = ReentrantLock()
-const _trace_refs = Vector{Any}()
+const _probprog_ref_lock = ReentrantLock()
+const _probprog_refs = IdDict()
 
-function _keepalive!(tr::ProbProgTrace)
-    lock(_trace_ref_lock)
+function _keepalive!(tr::Any)
+    lock(_probprog_ref_lock)
     try
-        push!(_trace_refs, tr)
+        _probprog_refs[tr] = tr
     finally
-        unlock(_trace_ref_lock)
+        unlock(_probprog_ref_lock)
     end
     return tr
 end
 
 get_choices(trace::ProbProgTrace) = trace.choices
 select(syms::Symbol...) = Set(syms)
-choicemap() = Constraint()
 
 function with_compiled_cache(f)
     cache = CompiledFnCache()

test/probprog/generate.jl

@@ -13,6 +13,19 @@ function model(rng, μ, σ, shape)
     return t
 end
 
+function two_normals(rng, μ, σ, shape)
+    x = ProbProg.sample(rng, normal, μ, σ, shape; symbol=:x, logpdf=normal_logpdf)
+    y = ProbProg.sample(rng, normal, x, σ, shape; symbol=:y, logpdf=normal_logpdf)
+    return y
+end
+
+function nested_model(rng, μ, σ, shape)
+    s = ProbProg.sample(rng, normal, μ, σ, shape; symbol=:s, logpdf=normal_logpdf)
+    t = ProbProg.sample(rng, two_normals, s, σ, shape; symbol=:t)
+    u = ProbProg.sample(rng, two_normals, t, σ, shape; symbol=:u)
+    return u
+end
+
 @testset "Generate" begin
     @testset "unconstrained" begin
         shape = (1000,)
@@ -31,15 +44,55 @@ end
         μ = Reactant.ConcreteRNumber(0.0)
         σ = Reactant.ConcreteRNumber(1.0)
 
-        constraint = Dict{Symbol,Any}(:s => (fill(0.1, shape),))
+        constraint = ProbProg.Constraint(:s => (fill(0.1, shape),))
 
         trace, weight = ProbProg.generate(rng, model, μ, σ, shape; constraint)
 
-        @test trace.choices[:s][1] == constraint[:s][1]
+        @test trace.choices[:s][1] == constraint[ProbProg.Address(:s)][1]
 
         expected_weight =
-            normal_logpdf(constraint[:s][1], 0.0, 1.0, shape) +
-            normal_logpdf(trace.choices[:t][1], constraint[:s][1], 1.0, shape)
+            normal_logpdf(constraint[ProbProg.Address(:s)][1], 0.0, 1.0, shape) +
+            normal_logpdf(
+                trace.choices[:t][1], constraint[ProbProg.Address(:s)][1], 1.0, shape
+            )
+        @test weight ≈ expected_weight atol = 1e-6
+    end
+
+    @testset "composite addresses" begin
+        shape = (10,)
+        seed = Reactant.to_rarray(UInt64[1, 4])
+        rng = ReactantRNG(seed)
+        μ = Reactant.ConcreteRNumber(0.0)
+        σ = Reactant.ConcreteRNumber(1.0)
+
+        constraint = ProbProg.Constraint(
+            :s => (fill(0.1, shape),),
+            :t => :x => (fill(0.2, shape),),
+            :u => :y => (fill(0.3, shape),),
+        )
+
+        trace, weight = ProbProg.generate(rng, nested_model, μ, σ, shape; constraint)
+
+        @test trace.choices[:s][1] == fill(0.1, shape)
+        @test trace.subtraces[:t].choices[:x][1] == fill(0.2, shape)
+        @test trace.subtraces[:u].choices[:y][1] == fill(0.3, shape)
+
+        s_weight = normal_logpdf(fill(0.1, shape), 0.0, 1.0, shape)
+        tx_weight = normal_logpdf(fill(0.2, shape), fill(0.1, shape), 1.0, shape)
+        ty_weight = normal_logpdf(
+            trace.subtraces[:t].choices[:y][1], fill(0.2, shape), 1.0, shape
+        )
+        ux_weight = normal_logpdf(
+            trace.subtraces[:u].choices[:x][1],
+            trace.subtraces[:t].choices[:y][1],
+            1.0,
+            shape,
+        )
+        uy_weight = normal_logpdf(
+            fill(0.3, shape), trace.subtraces[:u].choices[:x][1], 1.0, shape
+        )
+
+        expected_weight = s_weight + tx_weight + ty_weight + ux_weight + uy_weight
         @test weight ≈ expected_weight atol = 1e-6
     end
 
@@ -50,24 +103,24 @@ end
         μ = Reactant.ConcreteRNumber(0.0)
         σ = Reactant.ConcreteRNumber(1.0)
 
-        constraint1 = Dict{Symbol,Any}(:s => (fill(0.1, shape),))
+        constraint1 = ProbProg.Constraint(:s => (fill(0.1, shape),))
 
-        constrained_symbols = Set(keys(constraint1))
+        constrained_addresses = ProbProg._extract_addresses(constraint1)
 
         constraint_ptr1 = Reactant.ConcreteRNumber(
             reinterpret(UInt64, pointer_from_objref(constraint1))
         )
 
         wrapper_fn(constraint_ptr, rng, μ, σ) = ProbProg.generate_internal(
-            rng, model, μ, σ, shape; constraint_ptr, constrained_symbols
+            rng, model, μ, σ, shape; constraint_ptr, constrained_addresses
         )
 
         compiled_fn = @compile optimize = :probprog wrapper_fn(constraint_ptr1, rng, μ, σ)
 
         trace1, weight = compiled_fn(constraint_ptr1, rng, μ, σ)
         trace1 = unsafe_pointer_to_objref(Ptr{Any}(Array(trace1)[1]))
 
-        constraint2 = Dict{Symbol,Any}(:s => (fill(0.2, shape),))
+        constraint2 = ProbProg.Constraint(:s => (fill(0.2, shape),))
         constraint_ptr2 = Reactant.ConcreteRNumber(
             reinterpret(UInt64, pointer_from_objref(constraint2))
         )