From 01d8acfffc54d0954c7bd4330a641ce4a34d3ddb Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Wed, 28 Feb 2024 17:53:47 -0800
Subject: [PATCH] additional cleanup

---
 .../MLIR/Interfaces/EnzymeLogicReverse.cpp    |  2 +-
 .../MLIR/Interfaces/GradientUtilsReverse.cpp  | 59 -------------------
 .../MLIR/Interfaces/GradientUtilsReverse.h    |  7 ---
 3 files changed, 1 insertion(+), 67 deletions(-)

diff --git a/enzyme/Enzyme/MLIR/Interfaces/EnzymeLogicReverse.cpp b/enzyme/Enzyme/MLIR/Interfaces/EnzymeLogicReverse.cpp
index 37f692f957fb..06ccd344d35f 100644
--- a/enzyme/Enzyme/MLIR/Interfaces/EnzymeLogicReverse.cpp
+++ b/enzyme/Enzyme/MLIR/Interfaces/EnzymeLogicReverse.cpp
@@ -97,7 +97,7 @@ bool MEnzymeLogic::visitChildCustom(Operation *op, OpBuilder &builder,
     func::CallOp dCI =
         builder.create<func::CallOp>(op->getLoc(), srDiffe, resultTypes, args);
     for (int i = 0; i < (int)op->getNumOperands(); i++) {
-      gutils->mapInvertPointer(op->getOperand(i), dCI.getResult(i), builder);
+      gutils->setDiffe(op->getOperand(i), dCI.getResult(i), builder);
     }
 
     return true;
diff --git a/enzyme/Enzyme/MLIR/Interfaces/GradientUtilsReverse.cpp b/enzyme/Enzyme/MLIR/Interfaces/GradientUtilsReverse.cpp
index 0752b64f6f81..0e8c2752f101 100644
--- a/enzyme/Enzyme/MLIR/Interfaces/GradientUtilsReverse.cpp
+++ b/enzyme/Enzyme/MLIR/Interfaces/GradientUtilsReverse.cpp
@@ -41,14 +41,6 @@ mlir::enzyme::MGradientUtilsReverse::MGradientUtilsReverse(
                           ReturnActivity, ArgDiffeTypes_, originalToNewFn_,
                           originalToNewFnOps_, mode_, width, /*omp*/ false),
       symbolTable(symbolTable_) {
-
-  initInitializationBlock(invertedPointers_, ArgDiffeTypes_);
-}
-
-// for(auto x : v.getUsers()){x->dump();} DEBUG
-
-bool MGradientUtilsReverse::onlyUsedInParentBlock(Value v) {
-  return !v.isUsedOutsideOfBlock(v.getParentBlock());
 }
 
 Type mlir::enzyme::MGradientUtilsReverse::getIndexCacheType() {
@@ -129,57 +121,6 @@ void mlir::enzyme::MGradientUtilsReverse::addToDiffe(Value oldGradient,
   setDiffe(oldGradient, added, builder);
 }
 
-void mlir::enzyme::MGradientUtilsReverse::mapInvertPointer(
-    mlir::Value v, mlir::Value invertValue, OpBuilder &builder) {
-  assert(0);
-   /*
-  if (!invertedPointersGlobal.contains(v)) {
-    Value g = insertInitGradient(v, builder);
-    invertedPointersGlobal.map(v, g);
-  }
-  Value gradient = invertedPointersGlobal.lookupOrNull(v);
-  builder.create<enzyme::SetOp>(v.getLoc(), gradient, invertValue);
-  */
-}
-
-void MGradientUtilsReverse::initInitializationBlock(
-    IRMapping invertedPointers_, ArrayRef<DIFFE_TYPE> argDiffeTypes) {
-
-  OpBuilder initializationBuilder(
-      &*(this->newFunc.getFunctionBody().begin()),
-      this->newFunc.getFunctionBody().begin()->begin());
-
-  /*
-  for (const auto &[val, diffe_type] : llvm::zip(
-           this->oldFunc.getFunctionBody().getArguments(), argDiffeTypes)) {
-    if (diffe_type != DIFFE_TYPE::OUT_DIFF) {
-      continue;
-    }
-    auto iface = dyn_cast<AutoDiffTypeInterface>(val.getType());
-    if (!iface) {
-      llvm_unreachable(
-          "Type does not have an associated AutoDiffTypeInterface");
-    }
-    Value zero = iface.createNullValue(initializationBuilder, val.getLoc());
-    mapInvertPointer(val, zero, initializationBuilder);
-  }
-  for (auto const &x : invertedPointers_.getValueMap()) {
-    if (auto iface = dyn_cast<AutoDiffTypeInterface>(x.first.getType())) {
-      if (!iface.isMutable()) {
-        mapShadowValue(x.first, x.second,
-                       initializationBuilder); // This may create an unnecessary
-                                               // ShadowedGradient which could
-                                               // be avoidable TODO
-      } else {
-        mapInvertPointer(x.first, x.second, initializationBuilder);
-      }
-    } else {
-      llvm_unreachable("TODO not implemented");
-    }
-  }
-  */
-}
-
 void MGradientUtilsReverse::createReverseModeBlocks(Region &oldFunc,
                                                     Region &newFunc) {
   for (auto it = oldFunc.getBlocks().rbegin(); it != oldFunc.getBlocks().rend();
diff --git a/enzyme/Enzyme/MLIR/Interfaces/GradientUtilsReverse.h b/enzyme/Enzyme/MLIR/Interfaces/GradientUtilsReverse.h
index 24f7f332ed23..d3b2e818391f 100644
--- a/enzyme/Enzyme/MLIR/Interfaces/GradientUtilsReverse.h
+++ b/enzyme/Enzyme/MLIR/Interfaces/GradientUtilsReverse.h
@@ -40,8 +40,6 @@ class MGradientUtilsReverse : public MDiffeGradientUtils {
 
   void addToDiffe(mlir::Value oldGradient, mlir::Value addedGradient,
                   OpBuilder &builder);
-  void mapInvertPointer(mlir::Value v, mlir::Value invertValue,
-                        OpBuilder &builder);
 
   Type getIndexType();
   Value insertInit(Type t);
@@ -58,11 +56,6 @@ class MGradientUtilsReverse : public MDiffeGradientUtils {
   Type getIndexCacheType();
   Value initAndPushCache(Value v, OpBuilder &builder);
 
-  void initInitializationBlock(IRMapping invertedPointers_,
-                               ArrayRef<DIFFE_TYPE> argDiffeTypes);
-
-  bool onlyUsedInParentBlock(Value v);
-
   Operation *cloneWithNewOperands(OpBuilder &B, Operation *op);
 
   Value popCache(Value cache, OpBuilder &builder);