untraced call

Author: sbrantq

src/ProbProg.jl

@@ -222,6 +222,75 @@ function sample(
     return result
 end
 
+function call(f::Function, args::Vararg{Any,Nargs}) where {Nargs}
+    res = @jit optimize = :probprog call_internal(f, args...)
+    return res isa AbstractConcreteArray ? Array(res) : res
+end
+
+function call_internal(f::Function, args::Vararg{Any,Nargs}) where {Nargs}
+    argprefix::Symbol = gensym("callarg")
+    resprefix::Symbol = gensym("callresult")
+    resargprefix::Symbol = gensym("callresarg")
+
+    mlir_fn_res = invokelatest(
+        TracedUtils.make_mlir_fn,
+        f,
+        args,
+        (),
+        string(f),
+        false;
+        do_transpose=false,
+        args_in_result=:all,
+        argprefix,
+        resprefix,
+        resargprefix,
+    )
+    (; result, linear_args, in_tys, linear_results) = mlir_fn_res
+    fnwrap = mlir_fn_res.fnwrapped
+    func2 = mlir_fn_res.f
+
+    out_tys = [MLIR.IR.type(TracedUtils.get_mlir_data(res)) for res in linear_results]
+    fname = TracedUtils.get_attribute_by_name(func2, "sym_name")
+    fname = MLIR.IR.FlatSymbolRefAttribute(Base.String(fname))
+
+    batch_inputs = MLIR.IR.Value[]
+    for a in linear_args
+        idx, path = TracedUtils.get_argidx(a, argprefix)
+        if idx == 1 && fnwrap
+            TracedUtils.push_val!(batch_inputs, f, path[3:end])
+        else
+            if fnwrap
+                idx -= 1
+            end
+            TracedUtils.push_val!(batch_inputs, args[idx], path[3:end])
+        end
+    end
+
+    call_op = MLIR.Dialects.enzyme.untracedCall(batch_inputs; outputs=out_tys, fn=fname)
+
+    for (i, res) in enumerate(linear_results)
+        resv = MLIR.IR.result(call_op, i)
+        if TracedUtils.has_idx(res, resprefix)
+            path = TracedUtils.get_idx(res, resprefix)
+            TracedUtils.set!(result, path[2:end], resv)
+        elseif TracedUtils.has_idx(res, argprefix)
+            idx, path = TracedUtils.get_argidx(res, argprefix)
+            if idx == 1 && fnwrap
+                TracedUtils.set!(f, path[3:end], resv)
+            else
+                if fnwrap
+                    idx -= 1
+                end
+                TracedUtils.set!(args[idx], path[3:end], resv)
+            end
+        else
+            TracedUtils.set!(res, (), resv)
+        end
+    end
+
+    return result
+end
+
 function generate(f::Function, args::Vararg{Any,Nargs}; constraints=nothing) where {Nargs}
     trace = ProbProgTrace()
 

test/probprog/sample.jl

@@ -24,9 +24,9 @@ end
         seed = Reactant.to_rarray(UInt64[1, 4])
         μ = Reactant.ConcreteRNumber(0.0)
         σ = Reactant.ConcreteRNumber(1.0)
-        before = @code_hlo optimize = false ProbProg.generate_internal(one_sample, seed, μ, σ, shape)
+        before = @code_hlo optimize = false ProbProg.call_internal(one_sample, seed, μ, σ, shape)
         @test contains(repr(before), "enzyme.sample")
-        after = @code_hlo optimize = :probprog ProbProg.generate_internal(two_samples, seed, μ, σ, shape)
+        after = @code_hlo optimize = :probprog ProbProg.call_internal(two_samples, seed, μ, σ, shape)
         @test !contains(repr(after), "enzyme.sample")
     end
 
@@ -35,8 +35,8 @@ end
         seed = Reactant.to_rarray(UInt64[1, 4])
         μ = Reactant.ConcreteRNumber(0.0)
         σ = Reactant.ConcreteRNumber(1.0)
-        X = ProbProg.generate(one_sample, seed, μ, σ, shape)
-        Y = ProbProg.generate(two_samples, seed, μ, σ, shape)
+        X = ProbProg.call(one_sample, seed, μ, σ, shape)
+        Y = ProbProg.call(two_samples, seed, μ, σ, shape)
         @test !all(X .≈ Y)
     end
 end