[NaryReassociate] Check to avoid introducing poison when reusing SCEVs

[d0k]

Drop the poison flags if possible or skip the candidate if it's not. Otherwise we'd introduce poison in places where it previously wasn't, leading to miscompiles.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Benjamin Kramer (d0k)

Changes

Drop the poison flags if possible or skip the candidate if it's not. Otherwise we'd introduce poison in places where it previously wasn't, leading to miscompiles.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/NaryReassociate.cpp (+17-6)
• (modified) llvm/test/Transforms/NaryReassociate/nary-gep.ll (+40)

diff --git a/llvm/lib/Transforms/Scalar/NaryReassociate.cpp b/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
index 94e0b026eeef8..a10f909d5e5b6 100644
--- a/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
@@ -559,20 +559,31 @@ NaryReassociatePass::findClosestMatchingDominator(const SCEV *CandidateExpr,
     return nullptr;
 
   auto &Candidates = Pos->second;
+  const auto *FullExpr = SE->getSCEV(Dominatee);
   // Because we process the basic blocks in pre-order of the dominator tree, a
   // candidate that doesn't dominate the current instruction won't dominate any
   // future instruction either. Therefore, we pop it out of the stack. This
   // optimization makes the algorithm O(n).
   while (!Candidates.empty()) {
     // Candidates stores WeakTrackingVHs, so a candidate can be nullptr if it's
-    // removed
-    // during rewriting.
-    if (Value *Candidate = Candidates.back()) {
+    // removed during rewriting.
+    if (Value *Candidate = Candidates.pop_back_val()) {
       Instruction *CandidateInstruction = cast<Instruction>(Candidate);
-      if (DT->dominates(CandidateInstruction, Dominatee))
-        return CandidateInstruction;
+      if (!DT->dominates(CandidateInstruction, Dominatee))
+        continue;
+
+      // Make sure that the instruction is safe to reuse without introducing
+      // poison.
+      SmallVector<Instruction *> DropPoisonGeneratingInsts;
+      if (!SE->canReuseInstruction(FullExpr, CandidateInstruction,
+                                   DropPoisonGeneratingInsts))
+        continue;
+
+      for (Instruction *I : DropPoisonGeneratingInsts)
+        I->dropPoisonGeneratingAnnotations();
+
+      return CandidateInstruction;
     }
-    Candidates.pop_back();
   }
   return nullptr;
 }
diff --git a/llvm/test/Transforms/NaryReassociate/nary-gep.ll b/llvm/test/Transforms/NaryReassociate/nary-gep.ll
index 7e670773a7d85..d0ece1e11de5a 100644
--- a/llvm/test/Transforms/NaryReassociate/nary-gep.ll
+++ b/llvm/test/Transforms/NaryReassociate/nary-gep.ll
@@ -21,4 +21,44 @@ define void @no_sext_fat_pointer(ptr addrspace(2) %a, i32 %i, i32 %j) {
   ret void
 }
 
+define void @or_disjoint(ptr addrspace(2) %a, i32 %i, i32 %j, i32 %k) {
+; CHECK-LABEL: @or_disjoint(
+; CHECK-NEXT:    [[OR:%.*]] = or disjoint i32 [[I:%.*]], [[J:%.*]]
+; CHECK-NEXT:    [[V2:%.*]] = getelementptr float, ptr addrspace(2) [[A:%.*]], i32 [[OR]]
+; CHECK-NEXT:    call void @foo(ptr addrspace(2) [[V2]])
+; CHECK-NEXT:    [[ADD1:%.*]] = add nuw nsw i32 [[I]], [[J]]
+; CHECK-NEXT:    [[ADD2:%.*]] = add nuw nsw i32 [[ADD1]], [[K:%.*]]
+; CHECK-NEXT:    [[V3:%.*]] = getelementptr float, ptr addrspace(2) [[A]], i32 [[ADD2]]
+; CHECK-NEXT:    call void @foo(ptr addrspace(2) [[V3]])
+; CHECK-NEXT:    ret void
+;
+  %or = or disjoint i32 %i, %j
+  %v2 = getelementptr float, ptr addrspace(2) %a, i32 %or
+  call void @foo(ptr addrspace(2) %v2)
+  %add1 = add nuw nsw i32 %i, %j
+  %add2 = add nuw nsw i32 %add1, %k
+  %v3 = getelementptr float, ptr addrspace(2) %a, i32 %add2
+  call void @foo(ptr addrspace(2) %v3)
+  ret void
+}
+
+define void @drop_nuw_nsw(ptr addrspace(2) %a, i32 %i, i32 %j, i32 %k) {
+; CHECK-LABEL: @drop_nuw_nsw(
+; CHECK-NEXT:    [[ADD0:%.*]] = add i32 [[I:%.*]], [[J:%.*]]
+; CHECK-NEXT:    [[V2:%.*]] = getelementptr float, ptr addrspace(2) [[A:%.*]], i32 [[ADD0]]
+; CHECK-NEXT:    call void @foo(ptr addrspace(2) [[V2]])
+; CHECK-NEXT:    [[V3:%.*]] = getelementptr float, ptr addrspace(2) [[V2]], i32 [[K:%.*]]
+; CHECK-NEXT:    call void @foo(ptr addrspace(2) [[V3]])
+; CHECK-NEXT:    ret void
+;
+  %add0 = add nuw nsw i32 %i, %j
+  %v2 = getelementptr float, ptr addrspace(2) %a, i32 %add0
+  call void @foo(ptr addrspace(2) %v2)
+  %add1 = add i32 %i, %j
+  %add2 = add nuw nsw i32 %add1, %k
+  %v3 = getelementptr float, ptr addrspace(2) %a, i32 %add2
+  call void @foo(ptr addrspace(2) %v3)
+  ret void
+}
+
 declare void @foo(ptr addrspace(2))

[nikic]

Why are we using the SCEV of Dominatee here, rather than CandidateExpr? I'm not familiar with this pass, but from a cursory look what we're actually replacing is CandidateExpr, while Dominatee acts as a context instruction.

[d0k]

The pass is replacing one of the operands of Dominatee with CandidateInstruction (which has a SCEV of CandidateExpr). What I want to know is if CandidateInstruction is "more poisonous" than the original Dominatee.

[nikic]

But don't you want to know whether it is more poisonous than the operand being replaced, rather than Dominatee?

[d0k]

Hmm, I think you're right. Changed it.

[nikic]

LGTM

[fiigii]

This change seems problematic. Previously, we returned a candidate without pop_back the stack. With this change, the stack popped before returning the candidate. I have seen some problems.