generate op with constraints

Author: sbrantq

deps/ReactantExtra/API.cpp

@@ -353,6 +353,20 @@ enzymeActivityAttrGet(MlirContext ctx, int32_t val) {
                                               (mlir::enzyme::Activity)val));
 }
 
+extern "C" MLIR_CAPI_EXPORTED MlirAttribute enzymeConstraintAttrGet(
+    MlirContext ctx, uint64_t symbol, MlirAttribute values) {
+  mlir::Attribute vals = unwrap(values);
+  auto arr = llvm::dyn_cast<mlir::ArrayAttr>(vals);
+  if (!arr) {
+    ReactantThrowError(
+        "enzymeConstraintAttrGet: `values` must be an ArrayAttr");
+    return MlirAttribute{nullptr};
+  }
+  mlir::Attribute attr =
+      mlir::enzyme::ConstraintAttr::get(unwrap(ctx), symbol, arr);
+  return wrap(attr);
+}
+
 // Create profiler session and start profiling
 extern "C" tsl::ProfilerSession *
 CreateProfilerSession(uint32_t device_tracer_level,

src/ProbProg.jl

@@ -146,7 +146,8 @@ function sample(
             in_idx = nothing
             for (i, arg) in enumerate(linear_args)
                 if TracedUtils.has_idx(arg, argprefix) &&
-                    TracedUtils.get_idx(arg, argprefix) == TracedUtils.get_idx(res, argprefix)
+                    TracedUtils.get_idx(arg, argprefix) ==
+                   TracedUtils.get_idx(res, argprefix)
                     in_idx = i - 1
                     break
                 end
@@ -221,10 +222,12 @@ function sample(
     return result
 end
 
-function generate(f::Function, args::Vararg{Any,Nargs}) where {Nargs}
+function generate(f::Function, args::Vararg{Any,Nargs}; constraints=nothing) where {Nargs}
     trace = ProbProgTrace()
 
-    weight, res = @jit optimize = :probprog generate_internal(f, args...; trace)
+    weight, res = @jit sync = true optimize = :probprog generate_internal(
+        f, args...; trace, constraints
+    )
 
     trace.retval = res isa AbstractConcreteArray ? Array(res) : res
     trace.weight = Array(weight)[1]
@@ -233,7 +236,7 @@ function generate(f::Function, args::Vararg{Any,Nargs}) where {Nargs}
 end
 
 function generate_internal(
-    f::Function, args::Vararg{Any,Nargs}; trace::ProbProgTrace
+    f::Function, args::Vararg{Any,Nargs}; trace::ProbProgTrace, constraints=nothing
 ) where {Nargs}
     argprefix::Symbol = gensym("generatearg")
     resprefix::Symbol = gensym("generateresult")
@@ -276,9 +279,46 @@ function generate_internal(
 
     trace_addr = reinterpret(UInt64, pointer_from_objref(trace))
 
-    # Output: (weight, f's outputs...)
+    constraints_attr = nothing
+    if constraints !== nothing && !isempty(constraints)
+        constraint_attrs = MLIR.IR.Attribute[]
+
+        for (sym, constraint) in constraints
+            sym_addr = reinterpret(UInt64, pointer_from_objref(sym))
+
+            if !(constraint isa AbstractArray)
+                error(
+                    "Constraints must be an array (one element per traced output) of arrays"
+                )
+            end
+
+            sym_constraint_attrs = MLIR.IR.Attribute[]
+            for oc in constraint
+                if !(oc isa AbstractArray)
+                    error("Per-output constraints must be arrays")
+                end
+
+                push!(sym_constraint_attrs, MLIR.IR.DenseElementsAttribute(oc))
+            end
+
+            cattr_ptr = @ccall MLIR.API.mlir_c.enzymeConstraintAttrGet(
+                MLIR.IR.context()::MLIR.API.MlirContext,
+                sym_addr::UInt64,
+                MLIR.IR.Attribute(sym_constraint_attrs)::MLIR.API.MlirAttribute,
+            )::MLIR.API.MlirAttribute
+
+            push!(constraint_attrs, MLIR.IR.Attribute(cattr_ptr))
+        end
+
+        constraints_attr = MLIR.IR.Attribute(constraint_attrs)
+    end
+
     gen_op = MLIR.Dialects.enzyme.generate(
-        batch_inputs; outputs=out_tys, fn=fname, trace=trace_addr
+        batch_inputs;
+        outputs=out_tys,
+        fn=fname,
+        trace=trace_addr,
+        constraints=constraints_attr,
     )
 
     weight = TracedRArray(MLIR.IR.result(gen_op, 1))

test/probprog/generate.jl

@@ -40,4 +40,23 @@ end
         trace, weight = ProbProg.generate(model, seed, μ, σ, shape)
         @test mean(trace.retval) ≈ 0.0 atol = 0.05 rtol = 0.05
     end
+
+    @testset "constraints" begin
+        shape = (10,)
+        seed = Reactant.to_rarray(UInt64[1, 4])
+        μ = Reactant.ConcreteRNumber(0.0)
+        σ = Reactant.ConcreteRNumber(1.0)
+
+        s_constraint = fill(0.1, shape)
+        constraints = Dict(:s => [s_constraint])
+
+        trace, weight = ProbProg.generate(model, seed, μ, σ, shape; constraints)
+
+        @test trace.choices[:s] == s_constraint
+
+        expected_weight =
+            normal_logpdf(s_constraint, 0.0, 1.0, shape) +
+            normal_logpdf(trace.choices[:t], s_constraint, 1.0, shape)
+        @test weight ≈ expected_weight atol = 1e-6
+    end
 end