[SCEVExpander] Fix incorrect reuse of more poisonous instructions (PR63763)

[nikic]

Summary

SCEVExpander tries to reuse existing instruction with the same SCEV expression. However, doing this replacement blindly is not safe, because the instruction might be more poisonous.

What we were already doing is to drop poison-generating flags on the reused instruction. But this is not the only way that more poison can be introduced. The poison-generating flag might not be directly on the reused instruction, or the poison contribution might come from something like 0 * %var, which folds to 0 but can still introduce poison.

This patch fixes the issue in a principled way, by determining which values can contribute poison to the SCEV expression, and then checking whether any additional values can contribute poison to the instruction being reused. Poison-generating flags are dropped if doing that enables reuse.

This is a pretty big hammer and does cause some regressions in tests, but less than I would have expected. I wasn't able to come up with a less intrusive fix that still satisfies the correctness requirements.

Fixes https://github.com/llvm/llvm-project/issues/63763.
Fixes https://github.com/llvm/llvm-project/issues/63926.
Fixes https://github.com/llvm/llvm-project/issues/64333.
Fixes https://github.com/llvm/llvm-project/issues/63727.

[fhahn]

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

442

nit: should be -> must be?

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

1503

nit: this comment might be better placed after the check if, we handles the opposite case described.

1832

It's safe to ignore DropPoisonGeneratingInsts here, as the function is only used in isHighCostExpansionHelper, right?

It might be good to lock down the API to return a bool for getRelatedExistingExpansion (hasRelatedExistingExpansion?) to make sure the return value isn't used elsewhere in the future without dropping the flags?

[nikic]

Address review comments.

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

1832

Good idea! I've renamed the method and changed the return value to bool.

[fhahn]

LGTM, thanks!

llvm/include/llvm/Analysis/ScalarEvolution.h

public:
   bool loopHasNoAbnormalExits(const Loop *L) {
     return getLoopProperties(L).HasNoAbnormalExits;
   }
 
   /// Return true if this loop is finite by assumption.  That is,
   /// to be infinite, it must also be undefined.
   bool loopIsFiniteByAssumption(const Loop *L);
 
