workaround driver problems

enzyme/Enzyme/MLIR/Implementations/SCFAutoDiffOpInterfaceImpl.cpp

@@ -71,8 +71,7 @@ struct ForOpEnzymeOpsRemover
     auto forOp = cast<scf::ForOp>(op);
     scf::ForOp otherForOp; // where caches pops are
 
-    if (removeOpsWithinBlock(forOp.getBody(), rewriter).failed())
-      return failure();
+    (void)removeOpsWithinBlock(forOp.getBody(), rewriter);
 
     // Gradients whose values need to be passed as iteration variables.
     llvm::SetVector<Value> updatedGradients;

enzyme/Enzyme/MLIR/Passes/RemovalUtils.cpp

@@ -49,14 +49,14 @@ mlir::enzyme::CacheInfo::merge(mlir::enzyme::CacheInfo other,
 mlir::LogicalResult
 mlir::enzyme::removeOpsWithinBlock(mlir::Block *block,
                                    mlir::PatternRewriter &rewriter) {
-  bool valid = true;
+  bool matched = false;
 
   for (auto &it : *block) {
     mlir::Operation *op = ⁢
     if (auto iface = dyn_cast<mlir::enzyme::EnzymeOpsRemoverOpInterface>(op)) {
-      valid &= iface.removeEnzymeOps(rewriter).succeeded();
+      matched |= iface.removeEnzymeOps(rewriter).succeeded();
     }
   }
 
-  return success(valid);
+  return success(matched);
 }

enzyme/Enzyme/MLIR/Passes/RemoveUnusedEnzymeOps.cpp

@@ -29,6 +29,7 @@
 
 #include "mlir/IR/Dominance.h"
 #include "llvm/Support/raw_ostream.h"
+#include <cmath>
 
 using namespace mlir;
 using namespace enzyme;
@@ -307,7 +308,8 @@ static void applyPatterns(Operation *op) {
                   InitSimplify>(op->getContext());
 
   GreedyRewriteConfig config;
-  (void)applyPatternsAndFoldGreedily(op, std::move(patterns), config);
+  config.fold = true;
+  (void)applyPatternsGreedily(op, std::move(patterns), config);
 }
 
 struct RemoveUnusedEnzymeOpsPass
@@ -324,8 +326,11 @@ struct RemoveUnusedEnzymeOpsPass
     // of the interface will erase operations which are not only nested in the
     // currently matched operation (for example, the for op where the pops are
     // located). As such, we use the greedy driver with the option to run only
-    // on the pre-existing operations. This prevents the driver from running
-    // indefinitely.
+    // on the pre-existing and new operations. This prevents the driver from
+    // running indefinitely.
+    //
+    // TODO: improve this driver for the case where the newly created ops need
+    // ....  to be applied the pattern as well.
     GreedyRewriteConfig config;
     config.strictMode = GreedyRewriteStrictness::ExistingOps;
     (void)applyPatternsGreedily(op, std::move(patterns), config);