fix

EnzymeAD · Dec 4, 2024 · ecf0e9f · ecf0e9f
1 parent 0237c5c
commit ecf0e9f
Showing 1 changed file with 21 additions and 27 deletions.
diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -7137,7 +7137,7 @@ Value *GradientUtils::lookupM(Value *val, IRBuilder<> &BuilderM,
                 {
                   SCEVExpander OrigExp(
                       *OrigSE, ctx->getParent()->getParent()->getDataLayout(),
-                      "enzyme", /*PreserveLCSSA = */false);
+                      "enzyme", /*PreserveLCSSA = */ false);
 
                   OrigExp.setInsertPoint(
                       isOriginal(l1.header)->getTerminator());
@@ -7160,35 +7160,28 @@ Value *GradientUtils::lookupM(Value *val, IRBuilder<> &BuilderM,
                                     return OrigDT->dominates(A, B);
                                   });
                 for (auto a : InsertedInstructions) {
-                  Value *aV = a;
-                  if (auto PN = dyn_cast<PHINode>(a)) {
-                    if (PN->getNumIncomingValues() == 1) {
-                      aV = PN->getIncomingValue(0);
+                  if (isa<PHINode>(a)) {
+                    std::string str;
+                    raw_string_ostream ss(str);
+                    ss << "oldFunc: " << *oldFunc << "\n";
+                    ss << "newFunc: " << *newFunc << "\n";
+                    ss << "li: " << *li << "\n";
+                    ss << "start0: " << *start0 << "\n";
+                    ss << "Inserted a phi node (" << *a
+                       << ") during unwrap of SCEV: " << *ar1->getStart()
+                       << "\n";
+                    if (CustomErrorHandler) {
+                      CustomErrorHandler(str.c_str(), wrap(li),
+                                         ErrorType::InternalError, nullptr,
+                                         nullptr, nullptr);
                     } else {
-                      std::string str;
-                      raw_string_ostream ss(str);
-                      ss << "oldFunc: " << *oldFunc << "\n";
-                      ss << "newFunc: " << *newFunc << "\n";
-                      ss << "li: " << *li << "\n";
-                      ss << "start0: " << *start0 << "\n";
-                      ss << "Inserted a phi node (" << *a
-                         << ") during unwrap of SCEV: " << *ar1->getStart()
-                         << "\n";
-                      if (CustomErrorHandler) {
-                        CustomErrorHandler(str.c_str(), wrap(li),
-                                           ErrorType::InternalError, nullptr,
-                                           nullptr, nullptr);
-                      } else {
-                        EmitFailure("InsertedPHISCEV", li->getDebugLoc(), li,
-                                    ss.str());
-                      }
+                      EmitFailure("InsertedPHISCEV", li->getDebugLoc(), li,
+                                  ss.str());
                     }
                   }
-                  auto uwV = (aV == a)
-                                 ? unwrapM(aV, v, available,
-                                           UnwrapMode::AttemptSingleUnwrap,
-                                           /*scope*/ nullptr, /*cache*/ false)
-                                 : (Value *)available[aV];
+                  auto uwV =
+                      unwrapM(a, v, available, UnwrapMode::AttemptSingleUnwrap,
+                              /*scope*/ nullptr, /*cache*/ false);
                   auto uw = dyn_cast<Instruction>(uwV);
                   assert(uwV->getType() == a->getType());
 #ifndef NDEBUG
@@ -7200,6 +7193,7 @@ Value *GradientUtils::lookupM(Value *val, IRBuilder<> &BuilderM,
                       else if (auto inst = dyn_cast<Instruction>(op))
                         assert(inst->getParent()->getParent() == newFunc);
                     }
+                    assert(uw->getParent()->getParent() == newFunc);
                   }
 #endif
                   available[a] = uwV;