ProbProg: Making trace an operand

deps/ReactantExtra/API.cpp

@@ -367,6 +367,21 @@ enzymeActivityAttrGet(MlirContext ctx, int32_t val) {
                                               (mlir::enzyme::Activity)val));
 }
 
+extern "C" MLIR_CAPI_EXPORTED MlirType enzymeTraceTypeGet(MlirContext ctx) {
+  return wrap(mlir::enzyme::TraceType::get(unwrap(ctx)));
+}
+
+extern "C" MLIR_CAPI_EXPORTED MlirType
+enzymeConstraintTypeGet(MlirContext ctx) {
+  return wrap(mlir::enzyme::ConstraintType::get(unwrap(ctx)));
+}
+
+extern "C" MLIR_CAPI_EXPORTED MlirAttribute
+enzymeSymbolAttrGet(MlirContext ctx, uint64_t symbol) {
+  mlir::Attribute attr = mlir::enzyme::SymbolAttr::get(unwrap(ctx), symbol);
+  return wrap(attr);
+}
+
 // Create profiler session and start profiling
 extern "C" tsl::ProfilerSession *
 CreateProfilerSession(uint32_t device_tracer_level,

src/CompileOptions.jl

@@ -221,6 +221,8 @@ function CompileOptions(;
             :canonicalize,
             :just_batch,
             :none,
+            :probprog,
+            :probprog_no_lowering,
         ]
     end
 

src/Compiler.jl

@@ -1231,6 +1231,7 @@ end
 # TODO we want to be able to run the more advanced passes via transform dialect as an enzyme intermediate
 # However, this errs as we cannot attach the transform with to the funcop itself [as we run a functionpass].
 const enzyme_pass::String = "enzyme{postpasses=\"arith-raise{stablehlo=true},canonicalize,cse,canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math,canonicalize,cse,canonicalize\"}"
+const probprog_pass::String = "probprog{postpasses=\"arith-raise{stablehlo=true},canonicalize,cse,canonicalize\"}"
 
 function run_pass_pipeline!(mod, pass_pipeline, key=""; enable_verifier=true)
     pm = MLIR.IR.PassManager()
@@ -1641,6 +1642,7 @@ function compile_mlir!(
     blas_int_width = sizeof(BLAS.BlasInt) * 8
     lower_enzymexla_linalg_pass = "lower-enzymexla-linalg{backend=$backend \
                                    blas_int_width=$blas_int_width}"
+    lower_enzyme_probprog_pass = "lower-enzyme-probprog{backend=$backend}"
 
     legalize_chlo_to_stablehlo =
         if legalize_stablehlo_to_mhlo || compile_options.legalize_chlo_to_stablehlo
@@ -1807,6 +1809,122 @@ function compile_mlir!(
             ),
             "no_enzyme",
         )
+    elseif compile_options.optimization_passes === :probprog_no_lowering
+        run_pass_pipeline!(
+            mod,
+            join(
+                if compile_options.raise_first
+                    [
+                        "mark-func-memory-effects",
+                        opt_passes,
+                        kern,
+                        raise_passes,
+                        "enzyme-batch",
+                        opt_passes2,
+                        enzyme_pass,
+                        probprog_pass,
+                        opt_passes2,
+                        "canonicalize",
+                        "remove-unnecessary-enzyme-ops",
+                        "enzyme-simplify-math",
+                        (
+                            if compile_options.legalize_chlo_to_stablehlo
+                                ["func.func(chlo-legalize-to-stablehlo)"]
+                            else
+                                []
+                            end
+                        )...,
+                        opt_passes2,
+                    ]
+                else
+                    [
+                        "mark-func-memory-effects",
+                        opt_passes,
+                        "enzyme-batch",
+                        opt_passes2,
+                        enzyme_pass,
+                        probprog_pass,
+                        opt_passes2,
+                        "canonicalize",
+                        "remove-unnecessary-enzyme-ops",
+                        "enzyme-simplify-math",
+                        (
+                            if compile_options.legalize_chlo_to_stablehlo
+                                ["func.func(chlo-legalize-to-stablehlo)"]
+                            else
+                                []
+                            end
+                        )...,
+                        opt_passes2,
+                        kern,
+                        raise_passes,
+                    ]
+                end,
+                ",",
+            ),
+            "probprog_no_lowering",
+        )
+    elseif compile_options.optimization_passes === :probprog
+        run_pass_pipeline!(
+            mod,
+            join(
+                if compile_options.raise_first
+                    [
+                        "mark-func-memory-effects",
+                        opt_passes,
+                        kern,
+                        raise_passes,
+                        "enzyme-batch",
+                        opt_passes2,
+                        enzyme_pass,
+                        probprog_pass,
+                        opt_passes2,
+                        "canonicalize",
+                        "remove-unnecessary-enzyme-ops",
+                        "enzyme-simplify-math",
+                        (
+                            if compile_options.legalize_chlo_to_stablehlo
+                                ["func.func(chlo-legalize-to-stablehlo)"]
+                            else
+                                []
+                            end
+                        )...,
+                        opt_passes2,
+                        lower_enzymexla_linalg_pass,
+                        lower_enzyme_probprog_pass,
+                        jit,
+                    ]
+                else
+                    [
+                        "mark-func-memory-effects",
+                        opt_passes,
+                        "enzyme-batch",
+                        opt_passes2,
+                        enzyme_pass,
+                        probprog_pass,
+                        opt_passes2,
+                        "canonicalize",
+                        "remove-unnecessary-enzyme-ops",
+                        "enzyme-simplify-math",
+                        (
+                            if compile_options.legalize_chlo_to_stablehlo
+                                ["func.func(chlo-legalize-to-stablehlo)"]
+                            else
+                                []
+                            end
+                        )...,
+                        opt_passes2,
+                        kern,
+                        raise_passes,
+                        lower_enzymexla_linalg_pass,
+                        lower_enzyme_probprog_pass,
+                        jit,
+                    ]
+                end,
+                ",",
+            ),
+            "probprog",
+        )
     elseif compile_options.optimization_passes === :only_enzyme
         run_pass_pipeline!(
             mod,

src/Reactant.jl

@@ -191,6 +191,7 @@ include("Tracing.jl")
 include("Compiler.jl")
 
 include("Overlay.jl")
+include("probprog/ProbProg.jl")
 
 # Serialization
 include("serialization/Serialization.jl")

src/Types.jl

@@ -215,6 +215,7 @@ function ConcretePJRTArray(
 end
 
 Base.wait(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = foreach(wait, x.data)
+Base.isready(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = all(isready, x.data)
 XLA.client(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = XLA.client(x.data)
 function XLA.device(x::Union{ConcretePJRTArray,ConcretePJRTNumber})
     x.sharding isa Sharding.NoShardInfo && return XLA.device(only(x.data))
@@ -412,6 +413,7 @@ function ConcreteIFRTArray(
 end
 
 Base.wait(x::Union{ConcreteIFRTArray,ConcreteIFRTNumber}) = wait(x.data)
+Base.isready(x::Union{ConcreteIFRTArray,ConcreteIFRTNumber}) = isready(x.data)
 XLA.client(x::Union{ConcreteIFRTArray,ConcreteIFRTNumber}) = XLA.client(x.data)
 function XLA.device(x::Union{ConcreteIFRTArray,ConcreteIFRTNumber})
     return XLA.device(x.data)

src/probprog/Display.jl

@@ -0,0 +1,87 @@
+# Reference: https://github.com/probcomp/Gen.jl/blob/91d798f2d2f0c175b1be3dc6daf3a10a8acf5da3/src/choice_map.jl#L104
+function _show_pretty(io::IO, trace::ProbProgTrace, pre::Int, vert_bars::Tuple)
+    VERT = '\u2502'
+    PLUS = '\u251C'
+    HORZ = '\u2500'
+    LAST = '\u2514'
+
+    indent_vert = vcat(Char[' ' for _ in 1:pre], Char[VERT, '\n'])
+    indent = vcat(Char[' ' for _ in 1:pre], Char[PLUS, HORZ, HORZ, ' '])
+    indent_last = vcat(Char[' ' for _ in 1:pre], Char[LAST, HORZ, HORZ, ' '])
+
+    for i in vert_bars
+        indent_vert[i] = VERT
+        indent[i] = VERT
+        indent_last[i] = VERT
+    end
+
+    indent_vert_str = join(indent_vert)
+    indent_str = join(indent)
+    indent_last_str = join(indent_last)
+
+    sorted_choices = sort(collect(trace.choices); by=x -> x[1])
+    n = length(sorted_choices)
+
+    if trace.retval !== nothing
+        n += 1
+    end
+
+    if trace.weight !== nothing
+        n += 1
+    end
+
+    cur = 1
+
+    if trace.retval !== nothing
+        print(io, indent_vert_str)
+        print(io, (cur == n ? indent_last_str : indent_str) * "retval : $(trace.retval)\n")
+        cur += 1
+    end
+
+    if trace.weight !== nothing
+        print(io, indent_vert_str)
+        print(io, (cur == n ? indent_last_str : indent_str) * "weight : $(trace.weight)\n")
+        cur += 1
+    end
+
+    for (key, value) in sorted_choices
+        print(io, indent_vert_str)
+        print(io, (cur == n ? indent_last_str : indent_str) * "$(repr(key)) : $value\n")
+        cur += 1
+    end
+
+    sorted_subtraces = sort(collect(trace.subtraces); by=x -> x[1])
+    n += length(sorted_subtraces)
+
+    for (key, subtrace) in sorted_subtraces
+        print(io, indent_vert_str)
+        print(io, (cur == n ? indent_last_str : indent_str) * "subtrace on $(repr(key))\n")
+        _show_pretty(
+            io, subtrace, pre + 4, cur == n ? (vert_bars...,) : (vert_bars..., pre + 1)
+        )
+        cur += 1
+    end
+end
+
+function Base.show(io::IO, ::MIME"text/plain", trace::ProbProgTrace)
+    println(io, "ProbProgTrace:")
+    if isempty(trace.choices) && trace.retval === nothing && trace.weight === nothing
+        println(io, "  (empty)")
+    else
+        _show_pretty(io, trace, 0, ())
+    end
+end
+
+function Base.show(io::IO, trace::ProbProgTrace)
+    if get(io, :compact, false)
+        choices_count = length(trace.choices)
+        has_retval = trace.retval !== nothing
+        print(io, "ProbProgTrace($(choices_count) choices")
+        if has_retval
+            print(io, ", retval=$(trace.retval), weight=$(trace.weight)")
+        end
+        print(io, ")")
+    else
+        show(io, MIME"text/plain"(), trace)
+    end
+end