EnzymeAD · vimarsh6739 · May 13, 2026 · May 13, 2026 · May 14, 2026 · May 14, 2026
diff --git a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp
@@ -18,13 +18,10 @@
 #include "Interfaces/GradientUtilsReverse.h"
 
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/Support/LogicalResult.h"
 
-// TODO: We need a way to zero out a memref (which linalg.fill does), but
-// ideally we wouldn't depend on the linalg dialect.
-#include "mlir/Dialect/Linalg/IR/Linalg.h"
-
 using namespace mlir;
 using namespace mlir::enzyme;
 
@@ -203,6 +200,147 @@ struct SubViewOpInterfaceReverse
   }
 };
 
+struct AllocaScopeOpInterfaceReverse
+    : public ReverseAutoDiffOpInterface::ExternalModel<
+          AllocaScopeOpInterfaceReverse, memref::AllocaScopeOp> {
+  LogicalResult createReverseModeAdjoint(Operation *op, OpBuilder &builder,
+                                         MGradientUtilsReverse *gutils,
+                                         SmallVector<Value> caches) const {
+    auto scopeOp = cast<memref::AllocaScopeOp>(op);
+
+    SmallVector<bool> resultsActive(scopeOp->getNumResults(), false);
+    SmallVector<Value> incomingGradients;
+    incomingGradients.reserve(scopeOp->getNumResults());
+    for (OpResult result : scopeOp->getResults()) {
+      bool active = !gutils->isConstantValue(result);
+      resultsActive[result.getResultNumber()] = active;
+      if (active) {
+        incomingGradients.push_back(gutils->diffe(result, builder));
+        gutils->zeroDiffe(result, builder);
+      }
+    }
+
+    Region &scopeRegion = scopeOp.getBodyRegion();
+    SmallVector<Value> capturedInputs;
+    scopeRegion.walk([&](Operation *inner) {
+      for (Value operand : inner->getOperands()) {
+        Region *defRegion = operand.getParentRegion();
+        if (!defRegion || scopeRegion.isAncestor(defRegion))
+          continue;
+        if (gutils->isConstantValue(operand))
+          continue;
+        auto iface = dyn_cast<AutoDiffTypeInterface>(operand.getType());
+        if (!iface || iface.isMutable())
+          continue;
+        if (!llvm::is_contained(capturedInputs, operand))
+          capturedInputs.push_back(operand);
+      }
+    });
+
+    SmallVector<Value> capturedPre;
+    capturedPre.reserve(capturedInputs.size());
+    for (Value v : capturedInputs) {
+      capturedPre.push_back(gutils->diffe(v, builder));
+      gutils->zeroDiffe(v, builder);
+    }
+
+    auto newScope = cast<memref::AllocaScopeOp>(gutils->getNewFromOriginal(op));
+    newScope->moveBefore(builder.getInsertionBlock(),
+                         builder.getInsertionPoint());
+
+    Block &newBody = newScope.getBodyRegion().front();
+    OpBuilder bodyBuilder(newBody.getTerminator());
+
+    Block &oldBody = scopeOp.getBodyRegion().front();
+    bool valid = true;
+
+    for (Operation &innerOp : oldBody.getOperations()) {
+      for (Value result : innerOp.getResults()) {
+        if (!gutils->isConstantValue(result)) {
+          auto iface = dyn_cast<AutoDiffTypeInterface>(result.getType());
+          if (iface && !iface.isMutable())
+            gutils->zeroDiffe(result, bodyBuilder);
+        }
+      }
+    }
+
+    Operation *term = oldBody.getTerminator();
+    unsigned incomingIdx = 0;
+    for (auto indexedOperand : llvm::enumerate(term->getOperands())) {
+      if (!resultsActive[indexedOperand.index()])
+        continue;
+
+      Value operand = indexedOperand.value();
+      if (!gutils->isConstantValue(operand))
+        gutils->addToDiffe(operand, incomingGradients[incomingIdx],
+                           bodyBuilder);
+      incomingIdx++;
+    }
+
+    auto first = oldBody.rbegin();
+    first++;
+
+    for (auto it = first; it != oldBody.rend(); it++) {
+      valid &= gutils->Logic.visitChild(&*it, bodyBuilder, gutils).succeeded();
+    }
+
+    if (capturedInputs.empty())
+      return success(valid);
+
+    SmallVector<Value> contributions;
+    contributions.reserve(capturedInputs.size());
+    for (Value v : capturedInputs)
+      contributions.push_back(gutils->diffe(v, bodyBuilder));
+
+    unsigned numPrimal = newScope->getNumResults();
+    Operation *bodyTerm = newBody.getTerminator();
+    SmallVector<Value> retVals(bodyTerm->getOperands().begin(),
+                               bodyTerm->getOperands().end());
+    retVals.append(contributions.begin(), contributions.end());
+
+    SmallVector<Type> newResultTypes(newScope->getResultTypes().begin(),
+                                     newScope->getResultTypes().end());
+    for (Value v : capturedInputs)
+      newResultTypes.push_back(gutils->getShadowType(v.getType()));
+
+    OpBuilder scopeBuilder(newScope);
+    auto extScope = memref::AllocaScopeOp::create(scopeBuilder, op->getLoc(),
+                                                  newResultTypes);
+    extScope.getBodyRegion().takeBody(newScope.getBodyRegion());
+
+    Block &extBody = extScope.getBodyRegion().front();
+    Operation *movedTerm = extBody.getTerminator();
+    OpBuilder termBuilder(movedTerm);
+    memref::AllocaScopeReturnOp::create(termBuilder, movedTerm->getLoc(),
+                                        retVals);
+    gutils->erase(movedTerm);
+
+    for (unsigned i = 0; i < numPrimal; ++i) {
+      newScope->getResult(i).replaceAllUsesWith(extScope.getResult(i));
+      gutils->originalToNewFn.map(scopeOp->getResult(i), extScope.getResult(i));
+    }
+    gutils->erase(newScope);
+
+    builder.setInsertionPointAfter(extScope);
+    for (auto indexed : llvm::enumerate(capturedInputs)) {
+      Value v = indexed.value();
+      gutils->setDiffe(v, capturedPre[indexed.index()], builder);
+      gutils->addToDiffe(v, extScope.getResult(numPrimal + indexed.index()),
+                         builder);
+    }
+
+    return success(valid);
+  }
+
+  SmallVector<Value> cacheValues(Operation *op,
+                                 MGradientUtilsReverse *gutils) const {
+    return SmallVector<Value>();
+  }
+
+  void createShadowValues(Operation *op, OpBuilder &builder,
+                          MGradientUtilsReverse *gutils) const {}
+};
+
 class MemRefClonableTypeInterface
     : public ClonableTypeInterface::ExternalModel<MemRefClonableTypeInterface,
                                                   MemRefType> {
@@ -273,14 +411,33 @@ class MemRefAutoDiffTypeInterface
   LogicalResult zeroInPlace(Type self, OpBuilder &builder, Location loc,
                             Value val) const {
     auto MT = cast<MemRefType>(self);
-    if (auto iface = dyn_cast<AutoDiffTypeInterface>(MT.getElementType())) {
-      if (!iface.isMutable()) {
-        Value zero = iface.createNullValue(builder, loc);
-        linalg::FillOp::create(builder, loc, zero, val);
-      }
-    } else {
+    auto eltIface = dyn_cast<AutoDiffTypeInterface>(MT.getElementType());
+    if (!eltIface || eltIface.isMutable())
       return failure();
+    Value zero = eltIface.createNullValue(builder, loc);
+
+    if (MT.getRank() == 0) {
+      memref::StoreOp::create(builder, loc, zero, val, ValueRange{});
+      return success();
     }
+
+    Value c0 = arith::ConstantIndexOp::create(builder, loc, 0);
+    Value c1 = arith::ConstantIndexOp::create(builder, loc, 1);
+
+    SmallVector<Value> lbs(MT.getRank(), c0);
+    SmallVector<Value> steps(MT.getRank(), c1);
+    SmallVector<Value> ubs;
+    for (auto [i, d] : llvm::enumerate(MT.getShape())) {
+      ubs.push_back(
+          d == ShapedType::kDynamic
+              ? memref::DimOp::create(builder, loc, val, i).getResult()
+              : arith::ConstantIndexOp::create(builder, loc, d).getResult());
+    }
+
+    scf::ParallelOp::create(builder, loc, lbs, ubs, steps,
+                            [&](OpBuilder &b, Location l, ValueRange ivs) {
+                              memref::StoreOp::create(b, l, zero, val, ivs);
+                            });
     return success();
   }
 
@@ -299,5 +456,7 @@ void mlir::enzyme::registerMemRefDialectAutoDiffInterface(
     memref::LoadOp::attachInterface<LoadOpInterfaceReverse>(*context);
     memref::StoreOp::attachInterface<StoreOpInterfaceReverse>(*context);
     memref::SubViewOp::attachInterface<SubViewOpInterfaceReverse>(*context);
+    memref::AllocaScopeOp::attachInterface<AllocaScopeOpInterfaceReverse>(
+        *context);
   });
 }
