ProbProg: Raw MCMC (#2695)

* logpdf only mode

* mcmc op verify

Author: sbrantq

enzyme/Enzyme/MLIR/Dialect/EnzymeOps.td

@@ -683,6 +683,12 @@ def MCMCOp : Enzyme_Op<"mcmc", [DeclareOpInterfaceMethods<SymbolUserOpInterface>
   let description = [{
     Runs MCMC inference on selected addresses.
 
+    Two modes of operation:
+    1. Trace-based mode: `fn` and `original_trace` are provided. The model
+       function with `enzyme.sample` ops defines the density.
+    2. Custom logpdf mode: `logpdf_fn` and `initial_position` are provided.
+       The logpdf function maps position → scalar log-density directly.
+
     The `selection` attribute determines which addresses to sample via HMC/NUTS.
     All sample addresses are included in the trace tensor for consistency.
 
@@ -693,9 +699,9 @@ def MCMCOp : Enzyme_Op<"mcmc", [DeclareOpInterfaceMethods<SymbolUserOpInterface>
   }];
 
   let arguments = (ins
-    FlatSymbolRefAttr:$fn,
+    OptionalAttr<FlatSymbolRefAttr>:$fn,
     Variadic<AnyType>:$inputs,
-    AnyRankedTensor:$original_trace,
+    Optional<AnyRankedTensor>:$original_trace,
     AddressArrayAttr:$selection,
     AddressArrayAttr:$all_addresses,
 
@@ -711,6 +717,10 @@ def MCMCOp : Enzyme_Op<"mcmc", [DeclareOpInterfaceMethods<SymbolUserOpInterface>
     OptionalAttr<HMCConfigAttr>:$hmc_config,
     OptionalAttr<NUTSConfigAttr>:$nuts_config,
 
+    // Custom logpdf mode
+    OptionalAttr<FlatSymbolRefAttr>:$logpdf_fn,
+    Optional<AnyRankedTensor>:$initial_position,
+
     DefaultValuedStrAttr<StrAttr, "">:$name
   );
 
@@ -721,9 +731,13 @@ def MCMCOp : Enzyme_Op<"mcmc", [DeclareOpInterfaceMethods<SymbolUserOpInterface>
   );
 
   let assemblyFormat = [{
-    $fn `(` $inputs `)` `given` $original_trace
+    ($fn^)?
+    `(` $inputs `)`
+    (`given` $original_trace^)?
     (`inverse_mass_matrix` `=` $inverse_mass_matrix^)?
     (`step_size` `=` $step_size^)?
+    (`logpdf_fn` `=` $logpdf_fn^)?
+    (`initial_position` `=` $initial_position^)?
     attr-dict `:` functional-type(operands, results)
   }];
 

enzyme/Enzyme/MLIR/Dialect/Ops.cpp

@@ -1062,13 +1062,21 @@ LogicalResult MHOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
 //===----------------------------------------------------------------------===//
 
 LogicalResult MCMCOp::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 (auto fnAttr = getFnAttr()) {
+    auto global =
+        symbolTable.lookupNearestSymbolFrom<func::FuncOp>(*this, fnAttr);
+    if (!global)
+      return emitOpError("'")
+             << getFn().value() << "' does not reference a valid global funcOp";
+  }
+
+  if (auto logpdfAttr = getLogpdfFnAttr()) {
+    auto global =
+        symbolTable.lookupNearestSymbolFrom<func::FuncOp>(*this, logpdfAttr);
+    if (!global)
+      return emitOpError("'") << logpdfAttr.getValue()
+                              << "' does not reference a valid global funcOp";
+  }
 
   return success();
 }
@@ -1082,7 +1090,19 @@ LogicalResult MCMCOp::verify() {
         "Exactly one of hmc_config or nuts_config must be specified");
   }
 
-  // TODO: More verification
+  // TODO(#2695): More verification
+  if (!getFnAttr() && !getLogpdfFnAttr()) {
+    return emitOpError("one of `fn` or `logpdf_fn` must be specified");
+  }
+
+  if (getFnAttr() && getLogpdfFnAttr()) {
+    return emitOpError("specifying both `fn` and `logpdf_fn` is unsupported");
+  }
+
+  if (getLogpdfFnAttr() && !getInitialPosition()) {
+    return emitOpError(
+        "custom logpdf mode requires `initial_position` to be provided");
+  }
 
   return success();
 }