Convert to a static cast and add a more targeted test case

Author: vimarsh6739

enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp

@@ -225,13 +225,16 @@ 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();
-
+  // Reverse mode also uses invertedPointers for mutable values. For allocation
+  // results, forceAugmentedReturns() installs a placeholder shadow; the
+  // allocation derivative later calls setDiffe with the real shadow, so replace
+  // it here.
+  bool isMutable = cast<AutoDiffTypeInterface>(val.getType()).isMutable();
   if (mode == DerivativeMode::ForwardMode ||
       mode == DerivativeMode::ForwardModeSplit || isMutable) {
     if (isMutable) {
+      // Mutable arguments may already be mapped to a real user-provided shadow
+      // rather than a placeholder. Only replace placeholders here.
       auto existing = invertedPointers.lookupOrNull(val);
       if (existing && !existing.getDefiningOp<enzyme::PlaceholderOp>())
         return;

enzyme/test/MLIR/ReverseMode/alloca.mlir

@@ -1,4 +1,4 @@
-// RUN: %eopt --enzyme %s | FileCheck %s
+// RUN: %eopt --enzyme %s | FileCheck --implicit-check-not=enzyme.placeholder %s
 
 func.func @foo_flat(%x : f64) -> f64 {
   %buf = memref.alloca() : memref<f64>
@@ -59,4 +59,52 @@ func.func @dfoo_flat(%x: f64, %dout : f64) -> f64 {
 // CHECK:           %{{.+}} = "enzyme.get"(%[[GX]]) : (!enzyme.Gradient<f64>) -> f64
 // CHECK:           return %{{.+}} : f64
 
-// CHECK-NOT:       enzyme.placeholder
+
+// -----
+
+func.func @load_arg(%mem: memref<f64>) -> f64 {
+  %buf = memref.alloca() : memref<f64>
+  %x = memref.load %mem[] : memref<f64>
+  memref.store %x, %buf[] : memref<f64>
+  %y = memref.load %buf[] : memref<f64>
+  return %y : f64
+}
+
+func.func @dload_arg(%mem: memref<f64>, %dmem: memref<f64>, %dout : f64) {
+  enzyme.autodiff @load_arg(%mem, %dmem, %dout) {
+    activity = [#enzyme<activity enzyme_dup>],
+    ret_activity = [#enzyme<activity enzyme_activenoneed>]
+  } : (memref<f64>, memref<f64>, f64) -> ()
+  return
+}
+
+// CHECK-LABEL:   func.func private @diffeload_arg(
+// CHECK-SAME:      %[[MEM:[^,]+]]: memref<f64>,
+// CHECK-SAME:      %[[DMEM:[^,]+]]: memref<f64>,
+// CHECK-SAME:      %[[DOUT:[^)]+]]: f64) {
+// CHECK:           %[[DBUF:.+]] = memref.alloca() : memref<f64>
+// CHECK:           linalg.fill ins(%{{.+}} : f64) outs(%[[DBUF]] : memref<f64>)
+// CHECK:           %[[BUF:.+]] = memref.alloca() : memref<f64>
+// CHECK:           "enzyme.push"(%{{.+}}, %[[DMEM]]) : (!enzyme.Cache<memref<f64>>, memref<f64>) -> ()
+// CHECK:           memref.load %[[MEM]][] : memref<f64>
+// CHECK:           "enzyme.push"(%{{.+}}, %[[DBUF]]) : (!enzyme.Cache<memref<f64>>, memref<f64>) -> ()
+// CHECK:           memref.store %{{.+}}, %[[BUF]][] : memref<f64>
+// CHECK:           "enzyme.push"(%{{.+}}, %[[DBUF]]) : (!enzyme.Cache<memref<f64>>, memref<f64>) -> ()
+// CHECK:           memref.load %[[BUF]][] : memref<f64>
+// CHECK:           arith.addf %{{.+}}, %[[DOUT]] : f64
+// CHECK:           %[[POPL:.+]] = "enzyme.pop"(%{{.+}}) : (!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"(%{{.+}}) : (!enzyme.Cache<memref<f64>>) -> memref<f64>
+// CHECK:           memref.load %[[POPS]][] : memref<f64>
+// CHECK:           "enzyme.get"
+// CHECK:           arith.addf
+// CHECK:           "enzyme.set"
+// CHECK:           memref.store %{{.+}}, %[[POPS]][] : memref<f64>
+// CHECK:           "enzyme.get"
+// CHECK:           %[[POPD:.+]] = "enzyme.pop"(%{{.+}}) : (!enzyme.Cache<memref<f64>>) -> memref<f64>
+// CHECK:           memref.load %[[POPD]][] : memref<f64>
+// CHECK:           arith.addf
+// CHECK:           memref.store %{{.+}}, %[[POPD]][] : memref<f64>
+// CHECK:           return