[SCEVExpander] Recompute poison-generating flags on hoisting. PR57187

[mkazantsev]

Summary

Instruction being hoisted could have nuw/nsw flags inferred from the old
context, and we cannot simply move it to the new location keeping them
because we are going to introduce new uses to them that didn't exist before.

Example in https://github.com/llvm/llvm-project/issues/57187 shows how
this can produce branch by poison from initially well-defined program.

This patch forcefully recomputes poison-generating flag in the new context.

[mkazantsev]

Added context

[mkazantsev]

Rebased (test moved to X86).

[fhahn]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

1057

This seems too pessimistic and unnecessarily drops flags that are not inferred based on the context instruction, e.g. because the original IV already had nuw/nsw.

Would it be possible to preserve the flags valid at entry of the loop when hoisting?

[mkazantsev]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

1057

We can't do it here. But I guess this is mostly used in IndVars, and we can further strengthen no-wrap flags after this hoisting, if it happened. Didn't try it and don't know the consequences. Can it go as separate patch?

[fhahn]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

1057

After another thought, shouldn't moving the IV be fine by itself? It might generate poison at the earlier location, but that is fine, as long as it is not used before the original location. in a context where poison triggers UB.

Isn't the real issue the code that later replaces uses of the original IV with the hoisted IV? When updating uses, whether all flags can be retained would depend on whether the UB-on-poison user may execute in cases where the original IV increment didn't?

[fhahn]

llvm/test/Transforms/IndVarSimplify/X86/pr57187.ll

39

using an addresspace(1) pointer here has the unfortunate side effect that Alive2 cannot handle the code. It would probably be good to remove at least addrspace(1) here and maybe even better to just use a function call as use.

[mkazantsev]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

1057

Yes, I agree that moving is fine. The bug appears when we introduce a new use (in this case a trunc) to the hoisted IV in a point where the old IV didn't have any uses. I'll try to make a more targeted fix.

llvm/test/Transforms/IndVarSimplify/X86/pr57187.ll

39

Ok, I will update the test.

[mkazantsev]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

1057

Unfortunately, in practice it's undoable. This function may hoist multiple instructions, and then the notion that their flags were context-dependent gets lost. If we dropt flags immediately, we can't figure it out later. I propose to stick to this fix and then try to restore flags when possible.

[mkazantsev]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

1057

Actually we have all API available here, so let's just heal the flags.

[mkazantsev]

Healed poison flags. It also happened to improve some of the tests. :)

[mkazantsev]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

1066

This piece of code replicates similar piece of code in IndVars, I propose refectoring as a separate patch. Now we just need a conceptual agreement on how we fix it.

[mkazantsev]

Ping?

[mkazantsev]

Ping

[nikic]

I'm confused. Why would the nuw be legal in the "old context"? Wouldn't we already branch on poison there as well, because %iv.next is used in %loop.exit.cond is used in br (unconditionally)?

[fhahn]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

1057

Unfortunately, in practice it's undoable. This function may hoist multiple instructions, and then the notion that their flags were context-dependent gets lost

Yeah this is unfortunate indeed! I still think the right place to fix the flags would be where we actually introduce uses that cause UB on poison, but given the difficulty I guess our best option is doing it here. Could you add a comment explaining why this needs to be done here, even though hoisting in itself isn't a problem.

[nikic]

In D132022#3755012, @nikic wrote:

I'm confused. Why would the nuw be legal in the "old context"? Wouldn't we already branch on poison there as well, because %iv.next is used in %loop.exit.cond is used in br (unconditionally)?

Okay, the missing context here is that we create an additional IV with increment in the dead block and then perform a congruent IV replacement that hoists that increment and introduces uses in a live block.

[mkazantsev]

Ok, let's drop the flags right before introducing a new use.

[mkazantsev]

Made flag recomputation conditional to avoid unexpected impact. Because this logic may hoist multiple instruction, I don't think there is another place where it can easily done.

[nikic]

LGTM

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2024

posion -> poison

[fhahn]

LGTM thanks!

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

