add test for memref.alloca

Author: xys-syx

enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp

@@ -349,7 +349,7 @@ class MemRefAutoDiffTypeInterface
     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();
@@ -362,15 +362,16 @@ class MemRefAutoDiffTypeInterface
     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());
+      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);
-        });
+                            [&](OpBuilder &b, Location l, ValueRange ivs) {
+                              memref::StoreOp::create(b, l, zero, val, ivs);
+                            });
     return success();
   }
 

enzyme/test/MLIR/ReverseMode/alloca.mlir

@@ -0,0 +1,29 @@
+// RUN: %eopt --split-input-file --enzyme --canonicalize --remove-unnecessary-enzyme-ops --enzyme-simplify-math %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 {
+// 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-NOT:       enzyme.placeholder
+// The function ultimately returns the gradient w.r.t. x (= dout).
+// CHECK:           return %{{.*}} : f64