diff --git a/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp b/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp
@@ -223,8 +223,12 @@ void mlir::enzyme::MGradientUtils::setDiffe(mlir::Value val, mlir::Value toset,
   }
   assert(!isConstantValue(val));
 
+  bool isMutable = false;
+  if (auto iface = dyn_cast<AutoDiffTypeInterface>(val.getType()))
+    isMutable = iface.isMutable();
+
   if (mode == DerivativeMode::ForwardMode ||
-      mode == DerivativeMode::ForwardModeSplit) {
+      mode == DerivativeMode::ForwardModeSplit || isMutable) {
     setInvertedPointer(val, toset);
   }
   /*

diff --git a/enzyme/test/MLIR/ReverseMode/alloca.mlir b/enzyme/test/MLIR/ReverseMode/alloca.mlir
@@ -0,0 +1,62 @@
+// RUN: %eopt --enzyme %s | FileCheck %s
+
+func.func @foo_flat(%x : f64) -> f64 {
+  %buf = memref.alloca() : memref<f64>
+  memref.store %x, %buf[] : memref<f64>
+  %y = memref.load %buf[] : memref<f64>
+  return %y : f64
+}
+
+func.func @dfoo_flat(%x: f64, %dout : f64) -> f64 {
+  %dx = enzyme.autodiff @foo_flat(%x, %dout) {
+    activity = [#enzyme<activity enzyme_active>],
+    ret_activity = [#enzyme<activity enzyme_activenoneed>]
+  } : (f64, f64) -> (f64)
+  return %dx : f64
+}
+
+// CHECK-LABEL:   func.func private @diffefoo_flat(
+// CHECK-SAME:      %[[X:[^,]+]]: f64,
+// CHECK-SAME:      %[[DOUT:[^)]+]]: f64) -> f64 {
+
+// CHECK:           %[[GX:.+]] = "enzyme.init"() : () -> !enzyme.Gradient<f64>
+// CHECK:           "enzyme.set"(%[[GX]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
+// CHECK:           %[[CS:.+]] = "enzyme.init"() : () -> !enzyme.Cache<memref<f64>>
+// CHECK:           %[[CL:.+]] = "enzyme.init"() : () -> !enzyme.Cache<memref<f64>>
+// CHECK:           %[[GY:.+]] = "enzyme.init"() : () -> !enzyme.Gradient<f64>
+// CHECK:           "enzyme.set"(%[[GY]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
+
+// CHECK:           %[[DBUF:.+]] = memref.alloca() : memref<f64>
+// CHECK:           %[[ZERO_INIT:.+]] = arith.constant 0.000000e+00 : f64
+// CHECK:           memref.store %[[ZERO_INIT]], %[[DBUF]][] : memref<f64>
+
+// CHECK:           %[[BUF:.+]] = memref.alloca() : memref<f64>
+// CHECK:           "enzyme.push"(%[[CS]], %[[DBUF]]) : (!enzyme.Cache<memref<f64>>, memref<f64>) -> ()
+// CHECK:           memref.store %[[X]], %[[BUF]][] : memref<f64>
+// CHECK:           "enzyme.push"(%[[CL]], %[[DBUF]]) : (!enzyme.Cache<memref<f64>>, memref<f64>) -> ()
+// CHECK:           memref.load %[[BUF]][] : memref<f64>
+// CHECK:           cf.br ^bb1
+// CHECK:         ^bb1:
+
+// CHECK:           %{{.+}} = "enzyme.get"(%[[GY]]) : (!enzyme.Gradient<f64>) -> f64
+// CHECK:           arith.addf %{{.+}}, %[[DOUT]] : f64
+// CHECK:           "enzyme.set"(%[[GY]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
+
+// CHECK:           %{{.+}} = "enzyme.get"(%[[GY]]) : (!enzyme.Gradient<f64>) -> f64
+// CHECK:           %[[POPL:.+]] = "enzyme.pop"(%[[CL]]) : (!enzyme.Cache<memref<f64>>) -> memref<f64>
+// CHECK:           memref.load %[[POPL]][] : memref<f64>
+// CHECK:           arith.addf
+// CHECK:           memref.store %{{.+}}, %[[POPL]][] : memref<f64>
+
+// CHECK:           %[[POPS:.+]] = "enzyme.pop"(%[[CS]]) : (!enzyme.Cache<memref<f64>>) -> memref<f64>
+// CHECK:           memref.load %[[POPS]][] : memref<f64>
+// CHECK:           %{{.+}} = "enzyme.get"(%[[GX]]) : (!enzyme.Gradient<f64>) -> f64
+// CHECK:           arith.addf
+// CHECK:           "enzyme.set"(%[[GX]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
+// CHECK:           %[[ZERO_CLEAR:.+]] = arith.constant 0.000000e+00 : f64
+// CHECK:           memref.store %[[ZERO_CLEAR]], %[[POPS]][] : memref<f64>
+
+// CHECK:           %{{.+}} = "enzyme.get"(%[[GX]]) : (!enzyme.Gradient<f64>) -> f64
+// CHECK:           return %{{.+}} : f64
+
+// CHECK-NOT:       enzyme.placeholder
diff --git a/enzyme/test/MLIR/ReverseMode/alloca_scope.mlir b/enzyme/test/MLIR/ReverseMode/alloca_scope.mlir
@@ -0,0 +1,94 @@
+// RUN: %eopt --enzyme --canonicalize --remove-unnecessary-enzyme-ops --enzyme-simplify-math %s | FileCheck %s
+
+func.func private @scope_overwrite(%x: memref<f32>) {
+  memref.alloca_scope {
+    %val = memref.load %x[] : memref<f32>
+    %cos = math.cos %val : f32
+    memref.store %cos, %x[] : memref<f32>
+  }
+  return
+}
+
+func.func @dscope_overwrite(%x: memref<f32>, %dx: memref<f32>) {
+  enzyme.autodiff @scope_overwrite(%x, %dx) {
+    activity = [#enzyme<activity enzyme_dup>],
+    ret_activity = []
+  } : (memref<f32>, memref<f32>) -> ()
+  return
+}
+
+func.func private @scope_load_mul(%x: memref<f64>, %c: f64) -> f64 {
+  %r = memref.alloca_scope -> (f64) {
+    %v = memref.load %x[] : memref<f64>
+    %p = arith.mulf %v, %c : f64
+    memref.alloca_scope.return %p : f64
+  }
+  return %r : f64
+}
+
+func.func @dscope_load_mul(%x: memref<f64>, %dx: memref<f64>,
+                           %c: f64, %dout: f64) -> f64 {
+  %res = enzyme.autodiff @scope_load_mul(%x, %dx, %c, %dout) {
+    activity = [#enzyme<activity enzyme_dup>, #enzyme<activity enzyme_active>],
+    ret_activity = [#enzyme<activity enzyme_activenoneed>]
+  } : (memref<f64>, memref<f64>, f64, f64) -> f64
+  return %res : f64
+}
+
+// CHECK-LABEL:   func.func private @diffescope_overwrite(
+// CHECK-SAME:      %[[X:.*]]: memref<f32>,
+// CHECK-SAME:      %[[DX:.*]]: memref<f32>) {
+// CHECK:           %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:           %[[CDX0:.*]] = "enzyme.init"() : () -> !enzyme.Cache<memref<f32>>
+// CHECK:           %[[CVAL:.*]] = "enzyme.init"() : () -> !enzyme.Cache<f32>
+// CHECK:           %[[CDX1:.*]] = "enzyme.init"() : () -> !enzyme.Cache<memref<f32>>
+// CHECK:           memref.alloca_scope  {
+// CHECK:             "enzyme.push"(%[[CDX0]], %[[DX]])
+// CHECK:             %[[V:.*]] = memref.load %[[X]][] : memref<f32>
+// CHECK:             "enzyme.push"(%[[CVAL]], %[[V]])
+// CHECK:             %[[COS:.*]] = math.cos %[[V]] : f32
+// CHECK:             "enzyme.push"(%[[CDX1]], %[[DX]])
+// CHECK:             memref.store %[[COS]], %[[X]][] : memref<f32>
+// CHECK:           }
+// CHECK:           %[[DX_R1:.*]] = "enzyme.pop"(%[[CDX1]])
+// CHECK:           %[[DOUT:.*]] = memref.load %[[DX_R1]][] : memref<f32>
+// CHECK:           memref.store %[[ZERO]], %[[DX_R1]][] : memref<f32>
+// CHECK:           %[[VAL:.*]] = "enzyme.pop"(%[[CVAL]])
+// CHECK:           %[[SIN:.*]] = math.sin %[[VAL]] : f32
+// CHECK:           %[[NEG:.*]] = arith.negf %[[SIN]] : f32
+// CHECK:           %[[GRAD:.*]] = arith.mulf %[[DOUT]], %[[NEG]] : f32
+// CHECK:           %[[DX_R0:.*]] = "enzyme.pop"(%[[CDX0]])
+// CHECK:           %[[OLD:.*]] = memref.load %[[DX_R0]][] : memref<f32>
+// CHECK:           %[[NEW:.*]] = arith.addf %[[OLD]], %[[GRAD]] : f32
+// CHECK:           memref.store %[[NEW]], %[[DX_R0]][] : memref<f32>
+// CHECK:           return
+
+// CHECK-LABEL:   func.func private @diffescope_load_mul(
+// CHECK-SAME:      %[[X:[^,]+]]: memref<f64>,
+// CHECK-SAME:      %[[DX:[^,]+]]: memref<f64>,
+// CHECK-SAME:      %[[C:[^,]+]]: f64,
+// CHECK-SAME:      %[[DOUT:[^)]+]]: f64) -> f64 {
+// CHECK:           %[[CDX:.*]] = "enzyme.init"() : () -> !enzyme.Cache<memref<f64>>
+// CHECK:           %[[GC:.*]] = "enzyme.init"() : () -> !enzyme.Gradient<f64>
+// CHECK:           %[[CC:.*]] = "enzyme.init"() : () -> !enzyme.Cache<f64>
+// CHECK:           %[[CV:.*]] = "enzyme.init"() : () -> !enzyme.Cache<f64>
+// CHECK:           memref.alloca_scope  -> (f64) {
+// CHECK:             "enzyme.push"(%[[CDX]], %[[DX]])
+// CHECK:             %[[V:.*]] = memref.load %[[X]][] : memref<f64>
+// CHECK:             "enzyme.push"(%[[CV]], %[[V]])
+// CHECK:             "enzyme.push"(%[[CC]], %[[C]])
+// CHECK:             %[[P:.*]] = arith.mulf %[[V]], %[[C]] : f64
+// CHECK:             memref.alloca_scope.return %[[P]] : f64
+// CHECK:           }
+// CHECK:           memref.alloca_scope  {
+// CHECK:             %[[V_R:.*]] = "enzyme.pop"(%[[CV]])
+// CHECK:             %[[C_R:.*]] = "enzyme.pop"(%[[CC]])
+// CHECK-DAG:         %[[DC:.*]] = arith.mulf %[[DOUT]], %[[V_R]] : f64
+// CHECK-DAG:         %[[DV:.*]] = arith.mulf %[[DOUT]], %[[C_R]] : f64
+// CHECK:             %[[DX_R:.*]] = "enzyme.pop"(%[[CDX]])
+// CHECK:             %[[OLD:.*]] = memref.load %[[DX_R]][] : memref<f64>
+// CHECK:             %[[NEW:.*]] = arith.addf %[[OLD]], %[[DV]] : f64
+// CHECK:             memref.store %[[NEW]], %[[DX_R]][] : memref<f64>
+// CHECK:           }
+// CHECK:           %[[OUT:.*]] = "enzyme.get"(%[[GC]]) : (!enzyme.Gradient<f64>) -> f64
+// CHECK:           return %[[OUT]] : f64