[LIR][SCEVExpander] Restore original flags when aborting transform (#82362)

SCEVExpanderCleaner will currently remove instructions created by
SCEVExpander, but not restore poison generating flags that it may have
dropped. As such, running LIR can currently spuriously drop flags
without performing any transforms.

Fix this by keeping track of original instruction flags in SCEVExpander.

Fixes https://github.com/llvm/llvm-project/issues/82337.

1 files changed, 42 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
index fbe1dba..0f67cc3 100644
--- a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
+++ b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
@@ -43,6 +43,37 @@ cl::opt<unsigned> llvm::SCEVCheapExpansionBudget(
 
 using namespace PatternMatch;
 
+PoisonFlags::PoisonFlags(const Instruction *I) {
+  NUW = false;
+  NSW = false;
+  Exact = false;
+  Disjoint = false;
+  NNeg = false;
+  if (auto *OBO = dyn_cast<OverflowingBinaryOperator>(I)) {
+    NUW = OBO->hasNoUnsignedWrap();
+    NSW = OBO->hasNoSignedWrap();
+  }
+  if (auto *PEO = dyn_cast<PossiblyExactOperator>(I))
+    Exact = PEO->isExact();
+  if (auto *PDI = dyn_cast<PossiblyDisjointInst>(I))
+    Disjoint = PDI->isDisjoint();
+  if (auto *PNI = dyn_cast<PossiblyNonNegInst>(I))
+    NNeg = PNI->hasNonNeg();
+}
+
+void PoisonFlags::apply(Instruction *I) {
+  if (isa<OverflowingBinaryOperator>(I)) {
+    I->setHasNoUnsignedWrap(NUW);
+    I->setHasNoSignedWrap(NSW);
+  }
+  if (isa<PossiblyExactOperator>(I))
+    I->setIsExact(Exact);
+  if (auto *PDI = dyn_cast<PossiblyDisjointInst>(I))
+    PDI->setIsDisjoint(Disjoint);
+  if (auto *PNI = dyn_cast<PossiblyNonNegInst>(I))
+    PNI->setNonNeg(NNeg);
+}
+
 /// ReuseOrCreateCast - Arrange for there to be a cast of V to Ty at IP,
 /// reusing an existing cast if a suitable one (= dominating IP) exists, or
 /// creating a new one.
@@ -724,6 +755,7 @@ bool SCEVExpander::hoistIVInc(Instruction *IncV, Instruction *InsertPos,
   auto FixupPoisonFlags = [this](Instruction *I) {
     // Drop flags that are potentially inferred from old context and infer flags
     // in new context.
+    rememberFlags(I);
     I->dropPoisonGeneratingFlags();
     if (auto *OBO = dyn_cast<OverflowingBinaryOperator>(I))
       if (auto Flags = SE.getStrengthenedNoWrapFlagsFromBinOp(OBO)) {
@@ -1481,6 +1513,7 @@ Value *SCEVExpander::expand(const SCEV *S) {
     V = fixupLCSSAFormFor(V);
   } else {
     for (Instruction *I : DropPoisonGeneratingInsts) {
+      rememberFlags(I);
       I->dropPoisonGeneratingFlagsAndMetadata();
       // See if we can re-infer from first principles any of the flags we just
       // dropped.
@@ -1521,6 +1554,11 @@ void SCEVExpander::rememberInstruction(Value *I) {
   DoInsert(I);
 }
 
+void SCEVExpander::rememberFlags(Instruction *I) {
+  // If we already have flags for the instruction, keep the existing ones.
+  OrigFlags.try_emplace(I, PoisonFlags(I));
+}
+
 void SCEVExpander::replaceCongruentIVInc(
     PHINode *&Phi, PHINode *&OrigPhi, Loop *L, const DominatorTree *DT,
     SmallVectorImpl<WeakTrackingVH> &DeadInsts) {
@@ -2318,6 +2356,10 @@ void SCEVExpanderCleaner::cleanup() {
   if (ResultUsed)
     return;
 
+  // Restore original poison flags.
+  for (auto [I, Flags] : Expander.OrigFlags)
+    Flags.apply(I);
+
   auto InsertedInstructions = Expander.getAllInsertedInstructions();
 #ifndef NDEBUG
   SmallPtrSet<Instruction *, 8> InsertedSet(InsertedInstructions.begin(),