ProbProg: Simulate + Generate (with trace attributes)

enzyme/Enzyme/MLIR/Dialect/EnzymeOps.td

@@ -256,8 +256,114 @@ def BroadcastOp : Enzyme_Op<"broadcast"> {
 
 def SampleOp : Enzyme_Op<"sample",
     [DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
-  let summary = "Sample from a distribution. Arguments to the distribution are: a random number generator object, followed by arguments to the sample op itself";
-  let arguments = (ins FlatSymbolRefAttr:$fn, Variadic<AnyType>:$inputs, DefaultValuedStrAttr<StrAttr, "">:$name);
+  let summary = "Sample from a distribution";
+  let arguments = (ins
+    FlatSymbolRefAttr:$fn,
+    Variadic<AnyType>:$inputs,
+    OptionalAttr<FlatSymbolRefAttr>:$logpdf,
+    OptionalAttr<UI64Attr>:$symbol,
+    OptionalAttr<DenseI64ArrayAttr>:$traced_input_indices,
+    OptionalAttr<DenseI64ArrayAttr>:$traced_output_indices,
+    OptionalAttr<DenseI64ArrayAttr>:$alias_map,
+    DefaultValuedStrAttr<StrAttr, "">:$name
+  );
+  let results = (outs Variadic<AnyType>:$outputs);
+
+  let assemblyFormat = [{
+    $fn `(` $inputs `)` attr-dict `:` functional-type($inputs, results)
+  }];
+}
+
+def SimulateOp : Enzyme_Op<"simulate", [DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
+  let summary = "Simulate a probabilistic function to generate execution trace";
+  let description = [{
+  Simulate a probabilistic function to generate execution trace
+  by replacing all SampleOps with distribution calls and inserting
+  sampled values into the choice map.
+  }];
+
+  let arguments = (ins
+    FlatSymbolRefAttr:$fn,
+    Variadic<AnyType>:$inputs,
+    UI64Attr:$trace,
+    DefaultValuedStrAttr<StrAttr, "">:$name
+  );
+
+  let results = (outs Variadic<AnyType>:$outputs);
+
+  let assemblyFormat = [{
+    $fn `(` $inputs `)` attr-dict `:` functional-type($inputs, results)
+  }];
+}
+
+def AddSampleToTraceOp : Enzyme_Op<"addSampleToTrace"> {
+  let summary = "Add a sampled value into the execution trace";
+  let description = [{
+  Add a sampled value into the execution trace.
+  }];
+
+  let arguments = (ins
+    Variadic<AnyType>:$sample,
+    UI64Attr:$symbol,
+    UI64Attr:$trace,
+    DefaultValuedStrAttr<StrAttr, "">:$name
+  );
+
+  let results = (outs );
+
+  let assemblyFormat = [{
+    $sample attr-dict `:` functional-type($sample, results)
+  }];
+}
+
+def ConstraintAttr : Enzyme_Attr<"Constraint", "constraint"> {
+  let summary = "In probabilistic programming, mapping a symbol to a constraint value.";
+  let description = [{
+    This attribute represents a symbol/value pair inside the constraints dict
+    optionally attached to `enzyme.generate` op.  The `symbol` field should be
+    an ui64.  The `values` field is an array attribute containing constraint values
+    in the same order as the SampleOp's traced output indices.
+  }];
+
+  let parameters = (ins "uint64_t":$symbol, "ArrayAttr":$values);
+
+  let assemblyFormat = "`<` `symbol` `=` $symbol `,` `values` `=` $values `>`";
+}
+
+def UntracedCallOp : Enzyme_Op<"untracedCall"> {
+  let summary = "Call a probabilistic function without tracing";
+  let description = [{
+  Call a probabilistic function without tracing.
+  }];
+
+  let arguments = (ins
+    FlatSymbolRefAttr:$fn,
+    Variadic<AnyType>:$inputs,
+    DefaultValuedStrAttr<StrAttr, "">:$name
+  );
+
+  let results = (outs Variadic<AnyType>:$outputs);
+
+  let assemblyFormat = [{
+    $fn `(` $inputs `)` attr-dict `:` functional-type($inputs, results)
+  }];
+}
+
+def GenerateOp : Enzyme_Op<"generate", [DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
+  let summary = "Generate an execution trace and weight from a probabilistic function";
+  let description = [{
+  Generate an execution trace and weight from a probabilistic function. If a constraint dict
+  is provided, we will use the corresponding constraint values instead of generating new samples.
+  }];
+
+  let arguments = (ins
+    FlatSymbolRefAttr:$fn,
+    Variadic<AnyType>:$inputs,
+    UI64Attr:$trace,
+    OptionalAttr<ArrayAttr>:$constraints,
+    DefaultValuedStrAttr<StrAttr, "">:$name
+  );
+
   let results = (outs Variadic<AnyType>:$outputs);
 
   let assemblyFormat = [{

enzyme/Enzyme/MLIR/Dialect/Ops.cpp

@@ -439,22 +439,6 @@ void BroadcastOp::build(OpBuilder &builder, OperationState &result, Value input,
   build(builder, result, resultTy, input, shapeAttr);
 }
 
-//===----------------------------------------------------------------------===//
-// SampleOp
-//===----------------------------------------------------------------------===//
-
-LogicalResult SampleOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
-  // TODO: Verify that the result type is same as the type of the referenced
-  // func.func op.
-  auto global =
-      symbolTable.lookupNearestSymbolFrom<func::FuncOp>(*this, getFnAttr());
-  if (!global)
-    return emitOpError("'")
-           << getFn() << "' does not reference a valid global funcOp";
-
-  return success();
-}
-
 /**
  *
  * Modifies activities for the AutoDiffOp.
@@ -713,3 +697,60 @@ void AutoDiffOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                              MLIRContext *context) {
   patterns.add<ReverseRetOpt>(context);
 }
+
+//===----------------------------------------------------------------------===//
+// SampleOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult SampleOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  // TODO: Verify that the result type is same as the type of the referenced
+  // func.func op.
+  auto global =
+      symbolTable.lookupNearestSymbolFrom<func::FuncOp>(*this, getFnAttr());
+  if (!global)
+    return emitOpError("'")
+           << getFn() << "' does not reference a valid global funcOp";
+
+  if (getLogpdfAttr()) {
+    auto global = symbolTable.lookupNearestSymbolFrom<func::FuncOp>(
+        *this, getLogpdfAttr());
+    if (!global)
+      return emitOpError("'")
+             << getLogpdf().value() << "' does not reference a valid global "
+             << "funcOp";
+  }
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// GenerateOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult GenerateOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  // TODO: Verify that the result type is same as the type of the referenced
+  // func.func op.
+  auto global =
+      symbolTable.lookupNearestSymbolFrom<func::FuncOp>(*this, getFnAttr());
+  if (!global)
+    return emitOpError("'")
+           << getFn() << "' does not reference a valid global funcOp";
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// SimulateOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult SimulateOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  // TODO: Verify that the result type is same as the type of the referenced
+  // func.func op.
+  auto global =
+      symbolTable.lookupNearestSymbolFrom<func::FuncOp>(*this, getFnAttr());
+  if (!global)
+    return emitOpError("'")
+           << getFn() << "' does not reference a valid global funcOp";
+
+  return success();
+}

enzyme/Enzyme/MLIR/Interfaces/CMakeLists.txt

@@ -11,8 +11,10 @@ add_mlir_library(MLIREnzymeAutoDiffInterface
   CloneFunction.cpp
   GradientUtils.cpp
   GradientUtilsReverse.cpp
+  ProbProgUtils.cpp
   EnzymeLogic.cpp
   EnzymeLogicReverse.cpp
+  EnzymeLogicProbProg.cpp
 
   DEPENDS
   MLIRAutoDiffOpInterfaceIncGen

enzyme/Enzyme/MLIR/Interfaces/ProbProgUtils.cpp

@@ -0,0 +1,68 @@
+//===- ProbProgUtils.cpp - Utilities for probprog interfaces
+//--------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "Interfaces/ProbProgUtils.h"
+#include "Dialect/Ops.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
+
+// TODO: this shouldn't depend on specific dialects except Enzyme.
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+
+#include "CloneFunction.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Dominance.h"
+#include "llvm/ADT/BreadthFirstIterator.h"
+
+using namespace mlir;
+using namespace mlir::enzyme;
+
+MProbProgUtils *MProbProgUtils::CreateFromClone(FunctionOpInterface toeval,
+                                                MProbProgMode mode) {
+  if (toeval.getFunctionBody().empty()) {
+    llvm::errs() << toeval << "\n";
+    llvm_unreachable("Creating MProbProgUtils from empty function");
+  }
+
+  std::string suffix;
+
+  switch (mode) {
+  case MProbProgMode::Call:
+    suffix = "call";
+    break;
+  case MProbProgMode::Generate:
+    suffix = "generate";
+    break;
+  case MProbProgMode::Simulate:
+    suffix = "simulate";
+    break;
+  default:
+    llvm_unreachable("Invalid MProbProgMode\n");
+  }
+
+  OpBuilder builder(toeval.getContext());
+
+  auto NewF = cast<FunctionOpInterface>(toeval->cloneWithoutRegions());
+  SymbolTable::setSymbolName(NewF, toeval.getName().str() + "." + suffix);
+  NewF.setType(toeval.getFunctionType());
+
+  Operation *parent = toeval->getParentWithTrait<OpTrait::SymbolTable>();
+  SymbolTable table(parent);
+  table.insert(NewF);
+
+  IRMapping originalToNew;
+  std::map<Operation *, Operation *> originalToNewOps;
+  cloneInto(&toeval.getFunctionBody(), &NewF.getFunctionBody(), originalToNew,
+            originalToNewOps);
+
+  return new MProbProgUtils(NewF, toeval, originalToNew, originalToNewOps,
+                            mode);
+}

enzyme/Enzyme/MLIR/Interfaces/ProbProgUtils.h

@@ -0,0 +1,53 @@
+//===- ProbProgUtils.h - Utilities for probprog interfaces -------* C++
+//-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ENZYME_MLIR_INTERFACES_PROBPROG_UTILS_H
+#define ENZYME_MLIR_INTERFACES_PROBPROG_UTILS_H
+
+#include "mlir/IR/IRMapping.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
+
+#include "CloneFunction.h"
+#include "EnzymeLogic.h"
+
+#include <functional>
+
+namespace mlir {
+namespace enzyme {
+
+class MProbProgUtils {
+public:
+  FunctionOpInterface newFunc;
+
+  MProbProgMode mode;
+  FunctionOpInterface oldFunc;
+  IRMapping originalToNewFn;
+  std::map<Operation *, Operation *> originalToNewFnOps;
+
+private:
+  Block *initializationBlock;
+
+public:
+  MProbProgUtils(FunctionOpInterface newFunc_, FunctionOpInterface oldFunc_,
+                 IRMapping &originalToNewFn_,
+                 std::map<Operation *, Operation *> &originalToNewFnOps_,
+                 MProbProgMode mode_)
+      : newFunc(newFunc_), mode(mode_), oldFunc(oldFunc_),
+        originalToNewFn(originalToNewFn_),
+        originalToNewFnOps(originalToNewFnOps_),
+        initializationBlock(&*(newFunc.getFunctionBody().begin())) {}
+
+  static MProbProgUtils *CreateFromClone(FunctionOpInterface toeval,
+                                         MProbProgMode mode);
+};
+
+} // namespace enzyme
+} // namespace mlir
+
+#endif // ENZYME_MLIR_INTERFACES_PROBPROG_UTILS_H

enzyme/Enzyme/MLIR/Passes/CMakeLists.txt

@@ -5,6 +5,7 @@ add_mlir_doc(Passes EnzymePasses ./ -gen-pass-doc)
 
 add_mlir_dialect_library(MLIREnzymeTransforms
     EnzymeMLIRPass.cpp
+    ProbProgMLIRPass.cpp
     EnzymeBatchPass.cpp
     EnzymeWrapPass.cpp
     PrintActivityAnalysis.cpp

enzyme/Enzyme/MLIR/Passes/Passes.h

@@ -26,6 +26,8 @@ class DominanceInfo;
 namespace enzyme {
 std::unique_ptr<Pass> createDifferentiatePass();
 
+std::unique_ptr<Pass> createProbProgPass();
+
 std::unique_ptr<Pass> createBatchPass();
 
 std::unique_ptr<Pass> createDifferentiateWrapperPass();

enzyme/Enzyme/MLIR/Passes/Passes.td

@@ -39,6 +39,27 @@ def DifferentiatePass : Pass<"enzyme"> {
   let constructor = "mlir::enzyme::createDifferentiatePass()";
 }
 
+def ProbProgPass : Pass<"probprog"> {
+  let summary = "ProbProg Passes";
+  let dependentDialects = [
+    "arith::ArithDialect",
+    "complex::ComplexDialect",
+    "cf::ControlFlowDialect",
+    "tensor::TensorDialect",
+    "enzyme::EnzymeDialect",
+  ];
+  let options = [
+    Option<
+      /*C++ variable name=*/"postpasses",
+      /*CLI argument=*/"postpasses",
+      /*type=*/"std::string",
+      /*default=*/"",
+      /*description=*/"Optimization passes to apply to generated probabilistic programs"
+    >,
+  ];
+  let constructor = "mlir::enzyme::createProbProgPass()";
+}
+
 def BatchPass : Pass<"enzyme-batch"> {
   let summary = "Batch Passes";
   let dependentDialects = [