ProbProg: Making trace an operand

deps/ReactantExtra/API.cpp

@@ -364,6 +364,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,remove-unnecessary-enzyme-ops,enzyme-simplify-math,canonicalize,cse,canonicalize\"}"
 
 function run_pass_pipeline!(mod, pass_pipeline, key=""; enable_verifier=true)
     pm = MLIR.IR.PassManager()
@@ -1623,6 +1624,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}"
 
     if compile_options.optimization_passes === :all
         run_pass_pipeline!(
@@ -1824,6 +1826,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,