[llvm] [NaryReassociate] Check to avoid introducing poison when reusing SCEVs (PR #98156)

[Benjamin Kramer via llvm-commits]

https://github.com/d0k updated https://github.com/llvm/llvm-project/pull/98156

>From 619aa12d14e088fa06e1e24a03e25e3351f264b0 Mon Sep 17 00:00:00 2001
From: Benjamin Kramer <benny.kra at googlemail.com>
Date: Tue, 9 Jul 2024 15:15:23 +0200
Subject: [PATCH 1/2] [NaryReassociate] Check to avoid introducing poison when
 reusing SCEVs

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.
---
 .../lib/Transforms/Scalar/NaryReassociate.cpp | 23 ++++++++---
 .../Transforms/NaryReassociate/nary-gep.ll    | 40 +++++++++++++++++++
 2 files changed, 57 insertions(+), 6 deletions(-)

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))

>From bd5c200966c6910549220a8d37ec6ba3e5d7ce2a Mon Sep 17 00:00:00 2001
From: Benjamin Kramer <benny.kra at googlemail.com>
Date: Tue, 9 Jul 2024 16:07:00 +0200
Subject: [PATCH 2/2] Check CandidateExpr instead of Dominatee

---
 llvm/lib/Transforms/Scalar/NaryReassociate.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/NaryReassociate.cpp b/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
index a10f909d5e5b6..c00c71fcb0b43 100644
--- a/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
@@ -559,7 +559,6 @@ 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
@@ -575,7 +574,7 @@ NaryReassociatePass::findClosestMatchingDominator(const SCEV *CandidateExpr,
       // Make sure that the instruction is safe to reuse without introducing
       // poison.
       SmallVector<Instruction *> DropPoisonGeneratingInsts;
-      if (!SE->canReuseInstruction(FullExpr, CandidateInstruction,
+      if (!SE->canReuseInstruction(CandidateExpr, CandidateInstruction,
                                    DropPoisonGeneratingInsts))
         continue;