+  /// Return the set of Values that, if poison, will definitively result in S
+  /// being poison as well. The returned set may be incomplete, i.e. there can
+  /// be additional Values that also result in S being poison.
+  void getPoisonGeneratingValues(SmallPtrSetImpl<const Value *> &Result,
+                                 const SCEV *S);
+
   class FoldID {
     const SCEV *Op = nullptr;
     const Type *Ty = nullptr;
     unsigned short C;
 
   public:
     FoldID(SCEVTypes C, const SCEV *Op, const Type *Ty) : Op(Op), Ty(Ty), C(C) {
       assert(Op);

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

private:
   /// noop cast, doing what we can to share the casts.
   Value *InsertNoopCastOfTo(Value *V, Type *Ty);
 
   /// Expand a SCEVAddExpr with a pointer type into a GEP instead of using
   /// ptrtoint+arithmetic+inttoptr.
   Value *expandAddToGEP(const SCEV *Op, Type *Ty, Value *V);
 
   /// Find a previous Value in ExprValueMap for expand.
-  Value *FindValueInExprValueMap(const SCEV *S, const Instruction *InsertPt);
+  /// DropPoisonGeneratingInsts is populated with instructions for which
+  /// poison-generating flags must be dropped if the value is reused.
+  Value *FindValueInExprValueMap(
+      const SCEV *S, const Instruction *InsertPt,
+      SmallVectorImpl<Instruction *> &DropPoisonGeneratingInsts);
 
   Value *expand(const SCEV *S);
 
   /// Determine the most "relevant" loop for the given SCEV.
   const Loop *getRelevantLoop(const SCEV *);
 
   Value *expandMinMaxExpr(const SCEVNAryExpr *S, Intrinsic::ID IntrinID,
                           Twine Name, bool IsSequential = false);

llvm/lib/Analysis/ScalarEvolution.cpp

static bool impliesPoison(const SCEV *AssumedPoison, const SCEV *S) {
 
   // Make sure that no matter which SCEV in PC1.MaybePoison is actually poison,
   // it will also make S poison by being part of PC2.MaybePoison.
   return all_of(PC1.MaybePoison, [&](const SCEVUnknown *S) {
     return PC2.MaybePoison.contains(S);
   });
 }
 
+void ScalarEvolution::getPoisonGeneratingValues(
+    SmallPtrSetImpl<const Value *> &Result, const SCEV *S) {
+  SCEVPoisonCollector PC(/* LookThroughMaybePoisonBlocking */ false);
+  visitAll(S, PC);
+  for (const SCEVUnknown *SU : PC.MaybePoison)
+    Result.insert(SU->getValue());
+}
+
 const SCEV *
 ScalarEvolution::getSequentialMinMaxExpr(SCEVTypes Kind,
                                          SmallVectorImpl<const SCEV *> &Ops) {
   assert(SCEVSequentialMinMaxExpr::isSequentialMinMaxType(Kind) &&
          "Not a SCEVSequentialMinMaxExpr!");
   assert(!Ops.empty() && "Cannot get empty (u|s)(min|max)!");
   if (Ops.size() == 1)
     return Ops[0];

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

Value *SCEVExpander::expandCodeForImpl(const SCEV *SH, Type *Ty) {
   if (Ty) {
     assert(SE.getTypeSizeInBits(Ty) == SE.getTypeSizeInBits(SH->getType()) &&
            "non-trivial casts should be done with the SCEVs directly!");
     V = InsertNoopCastOfTo(V, Ty);
   }
   return V;
 }
 
-Value *SCEVExpander::FindValueInExprValueMap(const SCEV *S,
-                                             const Instruction *InsertPt) {
+static bool
+canReuseInstruction(ScalarEvolution &SE, const SCEV *S, Instruction *I,
+                    SmallVectorImpl<Instruction *> &DropPoisonGeneratingInsts) {
+  // If the instruction cannot be poison, it's always safe to reuse.
+  if (programUndefinedIfPoison(I))
+    return true;
+
+  // Otherwise, it is possible that I is more poisonous that S. Collect the
+  // poison-contributors of S, and then check whether I has any additional
+  // poison-contributors. Poison that is contributed through poison-generating
+  // flags is handled by dropping those flags instead.
+  SmallPtrSet<const Value *, 8> PoisonVals;
+  SE.getPoisonGeneratingValues(PoisonVals, S);
+
+  SmallVector<Value *> Worklist;
+  SmallPtrSet<Value *, 8> Visited;
+  Worklist.push_back(I);
+  while (!Worklist.empty()) {
+    Value *V = Worklist.pop_back_val();
+    if (!Visited.insert(V).second)
+      continue;
+
+    // Avoid walking large instruction graphs.
+    if (Visited.size() > 16)
+      return false;
+
+    // Either the value can't be poison, or the S would also be poison if it
+    // is.
+    if (PoisonVals.contains(V) || isGuaranteedNotToBePoison(V))
+      continue;
+
+    auto *I = dyn_cast<Instruction>(V);
+    if (!I)
+      return false;
+
+    // FIXME: Ignore vscale, even though it technically could be poison. Do this
+    // because SCEV currently assumes it can't be poison. Remove this special
+    // case once we proper model when vscale can be poison.
+    if (auto *II = dyn_cast<IntrinsicInst>(I);
+        II && II->getIntrinsicID() == Intrinsic::vscale)
+      continue;
+
+    if (canCreatePoison(cast<Operator>(I), /*ConsiderFlagsAndMetadata*/ false))
+      return false;
+
+    // If the instruction can't create poison, we can recurse to its operands.
+    if (I->hasPoisonGeneratingFlagsOrMetadata())
+      DropPoisonGeneratingInsts.push_back(I);
+
+    for (Value *Op : I->operands())
+      Worklist.push_back(Op);
+  }
+  return true;
+}
+
+Value *SCEVExpander::FindValueInExprValueMap(
+    const SCEV *S, const Instruction *InsertPt,
+    SmallVectorImpl<Instruction *> &DropPoisonGeneratingInsts) {
   // If the expansion is not in CanonicalMode, and the SCEV contains any
   // sub scAddRecExpr type SCEV, it is required to expand the SCEV literally.
   if (!CanonicalMode && SE.containsAddRecurrence(S))
     return nullptr;
 
   // If S is a constant, it may be worse to reuse an existing Value.
   if (isa<SCEVConstant>(S))
     return nullptr;
 
   for (Value *V : SE.getSCEVValues(S)) {
     Instruction *EntInst = dyn_cast<Instruction>(V);
     if (!EntInst)
       continue;
 
     // Choose a Value from the set which dominates the InsertPt.
     // InsertPt should be inside the Value's parent loop so as not to break
     // the LCSSA form.
     assert(EntInst->getFunction() == InsertPt->getFunction());
     if (S->getType() != V->getType() || !SE.DT.dominates(EntInst, InsertPt) ||
         !(SE.LI.getLoopFor(EntInst->getParent()) == nullptr ||
           SE.LI.getLoopFor(EntInst->getParent())->contains(InsertPt)))
       continue;
 
+    // Make sure reusing the instruction is poison-safe.
+    if (canReuseInstruction(SE, S, EntInst, DropPoisonGeneratingInsts))
       return V;
+    DropPoisonGeneratingInsts.clear();
   }
   return nullptr;
 }
 
 // The expansion of SCEV will either reuse a previous Value in ExprValueMap,
 // or expand the SCEV literally. Specifically, if the expansion is in LSRMode,
 // and the SCEV contains any sub scAddRecExpr type SCEV, it will be expanded
 // literally, to prevent LSR's transformed SCEV from being reverted. Otherwise,
Value *SCEVExpander::expand(const SCEV *S) {
   auto I = InsertedExpressions.find(std::make_pair(S, InsertPt));
   if (I != InsertedExpressions.end())
     return I->second;
 
   SCEVInsertPointGuard Guard(Builder, this);
   Builder.SetInsertPoint(InsertPt);
 
   // Expand the expression into instructions.
-  Value *V = FindValueInExprValueMap(S, InsertPt);
+  SmallVector<Instruction *> DropPoisonGeneratingInsts;
+  Value *V = FindValueInExprValueMap(S, InsertPt, DropPoisonGeneratingInsts);
   if (!V) {
     V = visit(S);
     V = fixupLCSSAFormFor(V);
   } else {
-    // If we're reusing an existing instruction, we are effectively CSEing two
-    // copies of the instruction (with potentially different flags).  As such,
-    // we need to drop any poison generating flags unless we can prove that
-    // said flags must be valid for all new users.
-    if (auto *I = dyn_cast<Instruction>(V))
-      if (I->hasPoisonGeneratingFlags() && !programUndefinedIfPoison(I))
-        I->dropPoisonGeneratingFlags();
+    for (Instruction *I : DropPoisonGeneratingInsts)
+      I->dropPoisonGeneratingFlagsAndMetadata();
   }
   // Remember the expanded value for this SCEV at this location.
   //
   // This is independent of PostIncLoops. The mapped value simply materializes
   // the expression at this insertion point. If the mapped value happened to be
   // a postinc expansion, it could be reused by a non-postinc user, but only if
   // its insertion point was already at the head of the loop.
   InsertedExpressions[std::make_pair(S, InsertPt)] = V;
for (BasicBlock *BB : ExitingBlocks) {
     if (SE.getSCEV(RHS) == S && SE.DT.dominates(RHS, At))
       return true;
   }
 
   // Use expand's logic which is used for reusing a previous Value in
   // ExprValueMap.  Note that we don't currently model the cost of
   // needing to drop poison generating flags on the instruction if we
   // want to reuse it.  We effectively assume that has zero cost.
-  return FindValueInExprValueMap(S, At) != nullptr;
+  SmallVector<Instruction *> DropPoisonGeneratingInsts;
+  return FindValueInExprValueMap(S, At, DropPoisonGeneratingInsts) != nullptr;
 }
 
 template<typename T> static InstructionCost costAndCollectOperands(
   const SCEVOperand &WorkItem, const TargetTransformInfo &TTI,
   TargetTransformInfo::TargetCostKind CostKind,
   SmallVectorImpl<SCEVOperand> &Worklist) {
 
   const T *S = cast<T>(WorkItem.S);

llvm/test/CodeGen/Thumb2/mve-blockplacement.ll

 ; CHECK-NEXT:    mov.w r9, #12
 ; CHECK-NEXT:    umull r1, r4, r1, r4
 ; CHECK-NEXT:    @ implicit-def: $r10
 ; CHECK-NEXT:    @ implicit-def: $r5
 ; CHECK-NEXT:    @ implicit-def: $r11
 ; CHECK-NEXT:    str r0, [sp, #12] @ 4-byte Spill
 ; CHECK-NEXT:    movs r1, #4
 ; CHECK-NEXT:    strd r2, r12, [sp, #4] @ 8-byte Folded Spill
-; CHECK-NEXT:    add.w r1, r1, r4, lsr #1
 ; CHECK-NEXT:    add.w r3, r3, r4, lsr #1
-; CHECK-NEXT:    bic r7, r1, #3
+; CHECK-NEXT:    add.w r1, r1, r4, lsr #1
+; CHECK-NEXT:    movw r4, #65532
+; CHECK-NEXT:    vdup.32 q6, r3
+; CHECK-NEXT:    movt r4, #32767
+; CHECK-NEXT:    and.w r7, r1, r4
 ; CHECK-NEXT:    adr r1, .LCPI1_0
+; CHECK-NEXT:    vdup.32 q7, r3
 ; CHECK-NEXT:    vldrw.u32 q0, [r1]
 ; CHECK-NEXT:    adr r1, .LCPI1_1
 ; CHECK-NEXT:    vldrw.u32 q5, [r1]
-; CHECK-NEXT:    vdup.32 q6, r3
 ; CHECK-NEXT:    vadd.i32 q4, q0, lr
-; CHECK-NEXT:    vdup.32 q7, r3
 ; CHECK-NEXT:    b .LBB1_4
 ; CHECK-NEXT:  .LBB1_2: @ %for.body6.preheader
 ; CHECK-NEXT:    @ in Loop: Header=BB1_4 Depth=1
 ; CHECK-NEXT:    mov r0, r11
 ; CHECK-NEXT:    cmn.w r11, #4
 ; CHECK-NEXT:    it le
 ; CHECK-NEXT:    mvnle r0, #3
 ; CHECK-NEXT:    movw r2, #18725

llvm/test/CodeGen/X86/AMX/amx-greedy-ra-spill-shape.ll

define void @foo(i32 %M, i32 %N, i32 %K, ptr %A, ptr %B_rcr4, ptr %C, i32 %c_row_from, i32 %c_row_to, i32 %c_row_tile, i32 %c_col_from, i32 %c_col_to, i32 %c_col_tile) {
   ; CHECK-NEXT:   MOV8mr %stack.0, 1, $noreg, 50, $noreg, [[MOVZX32rr16_]].sub_8bit :: (store (s512) into %stack.0 + 50, align 2, basealign 4)
   ; CHECK-NEXT:   [[SHR32ri:%[0-9]+]]:gr32 = SHR32ri [[SHR32ri]], 2, implicit-def dead $eflags
   ; CHECK-NEXT:   MOV32mr %stack.2, 1, $noreg, 0, $noreg, [[SHR32ri]] :: (store (s32) into %stack.2)
   ; CHECK-NEXT:   MOV8mr %stack.0, 1, $noreg, 50, $noreg, [[SHR32ri]].sub_8bit :: (store (s512) into %stack.0 + 50, align 2, basealign 4)
   ; CHECK-NEXT:   [[LEA64_32r:%[0-9]+]]:gr32 = LEA64_32r $noreg, 4, [[MOV89]], 0, $noreg
   ; CHECK-NEXT:   MOV16mr %stack.0, 1, $noreg, 20, $noreg, [[LEA64_32r]].sub_16bit :: (store (s512) into %stack.0 + 20, align 4)
   ; CHECK-NEXT:   PLDTILECFGV %stack.0, 1, $noreg, 0, $noreg, implicit-def dead $tmm0, implicit-def dead $tmm1, implicit-def dead $tmm2, implicit-def dead $tmm3, implicit-def dead $tmm4, implicit-def dead $tmm5, implicit-def dead $tmm6, implicit-def dead $tmm7 :: (load (s512) from %stack.0, align 4)
   ; CHECK-NEXT:   [[MOVSX64rr32_:%[0-9]+]]:gr64_nosp = MOVSX64rr32 [[COPY83]].sub_32bit
-  ; CHECK-NEXT:   [[COPY83]].sub_32bit:gr64_with_sub_8bit = nsw SUB32rr [[COPY83]].sub_32bit, [[SUB32rr]], implicit-def dead $eflags
+  ; CHECK-NEXT:   [[COPY83]].sub_32bit:gr64_with_sub_8bit = SUB32rr [[COPY83]].sub_32bit, [[SUB32rr]], implicit-def dead $eflags
   ; CHECK-NEXT:   undef %14.sub_32bit:gr64_with_sub_8bit = MOVZX32rr16 [[COPY83]].sub_16bit
   ; CHECK-NEXT:   ADD64mr %stack.1, 1, $noreg, 0, $noreg, %14, implicit-def dead $eflags :: (store (s64) into %stack.1)
   ; CHECK-NEXT:   undef %61.sub_32bit:gr64_with_sub_8bit = COPY %14.sub_32bit
   ; CHECK-NEXT:   %61.sub_32bit:gr64_with_sub_8bit = IMUL32rr %61.sub_32bit, [[COPY85]].sub_32bit, implicit-def dead $eflags
   ; CHECK-NEXT:   [[LEA64_32r1:%[0-9]+]]:gr32 = LEA64_32r $noreg, 4, [[COPY85]], 0, $noreg
   ; CHECK-NEXT:   [[MOVSX64rr32_1:%[0-9]+]]:gr64 = MOVSX64rr32 [[LEA64_32r1]]
   ; CHECK-NEXT:   MOV64mr %stack.3, 1, $noreg, 0, $noreg, [[MOVSX64rr32_1]] :: (store (s64) into %stack.3)
   ; CHECK-NEXT:   [[MOVSX64rr32_2:%[0-9]+]]:gr64_nosp = MOVSX64rr32 [[COPY85]].sub_32bit

llvm/test/Transforms/IndVarSimplify/pr63763.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
 ; RUN: opt -S -passes='print<scalar-evolution>,indvars' < %s 2>/dev/null | FileCheck %s
 
-; FIXME: This is a miscompile.
 ; We should use %invariant.op.us rather than %invariant.op for the exit
 ; value expansion. They have the same SCEV, but %invariant.op is more
 ; poisonous.
 define i32 @test(i1 %c) {
 ; CHECK-LABEL: define i32 @test
 ; CHECK-SAME: (i1 [[C:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[C]], i32 33, i32 0
 ; CHECK-NEXT:    [[SUB_I:%.*]] = add nsw i32 [[SEL]], -68
 ; CHECK-NEXT:    [[SHL_I:%.*]] = shl nuw nsw i32 1, [[SUB_I]]
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[SHL_I]], [[SEL]]
 ; CHECK-NEXT:    [[SEXT:%.*]] = shl i32 [[ADD]], 24
 ; CHECK-NEXT:    [[CONV2:%.*]] = ashr exact i32 [[SEXT]], 24
-; CHECK-NEXT:    [[INVARIANT_OP:%.*]] = sub i32 7, [[CONV2]]
+; CHECK-NEXT:    [[INVARIANT_OP:%.*]] = sub nsw i32 7, [[CONV2]]
 ; CHECK-NEXT:    call void @use(i32 [[INVARIANT_OP]])
+; CHECK-NEXT:    [[SEXT_US:%.*]] = shl i32 [[SEL]], 24
+; CHECK-NEXT:    [[CONV2_US:%.*]] = ashr exact i32 [[SEXT_US]], 24
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    br i1 true, label [[EXIT:%.*]], label [[LOOP]]
 ; CHECK:       exit:
-; CHECK-NEXT:    ret i32 [[INVARIANT_OP]]
+; CHECK-NEXT:    [[INVARIANT_OP_US:%.*]] = sub nsw i32 7, [[CONV2_US]]
+; CHECK-NEXT:    ret i32 [[INVARIANT_OP_US]]
 ;
 entry:
   %sel = select i1 %c, i32 33, i32 0
   %sub.i = add nsw i32 %sel, -68
   %shl.i = shl nuw nsw i32 1, %sub.i
   %add = add nsw i32 %shl.i, %sel
   %sext = shl i32 %add, 24
   %conv2 = ashr exact i32 %sext, 24

llvm/test/Transforms/LoopUnroll/runtime-exit-phi-scev-invalidation.ll

 ; CHECK-NEXT:    br label [[OUTER_MIDDLE]]
 ; CHECK:       outer.middle:
 ; CHECK-NEXT:    [[V_LCSSA1:%.*]] = phi i32 [ [[V_LCSSA1_PH]], [[OUTER_MIDDLE_UNR_LCSSA]] ], [ [[V_LCSSA1_PH2]], [[OUTER_MIDDLE_EPILOG_LCSSA]] ]
 ; CHECK-NEXT:    [[C_3:%.*]] = icmp ugt i32 [[V_LCSSA1]], 0
 ; CHECK-NEXT:    br i1 [[C_3]], label [[INNER_2_PREHEADER:%.*]], label [[EXIT:%.*]]
 ; CHECK:       inner.2.preheader:
 ; CHECK-NEXT:    br label [[INNER_2]]
 ; CHECK:       inner.2:
-; CHECK-NEXT:    [[OUTER_IV_NEXT]] = add nuw nsw i64 [[OUTER_IV]], 1
+; CHECK-NEXT:    [[OUTER_IV_NEXT]] = add i64 [[OUTER_IV]], 1
 ; CHECK-NEXT:    br label [[OUTER_HEADER]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ; CHECK:       exit.deopt.loopexit:
 ; CHECK-NEXT:    br label [[EXIT_DEOPT:%.*]]
 ; CHECK:       exit.deopt.loopexit3:
 ; CHECK-NEXT:    br label [[EXIT_DEOPT]]
 ; CHECK:       exit.deopt: