[llvm] 48ae614 - [InstCombine] Avoid infinite loop when negating phi nodes (#104581)

[via llvm-commits]

Author: Yingwei Zheng
Date: 2024-08-17T16:48:29+08:00
New Revision: 48ae61470104e9d7a8be5beb8739c24f52cc33c0

URL: https://github.com/llvm/llvm-project/commit/48ae61470104e9d7a8be5beb8739c24f52cc33c0
DIFF: https://github.com/llvm/llvm-project/commit/48ae61470104e9d7a8be5beb8739c24f52cc33c0.diff

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

Closes https://github.com/llvm/llvm-project/issues/96012

---------

Co-authored-by: Nikita Popov <github at npopov.com>

Added: 
    llvm/test/Transforms/InstCombine/pr96012.ll

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineInternal.h
    llvm/lib/Transforms/InstCombine/InstCombineNegator.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
index 16f1c3ba15eba2..a0f2399972b52f 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
+++ b/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;

diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineNegator.cpp b/llvm/lib/Transforms/InstCombine/InstCombineNegator.cpp
index 2dd1db6a4a757b..92293ef401465b 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineNegator.cpp
+++ b/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

diff  --git a/llvm/test/Transforms/InstCombine/pr96012.ll b/llvm/test/Transforms/InstCombine/pr96012.ll
new file mode 100644
index 00000000000000..1a1356f0742160
--- /dev/null
+++ b/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
+}