Register reverse mode AD for memref.alloca_scope

enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp

@@ -203,6 +203,81 @@ 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);
+      }
+    }
+
+    auto revScope =
+        memref::AllocaScopeOp::create(builder, op->getLoc(), TypeRange());
+    Block *revBody = builder.createBlock(&revScope.getBodyRegion());
+    OpBuilder bodyBuilder(revBody, revBody->end());
+    memref::AllocaScopeReturnOp::create(bodyBuilder, op->getLoc(),
+                                        ValueRange());
+    bodyBuilder.setInsertionPoint(revBody->getTerminator());
+
+    Block &oldBody = scopeOp.getBodyRegion().front();
+    bool valid = true;
+
+    // Values defined in the scoped region cannot be used outside it. Reset
+    // their adjoints before propagating gradients through the scoped body.
+    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();
+    }
+
+    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> {
@@ -299,5 +374,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);
   });
 }

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