bool isHighCostExpansion(const SCEV *Expr, Loop *L, unsigned Budget,
     assert(Cost <= ScaledBudget && "Should have returned from inner loop.");
     return false;
   }
 
   /// Return the induction variable increment's IV operand.
   Instruction *getIVIncOperand(Instruction *IncV, Instruction *InsertPos,
                                bool allowScale);
 
-  /// Utility for hoisting an IV increment.
-  bool hoistIVInc(Instruction *IncV, Instruction *InsertPos);
+  /// Utility for hoisting \p IncV (with all subexpressions requried for its
+  /// computation) before \p InsertPos. If \p RecomputePoisonFlags is set, drops
+  /// all poison-generating flags from instructions being hoisted and tries to
+  /// re-infer them in the new location. It should be used when we are going to
+  /// introduce a new use in the new position that didn't exist before, and may
+  /// trigger new UB in case of poison.
+  bool hoistIVInc(Instruction *IncV, Instruction *InsertPos,
+                  bool RecomputePoisonFlags = false);
 
   /// replace congruent phis with their most canonical representative. Return
   /// the number of phis eliminated.
   unsigned replaceCongruentIVs(Loop *L, const DominatorTree *DT,
                                SmallVectorImpl<WeakTrackingVH> &DeadInsts,
                                const TargetTransformInfo *TTI = nullptr);
 
   /// Return true if the given expression is safe to expand in the sense that

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

void SCEVExpander::fixupInsertPoints(Instruction *I) {
   for (auto *InsertPtGuard : InsertPointGuards)
     if (InsertPtGuard->GetInsertPoint() == It)
       InsertPtGuard->SetInsertPoint(NewInsertPt);
 }
 
 /// hoistStep - Attempt to hoist a simple IV increment above InsertPos to make
 /// it available to other uses in this loop. Recursively hoist any operands,
 /// until we reach a value that dominates InsertPos.
