fix reverse ad for memref.alloca

Author: xys-syx

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);
   }
   /*
@@ -240,11 +244,13 @@ void mlir::enzyme::MGradientUtils::setDiffe(mlir::Value val, mlir::Value toset,
 
 void mlir::enzyme::MGradientUtils::setInvertedPointer(Value val, Value toset) {
   assert(getShadowType(val.getType()) == toset.getType());
-  auto found = invertedPointers.lookupOrNull(val);
-  assert(found != nullptr);
-  auto placeholder = found.getDefiningOp<enzyme::PlaceholderOp>();
-  placeholder.replaceAllUsesWith(toset);
-  erase(placeholder);
+  
+  if (auto found = invertedPointers.lookupOrNull(val)) {
+    if (auto placeholder = found.getDefiningOp<enzyme::PlaceholderOp>()) {
+      placeholder.replaceAllUsesWith(toset);
+      erase(placeholder);
+    }
+  }
   invertedPointers.map(val, toset);
 }
 

enzyme/test/MLIR/ReverseMode/alloca.mlir

@@ -1,4 +1,4 @@
-// RUN: %eopt --split-input-file --enzyme --canonicalize --remove-unnecessary-enzyme-ops --enzyme-simplify-math %s | FileCheck %s
+// RUN: %eopt --enzyme %s | FileCheck %s
 
 func.func @foo_flat(%x : f64) -> f64 {
   %buf = memref.alloca() : memref<f64>
@@ -18,12 +18,45 @@ func.func @dfoo_flat(%x: f64, %dout : f64) -> f64 {
 // CHECK-LABEL:   func.func private @diffefoo_flat(
 // CHECK-SAME:      %[[X:[^,]+]]: f64,
 // CHECK-SAME:      %[[DOUT:[^)]+]]: f64) -> f64 {
-// A shadow memref.alloca must be created and zero-initialized so the
-// reverse-mode adjoint can accumulate gradients into it.
-// CHECK-DAG:       %[[ZERO:.*]] = arith.constant 0.000000e+00 : f64
-// CHECK-DAG:       %[[DBUF:.*]] = memref.alloca() : memref<f64>
-// CHECK:           memref.store %[[ZERO]], %[[DBUF]][] : memref<f64>
-// No leftover placeholders should survive the differentiation.
+
+// 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
-// The function ultimately returns the gradient w.r.t. x (= dout).
-// CHECK:           return %{{.*}} : f64

enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir

@@ -14,7 +14,7 @@ func.func @foo(%x : f64) -> f64{
 func.func @dfoo(%x: f64, %dout : f64) -> f64 {
   %dx = enzyme.autodiff @foo(%x, %dout) {
     activity = [#enzyme<activity enzyme_active>],
-    ret_activity = [#enzyme<acitivity enzyme_activenoneed>]
+    ret_activity = [#enzyme<activity enzyme_activenoneed>]
   } : (f64, f64) -> (f64)
   return %dx : f64
 }