[NaryReassociate] Check to avoid introducing poison when reusing SCEVs (llvm#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.

Author: d0k

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;
 }

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