mlir: Func call reverse diff

enzyme/Enzyme/MLIR/Implementations/FuncAutoDiffOpInterfaceImpl.cpp

@@ -29,9 +29,140 @@ namespace {
 #include "Implementations/FuncDerivatives.inc"
 } // namespace
 
+static std::optional<mlir::FunctionOpInterface>
+getContainingFunction(Operation *orig) {
+  Operation *parent;
+  while (parent = orig->getParentOp()) {
+    if (auto func = dyn_cast<mlir::FunctionOpInterface>(parent)) {
+      return std::optional(func);
+    }
+  }
+
+  return std::nullopt;
+}
+
+class AutoDiffCallRev
+    : public ReverseAutoDiffOpInterface::ExternalModel<AutoDiffCallRev,
+                                                       func::CallOp> {
+public:
+  LogicalResult createReverseModeAdjoint(Operation *orig, OpBuilder &builder,
+                                         MGradientUtilsReverse *gutils,
+                                         SmallVector<Value> caches) const {
+    DerivativeMode mode = DerivativeMode::ReverseModeGradient;
+
+    SymbolTable symbolTable = SymbolTable::getNearestSymbolTable(orig);
+
+    func::CallOp callOp = cast<func::CallOp>(orig);
+
+    auto parent = getContainingFunction(orig);
+    if (parent.has_value() &&
+        callOp.getCallee() == parent.value().getNameAttr()) {
+      // TODO: Recursion chains
+      orig->emitError() << "could not emit adjoint of recursive call at: "
+                        << *orig << "\n";
+      return failure();
+    }
+
+    Operation *callee = symbolTable.lookup(callOp.getCallee());
+    auto fn = cast<FunctionOpInterface>(callee);
+
+    auto narg = orig->getNumOperands();
+    auto nret = orig->getNumResults();
+
+    std::vector<DIFFE_TYPE> RetActivity;
+    for (auto res : callOp.getResults()) {
+      RetActivity.push_back(
+          gutils->isConstantValue(res) ? DIFFE_TYPE::CONSTANT
+          : res.getType().cast<AutoDiffTypeInterface>().isMutable()
+              ? DIFFE_TYPE::DUP_ARG
+              : DIFFE_TYPE::OUT_DIFF);
+    }
+
+    std::vector<DIFFE_TYPE> ArgActivity;
+    for (auto arg : callOp.getOperands()) {
+      ArgActivity.push_back(
+          gutils->isConstantValue(arg) ? DIFFE_TYPE::CONSTANT
+          : arg.getType().cast<AutoDiffTypeInterface>().isMutable()
+              ? DIFFE_TYPE::DUP_ARG
+              : DIFFE_TYPE::OUT_DIFF);
+    }
+
+    if (llvm::any_of(ArgActivity,
+                     [&](auto act) { return act == DIFFE_TYPE::DUP_ARG; }) ||
+        llvm::any_of(RetActivity,
+                     [&](auto act) { return act == DIFFE_TYPE::DUP_ARG; })) {
+      // NOTE: this current approach fails when the function is not read only.
+      //       i.e. it can modify its arguments.
+      orig->emitError() << "could not emit adjoint with mutable types in: "
+                        << *orig << "\n";
+      return failure();
+    }
+
+    std::vector<bool> volatile_args(narg, true);
+    std::vector<bool> returnShadow(narg, false);
+    std::vector<bool> returnPrimal(nret, false);
+
+    auto type_args = gutils->TA.getAnalyzedTypeInfo(fn);
+
+    bool freeMemory = true;
+    size_t width = 1;
+
+    auto revFn = gutils->Logic.CreateReverseDiff(
+        fn, RetActivity, ArgActivity, gutils->TA, returnPrimal, returnShadow,
+        mode, freeMemory, width, /*addedType*/ nullptr, type_args,
+        volatile_args, /*augmented*/ nullptr);
+
+    SmallVector<Value> revArguments;
+
+    for (auto cache : caches) {
+      revArguments.push_back(gutils->popCache(cache, builder));
+    }
+
+    for (auto result : callOp.getResults()) {
+      if (gutils->isConstantValue(result))
+        continue;
+      revArguments.push_back(gutils->diffe(result, builder));
+    }
+
+    auto revCallOp = builder.create<func::CallOp>(
+        orig->getLoc(), cast<func::FuncOp>(revFn), revArguments);
+
+    int revIndex = 0;
+    for (auto arg : callOp.getOperands()) {
+      if (gutils->isConstantValue(arg))
+        continue;
+      auto diffe = revCallOp.getResult(revIndex);
+      gutils->addToDiffe(arg, diffe, builder);
+      revIndex++;
+    }
+
+    return success();
+  }
+
+  SmallVector<Value> cacheValues(Operation *orig,
+                                 MGradientUtilsReverse *gutils) const {
+    SmallVector<Value> cachedArguments;
+
+    Operation *newOp = gutils->getNewFromOriginal(orig);
+    OpBuilder cacheBuilder(newOp);
+
+    for (auto arg : orig->getOperands()) {
+      Value cache = gutils->initAndPushCache(gutils->getNewFromOriginal(arg),
+                                             cacheBuilder);
+      cachedArguments.push_back(cache);
+    }
+
+    return cachedArguments;
+  }
+
+  void createShadowValues(Operation *op, OpBuilder &builder,
+                          MGradientUtilsReverse *gutils) const {}
+};
+
 void mlir::enzyme::registerFuncDialectAutoDiffInterface(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *context, func::FuncDialect *) {
     registerInterfaces(context);
+    func::CallOp::attachInterface<AutoDiffCallRev>(*context);
   });
 }

