only perform lookup for reverse builders

Author: tgymnich

enzyme/Enzyme/GradientUtils.cpp

@@ -746,18 +746,13 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
           return nullptr;
         }
 
-        Value *pidx;
+        Value *pidx = nullptr;
 
-        switch (mode) {
-        case DerivativeMode::ForwardMode:
+        if (isOriginalBlock(*BuilderM.GetInsertBlock())) {
           pidx = invertPointerM(dli->getOperand(0), BuilderM);
-          break;
-        case DerivativeMode::ReverseModePrimal:
-        case DerivativeMode::ReverseModeGradient:
-        case DerivativeMode::ReverseModeCombined:
+        } else {
           pidx =
               lookupM(invertPointerM(dli->getOperand(0), BuilderM), BuilderM);
-          break;
         }
 
         if (pidx == nullptr)

enzyme/Enzyme/GradientUtils.h

@@ -934,16 +934,13 @@ class GradientUtils : public CacheUtility {
     if (auto arg = dyn_cast<Argument>(ptr)) {
       assert(arg->getParent() == oldFunc);
     }
-    switch (mode) {
-    case DerivativeMode::ForwardMode:
+
+    if (isOriginalBlock(*BuilderM.GetInsertBlock())) {
       ptr = invertPointerM(ptr, BuilderM);
-      break;
-    case DerivativeMode::ReverseModePrimal:
-    case DerivativeMode::ReverseModeGradient:
-    case DerivativeMode::ReverseModeCombined:
+    } else {
       ptr = lookupM(invertPointerM(ptr, BuilderM), BuilderM);
-      break;
     }
+
     return BuilderM.CreateStore(newval, ptr);
   }