reorder

Author: sbrantq

src/ProbProg.jl

@@ -52,10 +52,12 @@ function __init__()
     return nothing
 end
 
-@noinline function generate!(f::Function, args::Vararg{Any,Nargs}) where {Nargs}
-    argprefix::Symbol = gensym("generatearg")
-    resprefix::Symbol = gensym("generateresult")
-    resargprefix::Symbol = gensym("generateresarg")
+@noinline function sample!(
+    f::Function, args::Vararg{Any,Nargs}; symbol::Symbol=gensym("sample")
+) where {Nargs}
+    argprefix::Symbol = gensym("samplearg")
+    resprefix::Symbol = gensym("sampleresult")
+    resargprefix::Symbol = gensym("sampleresarg")
 
     mlir_fn_res = invokelatest(
         TracedUtils.make_mlir_fn,
@@ -70,31 +72,45 @@ end
         resprefix,
         resargprefix,
     )
-    (; result, linear_args, in_tys, linear_results) = mlir_fn_res
+    (; result, linear_args, 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
+            idx -= fnwrap ? 1 : 0
             TracedUtils.push_val!(batch_inputs, args[idx], path[3:end])
         end
     end
 
-    gen_op = MLIR.Dialects.enzyme.generate(batch_inputs; outputs=out_tys, fn=fname)
+    out_tys = [MLIR.IR.type(TracedUtils.get_mlir_data(res)) for res in linear_results]
 
+    sym = TracedUtils.get_attribute_by_name(func2, "sym_name")
+    fn_attr = MLIR.IR.FlatSymbolRefAttribute(Base.String(sym))
+
+    traced_output_indices = Int[]
     for (i, res) in enumerate(linear_results)
-        resv = MLIR.IR.result(gen_op, i)
+        if TracedUtils.has_idx(res, resprefix)
+            push!(traced_output_indices, i - 1)
+        end
+    end
+
+    symbol_addr = reinterpret(UInt64, pointer_from_objref(symbol))
+
+    sample_op = MLIR.Dialects.enzyme.sample(
+        batch_inputs;
+        outputs=out_tys,
+        fn=fn_attr,
+        symbol=symbol_addr,
+        traced_output_indices=traced_output_indices,
+    )
+
+    for (i, res) in enumerate(linear_results)
+        resv = MLIR.IR.result(sample_op, i)
         if TracedUtils.has_idx(res, resprefix)
             path = TracedUtils.get_idx(res, resprefix)
             TracedUtils.set!(result, path[2:end], resv)
@@ -116,12 +132,10 @@ end
     return result
 end
 
-@noinline function sample!(
-    f::Function, args::Vararg{Any,Nargs}; symbol::Symbol=gensym("sample")
-) where {Nargs}
-    argprefix::Symbol = gensym("samplearg")
-    resprefix::Symbol = gensym("sampleresult")
-    resargprefix::Symbol = gensym("sampleresarg")
+@noinline function generate!(f::Function, args::Vararg{Any,Nargs}) where {Nargs}
+    argprefix::Symbol = gensym("generatearg")
+    resprefix::Symbol = gensym("generateresult")
+    resargprefix::Symbol = gensym("generateresarg")
 
     mlir_fn_res = invokelatest(
         TracedUtils.make_mlir_fn,
@@ -136,45 +150,31 @@ end
         resprefix,
         resargprefix,
     )
-    (; result, linear_args, linear_results) = mlir_fn_res
+    (; 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
-            idx -= fnwrap ? 1 : 0
+            if fnwrap
+                idx -= 1
+            end
             TracedUtils.push_val!(batch_inputs, args[idx], path[3:end])
         end
     end
 
-    out_tys = [MLIR.IR.type(TracedUtils.get_mlir_data(res)) for res in linear_results]
-
-    sym = TracedUtils.get_attribute_by_name(func2, "sym_name")
-    fn_attr = MLIR.IR.FlatSymbolRefAttribute(Base.String(sym))
-
-    traced_output_indices = Int[]
-    for (i, res) in enumerate(linear_results)
-        if TracedUtils.has_idx(res, resprefix)
-            push!(traced_output_indices, i - 1)
-        end
-    end
-
-    symbol_addr = reinterpret(UInt64, pointer_from_objref(symbol))
-
-    sample_op = MLIR.Dialects.enzyme.sample(
-        batch_inputs;
-        outputs=out_tys,
-        fn=fn_attr,
-        symbol=symbol_addr,
-        traced_output_indices=traced_output_indices,
-    )
+    gen_op = MLIR.Dialects.enzyme.generate(batch_inputs; outputs=out_tys, fn=fname)
 
     for (i, res) in enumerate(linear_results)
-        resv = MLIR.IR.result(sample_op, i)
+        resv = MLIR.IR.result(gen_op, i)
         if TracedUtils.has_idx(res, resprefix)
             path = TracedUtils.get_idx(res, resprefix)
             TracedUtils.set!(result, path[2:end], resv)
@@ -278,5 +278,4 @@ function print_trace(trace::Dict{Symbol,Any})
     end
     return println("### End of Trace ###")
 end
-
-end
+end