enzyme/Enzyme/MLIR/enzymemlir-opt.cpp

@@ -18,6 +18,7 @@
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
+#include "mlir/Dialect/Func/Extensions/InlinerExtension.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
@@ -72,6 +73,8 @@ int main(int argc, char **argv) {
 
   mlir::registerenzymePasses();
 
+  mlir::func::registerInlinerExtension(registry);
+
   // Register the standard passes we want.
   mlir::registerCSEPass();
   mlir::registerConvertAffineToStandardPass();

enzyme/test/MLIR/ReverseMode/func.mlir

@@ -0,0 +1,44 @@
+// RUN: %eopt --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops --enzyme-simplify-math %s | FileCheck %s
+
+module {
+    func.func @square(%arg0: f32) -> f32 {
+        %0 = arith.mulf %arg0, %arg0 : f32
+        return %0 : f32
+    }
+
+    func.func @inactive(%arg0: f32, %arg1: i32) -> (f32, i32) {
+        %0 = arith.constant 42.0 : f32
+        %1 = arith.constant 42 : i32
+        return %0, %1 : f32, i32
+    }
+
+    func.func @main(%arg0: f32) -> f32 {
+        %0 = func.call @square(%arg0) : (f32) -> f32
+        %1 = arith.constant 10 : i32
+        %2:2 = func.call @inactive(%arg0, %1) : (f32, i32) -> (f32, i32)
+        %3 = arith.addf %0, %2#0 : f32
+        return %3 : f32
+    }
+}
+
+// CHECK:  func.func @main(%arg0: f32, %arg1: f32) -> f32 {
+// CHECK-NEXT:    %c10_i32 = arith.constant 10 : i32
+// CHECK-NEXT:    %0 = call @square(%arg0) : (f32) -> f32
+// CHECK-NEXT:    %1:2 = call @inactive(%arg0, %c10_i32) : (f32, i32) -> (f32, i32)
+// CHECK-NEXT:    %2 = call @diffeinactive(%arg0, %c10_i32, %arg1) : (f32, i32, f32) -> f32
+// CHECK-NEXT:    %3 = call @diffesquare(%arg0, %arg1) : (f32, f32) -> f32
+// CHECK-NEXT:    %4 = arith.addf %2, %3 : f32
+// CHECK-NEXT:    return %4 : f32
+// CHECK-NEXT:  }
+
+// CHECK:  func.func private @diffeinactive(%arg0: f32, %arg1: i32, %arg2: f32) -> f32 {
+// CHECK-NEXT:    %cst = arith.constant 0.000000e+00 : f32
+// CHECK-NEXT:    return %cst : f32
+// CHECK-NEXT:  }
+
+// CHECK:  func.func private @diffesquare(%arg0: f32, %arg1: f32) -> f32 {
+// CHECK-NEXT:    %0 = arith.mulf %arg1, %arg0 : f32
+// CHECK-NEXT:    %1 = arith.mulf %arg1, %arg0 : f32
+// CHECK-NEXT:    %2 = arith.addf %0, %1 : f32
+// CHECK-NEXT:    return %2 : f32
+// CHECK-NEXT:  }