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

[via llvm-commits]

Author: Benjamin Kramer
Date: 2024-07-09T16:14:16+02:00
New Revision: d1fbdde9d400222b6db0e5509392007208ad1018

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

LOG: [NaryReassociate] Check to avoid introducing poison when reusing SCEVs (#98156)

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.

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/NaryReassociate.cpp
    llvm/test/Transforms/NaryReassociate/nary-gep.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/NaryReassociate.cpp b/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
index 94e0b026eeef8..c00c71fcb0b43 100644
--- a/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
@@ -565,14 +565,24 @@ NaryReassociatePass::findClosestMatchingDominator(const SCEV *CandidateExpr,
   // 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(CandidateExpr, 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))