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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

<details>
<summary>Changes</summary>

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

This workaround looks ugly. But I can't think of a better way to fix this.


---
Full diff: https://github.com/llvm/llvm-project/pull/104581.diff


3 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineInternal.h (+4-1) 
- (modified) llvm/lib/Transforms/InstCombine/InstCombineNegator.cpp (+8-3) 
- (added) llvm/test/Transforms/InstCombine/pr96012.ll (+24) 


``````````diff
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..bf4f277660879f 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, 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
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/104581