-bool SCEVExpander::hoistIVInc(Instruction *IncV, Instruction *InsertPos) {
+bool SCEVExpander::hoistIVInc(Instruction *IncV, Instruction *InsertPos,
+                              bool RecomputePoisonFlags) {
   if (SE.DT.dominates(IncV, InsertPos))
       return true;
 
   // InsertPos must itself dominate IncV so that IncV's new position satisfies
   // its existing users.
   if (isa<PHINode>(InsertPos) ||
       !SE.DT.dominates(InsertPos->getParent(), IncV->getParent()))
     return false;
for(;;) {
     IVIncs.push_back(IncV);
     IncV = Oper;
     if (SE.DT.dominates(IncV, InsertPos))
       break;
   }
   for (Instruction *I : llvm::reverse(IVIncs)) {
     fixupInsertPoints(I);
     I->moveBefore(InsertPos);
+    if (!RecomputePoisonFlags)
+      continue;
+    // Drop flags that are potentially inferred from old context and infer flags
+    // in new context.
+    I->dropPoisonGeneratingFlags();
+    if (auto *OBO = dyn_cast<OverflowingBinaryOperator>(I))
+      if (auto Flags = SE.getStrengthenedNoWrapFlagsFromBinOp(OBO)) {
+        auto *BO = cast<BinaryOperator>(I);
+        BO->setHasNoUnsignedWrap(
+            ScalarEvolution::maskFlags(*Flags, SCEV::FlagNUW) == SCEV::FlagNUW);
+        BO->setHasNoSignedWrap(
+            ScalarEvolution::maskFlags(*Flags, SCEV::FlagNSW) == SCEV::FlagNSW);
+      }
   }
   return true;
 }
 
 /// Determine if this cyclic phi is in a form that would have been generated by
 /// LSR. We don't care if the phi was actually expanded in this pass, as long
 /// as it is in a low-cost form, for example, no implied multiplication. This
 /// should match any patterns generated by getAddRecExprPHILiterally and
if (BasicBlock *LatchBlock = L->getLoopLatch()) {
         }
         // Replacing the congruent phi is sufficient because acyclic
         // redundancy elimination, CSE/GVN, should handle the
         // rest. However, once SCEV proves that a phi is congruent,
         // it's often the head of an IV user cycle that is isomorphic
         // with the original phi. It's worth eagerly cleaning up the
         // common case of a single IV increment so that DeleteDeadPHIs
         // can remove cycles that had postinc uses.
+        // Because we may potentially introduce a new use of OrigIV that didn't
+        // exist before at this point, its poison flags need readjustment.
         const SCEV *TruncExpr =
             SE.getTruncateOrNoop(SE.getSCEV(OrigInc), IsomorphicInc->getType());
         if (OrigInc != IsomorphicInc &&
             TruncExpr == SE.getSCEV(IsomorphicInc) &&
             SE.LI.replacementPreservesLCSSAForm(IsomorphicInc, OrigInc) &&
-            hoistIVInc(OrigInc, IsomorphicInc)) {
+            hoistIVInc(OrigInc, IsomorphicInc, /*RecomputePoisonFlags*/ true)) {
           SCEV_DEBUG_WITH_TYPE(
               DebugType, dbgs() << "INDVARS: Eliminated congruent iv.inc: "
                                 << *IsomorphicInc << '\n');
           Value *NewInc = OrigInc;
           if (OrigInc->getType() != IsomorphicInc->getType()) {
             Instruction *IP = nullptr;
             if (PHINode *PN = dyn_cast<PHINode>(OrigInc))
               IP = &*PN->getParent()->getFirstInsertionPt();

llvm/test/Transforms/IndVarSimplify/X86/pr57187.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -passes=indvars < %s | FileCheck %s
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-ni:1-p2:32:8:8:32-ni:2"
 target triple = "x86_64-unknown-linux-gnu"
 
-; FIXME: This test is demonstrating a miscompile. The original program is well-defined,
-; and after opt for any start != 0 branch by poisoned add nuw nsw is introduced.
+; Make sure we don't branch by poison here.
 define void @test(i32 %start) {
 ; CHECK-LABEL: @test(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = zext i32 [[START:%.*]] to i64
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    br label [[LOOP]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV_NEXT:%.*]], [[BACKEDGE:%.*]] ], [ [[TMP0]], [[ENTRY:%.*]] ]
-; CHECK-NEXT:    [[INDVARS_IV_NEXT]] = add nuw nsw i64 [[INDVARS_IV]], -1
+; CHECK-NEXT:    [[INDVARS_IV_NEXT]] = add nsw i64 [[INDVARS_IV]], -1
 ; CHECK-NEXT:    [[INDVARS:%.*]] = trunc i64 [[INDVARS_IV_NEXT]] to i32
 ; CHECK-NEXT:    [[LOOP_EXIT_COND:%.*]] = icmp slt i32 [[INDVARS]], 11
 ; CHECK-NEXT:    br i1 [[LOOP_EXIT_COND]], label [[EXIT:%.*]], label [[STUCK_PREHEADER:%.*]]
 ; CHECK:       stuck.preheader:
 ; CHECK-NEXT:    br label [[STUCK:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    [[IV_NEXT_LCSSA:%.*]] = phi i32 [ [[INDVARS]], [[LOOP]] ]
 ; CHECK-NEXT:    ret void
 ; CHECK:       stuck:
 ; CHECK-NEXT:    br i1 false, label [[BACKEDGE]], label [[STUCK]]
 ;
 entry:
   br label %loop
 
 backedge:                                         ; preds = %stuck
   br label %loop
 
 loop:                                             ; preds = %backedge, %entry
   %iv = phi i32 [ %start, %entry ], [ %iv.next, %backedge ]
   %iv.zext = zext i32 %iv to i64
   %gep = getelementptr inbounds i64, i64* undef, i64 %iv.zext
   %iv.next = add i32 %iv, -1
   %loop.exit.cond = icmp slt i32 %iv.next, 11
   br i1 %loop.exit.cond, label %exit, label %stuck
 
 exit:                                             ; preds = %loop
   %iv.next.lcssa = phi i32 [ %iv.next, %loop ]
   ret void
 
 stuck:                                            ; preds = %stuck, %loop
   br i1 false, label %backedge, label %stuck
 }