[InstCombine] Avoid infinite loop when negating phi nodes (#104581)

Closes #96012

---------

Co-authored-by: Nikita Popov <[email protected]>

Author: dtcxzyw

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

@@ -780,11 +780,14 @@ class Negator final {
   using BuilderTy = IRBuilder<TargetFolder, IRBuilderCallbackInserter>;
   BuilderTy Builder;
 
+  const DominatorTree &DT;
+
   const bool IsTrulyNegation;
 
   SmallDenseMap<Value *, Value *> NegationsCache;
 
-  Negator(LLVMContext &C, const DataLayout &DL, bool IsTrulyNegation);
+  Negator(LLVMContext &C, const DataLayout &DL, const DominatorTree &DT,
+          bool IsTrulyNegation);
 
 #if LLVM_ENABLE_STATS
   unsigned NumValuesVisitedInThisNegator = 0;

llvm/lib/Transforms/InstCombine/InstCombineNegator.cpp

@@ -90,13 +90,14 @@ static cl::opt<unsigned>
                     cl::desc("What is the maximal lookup depth when trying to "
                              "check for viability of negation sinking."));
 
-Negator::Negator(LLVMContext &C, const DataLayout &DL, bool IsTrulyNegation_)
+Negator::Negator(LLVMContext &C, const DataLayout &DL, const DominatorTree &DT_,
+                 bool IsTrulyNegation_)
     : Builder(C, TargetFolder(DL),
               IRBuilderCallbackInserter([&](Instruction *I) {
                 ++NegatorNumInstructionsCreatedTotal;
                 NewInstructions.push_back(I);
               })),
-      IsTrulyNegation(IsTrulyNegation_) {}
+      DT(DT_), IsTrulyNegation(IsTrulyNegation_) {}
 
 #if LLVM_ENABLE_STATS
 Negator::~Negator() {
@@ -309,6 +310,9 @@ std::array<Value *, 2> Negator::getSortedOperandsOfBinOp(Instruction *I) {
     auto *PHI = cast<PHINode>(I);
     SmallVector<Value *, 4> NegatedIncomingValues(PHI->getNumOperands());
     for (auto I : zip(PHI->incoming_values(), NegatedIncomingValues)) {
+      // Don't negate indvars to avoid infinite loops.
+      if (DT.dominates(PHI->getParent(), std::get<0>(I)))
+        return nullptr;
       if (!(std::get<1>(I) =
                 negate(std::get<0>(I), IsNSW, Depth + 1))) // Early return.
         return nullptr;
@@ -537,7 +541,8 @@ std::array<Value *, 2> Negator::getSortedOperandsOfBinOp(Instruction *I) {
   if (!NegatorEnabled || !DebugCounter::shouldExecute(NegatorCounter))
     return nullptr;
 
-  Negator N(Root->getContext(), IC.getDataLayout(), LHSIsZero);
+  Negator N(Root->getContext(), IC.getDataLayout(), IC.getDominatorTree(),
+            LHSIsZero);
   std::optional<Result> Res = N.run(Root, IsNSW);
   if (!Res) { // Negation failed.
     LLVM_DEBUG(dbgs() << "Negator: failed to sink negation into " << *Root

llvm/test/Transforms/InstCombine/pr96012.ll

@@ -0,0 +1,24 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=instcombine < %s | FileCheck %s
+
+target datalayout = "e-m:e-p:64:64-i64:64-i128:128-n32:64-S128"
+
+; Make sure we don't hang here.
+define i32 @pr96012() {
+; CHECK-LABEL: define i32 @pr96012() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br label %[[FOR_COND:.*]]
+; CHECK:       [[FOR_COND]]:
+; CHECK-NEXT:    br label %[[FOR_COND]]
+;
+entry:
+  br label %for.cond
+
+for.cond:
+  %indvar = phi i32 [ 1, %entry ], [ %shl, %for.cond ]
+  %conv1 = trunc i32 %indvar to i8
+  %neg = sub i8 0, %conv1
+  %conv2 = zext i8 %neg to i32
+  %shl = shl nuw i32 %conv2, 24
+  br label %for.cond
+}