[LIR][SCEVExpander] Restore original flags when aborting transform

llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h

@@ -41,13 +41,26 @@ struct SCEVOperand {
   const SCEV* S;
 };
 
+struct PoisonFlags {
+  unsigned NUW : 1;
+  unsigned NSW : 1;
+  unsigned Exact : 1;
+  unsigned Disjoint : 1;
+  unsigned NNeg : 1;
+
+  PoisonFlags(const Instruction *I);
+  void apply(Instruction *I);
+};
+
 /// This class uses information about analyze scalars to rewrite expressions
 /// in canonical form.
 ///
 /// Clients should create an instance of this class when rewriting is needed,
 /// and destroy it when finished to allow the release of the associated
 /// memory.
 class SCEVExpander : public SCEVVisitor<SCEVExpander, Value *> {
+  friend class SCEVExpanderCleaner;
+
   ScalarEvolution &SE;
   const DataLayout &DL;
 
@@ -70,6 +83,10 @@ class SCEVExpander : public SCEVVisitor<SCEVExpander, Value *> {
   /// InsertedValues/InsertedPostIncValues.
   SmallPtrSet<Value *, 16> ReusedValues;
 
+  /// Original flags of instructions for which they were modified. Used
+  /// by SCEVExpanderCleaner to undo changes.
+  DenseMap<AssertingVH<Instruction>, PoisonFlags> OrigFlags;
+
   // The induction variables generated.
   SmallVector<WeakVH, 2> InsertedIVs;
 
@@ -188,6 +205,7 @@ class SCEVExpander : public SCEVVisitor<SCEVExpander, Value *> {
     InsertedValues.clear();
     InsertedPostIncValues.clear();
     ReusedValues.clear();
+    OrigFlags.clear();
     ChainedPhis.clear();
     InsertedIVs.clear();
   }
@@ -491,6 +509,8 @@ class SCEVExpander : public SCEVVisitor<SCEVExpander, Value *> {
 
   void rememberInstruction(Value *I);
 
+  void rememberFlags(Instruction *I);
+
   bool isNormalAddRecExprPHI(PHINode *PN, Instruction *IncV, const Loop *L);
 
   bool isExpandedAddRecExprPHI(PHINode *PN, Instruction *IncV, const Loop *L);

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

llvm/test/Transforms/LoopIdiom/pr82337.ll

@@ -1,15 +1,15 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
 ; RUN: opt -S -passes=loop-idiom < %s | FileCheck %s
 
-; FIXME: The poison flags should be preserved, as no transform takes place.
+; The poison flags should be preserved, as no transform takes place.
 define void @test(ptr %p.end, ptr %p.start) {
 ; CHECK-LABEL: define void @test(
 ; CHECK-SAME: ptr [[P_END:%.*]], ptr [[P_START:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[P_END_INT:%.*]] = ptrtoint ptr [[P_END]] to i64
 ; CHECK-NEXT:    [[P_START_INT:%.*]] = ptrtoint ptr [[P_START]] to i64
-; CHECK-NEXT:    [[DIST:%.*]] = sub i64 [[P_END_INT]], [[P_START_INT]]
-; CHECK-NEXT:    [[LEN:%.*]] = lshr i64 [[DIST]], 5
+; CHECK-NEXT:    [[DIST:%.*]] = sub nuw i64 [[P_END_INT]], [[P_START_INT]]
+; CHECK-NEXT:    [[LEN:%.*]] = lshr exact i64 [[DIST]], 5
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[P_END]], [[P_START]]
 ; CHECK-NEXT:    br i1 [[CMP]], label [[EXIT:%.*]], label [[PREHEADER:%.*]]
 ; CHECK:       preheader: