simplify

EnzymeAD · Feb 29, 2024 · 8a74023 · 8a74023
1 parent 01d8acf
commit 8a74023
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 14 deletions.
diff --git a/enzyme/Enzyme/MLIR/Implementations/SCFAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/SCFAutoDiffOpInterfaceImpl.cpp
@@ -87,10 +87,10 @@ struct ForOpInterfaceReverse
  auto end = gutils->popCache(caches[1], builder);
  auto step = gutils->popCache(caches[2], builder);
 
- SmallVector<Value> nArgs;
  auto repFor = builder.create<scf::ForOp>(
- forOp.getLoc(), start, end, step, ArrayRef<Value>()); // TODO
- repFor.getRegion().begin()->erase();
+ forOp.getLoc(), start, end, step, ArrayRef<Value>());
+ // erase scf yield
+ repFor.getBody()->begin()->erase();
 
  for (auto &&[oldReg, newReg] : llvm::zip(op->getRegions(), repFor->getRegions())) {
 
@@ -99,7 +99,7 @@ struct ForOpInterfaceReverse
 
  auto loc = oBB->rbegin()->getLoc();
 
- auto idx = repFor.getRegion().begin()->getArgument(0);
+ auto idx = repFor.getInductionVar();
 
  auto lhs = builder.create<arith::AddIOp>(loc, idx, step);
 
@@ -130,8 +130,12 @@ struct ForOpInterfaceReverse
  builder.create<scf::YieldOp>(loc);
  };
 
- gutils->Logic.differentiate(gutils, oldReg, newReg, buildFuncReturnOp,
- nullptr);
+ for (auto &&[oBB, revBB] : llvm::zip(oldReg, newReg)) {
+ gutils->mapReverseModeBlocks.map(&oBB, &revBB);
+ gutils->Logic.visitChildren(&oBB, &revBB, gutils);
+ Block *newBB = gutils->getNewFromOriginal(&oBB);
+ gutils->Logic.handlePredecessors(&oBB, newBB, &revBB, gutils, buildFuncReturnOp);
+ }
  }
  }
 

diff --git a/enzyme/Enzyme/MLIR/Passes/RemoveUnusedEnzymeOps.cpp b/enzyme/Enzyme/MLIR/Passes/RemoveUnusedEnzymeOps.cpp
@@ -83,9 +83,15 @@ struct RemoveUnusedEnzymeOpsPass
  : public enzyme::RemoveUnusedEnzymeOpsPassBase<RemoveUnusedEnzymeOpsPass> {
  void runOnOperation() override {
 
+ SmallVector<enzyme::InitOp, 1> inits;
  getOperation()->walk([&](Operation *op) {
- DominanceInfo dInfo;
  if (auto initOp = dyn_cast<enzyme::InitOp>(op)) {
+ inits.push_back(initOp);
+ }
+ });
+
+ for (auto initOp : inits) {
+ DominanceInfo dInfo;
  Value v = initOp;
  if (auto type = dyn_cast<enzyme::GradientType>(initOp.getType())) {
  bool replaceable = true;
@@ -120,8 +126,8 @@ struct RemoveUnusedEnzymeOpsPass
  for (Operation *userGet : make_early_inc_range(v.getUsers())) {
  userGet->erase();
  }
- op->erase();
- return;
+ initOp->erase();
+ continue;
  }
  } else if (auto type = dyn_cast<enzyme::CacheType>(initOp.getType())) {
  bool replaceable = true;
@@ -177,13 +183,14 @@ struct RemoveUnusedEnzymeOpsPass
  for (Operation *user : make_early_inc_range(v.getUsers())) {
  user->erase();
  }
- op->erase();
+ initOp->erase();
+ continue;
  }
  }
- }
- });
- };
+ }
+ }
 };
+
 } // end anonymous namespace
 
 namespace mlir {

diff --git a/enzyme/test/MLIR/ReverseMode/pow.mlir b/enzyme/test/MLIR/ReverseMode/pow.mlir
@@ -1,4 +1,4 @@
-// RUN: %eopt --enzyme %s | FileCheck %s
+// RUN: %eopt --enzyme --remove-unnecessary-enzyme-ops -canonicalize %s | FileCheck %s
 
 module {
  func.func @ppow(%x: f64) -> f64 {